[zheyuf]

Stacked on #21: merge #21 first, then review the last 2 commits of this PR.

Description

With KV-cache block reuse on, the standalone DSpark drafter now keeps its context on a prefix-cache hit. Its draft pool used 64-token pages, so it could not pair with the target's 128-token prefix cache: on a hit the target reused the prefix while the drafter started from nothing. On #21 alone, greedy AL drops 9.4% from a cold request to the same request on a prefix-cache hit, and a sampled warm pass loses 6-8%.

1. Keep a DFlash generation request's drafter context across a missed iteration (pre-existing bug, independent of reuse). A generation request absent from one batch (e.g. suspended by the V2 scheduler) lost its drafter slot and came back at context length 0: it drafted from nothing and wrote its accepted K/V at positions 0..K through its block table. On the manager-bound pool the context stays in the request's pages, so a freed slot parks its read base, the returning request resumes it, and the length follows the target position of its first input token.
2. Reuse the drafter's KV with the target's prefix cache.
   • Page size by capability: without the non-causal context kernel, the DFlash TRTLLM ops serve power-of-two pages from 128 tokens (flashinfer's dynamic page-size kernels), so the draft pool takes the target's page size and pairs with it for joint block reuse. Only a pool the drafter reads is paired (not a private VANILLA / FA4 arena).
   • Claims: after a prefill of L tokens the drafter reads its context from L + 2 - W on (window W). A claim is valid where both pools hold it, or, with no draft page reused (a fresh mirror), up to L + 1 - W of the target's match. A fresh mirror whose prompt ends within one window of the shared prefix (its branch point, e.g. the end of a system prompt) recomputes from one window before that point and publishes the drafter's window there, so later requests branching off the same prefix claim it in both pools.
   • Publication: the paired pool publishes only pages the drafter wrote.
   • Host tier: an explicit host budget gives the paired pool as much as its GPU funding (at most half), a size that is the same on every rank.
   • Attention DP: the pad dummy also gets a cache in the paired draft pool. The V2 scheduler admits both pools together, so a dummy without one was never admitted and every rank skipped the forward.

Nothing to enable: with block reuse and KV cache manager V2, the drafter binds the draft pool at page_size=128 with reuse_addressable=True.

Results

GB300, MiniMax-M3.1-Flash + DSpark K7, on a wheel of the base (1aa2dfc) with #21 and this PR applied on top; one fresh server per cell. Reuse cells use a production-like recipe (TP4 EP1 bs30, per_conversation reuse, 140 GiB host tier). In the #21-alone column, the golden / default and agentic cells were measured on #21's current head; the other three cells come from earlier runs of #21.

	#21 alone	+ this PR
Golden semantics K7 (RS true + block), reuse off, 80 prompts c64, 2 passes	4.069 / 4.041	4.080 / 3.984
Default K7 (RS false), same	3.594	3.599 / 3.602
Sampled warm pass vs fresh, same 80 prompts, c64	-5.8% / -7.7%	+1.6%
Greedy c1, cold vs prefix-cache hit on one server (56 requests: repeats, two-turn, 16K / 64K shared prefix + new turn)	-9.4%	-0.9%
MLPerf agentic v6 (531 requests, 52 conversations, 48 in flight, MSL 1M), real AL	3.037	3.203 / 3.287 (+6.9%)
same: prefilled prompt tokens	561,484	535,300 / 535,289
Forced-AL step time (AL 2.6), reuse on, device median B1 / B8 / B64	8.079 / 14.374 / 28.080 ms	7.928 / 14.396 / 28.258 ms

• Agentic v6: 40 of the 52 conversations start with the same 30.6K-token system prompt. The first request that branches off it recomputes the drafter's window there once (4,397 tokens; the only draft-pool recompute in each run) and publishes it; later conversation starts claim the prompt in both pools (509 of the first 512 claims full).
• 10-minute soak (agentic replays + long-prompt sampled passes + 4K-token floods at once, 400K-token KV cap): 5 x 531 agentic requests, 13 sampled passes, 8 x 400 flood requests, 0 errors, no stall.
• Attention DP (DEP4, reuse on): 32-request sampled pass AL 3.671, 0 errors; greedy cold vs hit (34 requests) 3.831 -> 3.837.
• Step time is unchanged within the server-to-server spread (up to 6% at B64 between servers of one build).

Test Coverage

• test_paired_draft_reuse.py (new, GPU): two real KVCacheManagerV2s + KVCacheV2Scheduler: continuation hit, branch point (fresh-mirror recompute, then shared), commit gate with a negative control, attention-DP pad dummy admission.
• test_dflash_slot_resume.py (new, CPU): a request that misses an iteration resumes its parked context.
• test_kv_cache_estimation.py: page-size capability rule.
• GB300, on the stack: the KV-manager / scheduler / estimation / DFlash / DSpark / MSA / executor / disaggregation unit files, 1,676 passed; the 19 failures fail identically on [None][fix] MiniMax-M3.1 DSpark AL parity: non-causal draft blocks + opt-in golden sampling #21 alone (test_kv_cache_budget_split.py x8, test_estimation_temporarily_uses_inferred_pool_sizing x5, test_forward_step_carries_context_logits_request_to_runner x5, test_mla_dspark_auto_backend_resolves_to_cutedsl).
• Release Checks: as on [None][fix] MiniMax-M3.1 DSpark AL parity: non-causal draft blocks + opt-in golden sampling #21, cleanup_stale_prs.test.js and every pre-commit hook pass on the changed files; scripts/release_check.py fails only at its whole-package bandit scan, on the 10 issues the branch head already has.

🤖 Generated with Claude Code