From dc924cdf18652b8ffbc6d8ddf860d7835b02a70a Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Tue, 12 Aug 2025 04:12:03 +0000 Subject: [PATCH 1/3] [fix] Adjust default moe_max_num_tokens to fix OOM. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/modules/fused_moe/fused_moe_cutlass.py | 6 ++---- .../_torch/modules/fused_moe/fused_moe_vanilla.py | 8 ++------ .../_torch/modules/fused_moe/fused_moe_wide_ep.py | 5 ++--- tensorrt_llm/mapping.py | 4 ++++ 4 files changed, 10 insertions(+), 13 deletions(-) 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 e5a41fbf0941..fbb0996358c6 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -110,11 +110,9 @@ def __init__( assert len( self.initial_local_expert_ids) == self.expert_size_per_partition - max_num_tokens = model_config.max_num_tokens # The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled - if self.use_dp: - max_num_tokens *= model_config.mapping.world_size - self.moe_max_num_tokens = model_config.moe_max_num_tokens or max_num_tokens + max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size + self.moe_max_num_tokens = model_config.moe_max_num_tokens or model_config.max_num_tokens # The auxiliary CUDA stream and CUDA events are only used when MoE chunking is applied if self.moe_max_num_tokens < max_num_tokens: self.aux_stream = aux_stream_dict[ diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py index 3249bac979b2..29cc5f2fab4e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py @@ -81,13 +81,9 @@ def __init__( self.num_experts) self.expert_size_per_partition = self.expert_end - self.expert_start - max_num_tokens = model_config.max_num_tokens # The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled - if self.use_dp: - max_num_tokens *= model_config.mapping.world_size - self.moe_max_num_tokens = (model_config.moe_max_num_tokens - if model_config.moe_max_num_tokens - is not None else max_num_tokens) + max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size + self.moe_max_num_tokens = model_config.moe_max_num_tokens or max_num_tokens self._weights_created = False if not model_config.skip_create_weights_in_init: 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 c74c8966b68b..2862558df851 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 @@ -150,10 +150,9 @@ def __init__( assert len( self.initial_local_expert_ids) == self.expert_size_per_partition - max_num_tokens = model_config.max_num_tokens # The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled - max_num_tokens *= model_config.mapping.world_size - self.moe_max_num_tokens = model_config.moe_max_num_tokens if model_config.moe_max_num_tokens is not None else max_num_tokens + max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size + self.moe_max_num_tokens = model_config.moe_max_num_tokens or model_config.max_num_tokens # The auxiliary CUDA stream and CUDA events are only used when MoE chunking is applied if self.moe_max_num_tokens < max_num_tokens: self.aux_stream = aux_stream_dict[ diff --git a/tensorrt_llm/mapping.py b/tensorrt_llm/mapping.py index f78fe093f713..22824ea350d9 100644 --- a/tensorrt_llm/mapping.py +++ b/tensorrt_llm/mapping.py @@ -372,6 +372,10 @@ def node_rank(self): def local_rank(self): return self.rank % self.gpus_per_node + @property + def dp_size(self): + return self.tp_size if self.enable_attention_dp else 1 + def has_cp(self): return self.cp_size > 1 From 2bd9e773ac0b1ba6ad19bf99f0e4a7c20d7f7efd Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Tue, 12 Aug 2025 04:16:06 +0000 Subject: [PATCH 2/3] Update default_moe_max_num_tokens. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/modules/fused_moe/fused_moe_cutlass.py | 6 +++--- .../_torch/modules/fused_moe/fused_moe_deepgemm.py | 10 ++++++++++ .../_torch/modules/fused_moe/fused_moe_vanilla.py | 4 ++-- .../_torch/modules/fused_moe/fused_moe_wide_ep.py | 6 +++--- 4 files changed, 18 insertions(+), 8 deletions(-) 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 fbb0996358c6..3ad5b615917b 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -111,10 +111,10 @@ def __init__( self.initial_local_expert_ids) == self.expert_size_per_partition # The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled - max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size - self.moe_max_num_tokens = model_config.moe_max_num_tokens or model_config.max_num_tokens + moe_max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size + self.moe_max_num_tokens = model_config.moe_max_num_tokens or moe_max_num_tokens # The auxiliary CUDA stream and CUDA events are only used when MoE chunking is applied - if self.moe_max_num_tokens < max_num_tokens: + if self.moe_max_num_tokens < moe_max_num_tokens: self.aux_stream = aux_stream_dict[ AuxStreamType. MoeChunkingOverlap] if aux_stream_dict is not None else torch.cuda.Stream( 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 fcfdeb0fad28..0891aeae0ec7 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py @@ -327,6 +327,16 @@ def __init__( apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, ): + if model_config.moe_max_num_tokens is None: + moe_max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size + # The default moe_max_num_tokens is calculated from the following formula: + # max_isl = 8196, max_batch_size = 1024, mtp = 0 + # max_num_tokens = ((mtp+1)*max_batch_size+max_isl+128+63)//64*64 = 9344 + # moe_max_num_tokens = max_num_tokens * 2 = 18688 + # It can avoid OOM for 8k/1k cases. + default_moe_max_num_tokens = 18688 + if moe_max_num_tokens > default_moe_max_num_tokens: + model_config.moe_max_num_tokens = default_moe_max_num_tokens super().__init__( routing_method=routing_method, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py index 29cc5f2fab4e..ed6f11993b20 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py @@ -82,8 +82,8 @@ def __init__( self.expert_size_per_partition = self.expert_end - self.expert_start # The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled - max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size - self.moe_max_num_tokens = model_config.moe_max_num_tokens or max_num_tokens + moe_max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size + self.moe_max_num_tokens = model_config.moe_max_num_tokens or moe_max_num_tokens self._weights_created = False if not model_config.skip_create_weights_in_init: 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 2862558df851..4eb9d77606b4 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 @@ -151,10 +151,10 @@ def __init__( self.initial_local_expert_ids) == self.expert_size_per_partition # The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled - max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size - self.moe_max_num_tokens = model_config.moe_max_num_tokens or model_config.max_num_tokens + moe_max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size + self.moe_max_num_tokens = model_config.moe_max_num_tokens or moe_max_num_tokens # The auxiliary CUDA stream and CUDA events are only used when MoE chunking is applied - if self.moe_max_num_tokens < max_num_tokens: + if self.moe_max_num_tokens < moe_max_num_tokens: self.aux_stream = aux_stream_dict[ AuxStreamType. MoeChunkingOverlap] if aux_stream_dict is not None else torch.cuda.Stream( From 2ad8a28ba55d5aedadd0bd92d398af953c0f8d7a Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Tue, 12 Aug 2025 04:54:43 +0000 Subject: [PATCH 3/3] WAR frozen model_config. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py | 2 ++ 1 file changed, 2 insertions(+) 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 0891aeae0ec7..68bf255a1422 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py @@ -336,7 +336,9 @@ def __init__( # It can avoid OOM for 8k/1k cases. default_moe_max_num_tokens = 18688 if moe_max_num_tokens > default_moe_max_num_tokens: + model_config._frozen = False model_config.moe_max_num_tokens = default_moe_max_num_tokens + model_config._frozen = True super().__init__( routing_method=routing_method,