[Perf][DSv4.1] Fold the mHC post block into the delayed pre projection

[zyongye]

Purpose

DeepSeek-V4.1 runs an mHC post block and then the next sublayer's delayed pre block at every seam. Each seam therefore pays for three things: a post kernel, a split-k projection that re-reads the residual streams the post kernel just wrote, and the pre epilogue.

DSv4 already avoids the middle read: mhc_fused_tilelang computes the updated residual in registers and feeds the pre-norm GEMM from there. V4.1 could not reuse it, for the same reason it needed its own pre in the first place — it collapses each sublayer's input with the previous sublayer's pre-mix and hands the next one back, and the DSv4 fused wrapper neither accepts an incoming pre-mix nor returns one.

Approach

mhc_fused_post_pre_delayed_tilelang runs the existing fused kernel and then the delayed pre epilogue (MHC_PRE_NORM_KERNEL, use_pre_mix_in / save_pre_mix), returning (residual, post_mix, comb_mix, layer_input, next_pre_mix). Above the fused kernel's token range it falls back to exactly today's mhc_post_tilelang + split-k GEMM, so only the small-batch path changes shape.

Both V4.1 seams route through it, except one:

• Attention seam — fuses, unless Engram is injected. Engram reads the residual between the post and the pre, so the post cannot move into the projection there; that branch keeps the separate kernels.
• FFN seam — always fuses.
• DSpark draft reuses DeepseekV4DecoderLayer, so it inherits both.

The pre epilogue is JIT-compiled per split-k factor and warmed up from a token sweep that never produces the factors this GEMM picks, so MHCPreNormKernel.get_warmup_keys takes extra_splits and the decoder registers them. Without it the first decode step JITs inside the capture warmup.

Tuning

The tile, split-k and block size were inherited from DSv4 unmeasured (the code carried a TODO(gnovack): investigate autotuning these heuristics). Sweeping all three against the separate post + split-k GEMM, on GB300 at hc_mult=4, CUDA-graph device time:

tokens	1	2	4	8	16	24	32	48	64
inherited config, 5120	8.21	8.01	8.42	8.82	10.25	12.73	13.51	16.81	19.70
tuned, 5120	8.02	8.00	8.83	8.84	9.85	10.47	11.47	fallback	fallback
unfused, 5120	12.74	12.52	12.72	12.92	13.32	13.32	13.37	13.53	13.49

Two things were wrong with the inherited constants: the 16-token cutoff was early (the fused path still wins 1.2x at 32), and one fixed tile starves the projection of parallelism as the batch grows — it is 18% off at 24-32 tokens and slower than the unfused path past 48. The tuned heuristic is two bands (tile_n 2 below 16 and 6 through 32, n_thr 128), within ~1-2% of the per-shape optimum for one extra compiled variant. Same bands hold at hidden_size 7168.

The cutoff is 32, and it is set against a corrected baseline. vLLM configures the TF32 projection's split count through compute_mhc_pre_num_splits, which buckets a good estimate into {1, 4, 16} to bound startup specializations — and that bucketing is itself 1.14-1.19x off the optimum below 128 tokens (measured: the estimator says 80 splits at one token, the bucket gives 16, the optimum is 64; rounding the estimate down to a power of two reproduces the measured best at every shape I tried). Against a GEMM given the split its own estimator computes:

tokens	1	2	4	8	16	24	32	48
fused	8.44	8.23	8.83	9.03	9.85	10.61	11.29	13.55
unfused, split as configured today	13.13	12.74	12.76	13.33	13.35	13.55	13.56	13.35
unfused, split=64	11.09	11.09	11.29	11.66	11.50	11.69	11.62	11.72
fused vs tuned baseline	1.31x	1.35x	1.28x	1.29x	1.17x	1.10x	1.03x	0.86x

So the honest claim is 1.3x at decode widths, 1.17x at 16, break-even at 32. The cutoff sits at the break-even point rather than where a slow baseline would have put it (48, which is a 14% loss against a tuned GEMM). Fixing that split heuristic is worth 5-19% on every unfused seam — all prefill, all large-batch decode, and DSv4 as well as DSv4.1 — but it belongs in its own PR, not bundled here, since it changes a shared helper and a second model's hot path.

It is shared with the DSv4 wrapper, which uses the same kernel at the same shapes. The heuristic now also declines shapes the kernel cannot tile evenly (mix_size % tile_n, hidden_size % (n_splits * n_thr)) and falls back instead — previously those silently dropped whole tiles and k-slices.

Relationship to the other open mHC PRs

• [DSv4.1] Integrate Mega-mHC from DeepGEMM #56255 (Mega-mHC from DeepGEMM) — same two call sites, same 5-tuple refactor, same Engram carve-out. Its mhc_shifted_post_pre dispatches to DeepGEMM mega_mhc when device_capability_family(100) and hidden_size % 1024 == 0 and hc_mult == 4, and falls back to the unfused pair otherwise — which is what this PR replaces. The two are complementary and land cleanly together: benchmarked head to head below, the TileLang fusion is faster at small batch, Mega-mHC is faster at large batch. If both land, the right end state is one dispatcher picking between them by token count, and I am happy to write that as a follow-up (or fold this into [DSv4.1] Integrate Mega-mHC from DeepGEMM #56255, author's preference).
• [ROCm][DSV4.1][Perf] Fold the mHC post step into the delayed pre projection #56513 (ROCm) — same idea via AITER mhc_fused_post_pre_gemm_sqrsum, on _aiter_ops.py / mhc/aiter.py / layers/mhc.py / amd/model.py. No file overlap with this PR beyond the shared test file.
• [Perf][DSV4.1] Overlap shifted mHC statistics with attention and FFN #56611 — a different optimization (overlapping the coefficient GEMM with attention/FFN), env-gated.

Mega-mHC comparison

DeepGEMM 2.8.0 (66081d4c), GB300, hc_mult=4, hidden_size=5120, CUDA-graph device time per seam. Both paths are graph-capturable; mega_mhc lazily allocates per-stream split barriers and asserts that happens outside capture, so it must be warmed on the capture stream first (which vLLM's pre-capture warmup iteration does).

tokens	1	2	4	8	16	24	32	48	64	128
unfused	13.13	12.84	13.02	13.54	13.32	13.54	13.54	13.55	13.56	13.97
fused TileLang (this PR)	8.43	8.22	8.83	8.84	9.65	10.47	11.28	13.10	13.54	13.96
Mega-mHC (#56255)	10.19	10.46	10.07	10.67	10.47	10.47	10.88	11.09	11.26	11.82

This PR is 1.08–1.27x faster than Mega-mHC at 1–16 tokens; they cross at ~24; Mega-mHC is 1.18–1.20x faster at 48+. A torch.profiler kernel-time breakdown agrees (at 1 token: 4.32 µs fused kernel + 5.03 µs epilogue, against a single 10.44 µs sm100_mega_mhc_impl).

Numerically the two differ in one place: the fused TileLang kernel projects the FP32 post result it holds in registers, while Mega-mHC and the unfused path project its BF16 rounding. Against an FP32 reference at 8 tokens, max abs error: residual 9.77e-4 for all three; coefficients 9.25e-4 / 1.78e-3 (this PR) against 1.98e-4 / 3.46e-4 (Mega-mHC and unfused). This is inherited from the DSv4 kernel, not new — the stored residual streams are bit-identical either way.

Test Plan

.venv/bin/python -m pytest tests/kernels/test_mhc_kernels.py -v

New test_deepseek_v41_mhc_fused_post_pre_delayed (0/1/8/128 tokens x 4096/7168 hidden x carried/absent pre-mix) asserts the post-mapped residual and the collapsed layer input are bit-identical to mhc_post_tilelang + mhc_pre_delayed_tilelang, and the coefficients match an FP32-post torch reference to 1e-5 on the fused path. Plus a torch.library.opcheck case and a CUDA-graph capture/replay case, since decode replays this op from a graph. test_deepseek_v41_decoder_mixes_match_torch now covers both routes through the decoder.

Test Result

164 passed, 23 skipped in 53.80s     # tests/kernels/test_mhc_kernels.py, GB300

pre-commit run --all-files clean on the changed files (ruff, ruff-format, mypy-3.10).

Model evaluation

DeepSeek-V4.1-Flash, TP4 GB300, gsm8k 5-shot via mqa-run --api-type completions --temperature 0 --max-tokens 1024, n=1319, num_request_errors: 0 and 0 null answers in every run:

arm	runs	exact_match/flexible
this PR	3	0.9242, 0.9325, 0.9204
unfused baseline, same serve config	2	0.9348, 0.9249

The baseline arm's own two runs differ by 0.99 pt, so the arms' ranges overlap and the 0.42 pt difference in means sits inside the measured noise floor for this config (batching nondeterminism at concurrency 32). Dashboard history for this model is median 0.9314, p25–p75 0.9287–0.9333.

Second and third commits: aux hidden states for draft models

Once the post lives inside the next layer's pre, the standalone mhc_post_tilelang the model ran per aux-capture layer is recomputing a residual that already exists. The second commit takes mean(dim=1) from the fused call behind a capture_previous_aux flag and drops the duplicate post — the approach of #55575 (DeepSeek-V4), which only becomes available to V4.1 once V4.1 has the fused seam.

Two seams keep the standalone post: the Engram seam captures before the injection (matching the post it replaces — aux consumers read the stream before Engram touches it), and the last layer on a rank has no successor to fold into, so it takes its mean from the final post that runs anyway.

I did not port #55575's companion change that skips the _mtp_hidden_buffer copy when aux hidden states are present: that depends on its vllm/v1/worker/gpu/model_runner.py edit, which is model-agnostic and would collide with #55575 in a shared file. It should follow from that PR, not this one.

Third commit: fold the stream mean into the collapse

That leaves one pass still standing: mean(dim=1) over [T, hc, H]. The collapse in the pre epilogue already reads every one of those streams to weight them by the carried pre-mix, and the aux value is the unweighted mean of the same data — so the third commit accumulates it in that pass, behind a write_aux specialization so non-draft setups compile and run exactly what they did before. The buffer is allocated only when the caller asks.

Standalone mean(dim=1) is not cheap at scale (measured on GB300, [T, 4, 5120] bf16): 2.10 µs at 1 token, 17.9 µs at 1024, 72.1 µs at 4096, 141.8 µs at 8192 — per capture layer, per step. Folded, the aux-capture seam costs:

tokens	1	8	32	128	1024	4096
seam + separate mean	10.55	11.50	13.94	17.65	60.25	209.32
seam + folded mean	8.43	9.04	11.50	13.74	40.88	140.63
saved (µs)	2.13	2.46	2.45	3.90	19.38	68.70

Both epilogues carry the path — the one that applies RMSNorm and the plain one — so no caller falls back to a torch mean. Bit-exact against torch.mean on both, at 1/8/33/128/1024 tokens and 5120/7168 hidden, with every other output of the op unchanged when capturing is on. The Engram seam still uses a separate mean, since its aux is the pre-injection post and the epilogue runs after Engram.

DSpark verification

DeepSeek-V4.1-Flash, TP4 GB300, {"method":"dspark","num_speculative_tokens":3,"draft_sample_method":"probabilistic"} with real verification (not synthetic), 8 fixed prompts at temperature 0, same branch with and without the commit:

	drafts	draft tokens	accepted	rate	length	per position
neither commit	106	318	204	0.6415	2.9245	92 / 64 / 48
harvest	106	318	204	0.6415	2.9245	92 / 64 / 48
harvest + folded mean	106	318	204	0.6415	2.9245	92 / 64 / 48

Byte-identical generated text (sha256 of the concatenated completions matches) and identical acceptance counters, which is the end-to-end statement that the harvested tensor is the one the draft model used to get. test_deepseek_v41_capture_previous_aux pins the same thing at the unit level with atol=0, for both the fused and Engram seams.

That probe pins the tensor exactly but is only 8 prompts, so both arms were also run through GPQA-Diamond at full size — 198 questions x 4 epochs = 792 samples per arm, chat path with thinking on, temperature 1.0, top_p 0.95, max_tokens 57344, concurrency 128, same TP4 GB300 node:

arm	headline	answered-only	pass@4	request errors	truncated	null answers	empty@stop
DSpark k=5, real verification	0.9104	0.9138	0.9545	0	3	3	0
no spec decode	0.9129	0.9163	0.9495	0	3	3	0

Acceptance over that run, from the engine's own counters rather than a synthetic profile: 1,467,764 drafts, 7,338,820 draft tokens, 3,207,521 accepted — rate 0.4371, acceptance length 3.185 of a possible 6.

Both arms are this branch, so this is not a before/after for the PR; the before/after is the gsm8k table above. What it establishes is that the code these two commits touch carries 1.47M drafts at an accuracy indistinguishable from the no-draft arm: 2 questions of 792, where one question is 0.126 pt and the sign flips between exact_match and pass@4. Finished generations ran to a median of 2239 tokens with a p90 of 16232 and a max of 56320 against the 57344 cap, so the three truncations per arm are a genuine budget limit rather than a loop — identical in both arms, so they do not move the comparison.

AI assistance

This change was written with AI assistance (Claude Code). I reviewed every changed line, ran the tests and benchmarks above myself, and can defend the change end to end.

🤖 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 #88578 for commit e500862b3a1e.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88604 for commit f6baeecafbff.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88704 for commit 3432199c6dbe.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88749 for commit 9bf278d7b256.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88880 for commit 0f92779c48d4.