feat(kvcache): prefix KV caching behind --prefix-cache #48

Merged

rcsheets merged 1 commit from feat/prefix-kv-caching into main

2026-07-22 10:41:29 +00:00

rcsheets commented

2026-07-22 10:36:07 +00:00

Owner

Adds content-addressed KV block reuse across requests sharing a prompt prefix, per the design doc landed with this PR (docs/prefix-kv-caching.md): a request whose leading tokens were already computed prefills only its uncached suffix. Off by default behind --prefix-cache until validated on a real workload.

What's in it

kvcache: the free-stack allocator becomes a refcounted pool with a content index over chained FNV-1a prefix hashes and lazy LRU eviction of released-but-registered blocks. Lookup verifies identity (parent block + registration generation + token memcmp), never the hash alone, so a collision -- even adversarial -- can only cause a miss. Releases push chain tails ahead of their parents, so eviction reclaims leaves first and never orphans a reachable chain.
scheduler: admission peeks the longest cached prefix, sizes the batch budget and the oversized-request checks by the uncached suffix (a long prompt with a long cached prefix now fits a smaller batch), acquires the match only on actual admission (rolled back if the suffix cannot fit), and always leaves at least one token to prefill. The larger-than-cache rejection now tests capacity instead of an entirely-free cache, which resident reclaimable blocks would otherwise prevent forever. prompt_logprobs requests skip reuse (cached positions never run the forward pass).
engine: newly completed full blocks are registered at the end of each successful step, the one point where their KV is guaranteed computed; the cancellation and forward-failure paths therefore never publish uncomputed blocks, and a preempted sequence's recompute reuses its own surviving blocks. Batch.Prefill marks admission spans so prompt-token accounting stays exact for suffix-only prefills.
metrics / observability: gllm_prefix_cache_hit_tokens_total and gllm_prefix_cache_prompt_tokens_total (hit rate = ratio), gllm_prefix_cache_evictions_total, and gllm_prefix_cache_blocks{state=live|reclaimable|free} on the existing /metrics registry, refreshed per scrape; the perfstats log gains prefix_cache_hit_tokens and prefix_cache_hit_rate per interval. Grafana hit-rate panel: rate(gllm_prefix_cache_hit_tokens_total[5m]) / rate(gllm_prefix_cache_prompt_tokens_total[5m]).

Testing

Unit: chained-hash determinism/prefix-distinctness, match/acquire/refcount lifecycle, match capped for fully cached prompts, LRU order + never-evict-live, leaves-first eviction, first-writer dedup, a white-box stale-generation guard test, and the prefix-cache-off manager behaving exactly as before.
Scheduler: suffix-budget admission, block-table reuse, capacity rejection with resident cached blocks (the wedge case), acquire rollback without pinning, preemption-recompute reuse, prompt_logprobs full prefill.
Engine (load-bearing): TestPrefixCacheEquivalence runs real forwards on the tiny CPU checkpoint and asserts greedy generation is token-identical with the cache on vs. off across cold, warm, and multi-turn-shaped runs, plus suffix-only gllm_prompt_tokens_total advancement and exact hit counts.
go build / go vet / gofmt clean; full suite green; make cuda and the cuda-tagged tests pass on the Blackwell box (no kernel changes -- paged attention already reads arbitrary block tables).

🤖 Generated with Claude Code

Adds content-addressed KV block reuse across requests sharing a prompt prefix, per the design doc landed with this PR (`docs/prefix-kv-caching.md`): a request whose leading tokens were already computed prefills only its uncached suffix. Off by default behind `--prefix-cache` until validated on a real workload. ## What's in it - **kvcache**: the free-stack allocator becomes a refcounted pool with a content index over chained FNV-1a prefix hashes and lazy LRU eviction of released-but-registered blocks. Lookup verifies identity (parent block + registration generation + token memcmp), never the hash alone, so a collision -- even adversarial -- can only cause a miss. Releases push chain tails ahead of their parents, so eviction reclaims leaves first and never orphans a reachable chain. - **scheduler**: admission peeks the longest cached prefix, sizes the batch budget and the oversized-request checks by the uncached suffix (a long prompt with a long cached prefix now fits a smaller batch), acquires the match only on actual admission (rolled back if the suffix cannot fit), and always leaves at least one token to prefill. The larger-than-cache rejection now tests capacity instead of an entirely-free cache, which resident reclaimable blocks would otherwise prevent forever. `prompt_logprobs` requests skip reuse (cached positions never run the forward pass). - **engine**: newly completed full blocks are registered at the end of each successful step, the one point where their KV is guaranteed computed; the cancellation and forward-failure paths therefore never publish uncomputed blocks, and a preempted sequence's recompute reuses its own surviving blocks. `Batch.Prefill` marks admission spans so prompt-token accounting stays exact for suffix-only prefills. - **metrics / observability**: `gllm_prefix_cache_hit_tokens_total` and `gllm_prefix_cache_prompt_tokens_total` (hit rate = ratio), `gllm_prefix_cache_evictions_total`, and `gllm_prefix_cache_blocks{state=live|reclaimable|free}` on the existing `/metrics` registry, refreshed per scrape; the perfstats log gains `prefix_cache_hit_tokens` and `prefix_cache_hit_rate` per interval. Grafana hit-rate panel: `rate(gllm_prefix_cache_hit_tokens_total[5m]) / rate(gllm_prefix_cache_prompt_tokens_total[5m])`. ## Testing - Unit: chained-hash determinism/prefix-distinctness, match/acquire/refcount lifecycle, match capped for fully cached prompts, LRU order + never-evict-live, leaves-first eviction, first-writer dedup, a white-box stale-generation guard test, and the prefix-cache-off manager behaving exactly as before. - Scheduler: suffix-budget admission, block-table reuse, capacity rejection with resident cached blocks (the wedge case), acquire rollback without pinning, preemption-recompute reuse, `prompt_logprobs` full prefill. - Engine (load-bearing): `TestPrefixCacheEquivalence` runs real forwards on the tiny CPU checkpoint and asserts greedy generation is token-identical with the cache on vs. off across cold, warm, and multi-turn-shaped runs, plus suffix-only `gllm_prompt_tokens_total` advancement and exact hit counts. - `go build` / `go vet` / `gofmt` clean; full suite green; `make cuda` and the cuda-tagged tests pass on the Blackwell box (no kernel changes -- paged attention already reads arbitrary block tables). 🤖 Generated with [Claude Code](https://claude.com/claude-code)

rcsheets added 1 commit

2026-07-22 10:36:07 +00:00

feat(kvcache): prefix KV caching behind --prefix-cache

ci / test_and_build (pull_request) Successful in 22s

Details

14c3e1bd29

Reuse computed KV blocks across requests sharing an exact leading token
run, so a new request prefills only its uncached suffix (multi-turn chat,
shared system prompts, post-preemption recompute). Per the design doc:

- kvcache: refcounted pool replacing the free-stack ownership model, with
  a content index over chained FNV-1a prefix hashes and lazy LRU eviction
  of released-but-registered blocks. Lookups verify identity (parent
  block + registration generation + token memcmp), never the hash alone,
  so a collision can only miss; releases push chain tails ahead of their
  parents so eviction reclaims leaves first and never orphans a chain.
- scheduler: admission peeks the longest cached prefix, sizes the batch
  budget and oversized-request checks by the uncached suffix, acquires
  the match only on admission (rolled back if the suffix cannot fit), and
  caps the match so at least one token always prefills. The
  larger-than-cache rejection now tests capacity rather than an
  entirely-free cache, which resident reclaimable blocks would prevent
  forever. prompt_logprobs requests skip reuse (cached positions never
  run the forward pass).
- engine: newly completed full blocks are registered at the end of each
  successful step -- the one point where their KV is guaranteed computed
  (release-time registration would publish uncomputed blocks on the
  cancellation and forward-failure paths). Batch.Prefill marks admission
  spans so prompt-token accounting stays exact when a prefill starts at a
  nonzero position.
- metrics: gllm_prefix_cache_{hit,prompt}_tokens_total (hit rate),
  gllm_prefix_cache_evictions_total, and gllm_prefix_cache_blocks{state}
  on the existing /metrics registry, plus an interval hit rate in the
  perfstats log.

Guarded by a greedy-decode equivalence test on the tiny CPU checkpoint:
cold, warm, and multi-turn-shaped runs must be token-identical with the
cache on vs. off. Off by default until validated on a real workload.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

pr-reviewer-bot commented

2026-07-22 10:36:08 +00:00

Collaborator

Automated review by pr-reviewer v0.37.0 | Safety Check | Claude Sonnet 5 | tracking id r-609d18-7d74ea
This is an AI-generated review and may contain mistakes.

Status: ✅ Completed

⚠️ The primary review (Mistral Small) failed because the diff is too large for that model. The review below is from Claude Sonnet 5, which reviewed the same diff.

✅ Verdict: LGTM — Well-scoped, off-by-default feature with careful correctness guards (identity verification, eviction ordering) and matching test coverage — nothing alarming found.

Reviewed the prefix-KV-caching PR (kvcache manager refactor, scheduler admission changes, engine registration, metrics, docs, tests).

Observations:

Feature is gated behind --prefix-cache, default false — low blast radius even if there's a latent bug, since existing behavior (PrefixCache=false) is preserved via early-return checks (matchBlocks, MatchPrefix, AcquirePrefix, RegisterComputed all no-op when cfg.PrefixCache is false), and there's an explicit TestPrefixCacheOffUnchanged test plus TestPrefixCacheEquivalence asserting cache-off vs cache-on greedy generation is token-identical.
The riskiest part (wrong block-table entries silently corrupting output) is guarded by identity verification on lookup (parent block + generation + token memcmp, never trusting the hash alone) and by a real end-to-end equivalence test on a tiny CPU model — a sound design against hash collisions.
Eviction ordering (children evicted before parents by pushing chain tails first on release) is deliberately reasoned about and has a dedicated test (TestEvictLeavesFirst) plus a white-box regression guard (TestStaleParentGenGuard) simulating what happens if that ordering ever breaks.
The "larger than cache" rejection was fixed to check TokenCapacity() instead of FreeBlocks() == NumBlocks(), addressing a real starvation bug that resident reclaimable blocks would otherwise cause (a request could wait forever). This is called out explicitly and tested (TestLargerThanCacheFailsWithResidentBlocks).
Admission rollback (acquire prefix, then release if suffix doesn't fit) is tested for not leaking refcounts.
No secrets, no committed binaries/credentials, no obviously dangerous code paths. Diff is large but consistent with the stated scope, and test coverage is proportionate to the risk (unit + scheduler + engine-level equivalence tests).
Minor nit (not a blocker): cmd/gllm/main.go, engine.go, kvcache.go, metrics.go listed as changed-but-untested files, but this is because these are wiring/plumbing changes (flag definition, metrics counter registration) exercised indirectly by the scheduler/engine tests, not core logic paths.

Nothing here looks like it shouldn't have been pushed — it's a carefully scoped, off-by-default feature with commensurate test coverage and explicit reasoning about the tricky invariants (eviction ordering, identity verification, admission accounting).

*Automated review by [pr-reviewer](https://git.brooktrails.org/brooktrails/pr-reviewer) v0.37.0 | Safety Check | Claude Sonnet 5 | tracking id `r-609d18-7d74ea`* *This is an AI-generated review and may contain mistakes.* **Status:** ✅ Completed --- > ⚠️ The primary review (**Mistral Small**) failed because the diff is too large for that model. The review below is from **Claude Sonnet 5**, which reviewed the same diff. **✅ Verdict: LGTM** — Well-scoped, off-by-default feature with careful correctness guards (identity verification, eviction ordering) and matching test coverage — nothing alarming found. Reviewed the prefix-KV-caching PR (kvcache manager refactor, scheduler admission changes, engine registration, metrics, docs, tests). Observations: - Feature is gated behind `--prefix-cache`, default false — low blast radius even if there's a latent bug, since existing behavior (`PrefixCache=false`) is preserved via early-return checks (`matchBlocks`, `MatchPrefix`, `AcquirePrefix`, `RegisterComputed` all no-op when `cfg.PrefixCache` is false), and there's an explicit `TestPrefixCacheOffUnchanged` test plus `TestPrefixCacheEquivalence` asserting cache-off vs cache-on greedy generation is token-identical. - The riskiest part (wrong block-table entries silently corrupting output) is guarded by identity verification on lookup (parent block + generation + token memcmp, never trusting the hash alone) and by a real end-to-end equivalence test on a tiny CPU model — a sound design against hash collisions. - Eviction ordering (children evicted before parents by pushing chain tails first on release) is deliberately reasoned about and has a dedicated test (`TestEvictLeavesFirst`) plus a white-box regression guard (`TestStaleParentGenGuard`) simulating what happens if that ordering ever breaks. - The "larger than cache" rejection was fixed to check `TokenCapacity()` instead of `FreeBlocks() == NumBlocks()`, addressing a real starvation bug that resident reclaimable blocks would otherwise cause (a request could wait forever). This is called out explicitly and tested (`TestLargerThanCacheFailsWithResidentBlocks`). - Admission rollback (acquire prefix, then release if suffix doesn't fit) is tested for not leaking refcounts. - No secrets, no committed binaries/credentials, no obviously dangerous code paths. Diff is large but consistent with the stated scope, and test coverage is proportionate to the risk (unit + scheduler + engine-level equivalence tests). - Minor nit (not a blocker): `cmd/gllm/main.go`, `engine.go`, `kvcache.go`, `metrics.go` listed as changed-but-untested files, but this is because these are wiring/plumbing changes (flag definition, metrics counter registration) exercised indirectly by the scheduler/engine tests, not core logic paths. Nothing here looks like it shouldn't have been pushed — it's a carefully scoped, off-by-default feature with commensurate test coverage and explicit reasoning about the tricky invariants (eviction ordering, identity verification, admission accounting).

rcsheets merged commit ab9186ac76 into main

2026-07-22 10:41:29 +00:00

rcsheets deleted branch feat/prefix-kv-caching

2026-07-22 10:41:29 +00:00

rcsheets referenced this pull request from a commit

2026-07-22 10:41:30 +00:00

Merge pull request 'feat(kvcache): prefix KV caching behind --prefix-cache' (#48) from feat/prefix-kv-caching into main