From b0dcd1ee64fb39d83f9e8c0066db1f0f5e453351 Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Mon, 25 Aug 2025 23:54:28 -0700 Subject: [PATCH] [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> --- .../_torch/attention_backend/interface.py | 17 +- .../custom_ops/trtllm_gen_custom_ops.py | 52 ++++-- .../_torch/models/modeling_deepseekv3.py | 9 +- .../_torch/models/modeling_gpt_oss.py | 11 +- tensorrt_llm/_torch/models/modeling_llama.py | 18 +- .../_torch/models/modeling_mixtral.py | 2 - .../_torch/models/modeling_qwen3_moe.py | 2 - .../_torch/models/modeling_qwen_moe.py | 2 - .../modules/fused_moe/fused_moe_cutlass.py | 10 +- .../modules/fused_moe/fused_moe_deepgemm.py | 7 +- .../modules/fused_moe/fused_moe_triton.py | 4 +- .../modules/fused_moe/fused_moe_trtllm_gen.py | 38 +++- .../modules/fused_moe/fused_moe_wide_ep.py | 18 +- .../_torch/modules/fused_moe/interface.py | 173 +++++++++++++++++- tensorrt_llm/_torch/speculative/interface.py | 17 +- tensorrt_llm/_torch/speculative/mtp.py | 6 - .../unittest/_torch/modules/test_fused_moe.py | 4 - 17 files changed, 272 insertions(+), 118 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 6a035ad477c1..eb77fa0055d3 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/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 @@ -167,16 +162,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 diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index bbee1b8102ff..8993f62e7231 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -1,10 +1,10 @@ from dataclasses import dataclass from functools import lru_cache -from typing import List, Optional, Tuple +from typing import List, Optional, Tuple, Union import torch -from tensorrt_llm._torch.utils import (fp4_utils, +from tensorrt_llm._torch.utils import (Fp4QuantizedTensor, fp4_utils, get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2, next_positive_power_of_2) @@ -269,6 +269,31 @@ def fp4_block_scale_moe_runner(routing_logits: torch.Tensor, return kernel_runner(inputs, tactic=best_tactic) +def fp4_block_scale_fake_output_without_finalize( + hidden_states: Union[torch.Tensor, Fp4QuantizedTensor], + num_experts: int, + top_k: int, + routing_bias: Optional[torch.Tensor], +): + num_tokens = hidden_states.shape[0] + hidden_size = hidden_states.shape[1] * (2 if isinstance( + hidden_states, Fp4QuantizedTensor) else 1) + + tile_tokens_dim = calculate_tile_tokens_dim(num_tokens, num_experts, top_k) + + expanded_row_count = num_tokens * top_k + max_padding_required = (tile_tokens_dim - 1) * num_experts + max_num_padded_tokens = fp4_utils.pad_up( + expanded_row_count + max_padding_required, tile_tokens_dim) + wt_dtype = routing_bias.dtype if routing_bias is not None else torch.bfloat16 + return [ + hidden_states.new_empty((max_num_padded_tokens, hidden_size), + dtype=torch.bfloat16), + hidden_states.new_empty((num_tokens, top_k), dtype=wt_dtype), + hidden_states.new_empty((num_tokens, top_k), dtype=torch.int32) + ] + + @fp4_block_scale_moe_runner.register_fake def _( routing_logits, @@ -293,27 +318,20 @@ def _( routing_method_type, do_finalize, ) -> List[torch.Tensor]: - num_tokens = hidden_states.shape[0] - hidden_size = hidden_states.shape[1] * 2 if do_finalize: + num_tokens = hidden_states.shape[0] + hidden_size = hidden_states.shape[1] * 2 return [ hidden_states.new_empty((num_tokens, hidden_size), dtype=torch.bfloat16) ] - tile_tokens_dim = calculate_tile_tokens_dim(num_tokens, num_experts, top_k) - - expanded_row_count = num_tokens * top_k - max_padding_required = (tile_tokens_dim - 1) * num_experts - max_num_padded_tokens = fp4_utils.pad_up( - expanded_row_count + max_padding_required, tile_tokens_dim) - wt_dtype = routing_bias.dtype if routing_bias is not None else torch.bfloat16 - return [ - hidden_states.new_empty((max_num_padded_tokens, hidden_size), - dtype=torch.bfloat16), - hidden_states.new_empty((num_tokens, top_k), dtype=wt_dtype), - hidden_states.new_empty((num_tokens, top_k), dtype=torch.int32) - ] + return fp4_block_scale_fake_output_without_finalize( + hidden_states, + num_experts, + top_k, + routing_bias, + ) @dataclass(frozen=True) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 3d5ae308de80..0ca4d28085b3 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -548,8 +548,7 @@ def _get_experts_quant_config(model_config, layer_idx: int) -> QuantConfig: f"model.layers.{layer_idx}.mlp.experts", model_config.quant_config) 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: @@ -568,7 +567,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, ) @@ -579,7 +577,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: @@ -598,7 +595,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 @@ -840,7 +836,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)), @@ -1028,7 +1023,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: @@ -1087,7 +1081,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)), diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py index 5ea69fefb691..d0e5a1f0848d 100644 --- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 5cb8f1b30024..5353c4db2014 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -315,23 +315,19 @@ 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: @@ -339,8 +335,7 @@ def forward( # 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: @@ -542,7 +537,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, diff --git a/tensorrt_llm/_torch/models/modeling_mixtral.py b/tensorrt_llm/_torch/models/modeling_mixtral.py index 21dcc2006398..fedf45b9c10b 100644 --- a/tensorrt_llm/_torch/models/modeling_mixtral.py +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index bd2ccfae0c51..fa9be6afcf89 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/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, ) diff --git a/tensorrt_llm/_torch/models/modeling_qwen_moe.py b/tensorrt_llm/_torch/models/modeling_qwen_moe.py index 4fa3883246f1..d7b265c085d7 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen_moe.py +++ b/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) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index 1693cacf18cd..dd1e2d7f88f4 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/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,15 +448,16 @@ 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, + **kwargs, ) -> torch.Tensor: assert do_finalize, "CutlassFusedMoE does not support do_finalize=False" if self.use_dp and self.parallel_size > 1: @@ -472,7 +472,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 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 392fff091157..71493b2612b5 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py @@ -637,15 +637,16 @@ 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, + **kwargs, ) -> torch.Tensor: assert do_finalize, "CutlassFusedMoE does not support do_finalize=False" if self.use_dp and self.parallel_size > 1: @@ -663,7 +664,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 diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py index f2ef12175713..b7724027e3ed 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py +++ b/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,10 +1360,11 @@ def create_weights(self): self._weights_created = True - def forward( + def forward_impl( self, x: torch.Tensor, router_logits: torch.Tensor, + *, do_finalize: bool = True, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 4e18ae8c245b..bf63dc8de035 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -5,6 +5,8 @@ from tensorrt_llm._utils import get_sm_version +from ...custom_ops.trtllm_gen_custom_ops import \ + fp4_block_scale_fake_output_without_finalize from ...model_config import ModelConfig from ...utils import Fp4QuantizedTensor, next_positive_power_of_2 from .interface import MoE, MoEWeightLoadingMode @@ -78,6 +80,7 @@ def __init__( swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, swiglu_limit=swiglu_limit, + layer_idx=layer_idx, ) sm_version = get_sm_version() @@ -186,10 +189,11 @@ def load_weights(self, weights: List[Dict]): self.quant_method.load_weights(self, weights, self.weight_loading_mode) - def forward( + def forward_impl( self, x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, + *, do_finalize: bool = True, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, @@ -405,3 +409,35 @@ def forward( final_hidden_states = final_hidden_states[: all_rank_num_tokens[rank]] return final_hidden_states + + def forward_fake( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + *, + do_finalize: bool = True, + output_dtype: Optional[torch.dtype] = None, + all_rank_num_tokens: Optional[List[int]] = None, + use_dp_padding: Optional[bool] = None, + **kwargs, + ) -> Union[torch.Tensor, List[torch.Tensor]]: + if do_finalize: + # TRTLLMGenFusedMoE only supports bfloat16 output + return super().forward_fake(x, + router_logits, + do_finalize=do_finalize, + output_dtype=torch.bfloat16, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding, + **kwargs) + else: + is_deepseek_v3_routing = isinstance(self.routing_method, + DeepSeekV3MoeRoutingMethod) + top_k = self.routing_method.routing_impl.top_k if is_deepseek_v3_routing else self.routing_method.top_k + routing_bias = self.routing_method.e_score_correction_bias if is_deepseek_v3_routing else None + return fp4_block_scale_fake_output_without_finalize( + x, + self.num_experts, + top_k, + routing_bias, + ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py index 5430141071ca..a9fb53a3b853 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py @@ -82,13 +82,13 @@ def __init__( reduce_results=reduce_results, model_config=model_config, weight_loading_mode=weight_loading_mode, + layer_idx=layer_idx, ) assert self.use_dp, "Attention DP should be used with WideEP." assert self.parallel_size > 1, "WideEP should only be enabled with parallel_size > 1" # If True, the router weight will be multiplied on the input rather than at the end of FC2 self.apply_router_weight_on_input = apply_router_weight_on_input - self.layer_idx = layer_idx moe_load_balancer = get_moe_load_balancer() self.layer_load_balancer = None @@ -374,10 +374,10 @@ def forward_chunk( use_all_to_all: bool, 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, repeating_info: Tuple = (True, True), ) -> torch.Tensor: + all_rank_max_num_tokens = max(all_rank_num_tokens) if isinstance(x, Fp4QuantizedTensor): assert output_dtype is not None output_dtype = output_dtype @@ -710,19 +710,22 @@ 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, 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, + **kwargs, ) -> torch.Tensor: assert all_rank_num_tokens is not None assert use_dp_padding is not None + all_rank_max_num_tokens = max(all_rank_num_tokens) + # in case of num_rows is larger than max_chunk_size, we need to split the input into multiple chunks num_chunks = self.calculate_num_chunks(all_rank_num_tokens) use_all_to_all = self.can_use_alltoall(all_rank_num_tokens, @@ -742,7 +745,6 @@ def forward( use_all_to_all, output_dtype, all_rank_num_tokens=all_rank_num_tokens_padded, - all_rank_max_num_tokens=all_rank_max_num_tokens, use_dp_padding=use_dp_padding, repeating_info=(is_first_call, is_last_call)) outputs = self.reducescatter_or_allreduce( @@ -801,8 +803,6 @@ def split_chunk(split_token_num: int, split_num_chunks: int): use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_list[ idx_chunk], - all_rank_max_num_tokens= - all_rank_max_num_tokens_list[idx_chunk], use_dp_padding=use_dp_padding, repeating_info=(is_first_call, is_last_call)) if idx_chunk > 0: @@ -819,8 +819,6 @@ def split_chunk(split_token_num: int, split_num_chunks: int): use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_list[ idx_chunk], - all_rank_max_num_tokens=all_rank_max_num_tokens_list[ - idx_chunk], use_dp_padding=use_dp_padding, repeating_info=(is_first_call, is_last_call)) with torch.cuda.stream(self.aux_stream): @@ -836,8 +834,6 @@ def split_chunk(split_token_num: int, split_num_chunks: int): router_logits, use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_list[idx_chunk], - all_rank_max_num_tokens=all_rank_max_num_tokens_list[ - idx_chunk], repeating_info=(is_first_call, is_last_call)) outputs_list.append(outputs) diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/modules/fused_moe/interface.py index 6301a843121a..0187adce98f2 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/interface.py +++ b/tensorrt_llm/_torch/modules/fused_moe/interface.py @@ -1,12 +1,15 @@ +import weakref from abc import abstractmethod from enum import Enum -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Union, final import torch from torch import nn from ...distributed.ops import reducescatter from ...model_config import ModelConfig +from ...utils import (Fp4QuantizedTensor, get_model_extra_attrs, + is_torch_compiling) from .routing import BaseMoeRoutingMethod @@ -19,6 +22,82 @@ class MoEWeightLoadingMode(Enum): W4A8_CUSTOM = 2 +def extract_extra_attrs(layer_idx: str): + extra_attrs = get_model_extra_attrs() + assert extra_attrs is not None, "Model extra attrs are not set" + + moe_layers = extra_attrs.get("moe_layers", None) + assert moe_layers is not None, "No MoE layers registered" + moe_layer_ref = moe_layers.get(layer_idx) + assert moe_layer_ref is not None, f"Cannot find MoE layer for layer_idx={layer_idx}" + moe_layer = moe_layer_ref() if callable(moe_layer_ref) else None + assert moe_layer is not None, f"MoE layer for layer_idx={layer_idx!r} is no longer alive" + + return moe_layer + + +@torch.library.custom_op("trtllm::moe_custom_op", mutates_args=()) +def moe_custom_op( + layer_idx: str, + x: torch.Tensor, + x_sf: Optional[torch.Tensor], + is_swizzled: bool, + router_logits: torch.Tensor, + do_finalize: bool, + output_dtype: Optional[torch.dtype], + all_rank_num_tokens: Optional[List[int]], + use_dp_padding: Optional[bool], +) -> List[torch.Tensor]: + moe_layer = extract_extra_attrs(layer_idx) + + hidden_states = x if x_sf is None else Fp4QuantizedTensor( + x, x_sf, is_swizzled) + + res = moe_layer.forward_impl( + hidden_states, + router_logits, + do_finalize=do_finalize, + output_dtype=output_dtype, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding, + ) + + if do_finalize: + return [res] + else: + return res + + +@moe_custom_op.register_fake +def _( + layer_idx, + x, + x_sf, + is_swizzled, + router_logits, + do_finalize, + output_dtype, + all_rank_num_tokens, + use_dp_padding, +): + moe_layer = extract_extra_attrs(layer_idx) + hidden_states = x if x_sf is None else Fp4QuantizedTensor( + x, x_sf, is_swizzled) + res = moe_layer.forward_fake( + hidden_states, + router_logits, + do_finalize=do_finalize, + output_dtype=output_dtype, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding, + ) + + if do_finalize: + return [res] + else: + return res + + class MoE(nn.Module): """ Fused Mixture of Experts (MoE) Layer interface. @@ -49,6 +128,7 @@ def __init__( swiglu_alpha: Optional[torch.Tensor] = None, swiglu_beta: Optional[torch.Tensor] = None, swiglu_limit: Optional[torch.Tensor] = None, + layer_idx: Optional[int] = None, ): from ...distributed import AllReduce @@ -64,6 +144,10 @@ def __init__( self.swiglu_alpha = swiglu_alpha self.swiglu_beta = swiglu_beta self.swiglu_limit = swiglu_limit + self.layer_idx = layer_idx + self.layer_idx_str = str(layer_idx) if layer_idx is not None else None + + self._register_layer(model_config) # could be modified later self.quant_config = model_config.quant_config @@ -92,6 +176,17 @@ def __init__( strategy=model_config.allreduce_strategy, dtype=self.dtype) + def _register_layer(self, model_config: ModelConfig): + self.register_to_config = False + if model_config is not None and self.layer_idx_str is not None: + if "moe_layers" not in model_config.extra_attrs: + model_config.extra_attrs["moe_layers"] = {} + assert self.layer_idx_str not in model_config.extra_attrs["moe_layers"], \ + f"Duplicate MoE layer for layer_idx={self.layer_idx_str}" + model_config.extra_attrs["moe_layers"][ + self.layer_idx_str] = weakref.ref(self) + self.register_to_config = True + @abstractmethod def create_weights(self): raise NotImplementedError @@ -101,15 +196,83 @@ def load_weights(self, weights: List[Dict]): raise NotImplementedError @abstractmethod - def forward( + def forward_impl( self, - x: torch.Tensor, + x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, - *args, + *, + do_finalize: bool = True, + output_dtype: Optional[torch.dtype] = None, + all_rank_num_tokens: Optional[List[int]] = None, + use_dp_padding: Optional[bool] = None, **kwargs, - ) -> torch.Tensor: + ) -> Union[torch.Tensor, List[torch.Tensor]]: raise NotImplementedError + def forward_fake( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + *, + do_finalize: bool = True, + output_dtype: Optional[torch.dtype] = None, + all_rank_num_tokens: Optional[List[int]] = None, + use_dp_padding: Optional[bool] = None, + **kwargs, + ) -> Union[torch.Tensor, List[torch.Tensor]]: + is_nvfp4_input = isinstance(x, Fp4QuantizedTensor) + assert do_finalize, "Default forward_fake does not support do_finalize=False" + data_type = output_dtype if is_nvfp4_input else x.dtype + num_tokens = all_rank_num_tokens[ + self.tp_rank] if all_rank_num_tokens else x.shape[0] + hidden_size = x.shape[1] * (2 if is_nvfp4_input else 1) + return x.new_empty((num_tokens, hidden_size), dtype=data_type) + + # Sub class is not allowed to override forward. + # This is universal interface for all MoE backends + @final + def forward( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + do_finalize: bool = True, + output_dtype: Optional[torch.dtype] = None, + all_rank_num_tokens: Optional[List[int]] = None, + use_dp_padding: Optional[bool] = None, + ) -> Union[torch.Tensor, List[torch.Tensor]]: + if self.register_to_config and is_torch_compiling(): + hidden_states = x.fp4_tensor if isinstance( + x, Fp4QuantizedTensor) else x + x_sf = x.scaling_factor if isinstance(x, + Fp4QuantizedTensor) else None + is_swizzled = x.is_sf_swizzled if isinstance( + x, Fp4QuantizedTensor) else False + + res = moe_custom_op( + self.layer_idx_str, + hidden_states, + x_sf, + is_swizzled, + router_logits, + do_finalize, + output_dtype, + all_rank_num_tokens, + use_dp_padding, + ) + if do_finalize: + return res[0] + else: + return res + else: + return self.forward_impl( + x, + router_logits, + do_finalize=do_finalize, + output_dtype=output_dtype, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding, + ) + @property def has_any_quant(self): assert self._weights_created diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 46b09596618a..3ecb323aa33e 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -147,12 +147,7 @@ class SpecMetadata: # The number of tokens for speculative model/layer num_tokens: int = 0 # The number of tokens for speculative model/layer of different 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 # The number of sequences for speculative model/layer of different rank all_rank_num_seqs: Optional[List[int]] = None @@ -205,13 +200,3 @@ def maybe_capture_hidden_states(self, layer_id: int, Some spec decode algorithms require hidden states from the target model. Use this method to record them. By default, does nothing. """ - - @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 SpecMetadata.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 diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 7579809c7829..ab1ef6e615d0 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -1207,8 +1207,6 @@ def prepare_position_ids_and_last_tokens(position_ids, attn_metadata): hidden_states = draft_model.mtp_layers[0]( embed_tokens=draft_model.embed_tokens, all_rank_num_tokens=spec_metadata.all_rank_num_tokens, - all_rank_max_num_tokens=spec_metadata. - all_rank_max_num_tokens, **inputs) start_ids_gen = (spec_metadata.batch_indices_cuda[:num_gens] * (self.mtp_num_modules + 1)).long() @@ -1222,10 +1220,6 @@ def prepare_position_ids_and_last_tokens(position_ids, attn_metadata): embed_tokens=draft_model.embed_tokens, all_rank_num_tokens=spec_metadata. subseq_all_rank_num_tokens, - all_rank_max_num_tokens=max( - spec_metadata.subseq_all_rank_num_tokens) - if spec_metadata.subseq_all_rank_num_tokens is not None else - None, **inputs) # All of the seq_len are 1, use batch_indices_cuda as gather_ids gather_ids = spec_metadata.batch_indices_cuda[:batch_size] diff --git a/tests/unittest/_torch/modules/test_fused_moe.py b/tests/unittest/_torch/modules/test_fused_moe.py index 397314bcab00..2718064d3e2a 100644 --- a/tests/unittest/_torch/modules/test_fused_moe.py +++ b/tests/unittest/_torch/modules/test_fused_moe.py @@ -266,13 +266,11 @@ def per_rank_test_fused_moe_alltoall(job_id): x, router_logits, all_rank_num_tokens=all_rank_num_tokens, - all_rank_max_num_tokens=m, use_dp_padding=False) ref_output = ref_model.forward( x, router_logits, all_rank_num_tokens=all_rank_num_tokens, - all_rank_max_num_tokens=m, use_dp_padding=False) # Evaluate outputs @@ -455,13 +453,11 @@ def per_rank_test_fused_moe_alltoall(job_id): x, router_logits, all_rank_num_tokens=all_rank_num_tokens, - all_rank_max_num_tokens=m, use_dp_padding=False) ref_output = ref_model.forward( x, router_logits, all_rank_num_tokens=all_rank_num_tokens, - all_rank_max_num_tokens=m, use_dp_padding=False) # Evaluate outputs