[PCP][Spec Decode] Simplify PCP+MTP draft metadata handling; add DSpark support

[LucasWilkinson]

Simplify draft metadata handling for PCP+MTP

Suggested simplification for vllm-project#53427, offered as a PR into this PR's branch so it can be reviewed/merged (or cherry-picked) as a whole.

What changes

Build the draft prefill attention metadata inside AutoRegressiveSpeculator.propose(), gated on the drafter running replicated under PCP — mirroring DFlash, which always builds its own metadata in propose(). Building at that point (after prepare_prefill_inputs has filled the drafter's own buffers, and after the drafter's own cudagraph dispatch) makes most of the plumbing unnecessary:

• Drops pcp_manager._maybe_prepare_replicated_pcp_attn, which reached into speculator._build_draft_attn_metadata from outside the class.
• Reverts the query_start_loc_gpu/seq_lens parameters threaded through _build_draft_attn_metadata: at the build point the drafter's input_buffers already hold the global batch's values, so the base signature works unchanged.
• Reverts all DFlash changes (no longer needed once nothing routes PCP builds through it).
• Reduces the model_runner.py change to inlining the pre-hc hidden-state buffer restore (+6/−1); the propose call sites are untouched.
• Drops maybe_restore_pcp_hidden_state_buffer and maybe_restore_pcp_for_speculator (the latter is dead code: PCP requires MTP, which never produces aux hidden states).
• Replaces the mock-based tests of the free function with the two behavioral tests.

Net PR footprint becomes 5 files, +232/−45 vs main (from +374/−48).

Two correctness notes on the placement:

• The build uses the drafter's own prefill_batch_desc from its cudagraph dispatch rather than assuming the global batch's padding matches.
• No capture-side handling is needed because PCPManager.validate_config already rejects FULL cudagraph modes, and PIECEWISE metadata lives outside the graphs.

Validation (4x B300, zai-org/GLM-4.7-Flash unless noted)

• pytest tests/v1/worker/test_gpu_pcp_manager.py: 2 passed.
• pre-commit (ruff, mypy) clean on changed files.
• Qwen3-8B + eagle3 TP1 smoke (metadata-reuse path): outputs byte-identical to the pre-change baseline, acceptance identical (738 drafted / 261 accepted, per-pos [146, 78, 37]).
• GLM-4.7-Flash MTP3 PCP4 -cc.cudagraph_mode=NONE smoke: 7/8 prompts byte-identical to the prior validated run; the eighth is the known base-PCP greedy-degeneration prompt, where this run matches the spec-off PCP4 baseline byte-for-byte.
• GSM8K 200q 5-shot (temperature 0, concurrency 32, MTP3, cudagraph_mode=NONE):

Configuration	exact_match (strict)	Acceptance rate	Mean acceptance length	Per-position acceptance
TP4	0.835 ± 0.026	40.8–41.6%	2.23–2.25	0.924 / 0.263 / 0.038
PCP4	0.850 ± 0.025	40.6–41.2%	2.22–2.23	0.919 / 0.270 / 0.045

Accuracy within one standard error of each other; acceptance metrics aligned, consistent with the fixed-K results reported in vllm-project#53427.

DSpark support (second commit)

Because DFlash-family speculators already build all of their own attention metadata inside propose() — the pattern the first commit adopts for MTP — extending PCP speculative decoding to DSpark needs no metadata plumbing at all, just:

1. validate_config: admit dspark (including sparse-MLA targets with it; MTP remains dense-MLA only).
2. Restore aux hidden states to the global batch layout before propose() — the DSpark drafter reads aux layers, which are rank-local under PCP. (Inlined at the call site, mirroring the main hidden-states restore.)
3. Exempt draft layers from check_attention_cp_compatibility: the replicated drafter never runs CP-sharded attention, but its dense-attention layers resolve to backends without supports_pcp() (e.g. FLASHINFER). MTP passes this check only because its draft layers are MLA like the target.
4. Re-gather block tables for the global batch in DFlash's propose() under PCP — the runner gathered them for the rank-local sharded batch, which silently corrupts the drafter's context-KV slot mappings (signature: per-position acceptance [50, 0, 0, 0, ...], anchor-only, with outputs still correct since verification is target-side).

DSpark validation (8x B300)

GLM-5.2-FP8 (sparse-MLA DSA) + RedHatAI/GLM-5.2-speculator.dspark, k=8, cudagraph_mode=NONE, greedy:

Configuration	Acceptance rate	Per-position accepted
TP4	0.2716	115 / 78 / 56 / 38 / 26 / 18 / 12 / 9
TP4 x PCP2	0.2952	119 / 73 / 51 / 36 / 28 / 25 / 16 / 11

Acceptance aligned, per-position decay matches, outputs coherent (4/8 prompts byte-identical to TP4, remainder late FP-numerics divergences of the same character as base-model TP-vs-PCP differences). Note the drafter runs eagerly under PCP since PCP caps the cudagraph mode below FULL, which is what DFlash decode graphs require.

AI assistance (Claude) was used to prepare this change.

🤖 Generated with Claude Code

[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 for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream 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.

🚀