diff --git a/core/backend/options.go b/core/backend/options.go index 028ef3062..8fc47f231 100644 --- a/core/backend/options.go +++ b/core/backend/options.go @@ -223,13 +223,24 @@ func EffectiveContextSize(c config.ModelConfig) int { return DefaultContextSize } +// localGPU resolves the device that will run the model, for single-pass batch +// sizing. It is a package var so tests inject a deterministic device; production +// reads config.LocalGPU, whose detection is sync.Once-cached in xsysinfo — so the +// per-request call from the router's prompt trimmer (modelTokenTrim) stays cheap. +var localGPU = config.LocalGPU + // EffectiveBatchSize is the single-decode batch the backend will run with. // Score, embedding and rerank all process the whole input in one pass: score // decodes prompt+candidate (asserts n_tokens <= n_batch), and embedding/rerank -// pool over the full sequence in one physical batch (n_ubatch). So the batch -// is sized to the context — anything that fits the context fits one pass, +// pool over the full sequence in one physical batch (n_ubatch). Ideally the batch +// covers the whole context so any input that fits the context fits one pass, // avoiding both the GGML_ASSERT crash and the "input is too large to process" -// error. Explicit `batch:` always wins. +// error — BUT a full ctx-sized n_ubatch makes the per-device CUDA compute buffer +// multi-GiB (it scales ~ n_ubatch * n_ctx and can't be split across GPUs), so a +// large-context embedding model aborts on load with free VRAM to spare (#10485). +// So we cap the batch to the largest that fits the per-device VRAM headroom; an +// input longer than that cap is the accepted tradeoff (it can't be pooled in one +// pass, but the load no longer OOMs). Explicit `batch:` always wins. func EffectiveBatchSize(c config.ModelConfig) int { if c.Batch != 0 { return c.Batch @@ -238,7 +249,7 @@ func EffectiveBatchSize(c config.ModelConfig) int { c.HasUsecases(config.FLAG_EMBEDDINGS) || c.HasUsecases(config.FLAG_RERANK) if ctx := EffectiveContextSize(c); singlePass && ctx > DefaultBatchSize { - return ctx + return config.SinglePassBatchForContext(localGPU(), ctx) } return DefaultBatchSize } diff --git a/core/backend/options_internal_test.go b/core/backend/options_internal_test.go index 022d7b1d9..63b15d7c0 100644 --- a/core/backend/options_internal_test.go +++ b/core/backend/options_internal_test.go @@ -103,6 +103,19 @@ var _ = Describe("grpcModelOpts NBatch", func() { threads := 1 ctx := 4096 + // The single-pass batch is now VRAM-aware, so inject a deterministic GPU with + // ample per-device VRAM: at these small contexts the compute buffer fits + // easily, so EffectiveBatchSize returns the full context (the pre-#10485 + // behaviour these cases assert). Without injection the value would depend on + // the CI host's real (often unknown) VRAM. + const gib = uint64(1) << 30 + var origLocalGPU func() config.GPU + BeforeEach(func() { + origLocalGPU = localGPU + localGPU = func() config.GPU { return config.GPU{VRAM: 119 * gib} } + }) + AfterEach(func() { localGPU = origLocalGPU }) + It("defaults to 512 for an ordinary model", func() { cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}} opts := grpcModelOpts(cfg, "/tmp/models") @@ -162,6 +175,58 @@ var _ = Describe("grpcModelOpts NBatch", func() { }) }) +// Guards the VRAM-aware cap on the single-pass (embedding/score/rerank) batch: +// a large context must not turn n_ubatch into a multi-GiB compute buffer that +// aborts the load on a device with free VRAM (issue #10485). The GPU is injected +// via the localGPU package var so the cap is deterministic without a real device. +var _ = Describe("EffectiveBatchSize VRAM cap", func() { + const gib = uint64(1) << 30 + embeddings := config.FLAG_EMBEDDINGS + threads := 1 + + var origLocalGPU func() config.GPU + BeforeEach(func() { origLocalGPU = localGPU }) + AfterEach(func() { localGPU = origLocalGPU }) + + singlePassCfg := func(ctx int) config.ModelConfig { + return config.ModelConfig{ + Threads: &threads, + LLMConfig: config.LLMConfig{ContextSize: &ctx}, + KnownUsecases: &embeddings, + } + } + + It("caps a large embedding context to a batch below the context but at least the default", func() { + // Reproduces qwen3-embedding-4b: context 40960 on a modest 20 GiB card. + // Full-context n_ubatch=40960 aborts; the cap must fit the VRAM headroom. + localGPU = func() config.GPU { return config.GPU{VRAM: 20 * gib} } + batch := EffectiveBatchSize(singlePassCfg(40960)) + Expect(batch).To(BeNumerically(">=", DefaultBatchSize)) + Expect(batch).To(BeNumerically("<", 40960)) + }) + + It("keeps an explicit batch even with a large context and small VRAM", func() { + localGPU = func() config.GPU { return config.GPU{VRAM: 20 * gib} } + cfg := singlePassCfg(40960) + cfg.Batch = 512 + Expect(EffectiveBatchSize(cfg)).To(Equal(512)) + }) + + It("stays conservative (default, not the context) when per-device VRAM is unknown", func() { + // A detection gap (VRAM 0) must not fall back to the unbounded context — + // that's exactly what would OOM the load. + localGPU = func() config.GPU { return config.GPU{VRAM: 0} } + Expect(EffectiveBatchSize(singlePassCfg(40960))).To(Equal(DefaultBatchSize)) + }) + + It("returns the default batch for a non-single-pass model regardless of VRAM", func() { + localGPU = func() config.GPU { return config.GPU{VRAM: 20 * gib} } + ctx := 40960 + cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}} + Expect(EffectiveBatchSize(cfg)).To(Equal(DefaultBatchSize)) + }) +}) + // Guards the generic chat_template_kwargs forwarding: the model config map plus any // per-request metadata overrides are merged, coerced, and serialised into the // backend metadata blob that llama.cpp reads. Client metadata also overrides the diff --git a/core/config/hardware_defaults.go b/core/config/hardware_defaults.go index 81bc9fc7f..5ddf47d9e 100644 --- a/core/config/hardware_defaults.go +++ b/core/config/hardware_defaults.go @@ -149,6 +149,47 @@ func largeContextForDevice(g GPU, ctx int) bool { return extra > g.VRAM/blackwellBatchHeadroomDivisor } +// SinglePassBatchForContext caps the physical batch (n_batch / n_ubatch) for a +// single-pass load — embedding, score and rerank all decode/pool the whole input +// in ONE physical batch, so they want a batch >= the input length to avoid the +// GGML_ASSERT(n_tokens <= n_batch) abort and the "input is too large to process" +// error. The naive choice is batch == context, but n_ubatch == context turns the +// per-device CUDA compute buffer (which scales ~ n_ubatch * n_ctx and is NOT +// split across GPUs) into multi-GiB of scratch that must fit on a SINGLE card, so +// a large-context embedding model aborts on load (exitCode=-1) even with plenty +// of free VRAM — the same #10485 root cause the Blackwell batch boost guards +// against, which the single-pass path previously bypassed entirely. +// +// So instead of the full context we return the LARGEST batch whose compute buffer +// fits the per-device VRAM headroom (VRAM / blackwellBatchHeadroomDivisor), +// clamped to [DefaultPhysicalBatch, ctx]. The tradeoff: an input longer than the +// returned cap can no longer be pooled in a single pass — but a batch that OOMs +// the device processes nothing at all. +// +// g.VRAM must be the PER-DEVICE ceiling (smallest device on a multi-GPU host). +// VRAM 0 (unknown) stays conservative and returns DefaultPhysicalBatch rather +// than the unbounded context, so a detection gap can't OOM the load. +func SinglePassBatchForContext(g GPU, ctx int) int { + if ctx <= DefaultPhysicalBatch { + return DefaultPhysicalBatch + } + if g.VRAM == 0 { + return DefaultPhysicalBatch + } + perBatchCell := uint64(ctx) * computeBufferBytesPerCell + if perBatchCell == 0 { + return DefaultPhysicalBatch + } + batchCap := int(g.VRAM / blackwellBatchHeadroomDivisor / perBatchCell) + if batchCap < DefaultPhysicalBatch { + return DefaultPhysicalBatch + } + if batchCap > ctx { + return ctx + } + return batchCap +} + // IsManagedPhysicalBatch reports whether n is a value PhysicalBatch assigns. // Callers that re-tune a value chosen by an upstream host (the distributed // router correcting the frontend's guess) use this to avoid clobbering an @@ -254,6 +295,14 @@ var localGPU = func() GPU { } } +// LocalGPU exposes the locally-detected device descriptor to other packages +// (e.g. core/backend's single-pass batch sizing) so they resolve the same +// per-device VRAM this package's heuristics reason about. It goes through the +// injectable localGPU var, so a config-package test seam also affects callers. +func LocalGPU() GPU { + return localGPU() +} + // ApplyHardwareDefaults fills ModelConfig values that depend on the target GPU // and were left unset by the user. Currently: a larger physical batch on // Blackwell. Explicit config always wins (we only touch zero values). diff --git a/core/config/hardware_defaults_internal_test.go b/core/config/hardware_defaults_internal_test.go index d6878c86e..d69bca440 100644 --- a/core/config/hardware_defaults_internal_test.go +++ b/core/config/hardware_defaults_internal_test.go @@ -46,3 +46,38 @@ var _ = Describe("SetDefaults hardware defaults (single-instance)", func() { Expect(cfg.Batch).To(Equal(1024)) }) }) + +// SinglePassBatchForContext is the VRAM-aware cap for the single-pass +// (embedding/score/rerank) batch — the compute buffer scales ~ n_ubatch * n_ctx +// and must fit a single device, so a large context can't take the full context +// as its batch (issue #10485). +var _ = Describe("SinglePassBatchForContext", func() { + const gib = uint64(1) << 30 + + It("returns the default when the context is at or below the default batch", func() { + Expect(SinglePassBatchForContext(GPU{VRAM: 119 * gib}, DefaultPhysicalBatch)).To(Equal(DefaultPhysicalBatch)) + Expect(SinglePassBatchForContext(GPU{VRAM: 119 * gib}, 256)).To(Equal(DefaultPhysicalBatch)) + }) + + It("returns the full context when the compute buffer fits ample VRAM", func() { + // 4096 ctx on 119 GiB: the compute buffer is tiny, so the batch covers + // the whole context (single-pass pooling in one physical batch). + Expect(SinglePassBatchForContext(GPU{VRAM: 119 * gib}, 4096)).To(Equal(4096)) + }) + + It("caps below the context when a large context would overflow the VRAM headroom", func() { + batch := SinglePassBatchForContext(GPU{VRAM: 20 * gib}, 40960) + Expect(batch).To(BeNumerically(">=", DefaultPhysicalBatch)) + Expect(batch).To(BeNumerically("<", 40960)) + // The compute buffer for the capped batch must fit VRAM/headroom. + Expect(uint64(batch) * 40960 * computeBufferBytesPerCell).To(BeNumerically("<=", (20*gib)/blackwellBatchHeadroomDivisor)) + }) + + It("never caps below the default batch even when VRAM is very tight", func() { + Expect(SinglePassBatchForContext(GPU{VRAM: 1 * gib}, 40960)).To(Equal(DefaultPhysicalBatch)) + }) + + It("stays conservative (default) when per-device VRAM is unknown", func() { + Expect(SinglePassBatchForContext(GPU{VRAM: 0}, 40960)).To(Equal(DefaultPhysicalBatch)) + }) +})