[TRTLLM-7408][feat] Wrap MOE with custom op.

* Let all moe backend go through the same interface
* MOE is wrapped with custom op to improve full graph torch compile compatibility

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

Author: liji-nv

tensorrt_llm/_torch/attention_backend/interface.py

@@ -121,12 +121,7 @@ class AttentionMetadata:
         default_factory=AttentionRuntimeFeatures)
 
     # The number of tokens in each rank.
-    _all_rank_num_tokens: Optional[List[int]] = field(init=False,
-                                                      default=None,
-                                                      repr=False)
-    all_rank_num_tokens: Optional[List[int]]
-    # The max number of tokens among all ranks.
-    all_rank_max_num_tokens: Optional[int] = None
+    all_rank_num_tokens: Optional[List[int]] = None
 
     # These fields are set when changing seq_lens and _num_contexts to avoid computation
     # during execution. If the calculation happens during execution, torch compile treats it
@@ -163,16 +158,6 @@ def on_update(self):
         elif self._seq_lens is not None:
             self._num_tokens = self._seq_lens.sum().item()
 
-    @property
-    def all_rank_num_tokens(self) -> Optional[List[int]]:
-        return self._all_rank_num_tokens
-
-    @all_rank_num_tokens.setter
-    def all_rank_num_tokens(self, value: Optional[List[int]]):
-        value = value if value is not AttentionMetadata.all_rank_num_tokens else None
-        self._all_rank_num_tokens = value
-        self.all_rank_max_num_tokens = max(value) if value is not None else None
-
     @property
     def seq_lens(self) -> Optional[torch.Tensor]:
         return self._seq_lens

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -531,8 +531,7 @@ def _compute_shared_expert_tp_size(self, intermediate_size: int,
         return shared_tp_size, shared_output_scale
 
     def compute_routed_output(self, hidden_states, hidden_states_fp4,
-                              all_rank_num_tokens, all_rank_max_num_tokens,
-                              do_finalize):
+                              all_rank_num_tokens, do_finalize):
         # max-throughput
         use_dp_padding = False
         if self.use_dp and self.mapping.tp_size > 1:
@@ -551,7 +550,6 @@ def compute_routed_output(self, hidden_states, hidden_states_fp4,
             do_finalize=do_finalize,
             output_dtype=hidden_states.dtype,
             all_rank_num_tokens=all_rank_num_tokens,
-            all_rank_max_num_tokens=all_rank_max_num_tokens,
             use_dp_padding=use_dp_padding,
         )
 
@@ -562,7 +560,6 @@ def forward(
         hidden_states: torch.Tensor,
         hidden_states_fp4: Optional[Fp4QuantizedTensor] = None,
         all_rank_num_tokens: Optional[list[int]] = None,
-        all_rank_max_num_tokens: Optional[int] = None,
         final_all_reduce_params: Optional[AllReduceParams] = None,
         do_finalize: Optional[bool] = True,
     ) -> torch.Tensor:
@@ -581,7 +578,6 @@ def _compute_routed_output():
             routed_output = self.compute_routed_output(hidden_states,
                                                        hidden_states_fp4,
                                                        all_rank_num_tokens,
-                                                       all_rank_max_num_tokens,
                                                        do_finalize)
             return routed_output
 
@@ -804,7 +800,6 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize):
                 hidden_states,
                 hidden_states_fp4,
                 all_rank_num_tokens=attn_metadata.all_rank_num_tokens,
-                all_rank_max_num_tokens=attn_metadata.all_rank_max_num_tokens,
                 final_all_reduce_params=AllReduceParams(
                     enable_allreduce=not (self.fusion_config.POST_MOE_FUSION
                                           or self.mapping.tp_size == 1)),
@@ -992,7 +987,6 @@ def forward(
         embed_tokens: Embedding,
         attn_metadata: AttentionMetadata,
         all_rank_num_tokens: Optional[List[int]] = None,
-        all_rank_max_num_tokens: Optional[int] = None,
         **kwargs,
     ) -> torch.Tensor:
 
@@ -1051,7 +1045,6 @@ def norm_hidden():
         hidden_states = self.mlp(
             hidden_states,
             all_rank_num_tokens=all_rank_num_tokens,
-            all_rank_max_num_tokens=all_rank_max_num_tokens,
             final_all_reduce_params=AllReduceParams(
                 enable_allreduce=not (self.fusion_config.POST_MOE_FUSION
                                       or self.mapping.tp_size == 1)),

tensorrt_llm/_torch/models/modeling_gpt_oss.py

@@ -258,7 +258,6 @@ def forward_attn_dp(
 
         # Get attention_dp parameters
         all_rank_num_tokens = attn_metadata.all_rank_num_tokens
-        all_rank_max_num_tokens = attn_metadata.all_rank_max_num_tokens
 
         if self.mapping.tp_size > 1 and all_rank_num_tokens is not None:
             if (isinstance(self.experts, (TRTLLMGenFusedMoE, TritonFusedMoE))):
@@ -276,12 +275,10 @@ def forward_attn_dp(
 
         # Let CutlassFusedMoE handle allgather internally
         # Pass the normalized tensor (t) as input to experts, not x
-        expert_output = self.experts(
-            x=t,
-            router_logits=g,
-            all_rank_num_tokens=all_rank_num_tokens,
-            all_rank_max_num_tokens=all_rank_max_num_tokens,
-            use_dp_padding=False)
+        expert_output = self.experts(x=t,
+                                     router_logits=g,
+                                     all_rank_num_tokens=all_rank_num_tokens,
+                                     use_dp_padding=False)
 
         expert_output = expert_output.view(orig_shape)
         return expert_output, residual

tensorrt_llm/_torch/models/modeling_llama.py

@@ -309,32 +309,27 @@ def __init__(
         self.aux_stream = aux_stream
 
     def compute_routed_output(self, hidden_states, all_rank_num_tokens,
-                              all_rank_max_num_tokens,
                               cutlass_min_latency_mode):
         router_logits = self.router(hidden_states)
-        routed_output = self.experts(
-            hidden_states,
-            router_logits,
-            do_finalize=not cutlass_min_latency_mode,
-            all_rank_num_tokens=all_rank_num_tokens,
-            all_rank_max_num_tokens=all_rank_max_num_tokens,
-            use_dp_padding=False)
+        routed_output = self.experts(hidden_states,
+                                     router_logits,
+                                     do_finalize=not cutlass_min_latency_mode,
+                                     all_rank_num_tokens=all_rank_num_tokens,
+                                     use_dp_padding=False)
         return routed_output
 
     def forward(
         self,
         hidden_states: torch.Tensor,
         all_rank_num_tokens=None,
-        all_rank_max_num_tokens=None,
         final_all_reduce_params: Optional[AllReduceParams] = None,
         cutlass_min_latency_mode: Optional[bool] = False,
     ) -> torch.Tensor:
         # Only enable multi-stream for cuda graph since switch stream has extra host overhead
         # This design is mainly for low latency use case. Need to improve for max throughput use case.
         fn0 = lambda: self.shared_expert(hidden_states)
         fn1 = lambda: self.compute_routed_output(
-            hidden_states, all_rank_num_tokens, all_rank_max_num_tokens,
-            cutlass_min_latency_mode)
+            hidden_states, all_rank_num_tokens, cutlass_min_latency_mode)
         shared_output, routed_output = maybe_execute_in_parallel(
             fn0, fn1, self.moe_event[0], self.moe_event[1], self.aux_stream)
         if cutlass_min_latency_mode:
@@ -536,7 +531,6 @@ def forward(
         hidden_states = self.feed_forward(
             hidden_states,
             all_rank_num_tokens=attn_metadata.all_rank_num_tokens,
-            all_rank_max_num_tokens=attn_metadata.all_rank_max_num_tokens,
             final_all_reduce_params=AllReduceParams(
                 enable_allreduce=not self.disable_feed_forward_allreduce),
             cutlass_min_latency_mode=cutlass_min_latency_mode,

tensorrt_llm/_torch/models/modeling_mixtral.py

@@ -62,13 +62,11 @@ def forward(
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         all_rank_num_tokens = attn_metadata.all_rank_num_tokens
-        all_rank_max_num_tokens = attn_metadata.all_rank_max_num_tokens
         router_logits = self.gate(hidden_states)
         final_hidden_states = self.experts(
             hidden_states,
             router_logits,
             all_rank_num_tokens=all_rank_num_tokens,
-            all_rank_max_num_tokens=all_rank_max_num_tokens,
             use_dp_padding=False)
         return final_hidden_states
 

tensorrt_llm/_torch/models/modeling_qwen3_moe.py

@@ -127,7 +127,6 @@ def forward(
         hidden_states = hidden_states.view(-1, self.hidden_dim)
         use_dp_padding = False
         all_rank_num_tokens = attn_metadata.all_rank_num_tokens
-        all_rank_max_num_tokens = attn_metadata.all_rank_max_num_tokens
 
         if not do_finalize:
             assert not self.enable_attention_dp
@@ -144,7 +143,6 @@ def forward(
             hidden_states,
             router_logits,
             all_rank_num_tokens=all_rank_num_tokens,
-            all_rank_max_num_tokens=all_rank_max_num_tokens,
             use_dp_padding=use_dp_padding,
             do_finalize=do_finalize,
         )

tensorrt_llm/_torch/models/modeling_qwen_moe.py

@@ -84,13 +84,11 @@ def forward(
         hidden_states = hidden_states.view(-1, self.hidden_dim)
 
         all_rank_num_tokens = attn_metadata.all_rank_num_tokens
-        all_rank_max_num_tokens = attn_metadata.all_rank_max_num_tokens
         router_logits = self.gate(hidden_states)
         final_hidden_states = self.experts(
             hidden_states,
             router_logits,
             all_rank_num_tokens=all_rank_num_tokens,
-            all_rank_max_num_tokens=all_rank_max_num_tokens,
             use_dp_padding=False)
 
         shared_expert_output = self.shared_expert(hidden_states)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

@@ -85,6 +85,7 @@ def __init__(
             swiglu_alpha=swiglu_alpha,
             swiglu_beta=swiglu_beta,
             swiglu_limit=swiglu_limit,
+            layer_idx=layer_idx,
         )
 
         # Store original hidden size before any potential padding
@@ -96,8 +97,6 @@ def __init__(
             self.intermediate_size_per_partition = (
                 (self.intermediate_size_per_partition + 127) // 128) * 128
 
-        self.layer_idx = layer_idx
-
         self.num_slots = self.num_experts
         self.expert_size_per_partition = self.num_experts // self.ep_size
         self.initial_global_assignments = [
@@ -449,14 +448,13 @@ def split_chunk(self, split_token_num: int, split_num_chunks: int):
             split_num_chunks - val_mod)
         return split_chunk_size_list
 
-    def forward(
+    def forward_impl(
         self,
         x: Union[torch.Tensor, Fp4QuantizedTensor],
         router_logits: torch.Tensor,
         do_finalize: bool = True,  # used by other MoE backends
         output_dtype: Optional[torch.dtype] = None,
         all_rank_num_tokens: Optional[List[int]] = None,
-        all_rank_max_num_tokens: Optional[int] = None,
         use_dp_padding: Optional[bool] = None,
     ) -> torch.Tensor:
         assert do_finalize, "CutlassFusedMoE does not support do_finalize=False"
@@ -472,7 +470,7 @@ def forward(
                       1) // self.moe_max_num_tokens
 
         if use_dp_padding:
-            all_rank_num_tokens_padded = [all_rank_max_num_tokens
+            all_rank_num_tokens_padded = [max(all_rank_num_tokens)
                                           ] * len(all_rank_num_tokens)
         else:
             all_rank_num_tokens_padded = all_rank_num_tokens

tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py

@@ -637,14 +637,13 @@ def forward_chunk(
 
         return final_hidden_states
 
-    def forward(
+    def forward_impl(
         self,
         x: Union[torch.Tensor, Fp4QuantizedTensor],
         router_logits: torch.Tensor,
         do_finalize: bool = True,  # used by other MoE backends
         output_dtype: Optional[torch.dtype] = None,
         all_rank_num_tokens: Optional[List[int]] = None,
-        all_rank_max_num_tokens: Optional[int] = None,
         use_dp_padding: Optional[bool] = None,
     ) -> torch.Tensor:
         assert do_finalize, "CutlassFusedMoE does not support do_finalize=False"
@@ -663,7 +662,7 @@ def forward(
                           1) // self.moe_max_num_tokens
 
         if use_dp_padding:
-            all_rank_num_tokens_padded = [all_rank_max_num_tokens
+            all_rank_num_tokens_padded = [max(all_rank_num_tokens)
                                           ] * len(all_rank_num_tokens)
         else:
             all_rank_num_tokens_padded = all_rank_num_tokens

tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py

@@ -1287,6 +1287,7 @@ def __init__(
             reduce_results=reduce_results,
             model_config=model_config,
             weight_loading_mode=weight_loading_mode,
+            layer_idx=layer_idx,
         )
         if not IS_TRITON_KERNELS_AVAILABLE:
             raise ImportError("Triton kernels are not available.")
@@ -1359,7 +1360,7 @@ def create_weights(self):
 
         self._weights_created = True
 
-    def forward(
+    def forward_impl(
         self,
         x: torch.Tensor,
         router_logits: torch.Tensor,