[None][fix] : Fix OOM issue when dp padding is enabled (#8052)

Signed-off-by: peaceh <[email protected]>

Author: peaceh-nv

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

@@ -468,16 +468,17 @@ def forward_impl(
         else:
             num_rows = x.shape[0]
 
-        # in case of num_rows is larger than max_chunk_size, we need to split the input into multiple chunks
-        num_chunks = (num_rows + self.moe_max_num_tokens -
-                      1) // self.moe_max_num_tokens
-
         if use_dp_padding:
             all_rank_num_tokens_padded = [max(all_rank_num_tokens)
                                           ] * len(all_rank_num_tokens)
+            num_rows = sum(all_rank_num_tokens_padded)
         else:
             all_rank_num_tokens_padded = 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 = (num_rows + self.moe_max_num_tokens -
+                      1) // self.moe_max_num_tokens
+
         if num_chunks == 1:
             outputs = self.forward_chunk(
                 x,

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py

@@ -771,16 +771,16 @@ def forward_impl(
 
         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,
-                                               all_rank_max_num_tokens)
-
         if use_dp_padding:
             all_rank_num_tokens_padded = [all_rank_max_num_tokens
                                           ] * len(all_rank_num_tokens)
         else:
             all_rank_num_tokens_padded = 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_padded)
+        use_all_to_all = self.can_use_alltoall(all_rank_num_tokens_padded,
+                                               all_rank_max_num_tokens)
         if num_chunks == 1:
             is_first_call = self.repeat_idx == 0
             is_last_call = self.repeat_idx == self.repeat_count - 1