[Perf] Narrow DeepSeek V4 eager CUDA graph region

[WoosukKwon]

Summary

Narrow the DeepSeek V4 eager CUDA graph region so attention input preparation remains captured, and organize the flow so the execution stages are visible in one linear forward method.

The existing attention_impl eager break includes Q up-projection, fused Q normalization/RoPE/KV insertion, indexer input preparation, and MLA/indexer compression. This change moves those producer operations into forward and keeps only the sparse indexer operation plus forward_mqa in the eager region.

The final refactor removes the tuple-of-mostly-None plumbing, gives the input-projection helper a descriptive name and explicit return type, and reuses the indexer's existing forward method. It adds no net methods.

This branch is synced with main at fac808b36f.

Parallelism

This does not reduce stream parallelism relative to main:

• Input projections still use the default stream for fused_wqa_wkv plus three auxiliary streams for compressor KV score, indexer weights, and indexer compressor KV score.
• The next stage still overlaps Q projection/KV insertion on the default stream with the indexer on auxiliary stream 0 and the MLA compressor on auxiliary stream 1.
• Auxiliary stream 2 is reused by the indexer's internal Q/K overlap.

ROCm retains its existing sequential fallback.

Why this approach

An earlier DeepSeek V4 region-narrowing change was reverted after real-model output corruption. This version deliberately keeps forward_mqa eager together with the sparse indexer instead of capturing the attention backend. Persistent eager scratch workspaces also keep the tensors consumed by the eager region address-stable.

This mirrors the region boundary used by the DeepSeek V3.2 optimization while preserving the DeepSeek V4 attention execution boundary.

Duplicate-work check

I searched open vLLM PRs for DeepSeek V4 CUDA graph, DeepSeek V4 eager attention, Narrow DeepSeek V4, and eager CUDA graph region DeepSeek. No other open PR implements this DeepSeek V4 eager-region narrowing. #51425 is the analogous DeepSeek V3.2 change and does not modify DeepSeek V4.

Performance

DeepSeek-V4-Flash on 4x GB200, Rust frontend, FP8 KV cache, FP4 indexer cache, PIECEWISE breakable CUDA graphs, DeepGEMM mega-MoE, and MTP2.

Short-prefill TTFT uses 1,000 requests with 128 input tokens, 1 output token, and concurrency 1. Saturated throughput uses 4,096 requests with the same token lengths and concurrency 64. The comparison was run against main at b706fd162 before the final sync.

Metric	main	This PR	Change
Mean TTFT	41.45 ms	19.20 ms	-53.7%
P50 TTFT	41.18 ms	19.09 ms	-53.7%
Request throughput	430.05 req/s	433.29 req/s	+0.75%

A same-session A/B of the readability-only refactor measured 19.889 ms before and 19.742 ms after for mean TTFT (-0.74%); P50 was 19.663 ms and 19.674 ms respectively. This indicates no regression from the structural cleanup.

Validation

• pre-commit run --files vllm/models/deepseek_v4/attention.py vllm/models/deepseek_v4/nvidia/model.py vllm/models/deepseek_v4/amd/model.py
  • passed, including Ruff and Python 3.10 mypy
• .venv/bin/python -m pytest tests/v1/cudagraph/test_breakable_cudagraph.py -q
  • 14 passed
• Commit-time pre-commit and sign-off hooks
  • passed

Real-weight serving completed graph capture and all evaluation requests in both requested topologies. Both used MTP2 and otherwise identical serving/evaluator settings. The DP4/EP4 run used TP1, DP4, and expert parallel across all four GPUs.

Full GSM8K, 1,319 questions, 5-shot, seed 42, greedy decoding, 128 concurrent requests:

Configuration	Strict exact match	Flexible exact match
TP4/EP4 + MTP2	95.30% (1,257/1,319)	95.22% (1,256/1,319)
DP4/EP4 + MTP2	95.07% (1,254/1,319)	95.00% (1,253/1,319)

No evaluation requests failed. Raw samples and aggregate result JSON were retained for both runs.

AI assistance

This PR was developed with AI assistance. The human submitter reviewed the changed code and ran the validation and model evaluations listed above.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[WoosukKwon]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83167 for commit 3efc01ccee42.