[ROCm][DSv4.1] Fix TileLang mHC pre on 64-wide wavefronts and use the fused kernels on the AMD path

[JohnQinAMD]

Two commits: the first fixes the TileLang mHC pre kernels on 64-wide wavefronts, the second switches DeepSeek-V4.1's AMD path onto them. The second alone is a 91.1% -> 23.7% GSM8K regression, so they land together.

The bug

The three mhc_pre_big_fuse* kernels are warp-specialized: if tid < 32 puts the mix phase on warp 0 and the weighted sum with its fused RMSNorm on the rest. That is warp-uniform only at 32 lanes; on a 64-wide wavefront it splits a wave, TileLang hoists the barrier out of the branch, and the fused RMSNorm's cross-lane sumsq reduction then covers only part of the lanes. layer_input comes out a uniform sqrt(2) too large in every layer: measured ratio 1.41443 against sqrt(2) = 1.414214, implied sumsq factor 0.4998, constant at 1/8/64/256/1024 tokens.

Each phase needs a wave-aligned thread group starting at thread 0, and in one launch only one branch can start there. So the block size and split point become compile-time parameters instead of literals, and on a 64-wide wavefront the three kernel names are rebound at module scope to a wrapper that launches the two phases separately: mix_lanes = n_thr, then mix_lanes = 0.

Rebinding rather than editing call sites keeps tilelang.py and warmup.py out of the diff, and on CUDA those names still bind to the original kernel objects. CUDA compiles at (96, 32), which leaves the generated code byte-identical -- verified by dumping the kernel source at those values: same 23573 bytes, same 7 barriers, same md5. Restructuring the branch instead would not have been inert: rewriting the if/else as two complementary ifs takes the same arrangement from 7 barriers to 15. The second launch is free with cudagraphs, paid at capture rather than per replay.

Not latent: vllm/models/deepseek_v4/amd/model.py:754 uses MHCPreOp() with norm_weight=, so DeepSeek-V4 on ROCm reaches this kernel whenever the AITER gate fails.

Repro

vllm serve /path/to/DeepSeek-V4.1-Flash --tokenizer-mode deepseek_v41 \
  --tensor-parallel-size 4 --moe-backend aiter_triton_mxfp4_bf16 \
  --no-async-scheduling --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}'

# latency: the TPOT / TTFT / out tok/s columns below
vllm bench serve --backend vllm --model /path/to/DeepSeek-V4.1-Flash \
  --dataset-name random --random-input-len 1024 --random-output-len 1024 \
  --max-concurrency 1 --num-prompts 8 --ignore-eos

# the "without commit 1" row: revert the kernel commit only, keep the wiring
git revert --no-commit <commit 1>

For the GSM8K column: this checkpoint ships no chat template, so prompts come from its own encoding/encoding.py and sampling follows the model card. Per question, encode_messages([{"role": "user", "content": question}], thinking_mode="thinking", reasoning_effort=100), POSTed to /v1/completions with temperature=1.0, top_p=0.95, max_tokens=16384, scoring the text after </think> against the integer after #### over all 1319 test questions.

tests/evals/gsm8k/gsm8k_eval.py --num-shots 8 --num-questions 1319 also runs, but scores about 90.3% on either build: it is a base-model few-shot completion protocol, and this is the instruct checkpoint.

Measurements

4x MI355X (gfx950), ROCm 7.2.3, TP4. Latency 1k in / 1k out at concurrency 1. GSM8K is the full 1319-question set in the mode described above.

	TPOT	out tok/s	TTFT	GSM8K
torch reference	42.90 ms	23.21	206.0 ms	92.19 %
fused, without commit 1	15.80 ms	62.67	133.8 ms	23.65 %
fused	15.95 ms	62.18	133.1 ms	92.72 %

Sampled, one run each, so the 0.53 points between the two is not a claim of improvement.

Commit 2 also moves the DSpark draft path (dspark.py) onto the fused post kernel, and enabling DSpark is the only thing that exercises it. With --speculative-config '{"method":"dspark","model":"<same checkpoint>","num_speculative_tokens":5,"draft_sample_method":"probabilistic","enable_adaptive_verification":false}' -- the mtp.* weights ship in this checkpoint, and 5 is its dspark_block_size -- median TPOT is 8.91 ms at 111.2 output tok/s, mean acceptance length 2.33. The same DSpark config on the torch reference gives 22.63 ms, so the fused kernels are still worth 2.54x with speculative decoding on. On greedy 8-shot GSM8K (not the protocol in the table above) DSpark scores 90.90% against 89.76% for the same build without it, with 21 questions regressed and 36 improved, so no accuracy loss.

The fold is not bit-identical to the reference -- x * rsqrt * norm_weight is evaluated in fp32 and rounded once where the torch path rounds several times. Capturing the first 24 greedy tokens and logprobs for 120 prompts, the same build run twice already diverges on 33% of sequences (median logprob delta 2.37e-04) from batch composition alone; torch against fused gives 38% and 2.68e-04, so the fused deviation is at the level of the stack's own nondeterminism. At the kernel level the fold is equal or better than the provider it replaces (rms_norm priority is ['aiter', 'native']): NRMSE 0.0 / 4.3e-5 / 2.1e-5 / 3.0e-5 against AITER's 2.8e-7 / 7.9e-5 / 3.1e-5 / 3.1e-5 at 1/8/64/1024 tokens, against an exact fp64 evaluation of the same spec.

Notes

• No test is added here, deliberately rather than by oversight. tests/kernels/test_mhc_kernels.py::test_mhc_pre_tilelang runs on gfx950 and passes on the broken kernel, because it calls the op with nine positional arguments and never passes norm_weight, and the no-norm path is the correct one. The four *_applies_norm tests monkeypatch HAS_TILELANG_MHC = False, so they cover the dispatch rather than the kernel. The TileLang with-norm kernel therefore has no in-tree coverage on any platform, and CI will not catch a regression of this bug. Happy to add a fused_norm parametrization if you want one.
• Validated on the head of [Model] Support DeepSeek-V4.1-Flash #56214, because main cannot run DeepSeek-V4.1 on ROCm at all: rocm_aiter_sparse_attn_indexer takes 15 arguments there while the model passes 18. The code being changed is on main.
• Fusion is gated on model_config.dtype == bfloat16, the only dtype these kernels accept; other dtypes keep the torch reference with its separate norm.
• Not included: the AMD model registers no MHC JIT warmup keys, unlike nvidia/model.py:214-238, so a DSpark draft compile can land inside graph capture.

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

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @JohnQinAMD.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ChuanLi1101]

The wavefront split (tid < 32 on a 64-wide wave) is a real bug. Kernel fix + V4.1 wiring should stay in one PR — wiring alone is the ~24% GSM8K row.

Not mergeable yet:

Rebase; Mergify is already flagging conflicts.
Please add the with-norm_weight test you already described. Leaving CI green on the broken kernel is the actual merge risk.
_MHC_PRE_SPLIT_BLOCK = current_platform.is_cuda() runs at import time. Confirm that is safe if this module is imported before the platform is set.
Agree with the description on perf/accuracy: one-run TPOT, and 92.19 vs 92.72 is not an accuracy claim. CUDA (96, 32) staying byte-identical is the right NVIDIA bar. Happy to re-review after rebase + that test.

[Fangzhou-Ai]

Hi @JohnQinAMD The current default mHC kernels are from AITER. Can you compare you implementation with that? For what I can remember the tilelang mHC is not as fast as AITER mHC in DSV4 pro.

[Fangzhou-Ai]

adding a “bug" tag as this PR is fixing the gsm8k issue, perf-wise is neutral.

[JohnQinAMD]

Hi @JohnQinAMD The current default mHC kernels are from AITER. Can you compare you implementation with that? For what I can remember the tilelang mHC is not as fast as AITER mHC in DSV4 pro.

@Fangzhou-Ai I have tested the aiter mHC and the fixed-tilelang mHC backend on MI355X for dsv4.1 flash, the aiter backend is faster than tilelang backend, which is aligned with your mHC PR results #56503. As gfx950 currently use aiter mHC by default and gfx942 disabled tileLang but easy to port to aiter backed. This bug is currently not impact the vLLM users on ROCm platforms. I lean to close this PR for now. Any suggestions?

[Fangzhou-Ai]

instead of closing the PR, I would actually love to raise error/remove this path to make sure we are not falling back to a wrong implementation, @tjtanaa any suggestions?

[JohnQinAMD]

Close this pr temporally as this bug is not in active path.