Skip to content

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

Merged
mudler merged 2 commits into
masterfrom
fix/embedding-batch-vram-cap
Jul 6, 2026
Merged

fix(llama-cpp): cap single-pass embedding batch to fit VRAM#10695
mudler merged 2 commits into
masterfrom
fix/embedding-batch-vram-cap

Conversation

@localai-bot

Copy link
Copy Markdown
Collaborator

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 with
qwen3-embedding-4b (context=40960n_batch=40960 → abort) and
qwen3-embedding-0.6b (n_batch=8192). Pinning batch: 512 avoids it. The
failure 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.EffectiveBatchSize returned the full EffectiveContextSize as the
physical batch (n_batch / n_ubatch) with no upper bound and no VRAM awareness.
The CUDA compute buffer scales ~ n_ubatch * n_ctx and is allocated per device
(it can't be split across GPUs), so a full ctx-sized n_ubatch turns into multi-GiB
of 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 (EffectiveBatchSize didn'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 (no
duplication). EffectiveBatchSize calls it for the single-pass path:

  • Explicit c.Batch != 0 still short-circuits and wins.
  • Single-pass with ctx > DefaultBatchSizeSinglePassBatchForContext(localGPU(), ctx).
  • Unknown per-device VRAM (0) stays conservative (returns DefaultPhysicalBatch,
    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 localGPU seam in hardware_defaults.go). config.LocalGPU
wraps the sync.Once-cached xsysinfo detection, so the per-request call from the
router's prompt trimmer (modelTokenTrim) stays cheap. The exported
EffectiveBatchSize signature is unchanged, so both callers (grpcModelOpts and
route_model.go) need no edits.

The original single-pass intent is documented in a WHY comment: single-pass wants
batch >= input so pooling is one pass and avoids GGML_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 — new EffectiveBatchSize VRAM cap
    Describe: large embedding context (40960) on a small 20 GiB device → bounded batch
    < context and >= DefaultBatchSize; explicit batch: wins; unknown VRAM (0) →
    conservative (DefaultBatchSize, not context); non-single-pass → DefaultBatchSize.
    The existing grpcModelOpts NBatch cases now inject a deterministic high-VRAM GPU
    (they became VRAM-dependent).
  • core/config/hardware_defaults_internal_test.go — new SinglePassBatchForContext
    Describe: 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 DefaultPhysicalBatch on 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/ (and make lint new-from-master) — 0 issues.

Notes

  • The local pre-commit coverage gate (make test-coverage-check) was skipped at the
    maintainer's request; CI runs the full lint + coverage suite.
  • Per the repo's AI-assist policy, this commit carries an Assisted-by: trailer and
    intentionally has no Signed-off-by — a maintainer must add their own
    Signed-off-by to certify DCO.

mudler added 2 commits July 6, 2026 07:52
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>
@mudler mudler force-pushed the fix/embedding-batch-vram-cap branch from 9a0b4e6 to 6179bec Compare July 6, 2026 07:52
@mudler mudler merged commit 85f5267 into master Jul 6, 2026
60 checks passed
@mudler mudler deleted the fix/embedding-batch-vram-cap branch July 6, 2026 10:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants