[sysuTobo]

Summary

Implements #1238: keep prefix caching on with the DFlash/DSpark drafter by persisting the drafter's captured target hidden rows per sealed KV block.

• DFlashCaptureCache (new): fixed-budget, LRU, keyed by the KV block's SequenceHash, one 16-token block per slot (5 × 2560 bf16 = 25.6 KB/token on Qwen3-4B).
• PEGAINFER_SPEC_CAPTURE_CACHE_MB (default 128 MiB; explicit 0 = off, as is --no-prefix-cache), charged as a fixed term in DFlashMemoryReservation together with the per-request partial-block seal and the lane shift buffer.
• On a prefix hit the worker restores the matched blocks into pending_context; the suffix prefill appends the rest and the first draft round rebuilds the draft KV from the restored rows.
• A capture-cache miss keeps the target KV hit but leaves the request without drafter state, so the gate falls back to plain decode. A zero budget keeps the old spec-on => APC off behavior.
• The verify plan's page-list capacity now uses the by-reference bound (max_batch × pages-per-context), because shared prefix pages are listed once per request.

Review fixes (head 51a6c21a):

• Publish on finish: a retired request hands the lane its final block hashes and kv_position, and the lane drains the seal into the capture cache before dropping the state. Previously the next turn plain-decoded whenever the last served step crossed a block boundary.
• Bounded request-owned draft state: the reservation now exposes the per-request share of its per-token term (draft KV + projected-context scratch + pending capture, 56,320 B/token on Qwen3-4B). The lane charges each request's max_cache_len against the memory the pool set aside for draft state; a request that would exceed it takes the capture-miss path and plain-decodes. This keeps N requests sharing one prefix from each allocating a prompt of state the pool charged once.
• Long-prefix cap: a request with more full blocks than the capture cache has slots is not published (it could never restore whole and would evict every other conversation, including its own head). The seal still drains so it cannot overflow. At the 128 MiB default the longest restorable prefix is ~5.2k tokens.
• Capture misses are logged at debug level; they are the normal case once the LRU churns.

Verification

• cargo check -p pegainfer-qwen3
• cargo clippy -p pegainfer-qwen3 --all-targets -- -D warnings
• cargo test -p pegainfer-qwen3 --lib → 133 passed
• One L20 (46 GB, Qwen3-4B + deepseek-ai/dflash_qwen3_4b_block7, release), full gate file 10 passed / 1 ignored (168 s):
  • dflash_warm_prefix_keeps_drafting → 1 passed; cold and warm both 4 draft rounds / 27 accepted.
  • dflash_capture_cache_miss_falls_back_to_plain_decode → 1 passed; 400 cached tokens, drafts 4 → 4.
  • dflash_finish_publishes_last_sealed_block (new) → 1 passed; the drop publish reports target_blocks=50 from published=49, and the next turn drafts.
  • dflash_spec_levers_keep_losslessness → 1 passed; the no-prefix, gated, and trimmed children as before.
• CI on the current head is green (CPU clippy/unit, Qwen3 CUDA clippy sm_80, formatting, metadata, DCO, simulated frontend E2E).

Caveats

• The capture cache is a lane-owned LRU keyed by the block hash, not part of the KV block's refcount; moving the payload into the block state so both share one lifetime is the qwen3: keep APC on with DFlash/DSpark by persisting the drafter capture in the prefix state #1238 follow-up.
• The per-request draft state is now bounded against the reserved draft term, but it is still allocated outside the pool; billing it by reference in the pool itself remains qwen3: keep APC on with DFlash/DSpark by persisting the drafter capture in the prefix state #1238.
• The capture cache is all-or-nothing per hit; a miss costs the request's speculation (KV reuse remains). A joint cap that re-prefills the missing suffix is a follow-up.
• A turn whose final token lands exactly on a block boundary ends on the verify's bonus token, whose target hidden was never forwarded. That last block cannot be completed on finish, so the next turn's full-boundary match still plain-decodes; the joint cap above is what would recover it.

Refs #1238.

[github-actions]

This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged.

[FeathBow]

Requesting changes for the two points below; the mechanism itself holds up on this card.

Passed on an sm_89 card with Qwen3-4B: the gate file with the b16 drafter, the two new tests with the b16 and the DSpark block7 drafters (the DSpark numbers match your L20 run exactly), and a zero budget launch that keeps the prefix cache off. Two mutations of the restore, copying nothing and copying block i from slot i+1, both fail the warm gate on its accept rate floor, so the gate has the power it claims.

Publish on finish: the last step's sealed blocks never reach the cache, so the next turn of a conversation plain decodes whenever that step crossed a block boundary. Measured 2 of 12 two turn conversations, the same two with both drafters; details inline.

Draft state bound: with sharing, the per pool token billing no longer covers the lane. One seed request, then N warm requests on the same 4000 token prefix with 16 output tokens, b16 drafter, card total 49.1 GB:

N	main (prefix cache off with the drafter)	this PR
64	finished in 27 s, peak 41.7 GB	finished, peak 47.9 GB
96	finished in 42 s, peak 41.7 GB	no request finished or failed in 15 minutes, pinned at 48.5 GB, no error reported, killed by timeout

On main the 96 requests cannot all be resident (each owns 250 pool blocks of 9051) and run in waves; here they share 250 blocks, are admitted together, and each allocates about 220 MB of draft state, about 21 GB against a reservation near 9 GB. Details and two possible bounds inline.

Two smaller notes inline on the long prefix publish and the per miss warn.

[FeathBow]

Please also publish when a request finishes. The blocks sealed by its last step never get here, so the next turn's all or nothing lookup misses whenever that step crossed a block boundary. Twelve two turn conversations on this card (260 token prompt, 20 to 67 response tokens, then prompt + response + a short suffix): 10 draft on the second turn, 2 plain decode with a KV hit, the same two with both drafters, and both are the boundary cases (block 18 sealed at token 288 of 291, block 19 at token 304 of 304). The seal already holds those rows, and drop_request still has the request's KV handle before it drops it, so the drop command can carry the final hashes and kv_position and the lane can publish from the seal before discarding the state.

[FeathBow]

This TODO carries the default now. With the prefix cache forced off, every active request owned its context in the pool, so the per pool token billing covered what the lane allocates per request. With sharing it does not: 96 warm requests on a shared 4000 token prefix sat at 48.5 of 49.1 GB without finishing or failing until a 15 minute timeout killed the process, and 64 already reached 47.9 GB; the same 96 on main finish in 42 s at 41.7 GB. Before this ships as the default the lane needs a bound it can act on: count the bytes its request states hold against the reservation and treat a request that would exceed it like a capture miss, or keep the capture cache opt in until the draft state is billed by reference.

[FeathBow]

At 128 MiB the cache has 327 slots, so a request longer than that publishes over its own head: 64 warm requests on an 8000 token prefix all hit the KV cache and none restored, and every other conversation was evicted on the way. Skipping the publish when the request's block count exceeds the slot count, and saying in the doc that the budget caps the longest restorable prefix (about 5k tokens at the default), would keep this from surprising people.

[FeathBow]

Once the LRU churns this fires for most requests; debug level or a counter.

[sysuTobo]

Both blocking points and the two smaller notes are fixed in 51a6c21a.

Publish on finish. The drop path now carries the request's final block hashes + kv_position into the lane, which drains the seal into the capture cache before dropping the state. New gate dflash_finish_publishes_last_sealed_block: the first turn ends one token past a boundary, the drop publish reports target_blocks=50 from published=49, and the second turn drafts. The full gate file is 10 passed / 1 ignored (168 s); warm stays 4 rounds / 27 accepted on both sides, the 1 MiB miss gate stays 4 → 4, and the lever rows are unchanged.

One caveat for your 304/304 case: when the final token is the boundary token, it is the verify bonus, whose target hidden was never forwarded, so that last block cannot be completed on finish even with the drop publish. The recoverable case — the final step crosses a boundary and leaves a partial tail after it — is what the new gate pins. Recovering the exact-boundary case needs the joint cap (re-prefill one block); I added it to the #1238 follow-ups and the PR body.

Draft-state bound. The reservation now exposes the per-request share of its per-token term (56,320 B/token: draft KV + projected-context scratch + pending capture; the lane-level tail scratch is excluded). The executor passes the reserved draft term (pool_blocks × block_size × 65,536; 8,556 MiB on the L20 run) to the lane, which charges max_cache_len × 56,320 per request and sends an over-budget request down the capture-miss path (plain decode). The reservation's total is unchanged; the bound stays within what was already set aside.

Long-prefix publish. A request with more blocks than the cache has slots no longer inserts (it could never restore whole and would evict every other conversation, including its own head); the seal still drains so it cannot overflow. The docs now state that the budget caps the longest restorable prefix (~5.2k tokens at the 128 MiB default).

Per-miss warn. Moved to debug, with a debug ledger line per publish (request/target/published/captured/sealed/slots).

CI on the current head is green.

[sysuTobo]

Stage 2 (paged draft KV in the target's block space) is a larger follow-up, so it stays out of this PR. Full design, risks, and gates are in the issue: #1238 (comment)

Summary of the plan:

• Replace the capture cache with a per-layer draft KV pool addressed by the target's block IDs (a separate tensor over the same allocator/slot space, like SGLang's DFlash worker). Draft KV is 20,480 B/token (~14% of target KV on Qwen3-4B).
• Materialize immediately after target prefill/verify: fc + hidden_norm once, then per-layer K/V projection + K norm + RoPE, scattered to the target's page slots. Gated/plain verify writes too. No capture hash table, no first-round rebuild.
• Warm hit: the matched block IDs already carry draft KV; copy it into the request buffer (phase 1), or read the pool paged (phase 2). No capture lookup.
• Remove DFlashCaptureCache, the capture hashes/seal bookkeeping, PEGAINFER_SPEC_CAPTURE_CACHE_MB, and the capture/seal reservation terms; the per-request floor drops from 56,320 B/token toward ~25 KB/token.
• Lifetime/refcount follows the target block, so shared prefixes pay one physical draft KV copy and the ~5.2k-token capture cliff disappears.
• Risks: a cache_loc scatter writer, absolute positions, page-ID stability during a step, block-reuse pinning, TP coordination, and the exact-boundary bonus case.
• Gates: cold vs warm accept rate, first-step ITL A/B, 6k/8k-token prefixes, 64/96 shared-prefix requests within reservation, mixed/gated/chunked batches, eviction/reuse, TP1/TP2.

Stage 1 in this PR remains the capture-cache implementation under review.

[FeathBow]

Requesting changes once more, for one thing the round one fixes uncovered. The four points from round one are closed and measured on the same sm_89 card with both drafters.

Closed: the finish publish recovers the case I measured (the conversation whose last step sealed block 18 at token 288 of 291 now drafts on turn 2; 11 of the 12 conversations draft, and the new finish gate passes with both drafters); the draft state bound holds (the 96 shared prefix requests that wedged the card last round finish in 2.4 s at 40.8 GB, below main's 41.7 GB); the slot cap and the debug level are in; a zero budget still keeps the prefix cache off.

The remaining conversation is the boundary case you documented, and it is a KV side defect rather than a capture limitation. apply_speculative in the kvbm fork appends the whole committed span, bonus included, and then registers every block the appends completed, so a block completed by the bonus becomes matchable before the bonus has any KV. apply_decode registers before appending the new dangling token and apply_prefill leaves the sampled token dangling for exactly this reason. With the prefix cache forced off under the drafter nothing ever matched those blocks; with this PR the next turn does, and its prefill attends to a row that holds a rejected draft's KV or nothing. Ten turns per drafter that end exactly on a block boundary:

last token of turn 1	cases	turn 2 prefix match	turn 2 drafts
prefill's sampled token, never forwarded	2	one block short, as it should be	yes
verify bonus token, never forwarded	8	the full turn, including the block whose last row was never written	no

Turn 2 warm against the same prompt on a fresh engine: in 12 of the 16 bonus endings the first token's logprob moves by 0.05 to 0.94, while the four controls that sit one block short of a full match move it by at most 0.04. The capture miss is the safe half of this event; the plain decode that follows still runs on the corrupted prefix. The fix belongs in the fork's apply_speculative: register what the forward computed, then append the last token of the span without registering, the way apply_decode already does. After that every registered block has all its rows forwarded, the finish publish covers everything the KV side can match, and the caveat in the body, the doc and the gate comment can go.

One doc point inline: past the budget the whole decode batch stops speculating, because the step decision is all or nothing and a plain decode drops the restored states. 64 requests on a shared 4000 token prefix: 41 restored, 23 refused, 0 draft rounds for anyone; 4 requests: 44 rounds; main runs 704 rounds for the 64 in waves.

[FeathBow]

This caveat describes a KV side defect, and fixing it removes the caveat. apply_speculative in the kvbm fork appends the committed span, bonus included, and then registers every block the appends completed, so a block completed by the bonus is matchable before the bonus has any KV. apply_decode registers before appending the new dangling token and apply_prefill leaves the sampled token dangling for exactly this reason, and the box shows the asymmetry: a turn whose only generated token came from prefill matches one block short on the next turn and drafts, every turn that ends on a verify bonus at a boundary matches the full turn and plain decodes, and the warm first token's logprob moves by 0.05 to 0.94 against at most 0.04 for the controls. Registering before appending the span's last token, as the decode path does, makes the next turn match only rows that were forwarded; then this sentence, the one in the body and the gate comment stop being true and can go.

[FeathBow]

Per request this is right; per batch the step decision is all or nothing and the plain decode drops every restored state in the batch, so the first refused request ends speculation for everyone sharing the step. Measured with 64 requests on a shared 4000 token prefix: 41 restored, 23 refused here, 0 draft rounds for all 64 (the publish ledger shows only the seed); with 4 requests, 44 rounds. Worth one sentence in the doc next to the bound, with the ceiling it implies: the budget divided by max_cache_len times 56,320, which is 41 requests at 4k tokens on this card.