[Bugfix][Spec Decode] Mask prefix-cache-restored tokens out of the DFlash/DSpark draft context

[giorgiopiatti-caffeinated]

Purpose

DFlash/DSpark drafters build their context KV from the target's auxiliary hidden states, which only exist for tokens that flow through a target forward pass (precompute_and_store_context_kv, written per step for the scheduled tokens in DFlashSpeculator.propose). Tokens whose KV is restored at request (re)admission — automatic prefix caching hits, KV-connector restores, resumption after preemption — never run through the target, so their draft context KV slots are never written. The draft's attention nevertheless spans the full sequence (seq_lens = last_valid_pos + 1 + num_query_per_req in _prepare_dflash_inputs_kernel), so it reads uninitialized/stale KV for the whole restored region.

The impact scales with the cache-hit length: workloads where requests share a long common prefix degrade toward position-0-only acceptance with mean acceptance length ~1.0 (drafts are garbage from the first position), while MTP/EAGLE-style drafters on the same workload are unaffected because they draft from the last decode hidden state and build no context KV — which is what makes this easy to misattribute to the speculator checkpoint or quantization. Existing public DSpark benchmark recipes work around it by serving with --no-enable-prefix-caching, giving up prefix caching entirely.

This PR keeps prefix caching enabled and instead hides the restored tokens from the draft's attention:

• RequestState.num_cached_tokens (new): per request-slot count of tokens whose KV was restored at the most recent (re)admission, taken from the num_computed_tokens already passed to add_request (covers APC hits, connector restores, and resume-after-preemption uniformly).
• GPUModelRunner hands the buffer to the speculator via set_num_cached_tokens (hasattr-guarded; only DFlash/DSpark opt in).
• _prepare_dflash_inputs_kernel shortens the draft's seq_lens by the restored whole blocks, and a new Triton kernel (shift_draft_block_tables) left-shifts each request's draft block-table row in place by the same amount, so the draft attends only over slots that actually hold draft KV.

Design notes:

• The in-place shift is safe: input_block_tables are regathered from the persistent block tables every step, the shift runs after prepare_dflash_inputs (slot mappings index the unshifted table), and all consumers are ordered on the same stream. It is CUDA-graph compatible — graphs read the live buffer contents at replay.
• No position rewriting is needed: draft KV stores post-RoPE keys at absolute positions, and the shift is whole-block, so logical-to-physical block lookups agree between the (absolute) write path and the (shifted) read path.
• Requests without cache hits shift by zero and are bit-identical to current behavior; same for dense DFlash/DSpark checkpoints.
• Dummy runs skip the shift (their idx_mapping doesn't reference live requests) and the kernel clamps seq_lens to at least the query length as a guard.
• Known limitation: when the restored count is not block-aligned (e.g. a full-prompt cache hit, which recomputes only the last token), up to block_size - 1 restored slots remain visible to the draft. Output correctness is unaffected (rejection sampling), and the common APC case is block-aligned.

The draft loses the cached prefix from its visible context — bounded in practice by the drafter's training window anyway — in exchange for prefix caching and DFlash/DSpark composing at all. A follow-up could restore the full-context case by letting the draft KV cache group's blocks be reused with validity guarantees (the block-reuse machinery already exists; the gap is that cache_full_blocks is content-agnostic, so ingestion paths where the drafter never ran — e.g. KV-connector restores — can cache draft-group blocks that were never written). Happy to open that as a separate issue/RFC.

Test Plan

• New unit tests: tests/v1/spec_decode/test_dflash_prefix_cache_masking.py — shift kernel correctness for aligned/unaligned/zero shifts, per-request shifts through idx_mapping, and the overlapping in-place copy on long rows.
• End-to-end A/B with aiperf on GLM-5.2 (eigen-ai-labs/GLM-5.2-NVFP4, TP=4, B200) + RedHatAI/GLM-5.2-speculator.dspark, prefix caching enabled in both runs:

--speculative-config '{"model": "RedHatAI/GLM-5.2-speculator.dspark",
  "num_speculative_tokens": 7, "method": "dspark",
  "attention_backend": "FLASH_ATTN", "draft_sample_method": "probabilistic"}'

aiperf: --public-dataset spec-bench --extra-inputs temperature:0 top_p:1 --output-tokens-mean 256 --output-tokens-stddev 128 --concurrency 8 --request-count 200, isolated pod, metrics from vllm:spec_decode_num_accepted_tokens / vllm:spec_decode_num_draft_tokens / draft-step counters over the run window.

Test Result

SpecBench prompts are short, so cache hits only cover the shared chat-template/system prefix — even that is enough to measurably depress acceptance before the fix:

aiperf spec-bench (temp=0, 200/200 ok)	before (same nightly)	after (this PR)
Token acceptance (accepted / draft)	20.6%	27.6% (19,871 / 71,890)
Accepted draft tokens per step	1.44	1.93 (19,871 / 10,270)

Per-position accepted counts after the fix decay smoothly (pos0→pos6): 7,506 / 4,887 / 3,084 / 1,925 / 1,236 / 765 / 468 — versus an effectively position-0-only profile on cache-hit tokens before.

The improvement grows with the shared-prefix length of the workload. On an internal benchmark with a long shared system-prompt prefix (~20k tokens, near-total prefix-cache hit rate; same serve config, aiperf, temperature/top_p at serving defaults, concurrency 8, 200/200 ok):

internal long-shared-prefix benchmark	before (same nightly)	after (this PR)
Token acceptance (accepted / draft)	~0.6%	38.2% (40,689 / 106,575)
Mean acceptance length (vLLM log)	~1.02	3.41
Per-position acceptance rate (log)	pos0-only ~2–5%, pos1+ ≈ 0	83.5 / 57.2 / 38.9 / 27.8 / 16.8 / 11.1 / 6.0 %

i.e. speculative decoding goes from fully broken (drafts rejected from position 0) to a healthy per-position decay in line with the speculator's model card, without disabling prefix caching.

---

Essential Elements checklist

• Purpose of the PR explained
• Test plan included
• Test results included
• (docs update not applicable)

[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. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

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.

🚀

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @giorgiopiatti.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork