From eec674a69e7ea209e7a8038b00adc32510880341 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Tue, 26 May 2026 09:42:41 -0700 Subject: [PATCH 1/3] [TRTLLM-12214][perf] DeepGemmFusedMoE: fuse masked gather + finalize-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> --- .../modules/fused_moe/fused_moe_deepgemm.py | 181 +++++++++++++++--- 1 file changed, 158 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py index b051c9348cab..e5286f2b7887 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py @@ -233,6 +233,149 @@ def triton_masked_index_gather(output, input, start_offsets, row_indices): return +@triton.jit +def fused_gather_finalize_kernel( + output_ptr, + h3_ptr, + scales_ptr, + unpermuted_row_to_permuted_row_ptr, + token_to_expert_map_ptr, + expert_first_token_offset_ptr, + token_selected_experts_ptr, + num_rows, + experts_per_token, + col_size, + dim_size, + unpadded_dim_size, + num_experts_per_node, + start_expert_id, + BLOCK_SIZE: tl.constexpr, +): + """Fused gather + finalize kernel. + + Replaces masked_index_gather + finalizeMoeRoutingKernel by reading + directly from expert GEMM output (h3), applying routing weights, + and accumulating to the final output — eliminating the intermediate + permuted_data buffer entirely. + + Grid: (num_rows, cdiv(unpadded_dim_size, BLOCK_SIZE)) + """ + pid_row = tl.program_id(0) + pid_col = tl.program_id(1) + + # Hidden dimension offsets for this program + hidden_start = pid_col * BLOCK_SIZE + hidden_offsets = hidden_start + tl.arange(0, BLOCK_SIZE) + valid_hidden = hidden_offsets < unpadded_dim_size + + # Initialize accumulator in fp32 for precision + acc = tl.zeros((BLOCK_SIZE, ), dtype=tl.float32) + + # Iterate over topk experts for this token + for k in tl.range(0, experts_per_token): + # Check if expert is on this node + k_offset = pid_row * experts_per_token + k + expert_id_global = tl.load(token_selected_experts_ptr + k_offset) + expert_id_local = expert_id_global - start_expert_id + + if expert_id_local >= 0 and expert_id_local < num_experts_per_node: + # Get the expanded permuted row index + expanded_original_row = pid_row + k * num_rows + expanded_permuted_row = tl.load(unpermuted_row_to_permuted_row_ptr + + expanded_original_row) + + # Reverse the gather mapping: find h3 coordinates + # token_to_expert_map gives local expert index for this permuted row + local_expert_idx = tl.load(token_to_expert_map_ptr + + expanded_permuted_row) + expert_start = tl.load(expert_first_token_offset_ptr + + local_expert_idx) + col_idx = expanded_permuted_row - expert_start + + # Read from h3[local_expert_idx, col_idx, hidden_offsets] + h3_offset = (local_expert_idx * col_size * dim_size + + col_idx * dim_size + hidden_offsets) + h3_val = tl.load(h3_ptr + h3_offset, mask=valid_hidden, other=0.0) + + # Load routing weight and apply + scale = tl.load(scales_ptr + k_offset) + acc += h3_val.to(tl.float32) * scale + + # Store result + output_offset = pid_row * unpadded_dim_size + hidden_offsets + tl.store(output_ptr + output_offset, + acc.to(output_ptr.dtype.element_ty), + mask=valid_hidden) + + +@torch.no_grad() +def triton_fused_gather_finalize( + h3: torch.Tensor, + token_final_scales: torch.Tensor, + unpermuted_row_to_permuted_row: torch.Tensor, + token_to_expert_map: torch.Tensor, + expert_first_token_offset: torch.Tensor, + token_selected_experts: torch.Tensor, + num_rows: int, + hidden_size: int, + unpadded_hidden_size: int, + experts_per_token: int, + num_experts_per_node: int, + ep_rank: int, +) -> torch.Tensor: + """Fused gather + finalize: reads h3 directly, applies routing weights, + and accumulates to output, eliminating the intermediate permuted_data buffer. + + Args: + h3: Expert GEMM output [num_experts_local, col_size, hidden_size] + token_final_scales: Routing weights [num_rows, experts_per_token] + unpermuted_row_to_permuted_row: Mapping [num_rows * experts_per_token] + token_to_expert_map: Expanded token → local expert ID + expert_first_token_offset: Start offsets per expert [num_experts+1] + token_selected_experts: Global expert IDs [num_rows, experts_per_token] + num_rows: Number of original tokens + hidden_size: Padded hidden dimension (h3 stride) + unpadded_hidden_size: Actual output hidden dimension + experts_per_token: Top-K value + num_experts_per_node: Number of experts on this EP rank + ep_rank: Expert parallelism rank + + Returns: + output: [num_rows, unpadded_hidden_size] + """ + col_size = h3.shape[1] + dim_size = h3.shape[2] + start_expert_id = num_experts_per_node * ep_rank + + output = torch.empty( + (num_rows, unpadded_hidden_size), + dtype=h3.dtype, + device=h3.device, + ) + + BLOCK_SIZE = 1024 + grid = (num_rows, triton.cdiv(unpadded_hidden_size, BLOCK_SIZE)) + + fused_gather_finalize_kernel[grid]( + output, + h3, + token_final_scales, + unpermuted_row_to_permuted_row, + token_to_expert_map, + expert_first_token_offset, + token_selected_experts, + num_rows, + experts_per_token, + col_size, + dim_size, + unpadded_hidden_size, + num_experts_per_node, + start_expert_id, + BLOCK_SIZE=BLOCK_SIZE, + ) + return output + + @triton.jit def _preprocess_after_permute_kernel( expert_offsets_ptr, @@ -735,34 +878,26 @@ def run_moe( expected_m=expected_m, ) - # Gather and finalize - triton_masked_index_gather(permuted_data_tensor, h3, - expert_first_token_offset_tensor, - token_to_expert_map) - + # Fused gather + finalize: read h3 directly, apply routing weights, + # accumulate to output. Eliminates intermediate permuted_data buffer. topk = self.routing_method.top_k if token_selected_experts is not None: # For the deepgemmlowlatency, the topk has been viewed into 1 topk = token_selected_experts.shape[-1] - final_hidden_states = torch.ops.trtllm.moe_finalize_scale_op( - permuted_data_tensor, - None, # biases - token_final_scales, - unpermuted_row_to_permuted_row_tensor, - permuted_row_to_unpermuted_row_tensor, - token_selected_experts, - expert_first_token_offset_tensor, - False, # enable_alltoall - x.shape[0], # num_rows - x.shape[1], # (possibly padded) hidden_size - self.unpadded_hidden_size, # original hidden size - topk, - self.expert_size_per_partition, # num_experts_per_node - self.tp_size, - self.tp_rank, - self.ep_size, - self.ep_rank, + final_hidden_states = triton_fused_gather_finalize( + h3=h3, + token_final_scales=token_final_scales, + unpermuted_row_to_permuted_row=unpermuted_row_to_permuted_row_tensor, + token_to_expert_map=token_to_expert_map, + expert_first_token_offset=expert_first_token_offset_tensor, + token_selected_experts=token_selected_experts, + num_rows=x.shape[0], + hidden_size=x.shape[1], + unpadded_hidden_size=self.unpadded_hidden_size, + experts_per_token=topk, + num_experts_per_node=self.expert_size_per_partition, + ep_rank=self.ep_rank, ) return final_hidden_states From c5840a802b1ddfc12df6d2d6f65ca7e0525b5dc4 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Thu, 28 May 2026 15:12:10 -0700 Subject: [PATCH 2/3] [TRTLLM-12214][test] Add differential test for triton_fused_gather_finalize 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> --- .../test_deepgemm_fused_gather_finalize.py | 206 ++++++++++++++++++ 1 file changed, 206 insertions(+) create mode 100644 tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py diff --git a/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py b/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py new file mode 100644 index 000000000000..446187f7a845 --- /dev/null +++ b/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Differential correctness test for ``triton_fused_gather_finalize``. + +The DeepGemm post-GEMM pipeline used to be a pair of ops:: + + triton_masked_index_gather(permuted_data, h3, ...) # gather into buffer + torch.ops.trtllm.moe_finalize_scale_op(permuted_data, ...) # weight + reduce + +``triton_fused_gather_finalize`` fuses both: it reads the expert GEMM output +``h3`` directly via the reverse-permute map, applies the routing weights, and +accumulates to the final output, eliminating the intermediate ``permuted_data`` +buffer. + +This test feeds **the same, genuinely self-consistent permutation maps** (built +by the real ``moe_permute_op``) to both the old pair and the new fused kernel, +then asserts the two outputs match. The C++ ``finalizeMoeRoutingKernel`` and the +fused Triton kernel both accumulate ``k = 0..topk-1`` in order in an fp32 +accumulator with the same round-to-nearest fp32->bf16 final cast, so the result +should be bit-identical; the only residual difference is whether the +``acc += scale * value`` multiply-add contracts to an FMA identically across the +Triton and nvcc backends. The committed gate is therefore a tight tolerance, and +the test additionally reports whether the result happened to be bitwise equal. + +Run as:: + + pytest tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py -v +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm.*) +from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import ( + preprocess_after_permute, + triton_fused_gather_finalize, + triton_masked_index_gather, +) +from tensorrt_llm._utils import get_sm_version + +skip_unsupported = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() < 90, + reason="Requires CUDA SM90+ (DeepGemm MoE post-GEMM kernels)", +) + + +@dataclass(frozen=True) +class GatherFinalizeShape: + name: str + num_rows: int + hidden: int + num_experts: int + top_k: int + + +# Anchors mirror the PR-3 microbench sweep (study-plan note): the model anchor +# (32 x 8 x 4096 x 128), the single-token decode corner, and a medium batch. +SHAPES = [ + GatherFinalizeShape("anchor", num_rows=32, hidden=4096, num_experts=128, top_k=8), + GatherFinalizeShape("anchor_h7168", num_rows=32, hidden=7168, num_experts=128, top_k=8), + GatherFinalizeShape("single_token", num_rows=1, hidden=512, num_experts=64, top_k=4), + GatherFinalizeShape("medium_batch", num_rows=64, hidden=7168, num_experts=128, top_k=8), + GatherFinalizeShape("topk4", num_rows=16, hidden=4096, num_experts=64, top_k=4), +] + + +def _make_routing( + num_rows: int, num_experts: int, top_k: int, *, device: str, generator: torch.Generator +): + """Synthesize a valid (token_selected_experts, token_final_scales) pair. + + Each row selects ``top_k`` *distinct* experts (top-k never repeats an + expert), with positive routing weights that sum to one per row. + """ + logits = torch.randn( + (num_rows, num_experts), device=device, dtype=torch.float32, generator=generator + ) + topk_vals, topk_ids = logits.topk(top_k, dim=-1) + token_selected_experts = topk_ids.to(torch.int32) + token_final_scales = torch.softmax(topk_vals, dim=-1).to(torch.float32) + return token_selected_experts, token_final_scales + + +@skip_unsupported +@pytest.mark.parametrize("shape", SHAPES, ids=lambda s: s.name) +def test_fused_gather_finalize_matches_unfused(shape: GatherFinalizeShape) -> None: + device = "cuda" + gen = torch.Generator(device=device).manual_seed(1234) + + num_rows = shape.num_rows + hidden = shape.hidden + num_experts = shape.num_experts + top_k = shape.top_k + + # Single-rank configuration: every expert is local to this node. + tp_size, tp_rank, ep_size, ep_rank = 1, 0, 1, 0 + cluster_size, cluster_rank = 1, 0 + num_experts_per_node = num_experts + + token_selected_experts, token_final_scales = _make_routing( + num_rows, num_experts, top_k, device=device, generator=gen + ) + + # Driver activations for the permutation. Only the maps are consumed below; + # the expert GEMM output ``h3`` is synthesized independently. + x = torch.randn((num_rows, hidden), device=device, dtype=torch.bfloat16, generator=gen) + + # Real permutation maps (self-consistent by construction). + ( + permuted_row_to_unpermuted_row_tensor, + _permuted_token_selected_experts_tensor, + permuted_data_tensor, + expert_first_token_offset_tensor, + _permuted_token_final_scales_tensor, + unpermuted_row_to_permuted_row_tensor, + ) = torch.ops.trtllm.moe_permute_op( + x, + token_selected_experts, + token_final_scales, + None, # fc1_expert_weights + None, # fc2_expert_weights + None, # quant_scales + input_sf=None, + num_experts_on_rank=num_experts_per_node, + tp_size=tp_size, + tp_rank=tp_rank, + ep_size=ep_size, + ep_rank=ep_rank, + cluster_size=cluster_size, + cluster_rank=cluster_rank, + min_latency_mode=False, + use_fp8_block_scaling=True, + ) + + num_expanded = num_rows * top_k + assert permuted_data_tensor.shape[0] == num_expanded + + _masked_m, token_to_expert_map = preprocess_after_permute( + expert_first_token_offset_tensor, permuted_data_tensor + ) + + # Synthesize the expert GEMM output h3: [num_experts, max_tokens_per_expert, + # hidden]. At most ``num_rows`` tokens can route to any single expert (top-k + # picks distinct experts per token), so an m_max of align(num_rows, 128) + # safely bounds the per-expert column dimension, matching forward(). + m_max = ((num_rows + 127) // 128) * 128 + h3 = torch.randn( + (num_experts, m_max, hidden), device=device, dtype=torch.bfloat16, generator=gen + ) + + # ---- Old path: gather into permuted_data, then finalize-scale ---- + gather_out = torch.empty((num_expanded, hidden), device=device, dtype=h3.dtype) + triton_masked_index_gather( + gather_out, h3, expert_first_token_offset_tensor, token_to_expert_map + ) + out_unfused = torch.ops.trtllm.moe_finalize_scale_op( + gather_out, + None, # biases + token_final_scales, + unpermuted_row_to_permuted_row_tensor, + permuted_row_to_unpermuted_row_tensor, + token_selected_experts, + expert_first_token_offset_tensor, + False, # enable_alltoall + num_rows, + hidden, # (possibly padded) hidden_size + hidden, # unpadded hidden size (no padding in this test) + top_k, + num_experts_per_node, + tp_size, + tp_rank, + ep_size, + ep_rank, + ) + + # ---- New path: single fused kernel reading h3 directly ---- + out_fused = triton_fused_gather_finalize( + h3=h3, + token_final_scales=token_final_scales, + unpermuted_row_to_permuted_row=unpermuted_row_to_permuted_row_tensor, + token_to_expert_map=token_to_expert_map, + expert_first_token_offset=expert_first_token_offset_tensor, + token_selected_experts=token_selected_experts, + num_rows=num_rows, + hidden_size=hidden, + unpadded_hidden_size=hidden, + experts_per_token=top_k, + num_experts_per_node=num_experts_per_node, + ep_rank=ep_rank, + ) + + assert out_fused.shape == out_unfused.shape == (num_rows, hidden) + assert out_fused.dtype == out_unfused.dtype == torch.bfloat16 + + # Report bit-exactness (documents the PR claim) without gating on it. + bitwise_equal = torch.equal(out_fused, out_unfused) + max_abs_diff = (out_fused.float() - out_unfused.float()).abs().max().item() + print(f"[{shape.name}] bitwise_equal={bitwise_equal} max_abs_diff={max_abs_diff:.3e}") + + # Committed gate: tight tolerance (robust across FMA-contraction differences + # between the Triton and nvcc backends). + torch.testing.assert_close(out_fused, out_unfused, rtol=1e-2, atol=1e-2) From a32774669b57ef4b2413f00edb77fbb5ef6a8969 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 1 Jun 2026 10:54:18 -0700 Subject: [PATCH 3/3] [TRTLLM-12214][test] register fused gather+finalize unit test in CI test 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> --- tests/integration/test_lists/test-db/l0_b200.yml | 1 + tests/integration/test_lists/test-db/l0_b300.yml | 1 + tests/integration/test_lists/test-db/l0_h100.yml | 1 + 3 files changed, 3 insertions(+) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index def40f798183..3cbb3b7c0a55 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -105,6 +105,7 @@ l0_b200: - unittest/_torch/modules/test_moe_load_balancer.py - unittest/_torch/modules/test_moe_routing.py - unittest/_torch/modules/test_moe_host_sharer.py + - unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py # ------------- legacy MoE tests --------------- - unittest/_torch/modules/test_fused_moe.py # ------------- MoE: test_moe_backend (by backend) --------------- diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 063125788ed9..515dd979da6d 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -32,6 +32,7 @@ l0_b300: # ------------- MoE components tests --------------- - unittest/_torch/modules/test_moe_load_balancer.py - unittest/_torch/modules/test_moe_routing.py + - unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py # ------------- legacy MoE tests --------------- - unittest/_torch/modules/test_fused_moe.py # ------------- MoE: test_moe_backend (by backend) --------------- diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index e55c5ae8363e..3960fbc6e9c9 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -35,6 +35,7 @@ l0_h100: - unittest/_torch/modules/test_moe_load_balancer.py - unittest/_torch/modules/test_moe_routing.py - unittest/_torch/modules/test_moe_host_sharer.py + - unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py # ------------- legacy MoE tests --------------- - unittest/_torch/modules/test_fused_moe.py # ------------- MoE: test_moe_backend (by backend) ---------------