[https://nvbugs/5613089][fix] Fix the rank to access all_rank_chunk_size_list when chunked MoE is used (NVIDIA#8723)

Signed-off-by: Jinyang Yuan <[email protected]>

Author: jinyangyuan-nvidia

tensorrt_llm/_torch/autotuner.py

@@ -727,10 +727,10 @@ def choose_one(
         new_tuning_failure_occured = False
 
         for p in profiles:
-            tensors = self._prepare_input_tensors(p, inputs)
             is_cache_hit, *_ = self.profiling_cache.search_cache(
                 custom_op, runners, p.get_opt_shapes(), tuning_config)
             if not is_cache_hit:
+                tensors = self._prepare_input_tensors(p, inputs)
                 # Initialize runner and tactic as None in case of no valid tactic or runners are found
                 best_runner_id, best_tactic, min_time, has_tuning_failure_occured = self._profile_runners(
                     custom_op, runners, tensors, p, tuning_config, **kwargs)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

@@ -669,7 +669,7 @@ def forward_impl(
                 all_rank_num_tokens_list = [[
                     val[idx_chunk] for val in all_rank_chunk_size_list
                 ] for idx_chunk in range(num_chunks)]
-                chunk_size_list = all_rank_chunk_size_list[self.rank]
+                chunk_size_list = all_rank_chunk_size_list[self.parallel_rank]
             else:
                 all_rank_num_tokens_list = [None] * num_chunks
                 chunk_size_list = self.split_chunk(x.shape[0], num_chunks)
@@ -735,7 +735,7 @@ def _reducescatter_or_allreduce(x_, idx):
             outputs = torch.cat(outputs_list)
 
         if self.use_dp and self.parallel_size > 1:
-            rank = self.mapping.tp_rank
+            rank = self.parallel_rank
             outputs = outputs[:all_rank_num_tokens[rank]]
         self.repeat_idx = 0 if self.repeat_idx == self.repeat_count - 1 else self.repeat_idx + 1
         return outputs
@@ -765,7 +765,7 @@ def forward_fake(
             is_nvfp4_input = isinstance(x, Fp4QuantizedTensor)
             data_type = output_dtype if is_nvfp4_input else x.dtype
             num_tokens = all_rank_num_tokens[
-                self.mapping.tp_rank] if all_rank_num_tokens else x.shape[0]
+                self.parallel_rank] if all_rank_num_tokens else x.shape[0]
             hidden_size = x.shape[1] * (2 if is_nvfp4_input else 1)
             top_k = self.routing_method.experts_per_token
             return x.new_empty((num_tokens, top_k, hidden_size),

tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py

@@ -706,7 +706,7 @@ def forward_impl(
                 all_rank_num_tokens_list = [[
                     val[idx_chunk] for val in all_rank_chunk_size_list
                 ] for idx_chunk in range(num_chunks)]
-                chunk_size_list = all_rank_chunk_size_list[self.rank]
+                chunk_size_list = all_rank_chunk_size_list[self.parallel_rank]
             else:
                 all_rank_num_tokens_list = [None] * num_chunks
                 chunk_size_list = self.split_chunk(x.shape[0], num_chunks)
@@ -778,6 +778,6 @@ def _reducescatter_or_allreduce(x_, idx):
             outputs = torch.cat(outputs_list)
 
         if self.use_dp and self.parallel_size > 1:
-            rank = self.mapping.tp_rank
+            rank = self.parallel_rank
             outputs = outputs[:all_rank_num_tokens[rank]]
         return outputs

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

@@ -805,7 +805,7 @@ def forward_impl(
         self._load_balancer_done_set_cpu_stage(is_last_call)
 
         if use_dp_padding:
-            rank = self.mapping.tp_rank
+            rank = self.parallel_rank
             final_hidden_states = final_hidden_states[:
                                                       all_rank_num_tokens[rank]]
 

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py

@@ -762,7 +762,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int):
             ] for idx_chunk in range(num_chunks)]
             all_rank_max_num_tokens_list = split_chunk(all_rank_max_num_tokens,
                                                        num_chunks)
-            chunk_size_list = all_rank_chunk_size_list[self.rank]
+            chunk_size_list = all_rank_chunk_size_list[self.parallel_rank]
             if use_all_to_all:
                 all_rank_num_tokens_list = [[
                     1 if val == 0 else val for val in val_list
@@ -850,7 +850,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int):
                     self.event_dict[EventType.MoeChunkingOverlap].record()
                 self.event_dict[EventType.MoeChunkingOverlap].wait()
             outputs = torch.cat(outputs_list)
-        rank = self.mapping.tp_rank
+        rank = self.parallel_rank
         outputs = outputs[:all_rank_num_tokens[rank]]
         self.repeat_idx = 0 if self.repeat_idx == self.repeat_count - 1 else self.repeat_idx + 1
         return outputs

tensorrt_llm/_torch/modules/fused_moe/interface.py

@@ -184,6 +184,7 @@ def __init__(
 
         # All ranks participate in allreduce regardless of EP/TP combination
         self.mapping = model_config.mapping
+        self.parallel_rank = self.mapping.tp_rank
         self.parallel_size = self.mapping.tp_size
         self.intermediate_size_per_partition = intermediate_size // self.tp_size