Repository navigation
[ROCm][Perf] Use AITER tuned GEMM for the MoE router gate (backport #50535) - #21
Merged
Merged
Conversation
Tier 5 of the gate dispatch runs a bf16xbf16->fp32 GEMM, but it is reachable only through cuBLAS on Hopper and Blackwell. On ROCm the gate fell through to the generic F.linear fallback. Extend tier 5 with the AITER tuned GEMM. cuBLAS requires is_cuda() and AITER requires is_rocm(), so at most one is ever eligible and the two share the existing bf16 guard rather than adding a tier. Reaching it also requires leaving the weight in bf16: force_fp32_compute promoted the gate weight to fp32 whenever no specialized kernel was available, which on ROCm meant always. AITER only ships bf16-output tuned configs for the gate shapes, so the promotion would have kept the tier dead. The GEMM now takes the tuned bf16 kernel and casts the small num_tokens x num_experts output instead of upcasting the activation. Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com> Backport of upstream vllm-project/vllm PR vllm-project#50535 to moreh/v0.28.0. v0.28.0 keeps its six-tier dispatch docstring (it still lists the DSV3 kernel tier), so the AITER GEMM is documented as part of tier 5 there; the code change is identical. (cherry picked from commit 5fbeda20650903e9c21ec382b0affb804d2a2e76)
hgquannn
added a commit
that referenced
this pull request
Sep 21, 2026
…se branch Brings the changes tested on the tmp/v0.28.0 image series into moreh/v0.28.0. Content is exactly that of the fork PRs #16 to #22 (all against this branch or its tmp copy): #16 [ROCm][Perf] Fuse DSA indexer QK preprocessing with AITER (upstream vllm-project#53094) #17 [ROCm][Perf] Shard DSA prefill indexer rows across TP ranks (upstream vllm-project#56944) #18 [ROCm] Skip redundant sparse index remap on non-indexer layers (backport of upstream vllm-project/vllm vllm-project#51309) #19 [ROCm][MLA][Perf] Fuse decode QK-RoPE + Q-concat + KV-concat + KV-cache write for sparse MLA (backport of upstream vllm-project#47757) #20 [ROCm][Perf] Dispatch DSA prefill indexer logits to aiter's HIP kernel on gfx950 #21 [ROCm][Perf] Use AITER tuned GEMM for the MoE router gate (backport of upstream vllm-project#50535) #22 [ROCm][Bugfix] Add `positions` to unified_mla_attention_with_output_fake Verified as image vllm-forks:0.28.0-pr16-21-a4-9-260917 (this branch minus #22) built with aiter moreh/v0.1.19 + PR #4 + PR #9 on MI355X. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Context
Backport of upstream vllm-project#50535 (open, base
main, upstream commit5fbeda2065) tomoreh/v0.28.0, for the GLM-5.2 DSA test images built from this fork.The MoE router gate's bf16×bf16→fp32 GEMM tier was reachable only through cuBLAS on Hopper/Blackwell; on ROCm the gate fell through to
F.linear. This extends the tier with the AITER tuned GEMM (gfx950). It also keeps the gate weight in bf16 when the AITER tier is available, sinceforce_fp32_computewould otherwise promote it to fp32 and AITER only ships bf16-output tuned configs for the gate shapes; the smallnum_tokens x num_expertsoutput is cast instead. Upstream reports TPOT -3.05% / throughput +3.10% on MI355X TP4 for GLM-5.2-MXFP4.Changes
vllm/model_executor/layers/fused_moe/router/gate_linear.py:can_use_aiter_tuned_gemm,allow_aiter_router_gemm, bf16 weight when the AITER tier applies, the AITER branch in the bf16→fp32 tier, and therocm_aiter_router_gemmcustom op (opaque to torch.compile).tests/kernels/test_rocm_aiter_router_gemm.py: upstream test, unchanged.Backport difference: v0.28.0 still has the six-tier docstring (DSV3 kernel tier included), so the AITER GEMM is documented under tier 5 here instead of tier 4. The code is identical to upstream.
Test plan
python3 -m py_compileon both files: pass.vllm-forks:0.28.0-pr16-21-a4-9-260917(vllmtmp/v0.28.0= PR [ROCm][Perf] Fuse DSA indexer QK preprocessing with AITER (upstream #53094) #16 to [ROCm][Perf] Use AITER tuned GEMM for the MoE router gate (backport #50535) #21; aitertmp/v0.1.19-pr4-9) for GLM-5.2-MXFP4 TP8 serving on MI355X. Serving results to follow.Assisted by Claude Code; conflict resolution reviewed by hand.