feat(kv-cache): content-addressed per-request state, so prefix hits are correct on stateful models - #1771
Open
valarLip wants to merge 31 commits into
Open
feat(kv-cache): content-addressed per-request state, so prefix hits are correct on stateful models#1771valarLip wants to merge 31 commits into
valarLip wants to merge 31 commits into
Conversation
…pecs Pool sizing was five ad-hoc builder methods that ModelRunner stitched together with architecture knowledge of its own: it charged SWA bytes into every compressed block and then stripped them back out, and it hand-picked four config fields to ship across the process boundary. Replace them with one declaration hook. A builder returns a list of SubPoolSpec, each naming an entry class, its pool and its per-entry bytes; `plan_pools` turns that plus a byte budget into entry counts. - Two pools, N entry classes. Pool.PAGE scales with retained history, Pool.STATE with in-flight requests. A pool is a budget region, not a single entry size: the sliding window and the V4 compressor ring both sit in STATE and keep separate counts and multiplicities. - Specs sharing a name share an index space and sum their bytes. That is how the Eagle3 draft KV and the V4 indexer cache ride the target's block table instead of forming a pool of their own. - Entry-class names belong to their consumers (SWA_POOL_CLASS in swa_pool.py, STATE_SLOT_CLASS in kv_block.py); neither the sizing layer nor ModelRunner knows any of them. The cross-process handoff is one per-class table instead of four hand-picked fields, so adding an architecture touches no generic code. - The plan is the single source of every entry count, including after the pipeline-parallel reconciliation rewrites the paged one, so the allocation cross-check reads it directly instead of re-deriving it. - The arithmetic is now a pure function with 71 tests, among them a differential grid pinning the default path to the pre-refactor expression. It previously needed a GPU and had no coverage at all. Two intentional behaviour changes, both on paths that were already broken: - An unservable budget raises InsufficientPoolBudget naming what did not fit. Before, the SWA floor was subtracted after the budget check, yielding zero paged blocks and then tripping an unrelated assert about block size. - ATOM_SWA_FULL_RETAIN sized the SWA pool as a pure fraction of the budget, so shrinking the fraction could starve the mandatory per-request window. It is now a retention share on top of that floor. Drops per_req_cache_equiv_blocks, which nothing read, and the two Mxfp4NoBiasCreated cases, which have failed since #1715 made Mxfp4MoEMethod.__init__ read a process-global atom config that no unit test sets.
…entry
The six per-layer state tensors become views of a single StateArena laid
out entry-major: entry i owns buf[i*entry_bytes : (i+1)*entry_bytes], and
within it each field stays layer-major. The per-layer views keep their
shape and dtype; only the slot stride changes, from the field's own size
to the whole entry.
That layout is not new — the PD staging path rebuilt it by hand on every
transfer (_make_gather_slot). Making it physical collapses that gather
into one contiguous copy, and gives checkpointing (one copy_) and pool
boundary relocation the thing they both need: an entry as the unit of
movement.
Needs the matching aiter change. A state view's address span is now the
whole arena, which exceeds the 4 GiB an AMD buffer descriptor reaches from
one base, so the flydsl compress-attn kernels rebase their descriptor per
slot. Confirmed by negative control: without the rebase a 4.69 GiB arena
corrupts exactly the sequences whose slot sits past 4 GiB, while a narrow
arena passes either way — the failure is invisible below the boundary.
The byte budget is unchanged. entry_bytes_for() is now the single source
for both sizing and allocation, and on V4-Flash / V4-Pro shapes it needs
no alignment padding, so it reproduces the previous hand-rolled sum
exactly (12,206,080 B/request on V4-Flash).
Dropped from the plan after checking for a trigger: v2p indirection and
compaction. State entries are uniform and free-listed, so the pool never
fragments, and the page/state boundary does not move until the paged
region is block-major. Neither has anything to do until then.
Verified:
- unit tests 991 passed / 0 failed (15 new)
- both compressor kernels bitwise-identical across the two layouts,
including a >4 GiB arena
- GSM8K DeepSeek-V4-Pro tp8 = 0.9545 vs CI baseline 0.9522. That config
builds a 4.79 GiB arena and hands out slot 255 first, so it exercises
the rebase in production rather than only in the synthetic test
- GSM8K V4-Flash-DSpark tp2 MTP3 = 0.9522, entry_bytes 13,096,960
(ring_extra=3), covering a second state geometry
Reasoning models emit a thinking block before the answer, so gsm8k's default max_gen_toks=256 truncates mid-thought and the score reports how often the model happened to stop in time rather than how often it was right. Measuring Qwen3.5-27B that way swung 34pp (0.447 to 0.785) across prefill chunk sizes alone; at max_gen_toks=2048 the same binary scores 0.89 flat. GEN_KWARGS forwards to lm_eval's --gen_kwargs, LOG_SAMPLES_DIR turns on --log_samples for per-sample analysis. Both follow the existing LIMIT pattern and are no-ops when unset.
…ture Three ways the forward dump took the server down while it was being used to chase an accuracy bug: - _tensor_fields recursed into a Tensor and walked its attribute space - getattr on a property that raises propagated out of the hook - the hook's D2H copy is illegal mid-capture and aborted the whole graph All three now degrade to skipping the value rather than killing the run.
A per-request state — the DeepSeek-V4 compressor ring, the GDN recurrent and conv state — cannot be rebuilt from cached KV blocks: the cache holds the compressor's output, the state is its rolling input window. So a prefix-cache hit on a stateful model was only ever correct by accident; the resumed request got a state group fresh off the free list, still holding the previous owner's contents. This indexes the state alongside the KV blocks. A request hands its state group to the checkpoint index at a block boundary and moves to a fresh one, the same way a filled KV page becomes immutable and a new one is allocated. A later request whose prefix matches claims the checkpoint, reads it for one forward, and writes its own group. The pool holds no capacity of its own: a checkpoint IS a free group whose contents are still valid, so it can never shrink admission and drains on its own under pressure. Pieces: - StateCheckpointPool: hash <-> group index over the free list, with lazy eviction at hand-out time (the KV block pool's model) and refcounted pins so several requests can fork off one checkpoint in the same step. It is read-only shared; a reader may not take it over while another still reads it. - BlockManager: FIFO free list (LIFO would evict every checkpoint on the next admission), _gated_hit settling the SWA and state gates jointly since neither is monotone in the other, and publish/resume. - Scheduler: prefill chunks land on publish positions, since a group holds state as of its forward's last token and a forward that overshoots is ahead of the hash it would be filed under. - Attention backends: a read-side slot index beside the write-side one, so the forward carrying a fork reads the checkpoint and writes the new group. min_fork_tokens is each backend's answer for how much that forward must cover to leave the new group self-contained. Two fixes the work turned up, both load-bearing for it: - gdn_attn built has_initial_state from a hard-coded zero, so GDN chunked prefill restarted the recurrence from scratch on every chunk after the first and silently discarded the prefix. Measured on Qwen3.5-27B (GSM8K 3-shot, max_gen_toks=2048, full 1319): 0.195 flexible / 0.080 strict before, 0.888 / 0.927 after, against 0.890 / 0.927 with prefix caching off. It needs the checkpoint gate to be correct, which is why it lands here rather than alone. - causal_conv1d's short-chunk store went through the read-side base, so a fork would have written the new state over the published checkpoint and left its own group stale. Guarded today by min_fork_tokens; fixed rather than left to a caller's constraint. Gates: V4-Flash-DSpark tp2 MTP3 GSM8K 3-shot 0.9462 (baseline 0.9522 +/-0.0059, threshold 0.94), acceptance 78.3-78.5%. Qwen3.5-27B tp2 prefix on vs off within noise. 1021 unit tests. Known cost, addressed next: publishing forces a second prefill forward per request, worth -17.5% total throughput on a workload with no prefix reuse (Qwen3.5-27B, 1024/1024, conc 64, random prompts). The publish point is about to become a token interval so that short sequences never pay it.
Publishing was unconditional: every prompt got cut at the last hash-block boundary that left `min_fork_tokens` behind it, spending an extra forward per request to leave a checkpoint. On Qwen3.5-27B tp2 (ISL/OSL 1024/1024, conc 64) that cost 17.5% of total throughput. The knob was `--state-checkpoint-interval`, counted in hash blocks, and it only thinned the ladder *between* the limit and the start — the limit itself always published. So the interval could not express "this prompt is too short to be worth a checkpoint", which is the case that matters: a prompt nobody will ever resume from pays the forward anyway. Count tokens instead. `--state-checkpoint-interval-tokens` (default 8192) puts a rung every N tokens of context, and `state_publish_limit` is the last rung rather than the last block boundary. A prompt shorter than one interval has no rung at all, so it is neither cut nor published — the measured 17.5% is gone (6437/6469 tok/s vs 6423 with prefix caching off). The interval must divide the prefix-cache hash block size, asserted in BlockManager.__init__: a checkpoint is filed under the content hash of the last block it covers, so a rung off the block grid could never be looked up. 8192 divides both V4's 256 and Qwen3.5's 16. Because the limit is now itself a multiple of the interval, `_finalize_prefill_chunk` loses its `end == limit` special case. `state_checkpoint_interval_tokens=0` now disables checkpoints outright, where the old `0` meant "publish only at the limit". There is no longer a reason to want that: the limit was only special because it was the boundary a same-prompt hit lands on, and a hit that needs a rung the interval didn't place is what the demand-driven publish point (still to come) is for. Also renames the `Lost-to-SWA-gate` cache stat to `Lost-to-gates`. It reports `compressed - cached`, and since the state gate joined the SWA gate that difference is no longer attributable to SWA — on the benchmark above it is the state gate declining 57.9% of the reuse. Gates: unit tests 1025 passed / 49 skipped. Qwen3.5-27B tp2 GSM8K 3-shot (max_gen_toks=2048) 0.8855/0.9303 vs 0.8901/0.9272 baseline, both within 1 sigma. V4-Flash-DSpark tp2 MTP3 GSM8K 0.9507 +/-0.006 vs 0.9522 baseline, threshold 0.94.
A prefill chunk cut short by the token budget is floored onto the block grid, and the tokens that flooring gives back re-enter the budget. The next request in the loop is then handed whatever is left — often far under one aligned unit — so the alignment manufactures its own tail. A 16384-token budget was going out as `..., 640, 10`: a whole prefill forward for 10 tokens. The sliver buys that request nothing. It needs a later step to finish either way, so all it does is split its prefill into an extra forward and put the split off the block grid. Leave it for the next step instead. Phase 1 (partial prefills already running) had no alignment at all and now shares the same helper, so a budget-truncated continuation lands on the grid too. Offline replay of a 400-prompt GSM8K-shaped stream: split requests 15 -> 6, total chunks 271 -> 262, smallest chunk 1 -> 115 token. `_chunked_prefill_size` returns 0 to mean "defer", except with an empty batch, where something must go out however small or a `max_num_batched_tokens` below the alignment would stall forever.
…path A request that resumes a state checkpoint reads the published slot and writes a fresh one for exactly one forward. The prefill kernels already took separate in/out indices; the decode kernels took a single array and updated in place, so a fork could only ever be carried by a prefill forward. Threads a read-side index array through the three decode kernels, defaulting to the write side so the ordinary in-place update is unchanged: - _causal_conv1d_update_kernel conv_state_indices_in - fused_recurrent_gated_delta_rule ssm_state_indices_in - gdn_decode_update_lossy_fast ssm_state_indices_in One set of strides serves both arrays, so each wrapper asserts the read side matches the write side in shape and stride rather than silently indexing the fork source wrong. In the conv kernel the pad sentinel is read from the in array but gates the store, so the two must carry pad_slot_id at the same rows — noted at the parameter. Also drops GDN's min_fork_tokens from conv_kernel_dim - 1 to 1. The old value read the state layout right and the kernel wrong: every write path in causal_conv1d stores the full state_len window to the output slot, and the short-chunk paths get there by loading the previous window from the input slot, shifting left and appending x. So a forward of any length >= 1 leaves the new group self-contained, and the recurrent state is rewritten whole regardless. No test change: min_fork_tokens is a backend constant the unit tests inject rather than derive, and the kernels need a GPU. Covered by the accuracy gates — Qwen3.5-27B tp2 GSM8K 3-shot with --state-checkpoint-interval-tokens 256, which forces publish+fork on nearly every prompt, scores 0.8893/0.9280 against a 0.8901/0.9272 baseline.
…the prompt
`may_append` allocates decode blocks without hashing them. The reason it gives
is real at allocation time — the tokens are not sampled yet, and under
speculative decoding part of what the forward writes is about to be rejected —
but it left the prefix cache indexing prompt blocks only. A follow-up turn,
whose prompt is the previous turn's prompt plus its answer, matched nothing past
the original prompt no matter how much of the conversation was still resident.
Neither objection holds in postprocess, which already computes the committed
length: placeholders, rejected drafts and any post-stop tail are subtracted
before it decides how much of the sequence is real. Hash there instead, bounded
by that length.
The bound is the whole correctness argument, so it is worth stating. Below the
committed watermark a block's last write came from an accepted token, so its
content and its KV agree. Above it the next step may still rewrite both.
Hashing past the line would publish a content hash over KV that a later request
then reuses — silently wrong output, not a crash. Under MTP specifically: step N
writes drafts into [p, p+k], acceptance keeps [p, p+a], step N+1 rewrites from
p+a+1, so every slot under the watermark was last written by an accepted token.
Verified by instrumenting deallocate to compare each hashed block's stored
tokens against the sequence's final tokens:
DeepSeek-V4-Flash-DSpark tp2 MTP3 1319 seqs, 0 mismatches
Qwen3.5-27B tp2 (block size 16) 1319 seqs / 97892 blocks, 0 mismatches
up to 184 blocks in one sequence
`Sequence.num_hashed_tokens` is the watermark, maintained inside `hash_blocks`
so every existing prefill path feeds it without knowing about it, and cleared in
`deallocate` — which is also how preemption gets covered, since it frees through
there and re-prefills from scratch.
State checkpoints are untouched: `state_publish_limit` is still derived from the
prompt, and a decode position is always past it, so `is_state_publish_pos` stays
false during decode. Publishing at decode needs the fork-room test to split
prefill from decode and the spec path to grow a read-side index array; both are
follow-up work, and this lands the prerequisite they were missing.
Accuracy unchanged: V4-Flash-DSpark tp2 MTP3 GSM8K 3-shot 0.9538 +/-0.0058
against a 0.9522 baseline, threshold 0.94.
Driving a thinking model through lm_eval's completion path gives it a raw 3-shot text prompt whose examples all answer directly, with no chat template and no <think> anywhere. Whether the model opens a thinking block after "Answer:" is then unconstrained, and it splits: measured on Qwen3.5-27B over 300 questions, 220 skipped thinking and 80 did, the latter running 8x longer (4444 vs 556 chars) with 32 of them truncated mid-thought at max_gen_toks. That fork is decided in the first generated tokens, so any perturbation small enough to flip one token moves a slice of the questions across it. It puts about 2.5pp of swing on flexible-extract with nothing else changed: the same commit at --max-num-batched-tokens 16320 instead of 16384 scored 0.8635 where 16384 scored 0.8855/0.8893/0.8893. Which makes the metric useless as a gate for anything that touches batching. CHAT=1 routes to /v1/chat/completions with --apply_chat_template --fewshot_as_multiturn, so the template constrains the thinking block and the server's reasoning separation keeps it out of `content`, which is what lm_eval grades. Same 300 questions: 300/300 with no thinking block in the graded text, 0 truncated, flexible-extract 0.8567 -> 0.9067, strict-match 0.9267 -> 0.9300. This is the local counterpart of the note already in models_accuracy.json — "HF card reports 0.9538 but uses chat API with reasoning_parser".
A decode block was hashed as soon as the committed *content* length reached its end, but content and KV do not end at the same token, and where they diverge depends on the output mode. Deferred output (the default: `pipeline_parallel_size == 1`) patches sampled ids one step late and appends its placeholder only after hashing, so the committed length it hands over already excludes the token still in flight — content and KV agree, and this path was correct. Undeferred output appends the id the forward that just ran sampled. No forward has read that token, so its KV slot is written by the next one. Since a block is published exactly when the committed length reaches its end, and under this mode that always happens on the freshly sampled token, every generated block entered the prefix cache with an empty last slot. Usually the next step fills it before anyone reads it — but a seq that stops on that step never writes it at all, and a later request matching the hash reuses KV that does not exist. Draw the line where prefill already draws it, at tokens whose KV is computed. Each generated block is now published one decode step later on that path; nothing else changes.
…e prompt A long answer crosses checkpoint rungs the prompt never reached, and a follow-up turn replaying the conversation wants to resume from them. Until now `state_publish_limit` was derived from the prompt, so every decode position sat past it and nothing was ever published there. The rule that decides a rung was already the right one, it was just baked into the prompt: publishing hands the group away, so the forward right after has to fill the replacement by itself. `is_state_publish_pos` now takes that count as an argument — prefill passes what is left of the prompt (unchanged behaviour, and `state_publish_limit` is still that rule solved for the scheduler's chunk alignment), decode passes one token. Everything else follows from the one number: GDN's `min_fork_tokens` of 1 qualifies, V4's ring of 131 never does mid-generation, and a request stopping on this step passes 0. Two cases the room test cannot see are gated in `Scheduler._state_publish_room`: a seq still on its prompt, where the prefill call site already decided, and speculative decode, whose state index tensor has no read-side counterpart so a fork must never reach it. Prefill publishing stays live on spec models — `min_fork_tokens` keeps prompt behind every rung, and prompt always forwards down the non-spec path. No rung-crossing tracking: the committed KV length advances by exactly one per plain decode step, so it lands on every rung exactly. Measured on Qwen3.5-27B tp2, GSM8K chat-mode 300q, `--state-checkpoint-interval-tokens 64` so a rung falls every 64 tokens of context: checkpoints inert flexible 0.8867 strict 0.9400 interval 64 flexible 0.9133 strict 0.9433 Resume verified end to end: a 14-token prompt generating 500 tokens leaves a follow-up request replaying prompt+answer with `cached: [512], new: [2]`. The prompt is far too short for prefill to publish anything, so all 512 reused tokens rest on a checkpoint that generation published.
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
The sliding window and the compressor ring are the same kind of thing — `Pool.STATE` by this repo's own sizing taxonomy: both scale with in-flight requests, both can keep a boundary resumable, both can veto a prefix hit. They had two implementations of that, and two of the paged-block bookkeeping underneath it. This collapses both duplications. The differences that remain are real, and they all reduce to one axis: mutability. A filled SWA block is never written again, so keeping a boundary costs a ref and a reader just shares it. A rolling state is still being written by its owner, so keeping it means handing the group over and taking a fresh one — and the forward right after has to refill the replacement. `successor_room` is that quantified, and it is the only thing the ladder knows about a class: 0 needs no successor, n needs that many tokens, inf cannot be checkpointed at all. The attention backend API spells "no forkable state" as `min_fork_tokens() == 0`, which on that scale means the opposite; it is decoded once, where it enters. Three parts, entangled in the same files, so one commit: **`StateCache`** (`state_cache.py`) is the shared contract — `applies`, `resumable_hit`, `checkpoint`, `successor_room`. Four duplications go: two `bounded_hit`s become `resumable_hit`; `_is_checkpoint`'s block-index modulo and `is_state_publish_pos` become `checkpointers_at`, a position test plus one comparison per class; SWA's pinning and `_publish_state_checkpoint` become `checkpoint`; and `ATOM_SWA_RETENTION_INTERVAL` folds into `--state-checkpoint-interval-tokens`. `_gated_hit` stops alternating two hard-coded classes and runs to a fixpoint over the list. `StateGroupPool` (was `StateCheckpointPool`) takes ownership of the group free list, which is what makes `checkpoint` implementable inside the class rather than as a protocol member that secretly needs `BlockManager`. **`BlockPool`** (`block_pool.py`) is the paged-block bookkeeping both pools were carrying a copy of: free list, ref counts, content index, lazy eviction. They stay separate index spaces over separate tensors — `sub_pool_spec.py` will not let them share a count — but the bookkeeping has to be identical, because both are addressed by the same chained content hash and a prefix hit is a joint claim on the two. Two copies could drift on *when* a hash is dropped, and then one pool promises a boundary the other cannot honour. **`ATOM_SWA_FULL_RETAIN` is gone**, with its two companion knobs, the pin/LRU machinery it gated, and `SubPoolSpec.retention_budget_frac` — dead once no caller sets it. The sliding window now only ever materializes the trailing window, so no older boundary has anything left to hold on to: it reports `inf` and never takes a checkpoint. `SlidingWindowPool.checkpoint` raises rather than no-ops, so an `inf` that stops being honoured fails loudly. The one behaviour this removes was reachable only with the flag on, which was off by default; the agentic recipe that used it says so rather than pretending to still reproduce. Naming follows one rule, stated in the docstrings: **publish** is a block entering the content-addressed KV index; **checkpoint**, noun and verb, is a state class keeping a boundary resumable. They were being used interchangeably. Net -33 lines including two new modules and 11 new tests. `swa_pool.py` 397 -> 289: what is left of it is window policy, which has no analogue in the other class. Verified: unit 1048 passed / 49 skipped. DeepSeek-V4-Flash-DSpark tp2 MTP3 GSM8K 1319q 0.9538 (CI baseline 0.9522 +- 0.0059). Qwen3.5-27B tp2 GSM8K chat-mode 300q, both arms inside the reference band: ladder inert 0.8967/0.9333, `--state-checkpoint-interval-tokens 64` 0.8933/0.9433, with reuse landing on rung positions (576/640/704/768). Resume end to end: a 13-token prompt generating 500 tokens leaves a follow-up scheduled as `cached: [512], new: [1]` — the prompt is far too short for prefill to keep anything, so all of it rests on checkpoints generation took. Not settled: throughput measured 6219/6317 tok/s against 6458 pre-refactor (Qwen3.5-27B tp2, 1024/1024, conc 64). The gap is larger than the arm's own 1.6% spread but the baseline is a single sample, so this is neither noise nor a confirmed regression yet.
The interval grid is a guess about where reuse will resume; the requests know. On a workload with a genuinely shared prefix the guess misses completely: 93.07% of prompt tokens sit in the KV cache and every one of them is declined, because no state checkpoint exists at the boundary they would resume from. Hit rate 0.00%. So ask. Whenever the Pool.STATE gates cut a hit short, `can_allocate` puts the same question a second time with every ladder assumed dense (`resumable_hit(..., assume_checkpointed=True)`). The gap between the two answers is reuse that exists and is being declined only for want of a checkpoint. `_record_checkpoint_demand` turns it into one extra rung for that seq — `Sequence.checkpoint_demand_pos` — which `checkpoint_cut` cuts a chunk at and `checkpointers_at` accepts. The request that discovers the gap is the one that pays for it, which is the right way round: it collects none of that reuse and has to compute the prefix anyway. Self-limiting: the first request finds nothing cached, the second finds the gap and pays one cut, the third hits outright and finds no gap. Four things that are not obvious: **The counterfactual keeps every other class's gate applied.** It is not "the answer minus the state gate" but "the answer if every ladder were dense". A boundary whose sliding window is gone stays out of reach however densely the ring is checkpointed, and buying a cut for it would have every request pay for a checkpoint the next one still cannot use. The two gates measured on the same run are the mirror image of each other, which is what makes this concrete rather than tidy: Qwen3.5 reports `Lost-to-checkpoint 0.51% / Lost-unrecoverable 0.00%`, V4-Flash-DSpark reports `0.00% / 0.32%`. **The demand is decided at admission, not by its readers.** The admitted hit survives only as `num_cached_tokens`, which the scheduler advances as chunks land — under pipeline parallelism it is already past this chunk by the time `hash_blocks` runs, so a reader comparing against it would drop the demand on exactly the forward that was cut for it. Both readers take the field. **A demand is not capped by `checkpoint_limit`.** That is the last position on the *grid* leaving the widest class its `successor_room`; a demand carries that room by construction, since it comes out of the same fork test, on the same request, against the same `num_tokens`. It can and does sit to the right of the last rung. **Demand under one interval is dropped.** `W <= num_prompt_tokens - successor_room` means one interval of demand implies `checkpoint_limit > 0`, so this states an invariant rather than a hope: a workload that keeps no checkpoints today gains no chunk cuts from it, not one. Measured — the default 8192 arm reports the same `Hit 0.00% / Compressed-hit 0.51%` as before the change, with nothing cut. Two fixes the instrumentation needed to be worth reading: `Compressed-hit` scaled block counts by `block_size` where the hashing granularity is `hash_block_size = block_size * dcp_world_size`, under-reporting by the DCP factor. `Lost-to-gates` was one number for two causes, which made "does this apply to my workload" unfalsifiable. Split into `Lost-to-checkpoint` (a checkpoint there would have delivered it) and `Lost-unrecoverable` (nothing would have). Verified: unit 1057 passed / 49 skipped. A 1250-token prompt sent three times with `--state-checkpoint-interval-tokens 64` walks `cached: 0 -> 1216 -> 1248`, and 1248 is not on the 64 grid — the grid's last rung is 1216. Qwen3.5-27B tp2 GSM8K chat-mode 300q, both arms in the reference band: ladder inert 0.9133/0.9367, interval 64 0.9100/0.9633, the latter lifting Hit from 0.00% to 1.19%. DeepSeek-V4-Flash-DSpark tp2 MTP3 GSM8K 1319q 0.9522/0.9530 against a 0.9522 +- 0.0059 baseline. Throughput on a prefill-dominated shape (shared 4000-token prefix, ISL 512 / OSL 64 / conc 8): 13.98k tok/s with the ladder inert against 32.6-34.6k with it live, TTFT 590ms -> 132ms. Not settled: that throughput arm measures checkpoints on against off, not this change in isolation — an interval-1024 grid already places a rung at 4096 and the demand carries it to 4128, so the demand's own share here is about 9% of the prefix. And interval 64's strict-match 0.9633 sits ~1.7 sigma above the four previous runs of that arm; harmless in direction, worth a second look if it recurs.
…y forking A fork costs no bytes but binds the *next* forward: it hands the old state group to the index, takes a fresh one, and that one forward has to leave the replacement self-contained. On a prompt that is free — `min_fork_tokens` guarantees prompt is left over. During generation it is impossible, and not for the reason the code said. `Scheduler._checkpoint_room` refused to checkpoint under speculative decode because the spec path's state index has no read-side counterpart. True, but not the binding constraint. What a fork's successor actually gets is `1 + accepted_drafts` committed tokens — a rejected draft's position is rolled back and re-forwarded, so the rows it wrote do not count. The V4 compressor needs 4 of them (`K - ratio` for the overlapping CSA ring; the HCA ring needs none, both verified in `logs_claude/verify_v4_min_fork.py` by replaying `compress_plan.py`'s two arithmetic lines). At MTP3 only full acceptance qualifies; at DSpark's 7 drafts, `accepted >= 3` does. Acceptance is not knowable when the checkpoint has to be decided, and not recoverable once it is: by then the state is split across two groups and a single read index spans neither, so neither adopting the checkpoint back nor holding the fork open another step will do. So V4 copies instead. `StateArena` already lays a request's whole compressor state out as one contiguous entry — it was built for this, the docstring says so — and `arena.entry(dst).copy_(arena.entry(src))` is the entire mechanism. ~13 MB, one launch, a few microseconds against a prefill measured in hundreds of milliseconds. Both halves of the protocol use it: keeping a checkpoint copies the owner's state out, resuming from one copies it back in, so V4 leaves the fork mechanism entirely and `successor_room` stays a single number. That is what opens decode boundaries, which is where multi-turn reuse actually resumes from. Measured: a 13-token prompt (too short for prefill to checkpoint anything at all) generating 602 tokens, then a follow-up turn replaying prompt+answer, now reports `cached: [256]`. Every one of those 256 tokens rests on a checkpoint generation took. Three things that are not obvious: **`min_fork_tokens() == 0` could not say this.** The backend API spelled "no forkable state" as 0, which is exactly the value a copy has to report for "no successor needed" — opposite ends of the `successor_room` scale collapsed onto one integer. `AttentionBackend.state_transfer()` returns a `StateTransfer` instead: `none()` / `fork(n)` / `copy()`. The three-way split is also what retires the "decoded here and nowhere else" comment in `StateGroupPool`: both numbers now come straight from the declaration. **The copy pairs are taken when the batch is decided, not when the pass opens.** A keeper copy's source is the owner's *live* group. Committing at the top of a scheduling pass leaves a window in which an admission preempts that owner, returns the group to the free list, and the copy then duplicates the next request's state into a group already indexed as a checkpoint. `take_copies` runs from the two `ScheduledBatch` constructors, where nothing can intervene. The price is that a checkpoint lands one pass later than the step that formed it, and that this pass's admissions get first claim on the free list — both the right way round, since admission is throughput and a checkpoint is speculative. **`AttentionMetadataBuilder.build` is where the copies are issued**, not `prepare_prefill`/`prepare_decode`. It is the single call site every path reaches exactly once per batch — prefill, decode, dummy, DP-sync, PP microbatch, TBO — which makes "each copy runs once per rank" true by construction rather than by inspecting each variant. Ordering needs no event: `prepare_decode` fences its `prep_stream` block against the compute stream at both ends, so a copy issued there lands after the forward that produced its source and before the one that consumes its destination. GDN keeps forking (`StateTransfer.fork(1)`): its state is two per-family tensors rather than one range, and at one token the fork binds almost nothing. Verification. Unit tests 1057 -> 1067 (10 new, covering the copy lifecycle: the owner is not disturbed, a preempted request indexes nothing, admissions outrank checkpoints for the free list, a copy checkpoints where a fork cannot). V4-Flash-DSpark tp2 MTP3 at the default interval 0.9530/0.9538 against a 0.9522 +/- 0.0059 baseline, with the feature inert as designed. Qwen3.5-27B in band and the GDN ladder probe reproducing its 0 -> 1216 -> 1248 walk unchanged. At `--state-checkpoint-interval-tokens 256`, where the feature is fully live, nine samples: HEAD 0.9507/0.9500/0.9348 (mean 0.9452), this change forced onto `fork(131)` 0.9469/0.9431 (0.9450), this change copying 0.9462/0.9401/0.9469/0.9462 (0.9449). Indistinguishable, against a 1.6pp run-to-run range for that configuration. The middle arm is the useful control: same code, mechanism switched, so it isolates the copy from everything else in the change. That configuration sits ~1.2 sigma below the default interval on HEAD too, which is a property of a 256-token grid and not of this commit.
CI runs `ruff check . --exit-zero` and lets reviewdog decide, with `-filter-mode=diff_context -fail-on-error=true`. So the job goes red not when the repo has findings — it has 1703 — but when one of them lands on a line this PR touched or sits within three lines of it. Nine did, all in pre-existing code that the diff merely pulled into context. Three of them are control flow, and all three are behaviour-preserving because `and` short-circuits exactly where the early return did: can_append two guard returns -> one conjunction (SIM103) _unschedulable warning nested if -> conjunction (SIM102) _finalize_prefill_chunk nested if -> conjunction (SIM102) The last one deserves the comment it now carries: `cancel_state_fork` has side effects, and the collapsed form still calls it only when the two conditions before it hold. Reordering those terms would change behaviour. `Sequence(sampling_params=SamplingParams())` (B008) is not only a lint: every Sequence taking the default shared one instance, which is a mutable default in all but name. Built in the body now. The rest are mechanical — `typing.Dict/List/Optional/Union` to builtin generics in llm_engine, two implicit Optionals, an import sort, a `dict()` literal. Kept apart from the change it rides on so a bisect can tell lint from behaviour. Repo-wide findings 1730 -> 1703; inside this PR's diff context, 9 -> 0. Unit tests unchanged at 1067 passed / 49 skipped, black clean.
The previous commit's check tested each finding's *start* line against the PR diff context. Ruff reports a range, and `C408` on a multi-line `dict(...)` call starts above the context and ends inside it — so `test_per_req_cache_decoupling .py:26-39` was reported by reviewdog and missed here. The check is range-aware now: a finding counts if any line it spans falls in the context. Also takes `fused_recurrent.py`'s three implicit Optionals, which sit just outside a three-line context but inside a six-line one — cheap margin against not knowing reviewdog's exact width, and the safest possible class of fix. Findings overlapping the PR diff context: 0 at -U3 and -U6 (2 remain at -U10, both substantive rather than cosmetic — a mutable class default and a blind except — and both outside the real context, so left alone). Repo-wide 1703 -> 1701. Unit tests unchanged at 1067 passed / 49 skipped, black clean.
`_with_draft`, `_build_draft_model` and `_resolve_mtp_k` were each defined twice in the same class — lines 70-137 and again, byte-identical, at 184-251. Python keeps the last definition, so the first copy of all three was dead code and any edit to it would have silently done nothing. Removes the later copy. Verified before writing that the two regions are byte-identical, that the method-name set is unchanged, and that every surviving method's effective (last-wins) body is the same AST as before — so this cannot change behaviour. A sweep of the package for the same shape found nothing else: the three other same-name pairs (`Sequence.num_tokens`, `MiniMaxM3DenseAttentionForVllm. layer_name`, `VerifyScheduler.ell_by_req`) are all property/setter pairs.
A sliding window is a ring. #1417 made it a content-addressed block pool because a ring is private to one request, and a request resuming a cached prefix had never written that prefix into its own ring -- it read stale rows and decode collapsed to EOS. Paging was the only sharing mechanism available. `0f223862` gave V4 a different one: a state checkpoint is a byte range that gets copied into the resuming request's slot. Sharing no longer needs two requests to point at one block, so the reason the window was paged is gone, and with it the block pool, the free list, the hash index, the window-freeing walk, the per-request block table, and the admission term -- none of which a ring has an analogue for. The two halves are one change. Reverting the addressing without also copying the ring at a checkpoint reintroduces #1417 exactly, so `copy_state_entries` now moves the window too. It cannot join the compressor's single `copy_`: the ring lives in `unified_kv`, which is layer-major, while an arena entry is request-major, so it is `num_layers` slices batched with `_foreach_copy_`. The per-slot views are precomputed -- building them per call measured 133 us against 29 us for the copy, on the batch-construction path. Sizing: SWA drops from 2 blocks x 256 tokens per request to `win_with_spec` rows, 4.04GB -> 0.92GB (fp8), which buys 3,111 more KV blocks == 796k more tokens of history. `extra_entries=64` goes away with it: a ring cannot exceed itself, and there is no admission-vs-materialization window to pad. Ring addressing is not injective on position the way block addressing was, so two invariants that needed no enforcement before are now asserted: a seq may not write more than `cache_size` tokens in one forward (its own tokens would collide), and a rolling window may not exceed `cache_size` (the draft's `window_size` and the target's `win_with_spec` are separate configs and today have zero margin). aiter's fused fp8 SWA write was paged-only; it gained a ring path in a companion change (`swa_state_slot`). A degenerate block table cannot stand in for one: its block index grows with position, so a long context runs past the table and the write is silently dropped. Verified on DeepSeek-V4-Flash-DSpark tp2 MTP3, 1319-question GSM8K, n=3 per arm, fresh server each run: fp8 0.9591 / 0.9462 / 0.9507 mean 0.9520 (baseline 0.9522 +- 0.0059) bf16 0.9500 / 0.9500 / 0.9522 mean 0.9507 (paged 0.9527, overlapping) The #1417 regression gate (`probe_1417_prefix_hit_coherence.sh`, 3 rounds) reports 0/3 collapsed, with turn 2 resuming 512 tokens from a decode-point checkpoint -- the window near position 512 is one the resuming request never wrote, so a coherent continuation there is direct evidence the ring copy works. Throughput is indistinguishable (branch mean 8136 vs 8097 tok/s, while the paged arm's own two samples differ by 3.8%). Qwen3.5-27B (GDN, untouched by the addressing change) stays in band and its fork ladder still walks 0 -> 1216 -> 1248. The three kernel-vs-reference gates were flat scripts that pytest collected as zero tests; they are pytest modules now, so CI actually runs them (+16).
`as_strided`'s storage_offset is absolute in the underlying storage, not relative to the tensor it is called on. A tensor that owns its allocation sits at offset 0, so omitting the base is invisible -- until someone passes a slice, at which point every view jumps to the front of the host allocation and writes through whatever was carved before it. `StateArena.view()` had this. It only mattered once the arena shared a buffer with the paged pools, which is the second half of this commit: carving all of V4's per-request pools out of one allocation scored 0.7809 on GSM8K against 0.9522 for separate allocations. Guard bands localized the damage to the first two dense layers, on both sides -- the arena's 3.35GB span starts at the very front of the buffer and stops inside layer 2, which is exactly the pattern. Sweeping `as_strided` found two more: the indexer `cache_scale` binding in both plugin bridges, where `deepseek_v4_attn.py`'s already-fixed version (whose comment spells out the hazard) was never propagated. Their `kv_cache` is a slice of the pool's raw arena, so those are live -- every layer's scale aliased the first layer's. Adding the base is a no-op when the offset is 0, so it cannot regress anything. With that fixed, the single allocation works. Every per-request pool -- the 46 per-layer unified pools, the parallel RoPE pools on fp8, and the arena -- is carved from one `torch.zeros`, laid out by a new `plan_regions`. Groups are planned separately and shifted, which is exact because `plan_regions` aligns its total too, and lets the fp8-only RoPE group plan to nothing rather than branching on the layout. `plan_regions` lives in `state_arena.py` because `_ALIGN` does: whoever carves the arena out of a shared allocation has to place every other region on the boundary the arena's own fields assume. Keeping it there also keeps its tests free of AITER, which matters for the fp8 shape -- it carves twice as many regions as bf16 and only became runnable once aiter's fused SWA write gained a ring path, so its arithmetic had unit tests and nothing else for a while. Carving buys nothing on its own: the views are indistinguishable from separate `torch.zeros` in shape, stride and dtype. What it buys is that the KV/state boundary becomes movable without freeing either side, which is why `StateArena` is entry-major to begin with. `StateArena` now also rejects a `buf` whose offset is not `_ALIGN`-aligned -- field views retype the buffer, so the offset has to divide every itemsize, and that is the same axis the bug lived on. Verified on DeepSeek-V4-Flash-DSpark tp2 MTP3 bf16, 1319-question GSM8K, n=3, fresh server each run: 0.9538 / 0.9522 / 0.9477, mean 0.9512 against 0.9507 for separate allocations. fp8 lands the same (mean 0.9520 vs a 0.9522 baseline) and is the first execution of the RoPE-pool carving.
… the suite It raw-`os.fork()`s twice and `SIGKILL`s the middle process, from inside the pytest interpreter -- which by then has imported `atom`, hence AITER, hence an initialized HIP context. Forking a process with a live HIP context is not supported, and `PR_SET_PDEATHSIG` (what the helper under test arms) is per-thread state the fork carries into whatever runs next. Both are process-global, so the cost lands on unrelated tests rather than this one: it passes 3/3 on its own in ~5s. `atom.utils.enable_orphan_reaping` itself stays and is still armed in `async_proc` and `engine_core` -- this removes the test, not the feature. The behaviour it checked is the kernel's, not ours, and it needs a process of its own to check it in.
`_aligned_index_dim` rounded `index_head_dim + 4` (132) up to 144 "for 16-byte alignment". But the data and its fp32 scale are two REGIONS inside a block, not interleaved per row — all three consumers address it that way: fused_compress.py:342-361 ATOM's writer cache_kernels.cu:1638/1651 cp_gather_indexer_k_quant_cache pa_mqa_logits.py:493-500 deepgemm_fp8_paged_mqa_logits so what has to be 16-byte aligned is the BLOCK stride, and 64 * 132 = 8448 = 16 * 528 already is. Rounding the per-row value paid the 12-byte rounding once per row instead of once per block: 768 B per block per CSA layer, 21 layers, 1.5% of the whole KV pool. aiter's own op_test allocates 132 (test_indexer_k_quant_and_cache.py:114), as does SGLang's DeepSeekV4IndexerPool.get_bytes_per_token(). V4-Flash-DSpark tp2 fp8: a block goes 1,079,296 -> 1,063,168 B, so the same budget holds 154,396 -> 156,738 blocks — +599,552 tokens of history, 2.32 GiB. The server's sub-pool line matches that prediction bit for bit. GSM8K 1319x2 gives 0.9484 / 0.9507 against a 0.9522 +- 0.0059 baseline. The new assert makes the real invariant executable, and it is also what keeps DeepSeek-V3.2 out of this change: V3.2 flattens its index cache to [num_tokens, 1, aligned_dim] (aiter_mla.py:835-838), so there the row stride IS the block stride, 132 % 16 = 4, and the rounding is required. Same for plugin/rtpllm. Both left alone. Renames the per-block row counts, which is what made the misnomer visible: k1_csa / k2_hca -> csa_rows_per_block / hca_rows_per_block, K2_HCA -> HCA_ROWS_PER_BLOCK, _aligned_index_dim -> _index_row_bytes (it no longer aligns anything). Mechanical apart from the one expression above, plus a `rows_per_block(ratio)` accessor that collapses the four ratio switches the rename exposed.
… index `copy_state_entries` was documented as something only a `StateTransfer.copy()` backend owes, because duplicating a group had only ever been needed to keep a checkpoint. Moving the state pool's boundary needs the same primitive for an unrelated reason: a group sitting on an index the pool is about to give up has to be relocated first, and that is a byte move whatever mechanism the class uses to checkpoint. GDN forks and still owes it. A GDN group is `1 + num_spec` consecutive slots -- the extra ones hold the per-draft states a rejected speculation rolls back to -- so the whole span has to move together or a relocated request keeps drafting against another group's rollback states. That failure is silent, so it gets a test of its own. Both caches are layer-major with the slot as the second axis, so a group's rows are strided and there is no single range to copy; `_foreach_copy_` keeps a batch of pairs to one launch.
…inted halves The free list was one release-ordered queue holding both groups that carry a checkpoint and groups that carry nothing, which gets the eviction order wrong: a checkpoint handed back before a never-used group sits ahead of it in the queue and is spent first. Vacant groups now sit in their own container and are always drawn from before any checkpoint is touched, so a checkpoint can only be evicted once there is nothing free left to take. LRU within the checkpointed half is unchanged -- `claim` deliberately leaves the hash in place, so being resumed from still refreshes a group's position. Which half a free group belongs to is a function of `group_hash` rather than a third piece of state to keep in step, so `_set_hash` is the one place that moves a group across and no caller has to remember to. On top of that the group count stops being fixed for life: `extend` and `retire_top` move it for a pool whose share of the byte budget can change. Retiring is index-forced -- the bytes handed back are the ones the top group occupies -- but its cost is not. The top group's content is relocated and what gets spent is the least recently used checkpoint, wherever that one lives. Retiring by index alone would be anti-LRU: an index records the concurrency high-water mark at hand-out and is never refreshed by use, so the hottest checkpoint can be the one sitting on top. Allocating lowest-index-first is what keeps that cheap rather than correct. A high index is only reached at a high-water mark, so once a peak passes the top of the pool holds the least recently touched things and usually needs no move at all. Nothing calls `extend` or `retire_top` yet; wiring them to a moving KV/state boundary is the next step.
V4 keeps three things per pool: compressed history addressed by block, a sliding window per request, and the compressor's rolling state. They were three regions whose boundaries were startup constants, so none could ever give space back to another -- shrinking the block count moved the start of all 46 layers, and the state's bytes were reserved before the first request arrived. This states the layout as a single **row space** and lets each pool materialize it at its own row width. A row is one token's KV at one layer; what a row costs differs per plane (512 B packed fp8 NoPE, 128 B bf16 RoPE) but a row *index* means the same thing in both. That is a hard requirement rather than a convenience: `pa_sparse_prefill_fp8_opus` and `mla_decode_fwd_v4_nm` each take an NoPE pool, a RoPE pool and ONE shared index buffer. It is also why NoPE and RoPE cannot be two halves of one envelope -- an envelope of E bytes puts a block's row at `b*E/512 + r` in one plane and `b*E/128 + r` in the other, and those agree only for E == 0. Blocks grow from row 0, slots from the far end, and the gap between them is a host-side counter. Every per-layer view anchors at that layer's base row and runs to the end of the plane, so both regions are reached through one base pointer: moving the boundary re-carves nothing and re-captures nothing. Inside an envelope the layers are grouped by compress class rather than interleaved, because within a class every layer contributes the same number of rows. That regularity is what lets one index formula -- and so one index buffer -- serve every layer of a class, which is how the attention path already splits its buffers. A slot is the compressor state followed by every layer's windows. The two are allocated and given up together and no request can have one without the other, so they are one entry class priced as one; splitting them would only invite a split that cannot happen. The state is bytes rather than rows and does not divide by compress class, so it takes whole rows off the front of a slot and each plane materializes its share -- `plan_field_planes` enumerates every assignment of fields to planes and takes the one costing fewest rows, which is 2.4% slack on fp8 and none on bf16. A window's rows cannot be laid out contiguously per layer: the layer term is forced by the compress side, which cannot pad, and a window longer than a class's layer stride would overlap the next layer. So each class's window is cut into stride-sized runs and the runs interleaved, which `entry_rows_for` shows is not merely sufficient but the least any construction can use. The DSpark draft's window stops being a private plane whenever the pool's dtype can hold it. Its layer was always a dense layer of this row space and every slot always reserved its rows; only its bf16 block attention kept it out, so it now says which dtype it needs and takes the shared plane when that is what the pool is made of. Addressing had to widen twice on the way. A row index still fits int32 -- the index buffers are int32 by ABI -- but `row * row_width` does not: at 512 elements per row the plane passes 2^31 elements 2.8% of the way in, and a window row is at the far end by construction, so every window write and all but the first few thousand blocks wrapped silently. `pool_index.row_offset` is now the one place a row becomes an address. Folding the state into the slots widened its own reach the same way, from a 3.12 GB region to a stride running the length of a 152 GB plane -- 17x past a 32-bit product -- so the two kernels indexing it by slot widen too. aiter's gfx950 path already does this through `state_slot_byte_offset`. Verified on V4-Flash-DSpark tp2, bf16, GSM8K 1319 x3 with a restart each: 0.9507 / 0.9500 / 0.9515, mean 0.9507 against a 0.9507 baseline. DSpark (the only path exercising the merged draft window) 0.9515 at 65.2% acceptance. Capacity is neutral on bf16 -- a slot costs exactly what the two regions it replaces did -- and the win is that it is now on an elastic line at all. fp8 is NOT verified: three aiter kernels still turn a row number into an address in 32 bits, listed with evidence and repro in /app/logs_claude/aiter_32bit_addressing_deps.md. The full record of this work, including two GPU faults and what they taught, is in /app/logs_claude/phase_c_slot_geometry_record.md.
The draft layer's window KV is unquantized where the pool is packed, so it could not be rows of a plane and was given a plane of its own. That plane sits outside every slot, and `copy_state_entries` copies slots — so a checkpoint never carried it, and a request resuming a cached prefix drafted against whatever the previous occupant of its slot had left. Verification keeps the output right, so the only symptom is the acceptance rate quietly collapsing. fp8 is the only build that takes that path and it is still blocked on an aiter kernel, which is why nothing has ever run it. Make the window a state field instead. It takes bytes off the front of the slot like the compressor state and is read through a view of its plane retyped to its own width, where `head_dim` of that dtype is one row. `WindowParams` still describes it, counted in that view's rows, so `swa_write`, the DSpark gather and `copy_state_entries` are unchanged. Nothing branches on a dtype: the rows-per-window-row ratio is computed, and a draft whose KV matches the pool comes out at one and costs exactly what a shared dense class would have. `dspark_draft_dtype` becomes `window_kv_dtype`, which is the whole of what a future fused fp8 draft kernel has to change. Such a layer leaves the row space entirely (`ABSENT_RATIO`) rather than also reserving entry rows it would never read. Per slot on V4-Flash-DSpark: bf16 +399 arena rows - 399 entry rows, exactly neutral fp8 -7 rows, and three private planes gone: ~406 KB fp8 with an fp8 draft +0 rows; it fits the RoPE plane's existing slack Gate A (bf16 MTP3, which has no field windows, n=3): 0.9530 / 0.9507 / 0.9530, mean 0.9522 against a 0.9507 baseline. DSpark arm, where the field is live: 0.9507, acceptance 97,895/150,000 = 65.26% against 65.2% before. Acceptance is the discriminator, not accuracy: a mis-addressed window makes the draft read garbage while verification keeps the answers right. Sizing is byte-identical to the recorded numbers in both arms. A ratio other than one is only reachable on fp8, so e2e covers one only; tests/test_v4_pool_geometry.py covers 1, 2 and 8 at the geometry level.
`[Cache Stats]` is logged every hundredth request. A client measuring reuse over a handful of requests cannot wait for that interval, and cannot read a log line at all. Add `GET /debug/cache_stats` beside the MTP one, over the same utility-command path. The engine hands back counters rather than rates. Every rate this reports is a ratio of two of the counters, and DP ranks admit different numbers of tokens, so the mean of their rates is not the rate of their union — the merge has to be a sum, with the four rates derived once from it.
…p back Two changes to the paged pool, both about which block it spends next. The free list was one release-ordered queue holding blocks that still carry reusable content and blocks that carry nothing at all. Handing out the head therefore evicted a cached block while a never-used one waited behind it, on release order alone. Split it: vacant blocks go first and lowest id first, cached blocks after and least-recently-freed first. This is the mistake the state pool's free list already fixed, mirrored. The cached half is an insertion-ordered mapping rather than a deque because `claim` takes a *named* block off the free list on every prefix hit. Left in place, its entry would put the block back at its old position when freed again — LRU inverted for exactly the blocks being reused most. Removing it is O(1) here against O(n) from a deque, on a path that runs once per hit block. The state pool can afford the deque scan; it has 256 entries to this one's 150k. Second, `retire_top` and `extend` make the count movable, which a boundary between the block and slot regions needs. The highest id specifically, since the ids given up have to be the ones the boundary is about to cover: a cached block costs its content, a held one has to move and its new id is reported so the holder's block table can follow. It only fails when the top block is held and nothing is free to move it into — the same condition that would have blocked admitting that request. An entry in either container is stale when the block has been taken *or* has gained or lost a hash, because an id can sit in both at once and testing only that it is free hands a cached block out of the vacant half. Both conditions are the ones that decide which half it belongs to, so the test is the definition rather than a guard on top.
A demand rung is a boundary a request was seen to want and be refused. Two gates borrowed the interval to decide whether to act on one: the demand was dropped unless it sat at least an interval in, and the chunk cut was skipped unless the grid had already placed a rung on that prompt. Together they made the feature unreachable for any prompt shorter than the interval — the default 8192 against a 1024-token prompt left every reusable block declined, with the whole shortfall reported as `Lost-to-checkpoint` and nothing ever stored. The interval is a guess about where reuse will resume. A demand is reuse that was asked for and refused. The granularity of the guess is no reason to discard the evidence, so the grid keeps its rungs and a demand is kept wherever it lands. The position comes from the same forkable test as the hit, so it always leaves `successor_room` behind it: it is a position some class can really keep, grid or no grid. Both gates had to go together. Removing only the threshold would leave short prompts recording demands that `checkpoint_cut` then silently refuses to cut for — "recorded but never acted on", which reads as a bug in the recorder. Measured on DeepSeek-V4-Flash bf16 MTP3, four requests sharing a 1000-token prefix sent serially at the default interval: reuse goes 0 / 0 / 76.8% / 76.8%, where it was 0 throughout. The position it lands on is not a multiple of the interval, which is what says the demand and not the grid produced it. On a 1024/1024 conc-64 benchmark the hit rate goes 0% to 19.25%, at a cost of 38 extra chunk cuts across 384 requests — the mechanism is self-limiting, since a boundary somebody has already checkpointed presents no gap to the next request. Throughput over an interleaved A/B/A/B/A/B: 8022.9 against 7951.6 tok/s, inside each arm's own spread.
`Lost-to-checkpoint` says reuse was declined for want of a checkpoint. It does not say which of the four stages lost it, and the four want different fixes — so each is counted where it happens and reported together: demands_recorded the ladder saw a gap and asked for a rung chunks_cut_for_demand a prefill chunk was shortened to land on it checkpoints_kept the pool filed it checkpoints_dropped the pool had no free group to file it in checkpoints_evicted it was filed, then spent on an allocation The last two read identically in a hit rate and mean opposite things: dropped says the pool is too small for the rate of keeping, evicted that it is too small for how long a checkpoint has to last. `BlockManager.checkpoint_funnel` assembles both halves, since the ladder decides what to ask for and the pool decides what survives, and a reader needs them side by side. This pays for itself immediately. On a conc-64 benchmark the funnel reads 39 → 38 → 38 with dropped and evicted both zero, which refutes two plausible readings of the same 8% shortfall — that checkpoints were being churned out of a full pool, and that chunk cuts were being truncated by the token budget. Neither was happening; 384 requests simply only contain 39 distinct gaps.
A checkpoint holds state as of the forward's last token, so it can only be taken where the forward lands — `checkpointers_at` insists the position be exact, and it is right to. Prefill meets that by construction: `checkpoint_cut` shortens the chunk onto a rung. Generation cannot. A speculative step commits `1 + accepted` and steps over most rungs, so under MTP a decode checkpoint was kept only when the arithmetic happened to divide out. Bounding the drafts would land it there, and it was the first thing I tried. It costs more than it is worth: speculation thrown away near every rung, and `SpecStats` counts drafts offered as `mtp_k` regardless, so the acceptance rate would quietly under-report — the one number that tells us the draft path is healthy. The grid is not what the position has to satisfy. It exists to space checkpoints out, and a step landing on any hash-block boundary far enough past the last one spaces them just as well: a resumer finds a checkpoint by hash, never by arithmetic, so the position has to be findable, not predictable. So generation is held to the spacing and prefill stays on the grid, and `seq.last_checkpoint_pos` is what the spacing is measured from. The exactness invariant is untouched — a landing still has to be the block boundary the hash covers, which is what `pos % hash_block_size` says. What this buys scales with `interval / hash_block_size`: 23% per rung when the interval is one block, effectively certain at the 8192 default. The two rules coincide when the interval *is* the block, which is also the finest grid V4 admits, so a test or a probe at that setting measures nothing. The unit tests use `demand_config`, whose grid is four blocks wide, for exactly that reason; an end-to-end run at interval 256 confirmed only that nothing broke (512 of 512 cached tokens admitted), not that anything improved.
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
1 participant
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.
What this does
Makes the prefix cache correct — and then useful — on models with per-request
state (DeepSeek-V4's compressor ring, GDN's recurrent + conv state), by giving
that state the same content-addressed lifecycle the KV pages already have.
The correctness problem first. A per-request state cannot be rebuilt from cached
KV blocks: the cache holds the compressor's output, the state is its rolling
input window. So a prefix-cache hit on a stateful model was only ever right by
accident — the resumed request got a state group straight off the free list,
still holding the previous occupant's contents. Everything else here is either
what that fix needed underneath it, or what turns the now-correct path from
"safe" into "actually reused".
The shape of it
Sizing (
0da745ae,357d724d). Pool sizing was five ad-hoc builder methodsthat
ModelRunnerstitched together with architecture knowledge of its own.A backend now returns a list of
SubPoolSpec— entry class, pool, per-entrybytes — and
plan_poolsturns that plus a byte budget into entry counts. Twopools, N entry classes:
Pool.PAGEscales with retained history,Pool.STATEwith in-flight requests. Specs sharing a name share an index space, which is how
the Eagle3 draft KV and the V4 indexer cache ride the target's block table
instead of forming a pool of their own. Separately, a request's six per-layer
compressor tensors become views of one contiguous
StateArenaentry — thelayout the PD staging path was already rebuilding by hand on every transfer,
and the unit of movement checkpointing needs.
Checkpoints (
9756b4a8,60f6d8db). A request hands its state group to acontent-hash index at a block boundary and moves to a fresh one, the same way a
filled KV page becomes immutable. A later request whose prefix matches claims the
checkpoint, reads it for one forward, and writes its own. The index holds no
capacity of its own — a checkpoint is a free group whose contents are still
valid, the block pool's lazy-eviction model applied to state groups — so it can
never shrink admission. Rungs are spaced by
--state-checkpoint-interval-tokens(default 8192).
Reach (
e96f2404,e13bdb2b,a1f125ea). Three things had to be true beforegeneration could leave resume points behind: GDN's fork width had to be honest
(1, not
conv_kernel_dim - 1— everycausal_conv1dwrite path stores the fullwindow to the output slot), decode-generated blocks had to enter the prefix
cache at all, and the "can this rung be published" test had to stop being
derived from the prompt.
Two bugs found on the way, both pre-existing
92d836e3— decode blocks were hashed at the committed content length,which under undeferred output (
pipeline_parallel_size > 1) includes the tokenthe forward just sampled. No forward has read that token, so its KV slot is
written by the next one. Because a block is published exactly when the length
reaches its end, and under that mode it always does so on the fresh token,
every generated block entered the cache with an empty last slot. Usually the
next step fills it first — but a seq that stops on that step never writes it at
all, and a later request matching the hash reuses KV that does not exist.
358ba418— the checkpoint alignment could leave a budget remainder toosmall to be worth a step; a 16384-token budget was going out as
..., 640, 10.What it costs when nothing is reused
Publishing is capacity-neutral but costs the publisher a forward: its prompt
gets cut at the rung. The interval is what keeps that amortized rather than
per-request, and it counts tokens, so a prompt shorter than one interval
publishes nothing at all. That matters — on a workload of short, mutually
distinct prompts the hit rate is 0 by construction, so the feature has to be free
there. Publishing unconditionally at the last eligible boundary cost 17.5% of
total throughput for zero hits; at the token-counted default it is noise.
Verification
Unit — 1044 passed / 49 skipped. Every commit was checked out and run
standalone, so the branch is bisect-safe.
plan_poolswent from GPU-only anduntested to a pure function with 71 tests, including a differential grid pinning
the default path to the pre-refactor expression.
Accuracy — DeepSeek-V4-Pro tp8 GSM8K 0.9545 (CI baseline 0.9522);
V4-Flash-DSpark tp2 MTP3 0.9522. Both build a >4 GiB arena and hand out slot 255
first, exercising the descriptor rebase in production rather than only
synthetically.
Qwen3.5-27B tp2, GSM8K chat-mode 300q, with
--state-checkpoint-interval-tokens 64so a rung falls every 64 tokens and every request publishes many times:
Resume, end to end — a 14-token prompt generating 500 tokens leaves a
follow-up request replaying prompt+answer scheduled as
cached: [512], new: [2].The prompt is far too short for prefill to publish anything, so all 512 reused
tokens rest on a checkpoint that generation published.
Perf — Qwen3.5-27B tp2 ISL/OSL 1024/1024 conc 64: 6437–6469 tok/s with
prefix caching on vs 6422.7 with it off, i.e. within noise of each other.
Not in scope
Demand-driven publish points (record where a hit wanted a checkpoint and
publish there next time) and V4 ring carry-forward. The fork-room test is
self-consistent without the latter — V4's width of 131 simply never qualifies
mid-generation, no special case needed.
🤖 Generated with Claude Code