[TRTLLM-12214][perf] DeepGemmFusedMoE: fuse masked gather + finalize-scale into one Triton kernel - #14592
Conversation
…scale into one Triton kernel
The post-GEMM pipeline in DeepGemmFusedMoE was:
triton_masked_index_gather (writes permuted_data buffer, 25.6 MB at
bs=32 x topk=8 x hidden=7168 x bf16)
-> torch.ops.trtllm.moe_finalize_scale_op (reads permuted_data,
applies routing weights,
accumulates to output)
This commit fuses both steps into a single Triton kernel
triton_fused_gather_finalize that:
* reads h3 (the GEMM output) directly via the reverse-permute map
(permuted_row = unpermuted_to_permuted[token + k*num_rows], then
(expert_idx, col_idx) = (token_to_expert_map[permuted_row],
permuted_row - expert_first_token_offset[expert_idx]))
* applies token_final_scales in the same k=0..topk-1 order per token,
* accumulates in fp32 (matching the C++ finalizeMoeRoutingKernel
precision used by moe_finalize_scale_op).
The 25.6 MB intermediate buffer is eliminated; total HBM traffic for
this pair drops from ~80 MB to ~28.8 MB (-64%) at the model anchor.
Microbench on B200 (sm_100, do_bench warmup=50 rep=300), 40 cells
(num_source_tokens in {1,4,16,32,64} x topK in {4,8} x experts in
{64,128} x hidden in {4096,7168}):
anchor (N=32, K=8, E=128, H=4096): 27.71 us -> 16.38 us (-40.9%)
anchor (N=32, K=8, E=128, H=7168): 38.94 us -> 16.38 us (-57.9%)
small (N=1, K=4, E=64, H=4096): 20.29 us -> 10.27 us (-49.4%)
large (N=64, K=8, E=128, H=7168): 41.79 us -> 16.38 us (-60.8%)
aggregate: 40/40 win, median -54.9%, best -66.4%, worst -40.9%.
Zero regressing cells across the full sweep.
Full-bench context (Qwen3-235B-A22B-FP8 + EAGLE3 dyntree, 4xGB200 NVL72,
TP=4, bs=32, mtbench-32): replay median 35887 us -> 32587 us (-9.2%),
throughput 1447 -> 1583 tok/s (+9.4%), accuracy neutral
(acceptance-rate +0.0038, output-length -0.0031), all PR-1..PR-3 stacked.
Risk callouts for reviewers:
* moe_finalize_scale_op also handles enable_alltoall=True and other
paths the new kernel does not. DeepGemmFusedMoE.forward passes
enable_alltoall=False inline today; if alltoall is ever wired on
the DeepGemm path, this fusion must back out or fork.
* The 2D grid (num_rows, ceil_div(unpadded_hidden, 1024)) redundantly
loads topk metadata once per hidden-chunk program (7x at hidden=7168
/ BLOCK_SIZE=1024); L1 is expected to absorb this. A 1D-per-row
grid is the obvious fallback if NCU shows poor L1 hit rate.
* Single-file change, no C++ touch. Independent of PR-1 and PR-2 at
the code level.
Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
|
/bot run |
|
PR_Github #50371 [ run ] triggered by Bot. Commit: |
|
PR_Github #50371 [ run ] completed with state
|
…nalize Adds a differential correctness test that feeds the same self-consistent permutation maps (built by the real moe_permute_op) to both the old post-GEMM pair (triton_masked_index_gather + moe_finalize_scale_op) and the new fused triton_fused_gather_finalize kernel, then asserts the outputs match. Both paths accumulate k=0..topk-1 in order in an fp32 accumulator with the same round-to-nearest fp32->bf16 final cast, so the result is bit-identical modulo FMA-contraction differences between the Triton and nvcc backends. The committed gate is a tight tolerance (robust across GPU arch); the test also reports whether the result is bitwise equal. Covers the model anchor (32x8x4096x128), hidden=7168, single-token decode, medium-batch, and a top_k=4 shape. Gated to SM90+ (no deep_gemm GEMM is invoked, only the permute/gather/finalize kernels). Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
81e8dcf to
c5840a8
Compare
|
/bot run |
|
PR_Github #50900 [ run ] triggered by Bot. Commit: |
|
PR_Github #50900 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51105 [ run ] triggered by Bot. Commit: |
|
PR_Github #51105 [ run ] completed with state
|
…est lists The new differential-correctness test tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py was added in this PR but not picked up by CI because it lives in a new file not referenced by any test-db list. Register it in the SM90+ single-GPU lists (l0_h100, l0_b200, l0_b300) alongside the other MoE component tests so the bit-exactness check actually runs in pre-merge. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #51398 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR introduces a fused Triton kernel that combines MoE gather and finalize operations into a single operation. The kernel reads expert GEMM output directly, applies per-token routing weights, and writes final hidden states. The prior separate gather and finalize steps are removed from ChangesFused MoE gather-finalize
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/integration/test_lists/test-db/l0_h100.yml (1)
38-38: QA test-list follow-up is unnecessary for this change.This PR wires a unit test into existing pre-merge
test-dbsuites; notests/integration/defs/integration-test addition was made, sotests/integration/test_lists/qa/updates are optional/unnecessary here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/test_lists/test-db/l0_h100.yml` at line 38, The QA test-list follow-up is not needed because this change only wires an existing unit test (unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py) into the pre-merge test-db suite; remove any added or modified entries under tests/integration/test_lists/qa/ (or avoid creating a new QA integration-test def) and keep only the test-db list update (e.g., in tests/integration/test_lists/test-db/l0_h100.yml) so no qa/defs changes are committed.tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py (1)
71-85: ⚡ Quick winAdd an explicit return type to
_make_routing.Please annotate the helper return type (e.g.,
tuple[torch.Tensor, torch.Tensor]) to satisfy the repo’s typing rule for functions.As per coding guidelines: "Static type checking with mypy is opt-in by submodule; always annotate functions with return types (use
Noneif function does not return)".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py` around lines 71 - 85, Add an explicit return type annotation to the helper function _make_routing: change its signature to declare it returns tuple[torch.Tensor, torch.Tensor] (or Tuple[torch.Tensor, torch.Tensor] if you prefer typing.Tuple), so the signature references torch.Tensor for both elements; no functional changes required inside token_selected_experts and token_final_scales — just annotate the def _make_routing(...) -> tuple[torch.Tensor, torch.Tensor] (and add an import from typing if using Tuple).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/integration/test_lists/test-db/l0_h100.yml`:
- Line 38: The QA test-list follow-up is not needed because this change only
wires an existing unit test
(unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py) into
the pre-merge test-db suite; remove any added or modified entries under
tests/integration/test_lists/qa/ (or avoid creating a new QA integration-test
def) and keep only the test-db list update (e.g., in
tests/integration/test_lists/test-db/l0_h100.yml) so no qa/defs changes are
committed.
In
`@tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py`:
- Around line 71-85: Add an explicit return type annotation to the helper
function _make_routing: change its signature to declare it returns
tuple[torch.Tensor, torch.Tensor] (or Tuple[torch.Tensor, torch.Tensor] if you
prefer typing.Tuple), so the signature references torch.Tensor for both
elements; no functional changes required inside token_selected_experts and
token_final_scales — just annotate the def _make_routing(...) ->
tuple[torch.Tensor, torch.Tensor] (and add an import from typing if using
Tuple).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c564588f-ff41-4919-9832-39033e4d75b7
📒 Files selected for processing (5)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.pytests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_b300.ymltests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py
|
PR_Github #51398 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51442 [ run ] triggered by Bot. Commit: |
|
PR_Github #51442 [ run ] completed with state
|
|
/bot skip --comment "All failures in pipeline 40851 are irrelevant. This PR's new test passed." |
|
PR_Github #52110 Bot args parsing error: Failed to parse bot args |
|
/bot skip --comment "All failures in pipeline 40851 are irrelevant. This PRs new test passed." |
|
PR_Github #52114 [ skip ] triggered by Bot. Commit: |
|
PR_Github #52114 [ skip ] completed with state |
…scale into one Triton kernel (NVIDIA#14592) Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
3 participants
[TRTLLM-12214][perf] DeepGemmFusedMoE: fuse masked gather + finalize-scale into one Triton kernel
Summary
The post-GEMM pipeline in
DeepGemmFusedMoEwas:This PR fuses both steps into one Triton kernel
triton_fused_gather_finalizethat readsh3(the GEMM output) directly via the reverse-permute map:applies
token_final_scalesin the samek = 0..topk-1order per token, and accumulates in fp32 (matching the C++finalizeMoeRoutingKernelprecision used bymoe_finalize_scale_op).The 25.6 MB intermediate buffer is eliminated; total HBM traffic for this pair drops from ~80 MB to ~28.8 MB (−64%) at the model anchor.
Microbenchmark (B200, sm_100)
triton.testing.do_bench(warmup=50, rep=300)over 40 cells:num_source_tokens ∈ {1, 4, 16, 32, 64} × topK ∈ {4, 8} × experts ∈ {64, 128} × hidden ∈ {4096, 7168}.Baselineistriton_masked_index_gather + moe_finalize_scale_op(sum of the two timed individually).Patchedis the new fused kernel.N=1, K=4, E=64, H=4096N=32, K=8, E=128, H=4096N=32, K=8, E=128, H=7168N=64, K=8, E=128, H=7168N=1, K=8, E=128, H=7168(best)Zero regressing cells across the full sweep. The anchor measures −40.9% at H=4096 and −57.9% at H=7168.
Full-bench context
Qwen3-235B-A22B-FP8 + EAGLE3 dyntree, 4×GB200 NVL72, TP=4, EP=1, bs=32, mtbench-32, output_tokens=512,
--ignore_eos, all PR-1..PR-3 stacked:This is the largest single contribution of the three-PR series (~93 invocations per replay × ~16 µs saved per call).
Risk callouts
moe_finalize_scale_opalso handlesenable_alltoall=Trueand other paths the new kernel does not.DeepGemmFusedMoE.forwardpassesenable_alltoall=Falseinline today; if alltoall is ever wired onto the DeepGemm path, this fusion must back out or fork.(num_rows, ceil_div(unpadded_hidden, 1024))redundantly loads topk metadata once per hidden-chunk program (7× athidden=7168/BLOCK_SIZE=1024); L1 is expected to absorb this. A 1D-per-row grid is the obvious fallback if NCU shows poor L1 hit rate.h3addresses (proven by the reverse-permute math above), sametoken_final_scalesapplied in the samekorder, same fp32 accumulator precision as the C++finalizeMoeRoutingKernel. No new index math hazards.Files changed
tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py(single file)fused_gather_finalize_kernelJITtriton_fused_gather_finalizelaunchertriton_masked_index_gather(...) + torch.ops.trtllm.moe_finalize_scale_op(...)call pair with a single call to the new launcherNo C++ touch. Independent of PR-1 and PR-2 at the code level.
Test plan
Summary by CodeRabbit
Release Notes
New Features
Tests