[Perf][GDN] Share GDN metadata across cache groups for native MTP of any depth

[Leonccaa]

Purpose

Follow-up to #684 (see this comment). #684 made native MTP4 classify each decode batch once and build every GDN cache group's state rows in one fused launch. Before that, metadata was built once per group per step, at about 0.65 ms of host time per group on V100.

Nothing in that path depends on the draft depth except num_spec_state_tokens. The shared request metadata, the fused state rows (width num_spec + 1) and the sequence-length align column are all computed per row. Only the two gates required num_speculative_tokens == 4, so MTP1–MTP3 deployments, including our MTP3 Qwen3.8 serving stack, still build per group.

This PR:

• Accepts native MTP at any depth in the two gates (gdn_attn.py, mamba_hybrid.py). The VLLM_SM70_MTP4_SHARED_GDN_METADATA / VLLM_SM70_MTP4_FUSED_GDN_METADATA switches and their defaults are unchanged and still roll the paths back; the names keep their historical prefix, and comments and log lines now say MTP.
• Fixes the fused-metadata shadow oracle (VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW, debug-only), in a separate commit. The fused kernel writes a step's rows only up to num_actual_tokens, but the oracle required every entry of the shared spec_sequence_masks buffer after the live rows to be False, including rows no graph replay of this step reads. That buffer came from torch.empty, so the oracle passed or failed depending on leftover device memory; widening the tests to more depths made that visible. The mask is now allocated zeroed and the oracle checks only the padded rows the step writes and replays. Replay behavior is unchanged.

Test Plan

• pytest tests/v1/attention/test_gdn_metadata_builder.py tests/v1/spec_decode/test_dflash2.py on V100 (SM70). test_mtp_shared_gdn_metadata_matches_legacy and test_mtp_fused_gdn_group_metadata_matches_legacy_replay now cover MTP1 through MTP4 against separate per-group builds. The fused test runs every verifier query length at each depth, in both none and align cache modes, with the shadow oracle on. Run twice to check stability.
• End-to-end evidence for MTP3 comes from our serving stack (4x V100 TP4, AWQ g32, E4M3 KV, MTP3 greedy, prefix caching with align, grouped CPU offload, FULL_AND_PIECEWISE graphs), where we ran an equivalent implementation (details in the [Perf] Share Flash-Next decode fast paths with MTP4 by default #684 comment).

Test Result

• 293 passed on V100, on both runs. Before the oracle fix, one full run of the widened fused test failed 8 of its 20 cases (MTP1 and MTP3), and a stop-at-first-failure run failed an MTP2 case that had passed before, always with "left a live padded row".
• End to end on our stack with the shared and fused paths on for MTP3: C1 decode rose from 72.98 to 77.92 tok/s (with the spec-row selection change that [Perf] Share Flash-Next decode fast paths with MTP4 by default #684 also merged). Under target-only DCP2, which has six GDN cache groups, it rose from 71.21 to 75.76 tok/s. Token IDs and accepted/drafted counts matched the per-group reference exactly on our fixed prompts, and the offload, mixed-batch and 30K regressions passed. That stack is not on main yet (it needs [Model][SM70] Calibrated E4M3 KV cache for Qwen4Exp QSA with MTP #664 and [KV Offload] Grouped hybrid RAM caching, filesystem restoration and retention-sized Mamba pools (stacked on #617) #598), so this main-based branch itself was validated by the unit tests above rather than end to end.

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

🤖 Generated with Claude Code