[ROCm][DSV4.1][Perf] Use AITER mHC for the delayed pre block - #56503
Conversation
DeepSeek-V4.1-Flash on AMD still ran its mHC block as the eager Torch reference, so every sublayer seam spent ~141 kernel launches on a 4x4 comb matrix. On MI355X that dominates decode: a per-layer census over a captured trace attributes 85% of the layer's 350 launches to mHC, and the sinkhorn loop alone accounts for 37% of kernel time, ahead of the MoE. V4 already dispatches through MHCPreOp/MHCPostOp and reaches AITER, but V4.1 cannot reuse MHCPreOp: it collapses each sublayer's input with the *previous* sublayer's pre-mix, and AITER's mhc_pre neither returns its pre-mix nor accepts an incoming one. The gates are unaffected by that staggering, so drive AITER's two pre stages directly -- mhc_pre_big_fuse still computes the post/comb gates and every sinkhorn iteration -- recover the pre-mix from the same unreduced projection output with a small Triton epilogue, and collapse with hc_collapse_triton. Measured on gfx950 at hidden_size 5120, hc_mult 4, 20 sinkhorn iterations: 141 launches -> 4 per seam, worst relative error 1.2e-5 against the Torch reference over tokens 1..16384, with the collapse bit-exact. The model-entry broadcast seam projects a narrower operand and has no AITER equivalent, and the composed path folds no RMSNorm, so forward_hip declines both and they fall through to TileLang, then to the reference. Signed-off-by: Fangzhou Ai <fangzhou@semianalysis.ai>
mergify
Bot
added
deepseek
End-to-end ITL on MI355XConcurrency 1, TP4, fp4, 5-token DSpark MTP,
This is not a controlled A/B and should not be read as a 3.7x per-token speedup. The agentic trace is time-driven, so the faster run progressed further and sampled a different, longer segment of it:
The two ITL figures are therefore averages over different request populations. I am reporting them because the confound runs against the result -- this PR's run carried 22% longer contexts and 76% more output tokens per request, both of which should raise ITL, and ITL fell anyway -- but a clean number needs a fixed-request-count replay, which I have not run. The magnitude is at least mechanically consistent. This PR removes ~137 launches per seam across 80 seams, roughly 11,300 tiny kernels per step, which at ~1.7 us of GPU-side dispatch each accounts for the ~18.7 ms that disappeared. Graph capture removes CPU launch overhead but not per-kernel dispatch cost, which is why the win survives Provenance, since the install here is patched rather than a source build: the patch landed at 18:52:48, the server started at 18:55:15, and |
|
Follow-up on the ITL numbers above, since one confound I left open turns out to be closed. MTP acceptance is identical across the two runs: acceptance length 3.52 on both, acceptance rate 50.5% vs 50.3%. This matters because with 5-token DSpark MTP, ITL is a function of how many tokens each engine step emits, so a shift in acceptance would have moved ITL without step time changing at all. It did not, so the ITL delta is attributable to step time. It is also worth separating decode from end-to-end, because the two ratios differ and only one of them is the mHC seam:
At concurrency 1 this scenario is prefill-dominated -- mean input ~206K tokens against ~1.7K output, so ~122x more input than output tokens. Decode inter-token latency improves 3.71x, prefill 1.84x, and end-to-end output throughput lands between them at 2.65x. Quote 3.71x only for per-token responsiveness. The workload-mix caveat in the previous comment still stands and still cuts against the result: this PR's run carried 22% longer contexts, which should raise decode ITL, not lower it. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
… numerics test_mhc_pre_delayed_rocm_aiter_declines_unsupported checked that a fused norm is declined by comparing forward_hip's output against a Torch reference. That routes into mhc_pre_delayed_tilelang, whose fused-norm variant is independently broken on gfx950 with the pinned tilelang 0.1.10: it reduces the RMS sum of squares inside a thread-divergent branch and recovers half the sum, so layer_input comes out a factor of sqrt(2) high. The mix outputs are bit-identical, so only the folded norm is affected. That is a pre-existing defect in the TileLang path rather than anything this gate does, and the same kernel already fails test_deepseek_v41_mhc_pre_delayed[fused_norm=True] on main without this change. Assert the decision this gate actually owns -- that neither the broadcast seam nor a fused norm reaches AITER -- and add a positive control so the two declines cannot pass by the AITER path merely being unavailable. Signed-off-by: Fangzhou Ai <fangzhou@semianalysis.ai>
Re-validated against current main on MI355XRebuilt the stack on the 2026-09-11 ROCm nightly ( The gate reaches AITER for the shapes the model actually uses
Every token count from 1 to 128 takes the AITER path at the steady-state seam, and the model-entry broadcast is declined as intended. Controlled A/B of the block itselfThe end-to-end ITL numbers earlier in this thread are time-driven and, as noted there, not a controlled A/B. This is: one process, identical input tensors, same iteration count, only the kernel choice differs. Microseconds per call, median of 5 x 200 iterations, Under CUDA graphs, which is how decode runs here (
Accuracy against the Torch reference in the same run: max abs err 8.9e-08 at 1 token rising to 2.7e-05 at 128. The TileLang path reaches 1.6e-02 over the same sweep, so AITER is both faster and closer to the reference. One caveat worth stating. Timed eagerly instead, AITER is slower than TileLang (93.4 us vs 28.8 us at 1 token) because the Python wrapper in Accuracygsm8k, 1319/1319 problems, real block rejection (
That is within one standard error of the pre-mHC baseline measured on the same recipe (0.9704). Note on the test change in the new commit
This is pre-existing on main and unrelated to this PR, which adds no TileLang path. Verified by running the kernel suite on main + #56214 with this PR reverted: the same 15 So the new commit asserts the decision this gate actually owns -- that neither the broadcast seam nor a fused norm reaches AITER -- plus a positive control so the two declines cannot pass merely because the AITER path is unavailable. The fallback's numerics stay covered by Kernel suite
|
|
/ci run |
|
✅ Triggered Buildkite CI #88433 for commit |
|
/ci run |
|
✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/88433 |
github-actions
Bot
added
the
ready
The deepseekv41-flash-0909 tag predates vllm-project/vllm#56503, which moves the mHC delayed pre block off the eager Torch reference and onto AITER. That block is 85% of the decoder's kernel launches for this model, so the release tag leaves most of the decode cost unaddressed. Tag is nightly-TBD until the first nightly containing that commit is published; it is then pinned to nightly-<sha> like the other ROCm entries.
The nightly containing vllm-project/vllm#56503 is out: nightly-eed1f3d0c6043bd494424a22443ee198dd56f657 (sha256:960228cf..., published 2026-09-12). Replace the nightly-TBD placeholder in the config, recipe header, and both docs, and record the upstream recipe PR that moves vllm-project/recipes to the same image. Co-authored-by: Cursor <cursoragent@cursor.com>
The deepseekv41-flash-0909 tag predates vllm-project/vllm#56503, which moves the mHC delayed pre block off the eager Torch reference and onto AITER. That block is 85% of the decoder's kernel launches for this model, so the release tag leaves most of the decode cost unaddressed. Tag is nightly-TBD until the first nightly containing that commit is published; it is then pinned to nightly-<sha> like the other ROCm entries.
The nightly containing vllm-project/vllm#56503 is out: nightly-eed1f3d0c6043bd494424a22443ee198dd56f657 (sha256:960228cf..., published 2026-09-12). Replace the nightly-TBD placeholder in the config, recipe header, and both docs, and record the upstream recipe PR that moves vllm-project/recipes to the same image. Co-authored-by: Cursor <cursoragent@cursor.com>
…oject#56503) Signed-off-by: Fangzhou Ai <fangzhou@semianalysis.ai> Co-authored-by: Fangzhou Ai <fangzhou@semianalysis.ai>
…oject#56503) Signed-off-by: Fangzhou Ai <fangzhou@semianalysis.ai> Co-authored-by: Fangzhou Ai <fangzhou@semianalysis.ai>
2 participants
Purpose
DeepSeek-V4.1-Flash on AMD still runs its mHC block as the eager Torch reference (
vllm/models/deepseek_v4_1/amd/model.pyimportsmhc_pre_delayed_torch/mhc_post_torchdirectly). DeepSeek-V4 was moved onto theMHCPreOp/MHCPostOpdispatch layer and reaches AITER (#41946, #43950, #52737), but V4.1 was never migrated, so on MI355X it pays ~141 kernel launches per sublayer seam to Sinkhorn-normalize a 4x4 matrix.That cost dominates decode. A per-layer census over a captured decode trace on MI355X (gfx950, TP4, DeepSeek-V4.1-Flash):
85% of the layer's 350 launches are mHC. Attributing GPU time through correlation IDs, the Sinkhorn roll-up is 37.5% of kernel time across 6,966 launches per decode step — ahead of
vllm::moe_forward_sharedat 15.4% — andaten::sumon[6, 4, 4]inputs alone is 21.5%, from 3,120 calls per step on 96-element tensors. Decode ITL on this machine degrades only 1.9x between concurrency 1 and 32 while per-GPU output throughput rises 13x, so the device is far from saturated and the per-step cost is dominated by something fixed rather than by bandwidth or compute.Why V4.1 cannot just reuse
MHCPreOpV4.1 collapses each sublayer's input with the previous sublayer's pre-mix, and returns the current one for the next seam. AITER's
mhc_preneither returns its pre-mix nor accepts an incoming one, so it cannot serve that shape as-is.The staggering does not change the gates, though. Comparing the delayed and non-delayed formulations on identical inputs:
So the expensive part — the projection, the RMS, the gates and all 20 Sinkhorn iterations, i.e. 141 of the 142 kernels — is bit-identical between the two. Only the collapse differs.
Approach
AITER exposes its two pre stages separately, and
mhc_pre_big_fuseconsumes the unreduced split-k projection output. This PR drives those stages directly:mhc_pre_gemm_sqrsum— projection againstfnplus row square-sums (1 kernel).mhc_pre_big_fuse— post/comb gates and every Sinkhorn iteration (1 kernel). Its collapse applies the current pre-mix, which the delayed form does not want, so it is discarded.mhc_pre_mix_tritonepilogue — recovers this seam's pre-mix from the same GEMM output and square-sums thatmhc_pre_big_fusealready consumed, rather than repeating the projection (1 kernel).hc_collapse_triton— the delayed collapse, with the carried coefficients (1 kernel).New
MHCPreDelayedOpdispatches AITER -> TileLang -> reference, mirroringMHCPreOp. Two cases decline AITER and fall through to TileLang: the model-entry broadcast seam (it projects a narrower operand,x, for which there is no AITER equivalent) and a fused RMSNorm (the composed path folds none).hidden_size5120 is also outside_AITER_MHC_FUSED_RMSNORM_SIZES, so the fused-RMSNorm variant would not apply here regardless.Result
Measured on gfx950,
hidden_size=5120,hc_mult=4, 20 Sinkhorn iterations, againstmhc_pre_delayed_torch:35x fewer launches per seam, worst relative error 1.2e-5, and the collapse is bit-exact because it is the same FP32 multiply-and-sum in both paths.
num_tokensspans AITER's split-k choices (80 down to 2), which matters because the pre-mix is recovered from that unreduced output.The model change also routes the three eager
(pre_mix.unsqueeze(-1) * x.float()).sum(dim=1)collapses (including the DSpark head) throughhc_collapse_triton, which is bit-exact for that reduction, and themhc_postcall sites throughMHCPostOp.Test Plan
tests/kernels/test_mhc_kernels.py::test_mhc_pre_delayed_rocm_aiter— AITER vs the delayed reference over tokens 1/2/7/128/1024, both seam variants (carried and model-entry pre-mix), asserting the collapse bit-exact.tests/kernels/test_mhc_kernels.py::test_mhc_pre_delayed_rocm_aiter_declines_unsupported— the broadcast seam and a fused norm must not silently take the AITER path. Asserted by recording which custom op runs, with a positive control so the two declines cannot pass merely because the AITER path is unavailable. (The earlier version checked this by comparing the fallback's output against a Torch reference, which fails for reasons that have nothing to do with this PR:mhc_pre_big_fuse_with_norm_tilelangfrom [Model] Support DeepSeek-V4.1-Flash #56214 is broken on gfx950 with the pinned tilelang 0.1.10, reducing the RMS sum of squares inside a thread-divergent branch and recovering half of it, solayer_inputcomes out sqrt(2) high. Confirmed pre-existing by running the suite on main + [Model] Support DeepSeek-V4.1-Flash #56214 with this PR reverted: the same 15test_deepseek_v41_mhc_pre_delayed[fused_norm=True]cases fail there, and applying this PR leaves the count unchanged. It does not affect serving, since the AMD decoder appliesattn_norm/ffn_normseparately and never passesnorm_weightinto the pre block.)End-to-end on MI355X (8x gfx950, TP4, DeepSeek-V4.1-Flash), gsm8k 5-shot via lm-eval:
Within one standard error on both filters, and well clear of the 0.91 gate. The run also exercises the TileLang fallback:
hc_prenorm_gemm_tilelangcompiles 8 times on this branch and never on the baseline, which is the model-entry broadcast seam declining AITER as intended.Re-run on the 2026-09-11 nightly (
0.28.1rc1.dev681+ge7edf17ce+ [Model] Support DeepSeek-V4.1-Flash #56214), which is main as it stands rather than the image the row above came from: 0.9719 ± 0.0045 strict-match, 0.9712 ± 0.0046 flexible-extract, 1319/1319 problems.Controlled A/B of the block itself, same process and identical input tensors so only the kernel choice differs, under CUDA graphs (
CUDAGraphMode.FULL_AND_PIECEWISE, which is how decode runs here). Microseconds per call, median of 5 x 200 iterations:Timed eagerly instead, the AITER path is slower than TileLang (93.4 us vs 28.8 us at 1 token): the Python wrapper in
_aiter_ops.mhc_pre_delayedcosts roughly 75 us of host time per call that a captured graph removes. It remains 7.6x-8.2x faster than the Torch reference main runs today, so eager is not a regression against the baseline this PR replaces, but trimming that host overhead is a worthwhile follow-up forenforce_eagerusers.The end-to-end ITL figures in the thread below are from a time-driven agentic trace and are not a controlled A/B; the table above is the controlled measurement.
Notes / follow-ups
mhc_pre_big_fusestill writes a collapse we discard. That redundant[tokens, hidden_size]bf16 store is the price of composing rather than having a native delayed kernel; it is negligible at decode shapes and ~1-2% of a prefill step. A nativemhc_pre_delayedin AITER (have themhc_preepilogue returnpre_mixand accept an incoming one —pre_mixis already computed inside the kernel as a slice ofmixes, so no new math) would remove both it and the Triton epilogue, taking the seam to 2 launches. Happy to take that up on the AITER side as a follow-up.aiter/jit/module_fused_ar_mhc.socould additionally fuse the TP all-reduce into the mHC seam (4 collective launches per layer today). AITER already exposesfused_allreduce_mhc_post_{only,one_stage,split}and vLLM calls none of them. Out of scope here; tracked in the RFC below.using torch solution:0for the indexer shapes (N:128,K:512andN:32,K:5120), i.e. no tuned config on gfx950.Part of a broader set of ROCm gaps for this model, collected in #56506.