[Perf][DSpark] Collapse DeepSeek-V4.1 draft states before SP all-gather

[Juntian777]

Purpose

Under sequence parallelism, the DeepSeek-V4.1 DSpark draft gathers BF16 [T_local, 4, H] residual streams and FP32 [T_local, 4] pre-mix coefficients before collapsing them into [T, H]. The draft returns only the collapsed pre-norm head states; it does not export a full-HC MTP buffer at this boundary.

Collapse each token's streams locally, then gather only [T_local, H] and trim SP padding. This reduces the hidden-state gather payload by 4x for hc_mult=4 and removes the pre-mix gather, reducing this tail from two collectives to one. The target model already follows this ordering when it does not need to export full HC states. Normalization remains in compute_logits.

This is a communication optimization, not an output-correctness fix. SP currently requires PP=1, EP enabled, TP>1, and either a MegaMoE backend or DP>1.

Duplicate-work check

No existing issue was identified for this optimization; issue-number duplicate checks are not applicable. Open-PR searches covered dspark collapse, DeepSeek sequence parallel, "hc_collapse_triton", and "dspark" "gather" in:title. No matching fix was found. Reviewed related PRs: #50658 projects Kimi-K3 target auxiliary states; #56438 handles V4.1 pipeline/SP stage boundaries; #51244 adds V4 kernel fusions; #56513 changes the ROCm delayed mHC projection. None applies this NVIDIA V4.1 draft-tail reorder.

Test Plan

• Extend the existing mHC kernel suite with a single-GPU dummy test comparing concatenate-then-collapse against collapse-each-shard-then-concatenate. Cover TP4/TP8, T=1/7/32, padding, and ranks with no valid tokens. Require exact BF16 equality (atol=0, rtol=0).
• Run the existing collapse correctness, custom-op/compile, and draft-head tests.
• Use a standalone measurement harness (not included in the PR diff) to benchmark the exact tail (mhc_post_tilelang + collapse + production sp_all_gather) before and after, using random tensors and no model weights. Check eager and captured outputs bitwise before timing. Include an NCCL-only control.

.venv/bin/python -m pytest tests/kernels/test_mhc_kernels.py \
  -k 'hc_collapse or v41_dspark_head' -v

# Retrieve the standalone harness from the measurement commit.
git show 16dc4cc4d05f86cf8fd0c7413e327395337d2371:benchmarks/kernels/benchmark_dspark_sp_collapse.py \
  > /tmp/benchmark_dspark_sp_collapse.py

.venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=4 \
  /tmp/benchmark_dspark_sp_collapse.py

.venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=4 \
  /tmp/benchmark_dspark_sp_collapse.py --nccl-only

.venv/bin/pre-commit run --files \
  vllm/models/deepseek_v41/nvidia/dspark.py \
  tests/kernels/test_mhc_kernels.py

git diff --check

Test Result

• 25 passed, 177 deselected. Includes all six new single-GPU shard-order cases; 14 existing Torch deprecation warnings.
• TP4: all six shapes passed bitwise checks in eager execution and CUDA graph replay, for both default SP dispatch and NCCL-only. Repeated the default-dispatch run with the same result.
• Changed-file pre-commit hooks and whitespace check passed.
• Full-model evaluation and end-to-end serving benchmarks were not run. Validation is limited to tail-output equivalence and the scoped microbenchmark; no acceptance-rate, model-quality, or serving-throughput claim is made.

Scoped tail performance

One node, 4x NVIDIA GB200 connected by NV18 links. PyTorch 2.13.0+cu132, CUDA 13.2, NCCL 2.29.7, Triton 3.7.1, TileLang 0.1.12, FlashInfer 0.6.18. BF16 streams/output, FP32 mixes, H=5120, hc_mult=4. Measurement commit: 16dc4cc4d05f86cf8fd0c7413e327395337d2371 (the standalone benchmark was subsequently removed from the PR; implementation and tests are unchanged), based on 00972dfd72988942138a7a6089eaee08580210b8.

CUDA graph internal events exclude CPU launch gaps. Each measurement flushes 128 MiB of L2 and uses a device-side all-reduce rendezvous before replay, outside the timed interval. Four pointer-distinct workspaces, 20 warmup replays, 50 samples per variant, alternating A/B order. Report the median of each sample's slowest rank. Compilation, allocation and input generation are outside timing; all production gather staging/copies remain inside. The benchmark also reports logical receive bytes per rank and payload/time; this is not a measurement of physical link bandwidth.

All times below are microseconds for the entire tail. Default-dispatch figures are from the repeat run.

Global T	Local padded T	Default before	Default after	Speedup	NCCL before	NCCL after	Speedup
1	1	24.61	18.88	1.30x	42.78	31.68	1.35x
7	2	26.82	18.94	1.42x	42.82	31.81	1.35x
32	8	43.14	19.60	2.20x	45.50	32.91	1.38x
128	32	108.69	20.78	5.23x	62.38	35.81	1.74x
512	128	163.89	28.80	5.69x	96.96	52.93	1.83x
2048	512	168.43	88.14	1.91x	180.82	90.94	1.99x

The first default-dispatch run measured 1.29x/1.40x/2.10x/5.13x/5.76x/1.88x for the same ascending token counts. These are microbenchmark results, not end-to-end speedups.

Backend selection is part of the measured production wrapper: at T<=128, both versions' state gathers use MNNVL Lamport; at T=512, the original full-state gather falls back while the smaller collapsed gather still uses MNNVL; at T=2048, both state gathers fall back. The original pre-mix gather uses MNNVL throughout. The NCCL-only control disables custom gathers. Real distributed measurements cover TP4; TP8 is covered by single-GPU sharding-equivalence tests only.

AI assistance was used (OpenAI Codex) for implementation, tests, benchmarking, and this description.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[WoosukKwon]

/ci run

[github-actions]

✅ @Juntian777, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[github-actions]

❌ This PR is 34 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[WoosukKwon]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89043 for commit 491d037907a5.