fix(llama-cpp): cap single-pass embedding batch to fit VRAM (#10695)

* fix(llama-cpp): cap single-pass embedding batch to fit VRAM

Embedding/score/rerank all decode or pool the whole input in one physical
batch, so EffectiveBatchSize sized the batch to the full context window. For
a large context that makes n_ubatch huge, and the per-device CUDA compute
buffer (forward-graph scratch, ~n_ubatch * n_ctx, NOT split across GPUs)
balloons into multi-GiB: a large-context embedding model then aborts on load
(exitCode=-1) even with plenty of free VRAM. Reproduced with qwen3-embedding-4b
(context 40960 -> n_batch 40960 -> abort) and qwen3-embedding-0.6b
(n_batch 8192); pinning batch:512 avoided it.

This is the same root cause as issue #10485 (a large context turns the batch
into multi-GiB of scratch that must fit on a SINGLE card), but the single-pass
path bypassed the VRAM headroom guard the config layer already had — it
returned the unbounded context as the batch with no GPU awareness.

Make the single-pass batch VRAM-aware: cap it to the largest batch whose
compute buffer fits the per-device VRAM headroom, clamped to
[DefaultPhysicalBatch, ctx], reusing the existing computeBufferBytesPerCell and
headroom-divisor math (no duplication). Unknown per-device VRAM (0) stays
conservative (DefaultPhysicalBatch, not the context) so a detection gap can't
OOM. The GPU is resolved through an injectable package var (config.LocalGPU,
backed by sync.Once-cached xsysinfo detection) so the per-request router call
stays cheap and tests inject a deterministic device. Explicit batch: still
wins. An input longer than the cap can no longer be pooled in one pass — the
accepted tradeoff, since a batch that OOMs the device processes nothing.

Assisted-by: Claude:claude-opus-4-8 golangci-lint go-test
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(config): single-pass batch follows context on unknown VRAM

The single-pass (embedding/score/rerank) batch cap must only shrink the batch
when the per-device VRAM ceiling is KNOWN. On unknown VRAM (CPU-only or a GPU
detection gap) SinglePassBatchForContext returned DefaultPhysicalBatch, which
under-sized the batch below the context — over-trimming score/embed/rerank
inputs (the modelTokenTrim middleware regression) with no OOM benefit on CPU
where the compute buffer lives in system RAM. Return the full context instead,
preserving the original single-pass behavior; the VRAM cap stays a downward
safety that only engages when VRAM is known.

Assisted-by: Claude:claude-opus-4-8 [go-test go-vet]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@localai.io>

Author: localai-bot

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
 }

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,61 @@ 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("returns the full context when per-device VRAM is unknown", func() {
+		// Unknown VRAM (CPU / detection gap) preserves the original single-pass
+		// behavior: batch follows context. The VRAM cap is a downward safety that
+		// only engages when the per-device ceiling is known — clamping here would
+		// re-break single-pass pooling and over-trim inputs, with no OOM benefit on
+		// CPU where the compute buffer lives in system RAM.
+		localGPU = func() config.GPU { return config.GPU{VRAM: 0} }
+		Expect(EffectiveBatchSize(singlePassCfg(40960))).To(Equal(40960))
+	})
+
+	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

core/config/hardware_defaults.go

@@ -149,6 +149,51 @@ 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 — CPU-only or a detection gap) returns the full context,
+// preserving the original single-pass behavior (batch follows context): the cap
+// is a DOWNWARD safety that only engages when the per-device ceiling is known.
+// Returning a smaller batch on unknown VRAM would re-break single-pass pooling
+// (n_tokens > n_batch) and over-trim score/embed/rerank inputs, with no OOM
+// benefit on CPU where the buffer lives in system RAM.
+func SinglePassBatchForContext(g GPU, ctx int) int {
+	if ctx <= DefaultPhysicalBatch {
+		return DefaultPhysicalBatch
+	}
+	if g.VRAM == 0 {
+		return ctx
+	}
+	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 +299,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).

core/config/hardware_defaults_internal_test.go

@@ -46,3 +46,41 @@ 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("returns the full context (unclamped) when per-device VRAM is unknown", func() {
+		// Unknown VRAM (CPU / detection gap) preserves the original single-pass
+		// behavior — the cap is a downward safety that only engages when VRAM is
+		// known. Clamping here would over-trim score/embed/rerank inputs.
+		Expect(SinglePassBatchForContext(GPU{VRAM: 0}, 40960)).To(Equal(40960))
+	})
+})