[ShuoleiWang]

Purpose

Without sequence parallelism, every TP rank computes DeepSeek-V4.1's fused_wqa_wkv (5120 → 1280 + 512) for all tokens, so all ranks do the same work. The checkpoint stores this weight as FP8 with UE8M0 block scales, which vLLM runs through its MXFP8 path. On SM90 that path is a W8A16 Marlin GEMM: on H20 it takes 1.79 ms per layer for an 8192-token prefill chunk, 72 ms per chunk over the 40 layers.

This PR splits the projection by token for batches of at least 2048 tokens. Each rank projects its token shard (sp_shard), and sp_all_gather rebuilds qr_kv. The change is confined to DeepseekV4Attention._fused_wqa_wkv_gemm:

• _run_parallel_input_projections is unchanged, so the multi-stream decision is the same as on main.
• The all-gather is issued on the current stream, right after the GEMM that produces its input.

_shard_fused_wqa_wkv in the NVIDIA decoder layer enables the split only when all of the following hold:

• SM90. On SM100, fused_wqa_wkv runs as a native MXFP8 GEMM, far cheaper than W8A16 Marlin, so the gather would be a much larger share of the cost. I have not measured SM100, and the H20 results below do not carry over to it.
• No sequence parallelism. There the hidden states are already sharded, and [Perf][DSV4.1] Project sequence shards before all-gather #57257 splits this projection for that case.
• No LoRA. The LoRA wrapper maps every token of the batch.
• No micro-batching. Two micro-batches would share the custom all-gather buffer.
• No batch invariance. Whether a token takes the split depends on the batch size.
• The custom all-reduce is enabled and fully connected. This serves only as a proxy for NVLink between all TP ranks. From 6144 tokens at TP8, the shard exceeds the custom all-gather's 2 MiB limit, so the gather itself runs on NCCL.

TP4 and TP2 take the same gate but were not measured. On Hopper they need pipeline parallelism to hold the checkpoint.

The compressor's fused_wkv_wgate and the indexer's weights_proj stay replicated, and _run_parallel_input_projections, which can run them on separate streams, is unchanged. They appear in only a few layers:

• the compressor in the 4 KV-source layers (2, 8, 14, 20);
• the indexer in the 8 index-source layers, which are those four plus 24, 28, 32 and 36.

At 8192 tokens, splitting them as well would save at most another 457 µs in a KV-source layer (last column of the second table; layer 20's compressor is half as wide). In the four indexer-only layers the projection is only 5120 → 32. This PR saves about 1.45 ms in each of the 40 layers, so it limits the split to this dominant projection.

Why 2048 tokens. On H20 the split is faster from 48 tokens (first table below). H100 and H200 have about 6.7 times H20's dense throughput with the same 900 GB/s NVLink, so the replicated GEMM costs much less there while the gather costs the same.

I have no H100 or H200 to measure on, so I scaled the H20 numbers pessimistically:

• a GEMM of 1024 rows or more runs 6.7 times faster than on H20;
• one of 512 rows runs 3.7 times faster, and one of 256 rows or fewer only twice as fast;
• the gather runs no faster.

Under these assumptions the break-even is about 1450 tokens, and 2048 is the next power of two. A measurement on H100 or H200 would confirm or retune it.

Related work. I found no upstream PR for the case without sequence parallelism. The fork commit bytedance-iaas/vllm@383776ac20 splits fused_wqa_wkv the same way, but behind an environment variable that is off by default, and not when speculative decoding is on.

Test Plan

Kernel benchmark. Setup:

• Hardware: 8× H20 (SM90, 96 GB, NVLink), TP8.
• Software: torch 2.13.0+cu130, NCCL 2.29.7, vLLM nightly 0.30.1rc1.dev85+g8b84e1506 with this diff applied as a Python overlay.
• What it runs: fused_wqa_wkv is built as the model builds it, with random weights and the checkpoint's quantization, so SM90 runs Marlin. The second table times _run_parallel_input_projections from this branch, which calls _fused_wqa_wkv_gemm, with the model's projection shapes and three aux streams.

Timing follows .agents/skills/kernel-microbenchmark:

• One CUDA graph per candidate.
• Before every measured replay, an L2 flush (twice the L2 size) and a device-side barrier, both outside the timed interval.
• The time of a replay is the CUPTI span of its kernels. Each sample is the maximum over the 8 ranks, and the tables report the median of 50 samples.
• Candidates alternate within every round, and correctness is checked before timing.

The benchmark script is not part of this PR; I can add it under benchmarks/kernels/ if that helps.

Serving. TP8 without expert parallelism. DSpark uses synthetic acceptance with a mean acceptance length of 3.51, the same distribution in both arms:

vllm serve deepseek-ai/DeepSeek-V4.1-Flash --revision dba1be0a40aa45a94ad051997016db3960a90277 \
  --tensor-parallel-size 8 \
  --language-model-only --tokenizer-mode deepseek_v41 --max-model-len 262144 \
  --max-num-batched-tokens 8192 --max-num-seqs 128 --gpu-memory-utilization 0.90 \
  --no-enable-prefix-caching --engram-config '{"cpu_offload":false}' \
  --speculative-config '{"method":"dspark","num_speculative_tokens":5,"draft_sample_method":"probabilistic","rejection_sample_method":"synthetic","synthetic_acceptance_length":3.51,"enable_adaptive_verification":false}'

One vllm bench serve call per point:

Point	ISL	OSL	C	N
8K / 32K / 128K prompt TTFT	8192 / 32768 / 131072	1	1	5 / 3 / 2
1K in / 1K out, C1	1024	1024	1	6
1K in / 1K out, C64	1024	1024	64	192
8K in / 1K out, C16	8192	1024	16	48
8K in / 1K out, C64	8192	1024	64	128

vllm bench serve --backend openai --endpoint /v1/completions --dataset-name random \
  --random-input-len <ISL> --random-output-len <OSL> --random-range-ratio 0 --ignore-eos \
  --max-concurrency <C> --num-prompts <N> --num-warmups <C> --seed 1

Each arm ran on a fresh server, in the order this PR (B1), main (A1), main (A2), this PR (B2).

Accuracy. Full GSM8K (1,319 questions) through /v1/chat/completions, with real DSpark verification (rejection_sample_method: block) instead of synthetic acceptance:

• 5-shot multi-turn chat with thinking off (chat_template_kwargs: {"thinking": false});
• temperature 0, seed 123, max_tokens 1024, 32 concurrent requests;
• each reply is scored by the number after ####, or else by its last number.

Test Result

fused_wqa_wkv, split vs replicated (µs per call, TP8, every batch size split so the crossover is visible):

Tokens	Rows per rank	Replicated	Split + all-gather	Speedup
6	1	14.1	23.2	0.61x
12	2	14.6	23.1	0.63x
24	3	19.2	21.1	0.91x
48	6	25.9	20.3	1.28x
96	12	49.6	21.5	2.31x
384	48	99.7	35.8	2.78x
1024	128	239.4	66.8	3.58x
2048	256	458.5	99.7	4.60x
4096	512	903.6	173.3	5.22x
8192	1024	1794.7	346.9	5.17x
16384	2048	3570.8	635.4	5.62x

Up to 4096 tokens the gather is the custom all-reduce's all-gather; from 6144 tokens it is NCCL. From 4096 tokens on, the replicated Marlin GEMM runs at 83–84 TFLOPS.

Input projections of one attention layer, main vs this PR (µs; only this stage is timed, not the whole layer):

Tokens	Layer without compressor or indexer (32 layers): main	This PR	Speedup	Layer with compressor and indexer (layers 2, 8, 14): main	This PR	Speedup	If compressor and indexer were also split
1024	236.7	236.5	same code	337.6	337.7	same code	337.5
1536	351.1	348.8	same code	478.5	480.6	same code	481.1
2048	458.1	99.9	4.59x	660.4	301.8	2.19x	198.3
4096	903.0	172.7	5.23x	1276.2	545.9	2.34x	343.5
8192	1794.2	346.0	5.19x	2473.9	1028.4	2.41x	571.6
16384	3569.4	635.6	5.62x	4902.1	1976.6	2.48x	1052.3

• Estimate per 8K prefill chunk: this stage gets about 1.45 ms faster in every layer, about 58 ms over the 40 layers. The serving results below are the measurement.
• Correctness: qr_kv differs from the replicated projection by at most 0.031 across two sweeps with different inputs. That is bf16 rounding from Marlin's M-dependent tiling. kv_score and indexer_weights are bitwise identical to main.

Serving and accuracy. The serving numbers are a sanity check with two servers per arm, not a precise measurement. Lower TTFT and higher output throughput are better, and Change compares the means of the two arms.

Point	Metric	main A1	main A2	this PR B1	this PR B2	Change
TTFT, 8K prompt, C1	ms	1117.8	1117.8	1066.5	1075.2	−4.2%
TTFT, 32K prompt, C1	ms	4535.6	4559.3	4320.5	4322.3	−5.0%
TTFT, 128K prompt, C1	ms	19648	19680	18859	18807	−4.2%
8K in / 1K out, C16	output tok/s	601.1	602.1	621.4	620.8	+3.2%
8K in / 1K out, C64	output tok/s	709.9	709.4	739.0	739.4	+4.2%
1K in / 1K out, C64	output tok/s	1983.8	1987.6	2014.9	2010.1	+1.4%
1K in / 1K out, C1	output tok/s	335.2	335.2	338.0	334.1	+0.2%

At 1K in / 1K out, C1, no batch reaches 2048 tokens, so both arms run main's kernels. The two servers of this PR differ by 1.2% there, more than the gap between the arms.

GSM8K, 5-shot chat, greedy (1,319 questions)	main	this PR
Accuracy	96.59% (1274)	96.97% (1279)
DSpark acceptance length	3.867	3.863

In three runs of main on the same nightly, one of them the run above, GSM8K accuracy ranged from 96.36% to 96.97%.

AI assistance: Claude Code helped with the implementation and ran the benchmarks. I reviewed every changed line and checked the results above.

🤖 Generated with Claude Code

[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.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94982 for commit df929699b597.

[ShuoleiWang]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #94982.