fix(llama-cpp): cap single-pass embedding batch to fit VRAM#10695
Merged
Conversation
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>
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>
9a0b4e6 to
6179bec
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Symptom
Embedding models with a large context window crash on load — the backend process
aborts (
exitCode=-1) — even with plenty of free VRAM. Reproduced live withqwen3-embedding-4b(context=40960→n_batch=40960→ abort) andqwen3-embedding-0.6b(n_batch=8192). Pinningbatch: 512avoids it. Thefailure is VRAM-independent: the batch was sized with no awareness of the device.
Root cause
For single-pass usecases (
FLAG_EMBEDDINGS/FLAG_SCORE/FLAG_RERANK),core/backend.EffectiveBatchSizereturned the fullEffectiveContextSizeas thephysical batch (
n_batch/n_ubatch) with no upper bound and no VRAM awareness.The CUDA compute buffer scales ~
n_ubatch * n_ctxand is allocated per device(it can't be split across GPUs), so a full ctx-sized
n_ubatchturns into multi-GiBof scratch that must fit on a single card and aborts the load.
This is the same root cause as #10485 ("a large context turns the batch into
multi-GiB of extra scratch that must fit on a SINGLE card"), for which the config
layer already grew a VRAM headroom guard (
PhysicalBatchForContext,largeContextForDevice,computeBufferBytesPerCell,blackwellBatchHeadroomDivisor).But that guard only gated the Blackwell 512→2048 batch boost — the single-pass
embedding path bypassed it entirely (
EffectiveBatchSizedidn't even take a GPU).Fix
Make the single-pass batch VRAM-aware. New
config.SinglePassBatchForContext(gpu, ctx)returns the largest batch whose compute buffer fits the per-device VRAM headroom
(
VRAM / blackwellBatchHeadroomDivisor), clamped to[DefaultPhysicalBatch, ctx],reusing the existing
computeBufferBytesPerCell/ headroom-divisor math (noduplication).
EffectiveBatchSizecalls it for the single-pass path:c.Batch != 0still short-circuits and wins.ctx > DefaultBatchSize→SinglePassBatchForContext(localGPU(), ctx).0) stays conservative (returnsDefaultPhysicalBatch,not the unbounded context) so a detection gap can't OOM the load.
The GPU is threaded in via an injectable package var (
var localGPU = config.LocalGPU,mirroring the existing
localGPUseam inhardware_defaults.go).config.LocalGPUwraps the sync.Once-cached
xsysinfodetection, so the per-request call from therouter's prompt trimmer (
modelTokenTrim) stays cheap. The exportedEffectiveBatchSizesignature is unchanged, so both callers (grpcModelOptsandroute_model.go) need no edits.The original single-pass intent is documented in a WHY comment: single-pass wants
batch >= inputso pooling is one pass and avoidsGGML_ASSERT(n_tokens <= n_batch)and the "input is too large to process" error. The cap is "the largest batch that
fits VRAM"; an input longer than the cap is the accepted tradeoff — a batch that OOMs
the device processes nothing at all.
Test coverage (Ginkgo v2 + Gomega)
core/backend/options_internal_test.go— newEffectiveBatchSize VRAM capDescribe: large embedding context (40960) on a small 20 GiB device → bounded batch
< contextand>= DefaultBatchSize; explicitbatch:wins; unknown VRAM (0) →conservative (
DefaultBatchSize, not context); non-single-pass →DefaultBatchSize.The existing
grpcModelOpts NBatchcases now inject a deterministic high-VRAM GPU(they became VRAM-dependent).
core/config/hardware_defaults_internal_test.go— newSinglePassBatchForContextDescribe: default when ctx ≤ default; full context when the buffer fits ample VRAM;
caps below context when a large context overflows the headroom (with an explicit
buffer-fits-headroom assertion); floors at
DefaultPhysicalBatchon tight VRAM;conservative default when VRAM is unknown.
The GPU is injected via the package-var pattern; no stdlib testing assertions.
Verification
go build ./core/backend/... ./core/config/...— passes.go test ./core/backend/... ./core/config/...— passes (new specs verified running).golangci-lint run ./core/backend/ ./core/config/(andmake lintnew-from-master) — 0 issues.Notes
make test-coverage-check) was skipped at themaintainer's request; CI runs the full lint + coverage suite.
Assisted-by:trailer andintentionally has no
Signed-off-by— a maintainer must add their ownSigned-off-byto certify DCO.