feat(llm): --low-vram mode with on-demand reclaim on embed VRAM pressure

[brettdavies]

--low-vram: share a GPU with another large model

Why

qmd query (and any other path that runs the full pipeline) loads three GGUF models into a single process: embed (~320 MB), generate (~2 GB), rerank (~2.3 GB). Once all three are resident, peak VRAM sits around 5.4 GB. On a GPU that's already shared with another model (say a 24 GB card running a ~20 GB Ollama), there isn't enough free VRAM left, and rerank context creation fails with Failed to create any rerank context. That's the failure mode tracked in #275 (closed during the v2.5.1 backlog cleanup; reproducible on hardware ranging from a 2 GB GTX 960M to a 6 GB RTX 3060 per reporters there, plus the 24 GB / Ollama-coexistence case below).

The three pipeline stages (expand → embed → search → rerank) are inherently sequential, so the win is straightforward: keep the tiny embed model resident, dispose the heavy generate and rerank models after each use, reload them on demand. Peak drops to ~2.6 GB. The cost is per-stage load latency.

Measured on an RTX 3090 Ti running alongside Ollama Gemma 4 26B (~20.3 GB VRAM):

Mode	Query time	Peak qmd VRAM	Coexists with 20 GB Ollama?
Bare qmd query (cold)	33 s	~3 GB (churning)	partially; CPU offload thrash + occasional rerank failures
qmd query (warm, default)	3.0 s	5.4 GB	no (25.7 GB total > 24 GB free)
qmd query --low-vram	5.6 s	2.6 GB	yes (22.6 GB total, headroom)

22K-file collection, hybrid query, full pipeline (expand + vec + rerank).

What this PR does

Adds lowVram to LlamaCppConfig, surfaced as a global --low-vram CLI flag (any subcommand) and a QMD_LOW_VRAM=1 env var. Naming follows the existing engine-knob pattern (QMD_FORCE_CPU, QMD_LLAMA_GPU).

When enabled, LlamaCpp:

1. Disposes the generate model in a finally block after each expandQuery call.
2. Disposes the rerank model and ranking contexts in a finally block after each rerank call.
3. Serializes overlapping expandQuery calls through an internal promise chain so a dispose call can never race with another caller's in-flight use of the same model. Same for rerank.
4. Leaves embed and embedBatch running in parallel as before: the embed model stays resident.
5. Adaptive context sizing. Before allocating rerank contexts, probes getVramState(). Free VRAM at or above 1.5 GB keeps the 4096-token default; below that shrinks to 2048 tokens (~568 MB allocation cost instead of ~1146 MB). The QMD_RERANK_CONTEXT_SIZE env override still wins when set. CPU mode keeps the 4096 default. Document truncation in rerankImpl uses the actually allocated size, not a static constant, so chunks are never under-truncated when the context was shrunk under pressure.
6. Two-pass on-demand reclaim. If embed, expandQuery, or rerank fails with an insufficient-VRAM error from node-llama-cpp, the wrapper drains the unrelated chains, disposes everything that isn't needed for the failing operation (generate model, rerank model, embed contexts), and retries once. If the retry also fails for VRAM AND the failing operation is not embed, a second pass drops the embed model itself (~568 MB) and retries one more time. Outside lowVram mode this is a pass-through, so default-path callers see zero overhead. Detail in the reclaim section below.

The two chains are independent: expandQuery and rerank against their separate heavy models can run in parallel against each other.

Because the flag is global and the constructor reads QMD_LOW_VRAM, it works everywhere LlamaCpp is constructed (qmd query, qmd embed, qmd mcp --http --daemon, qmd vsearch, etc.) without per-subcommand plumbing.

Two-pass reclaim

The reclaim path addresses the case where the catch-after-throw approach by itself does not free enough VRAM in one shot. The original observation that the lib's own pre-allocation VRAM check (InsufficientMemoryError: "A context size of N is too large for the available VRAM") is the right signal to react to still applies — we do not race to predict it. What is new is what we do after catching it.

First pass, on any insufficient-VRAM error:

1. Drain the other operations' serialization chains so their finally-dispose has run.
2. Dispose the heavy models we don't need for the failing operation: generate (for non-expandQuery callers), rerank (for non-rerank callers).
3. Dispose embed contexts (not the model) when the failing operation is not embed itself. Contexts are ~143 MB each at the embed default; dropping them is cheap. The model stays resident because reload cost dominates the marginal VRAM win on cards with reasonable headroom.
4. Retry the operation once.

If the retry also fails with an insufficient-VRAM error AND the failing operation is not embed, a second pass evicts the embed model itself (~568 MB) and retries one final time. Embed reclaim short-circuits the second pass — it cannot make progress without the model resident. Worst-case attempt count is capped at three (initial + first-pass retry + second-pass retry); no infinite retry loop is possible.

Each dispose runs best-effort: a failure is logged with the dispose label, but the remaining disposes still run. A retry against a partial reclaim is still more likely to succeed than no retry at all, and the original VRAM error is more useful than a downstream dispose error.

Per-context VRAM cost model in computeParallelism now derives from the actually chosen rerank context size (~0.28 MB/token, measured at both 2048 and 4096) instead of the previous 1000 MB hardcode that assumed 2048. Parallelism estimates scale correctly across both sizes.

isInsufficientVramError matches node-llama-cpp's "too large for the available VRAM" template (single-sequence and multi-sequence variants), our own "Failed to create any (embedding|rerank) context" wrapper messages, generic "out of memory" (CUDA / GGML), and "insufficient VRAM" safety-check variants. Conservative on purpose — too broad and we'd evict on unrelated failures; too narrow and we'd miss the case the feature targets. Pinned against a curated list of real lib strings plus a red-team negative list (model-not-found, unsupported-architecture, network errors, substring-greediness cases) so a future lib message change is caught as a unit-test failure before it ships.

Test coverage

File	What it covers
test/llm-low-vram.test.ts	Concurrency: lowVram serialization prevents dispose-mid-use races, two-pass reclaim event ordering for all three callers, second-pass cap (max 3 attempts), best-effort dispose continues when a helper throws, non-VRAM errors short-circuit, post-first-pass non-VRAM errors are terminal, in-flight chains drain before dispose.
test/llm-adaptive-rerank.test.ts	Adaptive sizing: env override wins, CPU mode skips probe, boundary behavior at threshold, Ollama-hogged scenario picks 2048, probe failure falls back to default, rerankImpl uses the chosen size, computeParallelism budget reflects the chosen size, lifecycle reset on dispose.
test/llm-vram-error-pinning.test.ts	Tripwire: positives the regex must catch (InsufficientMemoryError template variants, synthetic codebase strings, CUDA / GGML OOM), red-team negatives that must NOT match (model-not-found, wrong-architecture, network errors, substring greediness), input-shape handling.
test/llm-reclaim-integration.test.ts	Opt-in real-GPU integration (QMD_RECLAIM_INTEGRATION=1). Oversized context (1 M tokens) forces a real allocation failure, asserts the 3-attempt cap holds and isInsufficientVramError classifies the live InsufficientMemoryError. Missing-model-file case asserts non-VRAM failures do not retry.

The integration test is gated because it downloads the rerank model on first run and exercises real GPU allocation; it cannot meaningfully execute in CI. All synthetic tests run unconditionally.

Architecture

d21ddd3  Brett  test(llm): single-step assertions to satisfy the anti-slop lint fence
9af2ed0  Brett  test(llm): isolate low-vram suite from ambient CI and QMD_LOW_VRAM
515d3af  Brett  test(llm): opt-in real-GPU integration test for the two-pass reclaim path
493dc3c  Brett  feat(llm): adaptive rerank context size with correct parallelism cost
a7a3230  Brett  feat(llm): two-pass VRAM reclaim with embed-model escalation
32cfefb  Brett  test(llm): pin isInsufficientVramError against real node-llama-cpp error strings
249d592  Brett  feat(llm): extend reclaim to expandQuery and rerank
868972e  Brett  feat(llm): reclaim heavies on embed VRAM pressure
e490614  Brett  feat(llm): add low-vram mode that loads one heavy model at a time

Changelog

Added

• --low-vram global flag and QMD_LOW_VRAM=1 env var. Disposes the generate and rerank models after each use while keeping the embed model resident. Peak VRAM drops from ~5.4 GB to ~2.6 GB at the cost of per-stage load latency. Intended for shared GPUs (Ollama coexistence) and small-VRAM cards where loading all three models at once would fail.
• Two-pass on-demand reclaim when embed, expandQuery, or rerank hits an insufficient-VRAM error from node-llama-cpp. First pass disposes the other models and the embed contexts; second pass evicts the embed model itself. Worst-case three attempts. Off outside --low-vram.
• Adaptive rerank context size. Picks 2048 instead of the 4096 default when free GPU VRAM is below ~1.5 GB. The QMD_RERANK_CONTEXT_SIZE env override still wins.
• isInsufficientVramError exported from src/llm.ts for direct unit-testing.

Changed

• computeParallelism budgets each rerank context against its actual chosen size (~0.28 MB/token) instead of a fixed 1000 MB-per-context estimate. Parallelism scales correctly across both 2048 and 4096 sizes.
• Long-document truncation in rerankImpl now uses the actually allocated context size, so chunks are never under-truncated when the context was shrunk under VRAM pressure.

Fixed

• In --low-vram mode, a VRAM failure creating the first rerank context propagates to withReclaim (which evicts heavies and retries) instead of falling into the warn-and-skip path fix(llm): surface the real error when rerank context creation fails #866 added for the default mode. The original error is rethrown as-is; isInsufficientVramError recognizes the underlying node-llama-cpp messages directly (pinned by test/llm-vram-error-pinning.test.ts). Default-mode behavior is unchanged: reranking still degrades gracefully with a warning.

Type of Change

• feat: New feature (non-breaking change which adds functionality)
• fix: Bug fix
• BREAKING CHANGE: Breaking API change

Related Issues/Stories

• Story: n/a
• Issue: closes Low-VRAM GPUs: evict idle models before loading reranker #275
• Architecture: n/a
• Related PRs: continues the lowVram-coexistence direction; independent of feat(serve): qmd serve shared model server + RemoteQMD client (supersedes #511) #663 (qmd serve): the two merge in either order, and the combined serve+low-vram shape lives in feat: low-vram engine + qmd serve combined (production shape of #662 + #663) #927.

Testing

• Unit tests added/updated
• Integration tests added/updated (opt-in real-GPU)
• Manual testing completed
• All tests passing

Test Summary:

• test/llm-low-vram.test.ts: 23 cases (concurrency, two-pass reclaim, best-effort, non-VRAM termination).
• test/llm-adaptive-rerank.test.ts: 18 cases (env override, CPU mode, boundary, lifecycle).
• test/llm-vram-error-pinning.test.ts: 21 cases (positive catalog, red-team negatives, input shapes).
• test/llm-reclaim-integration.test.ts: 2 opt-in real-GPU cases gated by QMD_RECLAIM_INTEGRATION=1. Verified locally: 2 pass / 0 fail in 794 ms with the rerank model pre-cached.

Files Modified

Modified:

• src/llm.ts — adaptive context sizing, two-pass reclaim, disposeEmbedContexts / disposeEmbedModel / bestEffort helpers, dead-flash-attention-retry removal, isInsufficientVramError export, rerankContextSize lifecycle reset on the three dispose sites.

Created:

• test/llm-adaptive-rerank.test.ts
• test/llm-vram-error-pinning.test.ts
• test/llm-reclaim-integration.test.ts

Renamed:

• None.

Deleted:

• None.

Breaking Changes

• No breaking changes.

Deployment Notes

• No special deployment steps required.

[brettdavies]

Hi @tobi, friendly nudge on this one when you have a review cycle.

Single opt-in flag with no behaviour change in the default path: LlamaCppConfig.lowVram defaults to false, so existing callers, indexes, and the qmd query / qmd embed / qmd mcp paths run identically when the flag isn't set. The lowVram code paths are only entered when --low-vram / QMD_LOW_VRAM=1 is explicitly passed.

The win is concrete and matches the case in the body: on a 24 GB GPU sharing space with a ~20 GB Ollama, qmd query currently fails the rerank stage with Failed to create any rerank context. With --low-vram, peak VRAM drops from 5.4 GB to 2.6 GB and the full pipeline runs (3.0s → 5.6s; table in the body has measurements). Same fix applies to the smaller-card scenarios reported on #275.

Rebased onto current main (commit 29073bc); tsc --noEmit clean; six new concurrency tests in test/llm-low-vram.test.ts pass against a stubbed LlamaCpp so they run without a GPU and cover the dispose-mid-call race the chain is meant to prevent. The same engine flag underpins the qmd serve --low-vram passthrough in #663, but this PR stands alone.

Happy to rename the flag (the body's open question on --low-vram vs --sequential) or adjust scope.

[brettdavies]

Rebased onto current main (post-v2.8.3). Scope is unchanged (engine-only, 8 commits), but two things are worth calling out:

• Main made rerank-context creation fail soft while this PR was open: fix(llm): surface the real error when rerank context creation fails #866 surfaces the real error, warns, and skips reranking instead of throwing. That graceful skip is exactly wrong under --low-vram, where the whole point is that a VRAM failure triggers eviction and a retry. The rebased branch keeps both behaviors: with lowVram set and isInsufficientVramError(error) true, ensureRerankContexts rethrows the original error so withReclaim can evict heavies and retry; in every other case main's warn-and-skip stands. The synthetic Failed to create any rerank context message no longer appears on this path: the real error propagates, and isInsufficientVramError already recognizes the underlying node-llama-cpp messages (pinned by the error-pinning test in this PR).
• feat(serve): qmd serve shared model server + RemoteQMD client (supersedes #511) #663 no longer stacks on this PR. It's been split back to serve-only, per the discussion there, so the two merge independently in either order. The combined engine+serve shape lives in feat: low-vram engine + qmd serve combined (production shape of #662 + #663) #927.

Everything else rebased clean, including onto main's node-llama-cpp 3.20 bump and the async sequence-dispose changes.