[Bugfix][Core][Spec Decode] Enable prefix caching for the Kimi K3 DFlash drafter

[rchalamala]

Purpose

A Kimi K3 DFlash lane does not work in vLLM today. With --enable-prefix-caching it dies on the first request:

AssertionError: SlidingWindowManager does not support fine-grained (partial) cache hits

Without prefix caching it survives, but every turn of a multi-turn agentic workload re-reads a ~33,000-token prompt, which is not a usable lane either. This PR is the set of changes needed to make it run.

The root cause of the assertion is that the DFlash drafter declares all of its layers sliding_attention. That builds a SlidingWindowSpec group whose manager cannot serve the fine-grained partial hits the hybrid coordinator hands it. Booking those layers as FullAttentionSpec for KV-accounting purposes — while still enforcing the window at compute time — removes the group that cannot take the hits, and the rest of the changes are the fallout and adjacent defects found on the way.

Based on kimi-k3 rather than main, since the K3 model wiring this lane needs lands there (vllm-project#50000). Only commit 7 touches K3 model code; the other seven are model-independent and apply to main unchanged.

Prefix caching

1. Bound the rejection sampler per-block argmax gather. tl.argmax is reduced over the padded lane count and the resulting index gathers from a companion tensor sized to the real block count, unclamped and unmasked, at two sites in rejection_sampler_utils.py. A padded lane cannot win on value, but it wins when a real lane holds NaN, since Triton's argmax combine is false in both comparisons for an unordered operand. Needed for the lane to survive V2 warmup at depth 16.
2. Reconcile hybrid prefix-cache hits across every KV cache group. Disables fine-grained hits when a group needs block alignment, and trims every downward-closed group to the reconciled hit rather than only attention_groups[0] — commit 3 creates a second full-attention group, and the truncation assumed there was one. Declares downward-closure as a manager property instead of repeating an isinstance check per site, and shares the predicate and the truncation with the Mooncake store mirror so the two cannot drift.
3. Let an all-sliding DFlash drafter book its KV as full attention. The change that actually buys the cache. sliding_window is preserved on the converted spec because FlashAttentionMetadataBuilder.build reads the window off the spec; dropping it would silently enforce no window at all. Derived from the drafter's own layer types and the caching state rather than gated on an environment variable.
4. Reconcile the connector's dense reference across full-attention groups. With two dense groups at different block sizes, the eviction guard and the reported hit length are different questions: max over dense hits for the guard, since a group is only evidence of eviction if it hit deeper than every dense group; min for the reported length, since it must be a length every returned block list covers. Both collapse to the previous behaviour with a single full-attention group.

Speculative decoding correctness

5. Reserve spec-decode lookahead blocks in V2 warmup, matching what KVCacheManager.allocate_slots actually reserves. The reservation now derives from VllmConfig.num_lookahead_tokens, the same property the scheduler reads, rather than from a second copy of the rule.
6. Require all requests to be decoding before dispatching a uniform-decode CUDA graph. A chunked prefill whose final chunk equals 1 + num_speculative_tokens otherwise gets a decode graph replayed over prompt tokens — wrong output, silently.
7. Tap the pre-norm AttnRes mixture as the K3 DFlash aux hidden state.

Diagnostics

8. Warn when --block-size is silently discarded by backend alignment. Omitting the flag does not fall back to the block-size default; it lands on the MLA floor of 128, so any value at or below 128 is indistinguishable from passing nothing. This cost us a day of chasing an acceptance difference that was a serving flag.

One environment variable remains, VLLM_KIMI_K3_AUX_ATTN_RES_STREAM, gating commit 7 and defaulting off. It changes model behaviour, which probably belongs in SpeculativeConfig; it is a flag here so the change could be A/B'd on a single container. Happy to convert it before merge.

Test Plan

Unit tests, all new or extended by this PR:

pytest tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py \
       tests/v1/spec_decode/test_dflash_draft_full_attention_spec.py \
       tests/v1/worker/test_gpu_warmup_blocks.py \
       tests/v1/worker/test_uniform_decode_token_count.py \
       tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py \
       tests/models/kimi_k3/test_aux_attn_res_stream.py

The warmup test validates _reserved_block_count against the real KVCacheManager.allocate_slots rather than against a hand-computed number, so it cannot drift from the allocator it is predicting. It does so for a full-attention group alongside one Mamba group per cache mode — "none", "all" and "align".

End to end: Kimi K3, 8x B300, TP8, DFlash drafter at k=16, --enable-prefix-caching, V2 model runner, marlin MoE. Multi-turn agentic trajectory replay with ~33,000-token prompts, swept at 1/4/8/16 concurrent users for 120 s per level. Patch liveness confirmed from the boot log rather than from the source tree, since a patched file on disk is not evidence the running process carries it.

Test Result

Unit: 98 tests, of which 95 were run against this tree and pass. The remaining 3 are parametrizations of a shared fixture that fetches facebook/opt-125m, which the offline test environment cannot reach; they are untouched by this PR. The 26 warmup cases were run on a CUDA box, the rest on CPU.

Boot: clean with prefix caching enabled. Engine init 251.6 s, KV pool 1,870,846 tokens, attention block size 1920. All eight workers log the drafter's full-attention booking, and commit 2's "disabling fine-grained hits" warning never fires — the drafter no longer forms a sliding-window group, which is the mechanism being fixed.

The assertion is gone. Before this PR the same configuration dies on request one.

Caching, same container, only the gate changed:

	before	after
prefix cache hit rate	46.5%	95.5%
TTFT p50	3.45 s	0.91 s
requests / 120 s	22	50
completion throughput	146	314 tok/s/GPU

Sweep on the agentic workload, zero failed requests at every level:

concurrent users	requests	failed	prefix cache hit rate	TTFT p50	decode p50
1	76	0	93.5%	1.05 s	161.5 tok/s/user
4	101	0	91.6%	1.80 s	59.8
8	144	0	90.1%	2.18 s	25.6
16	127	0	83.0%	2.37 s	14.2

Hit rate holds above 83% out to 16 users and TTFT stays near a second on a 33k prompt. Reproduced on three machines, the third built from a written recipe rather than by hand.

What this does not claim. Draft acceptance is not claimed to improve. Acceptance moves with an unrelated serving flag (--block-size, via the MLA alignment floor — see commit 8), and no contribution from these changes has been isolated. Commit 7's +1.903 accept length was measured with prefix caching off; with caching on it is +0.083, inside the noise, and it is included for the mechanism rather than the number.

DFlash also still trails an unspeculated lane at 8 concurrent users on this workload. That is the compute cost of running sixteen sequential drafter forwards, not a caching problem, and nothing here addresses it.

Known gaps.

• Commit 4 is not exercised end to end, as we have no KV-connector deployment. It is covered by a unit test that pins the returned pair as internally consistent, which is a weaker claim than a live connector.
• Commit 1's tests are guards rather than regression tests: the original failure was found by reading the kernel and never deliberately reproduced.
• Commit 7's unit test covers which weights the tap selects, including the pipeline-parallel fallback. It does not assert the mixture against what the drafter was trained on; the evidence for that is the measured accept-length A/B, not a test.
• Withdrawn: an earlier revision of this description said a correctness eval (AIME 2025) was running and that I would post the result. That run did not complete and no eval has been run against this branch. Treat correctness here as unvalidated beyond the unit tests and the serving sweep above.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use /ci run or /ci retry. New commits do not start CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀