[TRTLLM-12214][perf] DeepGemmFusedMoE: fuse masked gather + finalize-scale into one Triton kernel

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

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) ---------------

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) ---------------

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) ---------------