Repository: NVIDIA/TensorRT-LLM
Pull request #3370: refactor: remove ParallelConfig in tensorrt_llm._torch.distributed module
Author: login: QiJune
Created: 2025-04-08T11:23:36Z | Merged: 2025-04-11T22:34:21Z

To avoid unnecessary cyclical conversion "Mapping -> ParallelConfig -> Mapping"


Diff:
diff --git a/examples/pytorch/out_of_tree_example/modeling_opt.py b/examples/pytorch/out_of_tree_example/modeling_opt.py
index c3a68aa041b0..7d0786f26492 100644
--- a/examples/pytorch/out_of_tree_example/modeling_opt.py
+++ b/examples/pytorch/out_of_tree_example/modeling_opt.py
@@ -7,7 +7,6 @@
 from transformers.activations import ACT2FN
 
 from tensorrt_llm._torch.attention_backend import AttentionMetadata
-from tensorrt_llm._torch.distributed import ParallelConfig, TensorParallelMode
 from tensorrt_llm._torch.model_config import ModelConfig
 from tensorrt_llm._torch.models.modeling_utils import (DecoderModel,
                                                        DecoderModelForCausalLM,
@@ -16,7 +15,7 @@
 from tensorrt_llm._torch.modules.attention import Attention
 from tensorrt_llm._torch.modules.decoder_layer import DecoderLayer
 from tensorrt_llm._torch.modules.embedding import Embedding
-from tensorrt_llm._torch.modules.linear import Linear
+from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
 
 
 class LayerNorm(nn.LayerNorm):
@@ -70,10 +69,8 @@ def __init__(
             config.ffn_dim,
             bias=config.enable_bias,
             dtype=config.torch_dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=model_config.mapping.tp_rank,
-                tensor_parallel_size=model_config.mapping.tp_size,
-                tensor_parallel_mode=TensorParallelMode.COLUMN),
+            mapping=model_config.mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
             quant_config=model_config.get_quant_config(),
         )
         self.fc2 = Linear(
@@ -81,10 +78,8 @@ def __init__(
             config.hidden_size,
             bias=config.enable_bias,
             dtype=config.torch_dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=model_config.mapping.tp_rank,
-                tensor_parallel_size=model_config.mapping.tp_size,
-                tensor_parallel_mode=TensorParallelMode.ROW),
+            mapping=model_config.mapping,
+            tensor_parallel_mode=TensorParallelMode.ROW,
             quant_config=model_config.get_quant_config(),
         )
         self.final_layer_norm = LayerNorm(
@@ -150,10 +145,8 @@ def __init__(self, model_config: ModelConfig[OPTConfig]):
             config.vocab_size,
             config.word_embed_proj_dim,
             dtype=config.torch_dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=model_config.mapping.tp_rank,
-                tensor_parallel_size=model_config.mapping.tp_size,
-                tensor_parallel_mode=TensorParallelMode.COLUMN))
+            mapping=model_config.mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN)
         self.embed_positions = nn.Embedding(config.max_position_embeddings + 2,
                                             config.hidden_size,
                                             dtype=config.torch_dtype)
diff --git a/tensorrt_llm/_torch/attention_backend/star_flashinfer.py b/tensorrt_llm/_torch/attention_backend/star_flashinfer.py
index 29ca290258c7..d3911a7d56c1 100644
--- a/tensorrt_llm/_torch/attention_backend/star_flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/star_flashinfer.py
@@ -9,7 +9,6 @@
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
 from ..distributed import allgather
-from ..modules.linear import ParallelConfig
 from .flashinfer import FlashInferAttentionMetadata, PlanParams
 from .interface import AttentionBackend, AttentionMask, PredefinedAttentionMask
 from .vanilla import VanillaAttention
@@ -440,12 +439,10 @@ def reduce_results(output, lse):
             out_tmp = output
             lse = lse.unsqueeze(-1) / np.log2(np.e)  # [b * s, nheads, 1]
             if metadata.mapping.cp_size != 1:
-                parallel_cfg = ParallelConfig(
-                    tensor_parallel_size=metadata.mapping.cp_size,
-                    tensor_parallel_rank=metadata.mapping.cp_rank,
-                    pipeline_parallel_size=metadata.mapping.pp_size)
-                output_tensor = allgather(output, parallel_cfg, gather_dim=0)
-                lse_tensor = allgather(lse, parallel_cfg, gather_dim=0)
+                output_tensor = allgather(output,
+                                          metadata.mapping,
+                                          gather_dim=0)
+                lse_tensor = allgather(lse, metadata.mapping, gather_dim=0)
                 output_tensor = output_tensor.to(torch.float32)
             else:
                 lse_tensor = lse
diff --git a/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py b/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py
index 49d8219dae5e..cada12f2100c 100644
--- a/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py
+++ b/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py
@@ -4,18 +4,19 @@
 
 # use trtllm distributed ops to improve TP performance if possible
 try:
+    from ....mapping import Mapping
     from ...distributed import AllReduce, allgather
-    from ...modules.linear import AllReduceFusionOp, AllReduceParams, ParallelConfig
+    from ...modules.linear import AllReduceFusionOp, AllReduceParams
 
     def trtllm_allgather(tensor, dim):
         rank, world_size = get_rank_world_size()
-        p_config = ParallelConfig(tensor_parallel_size=world_size, tensor_parallel_rank=rank)
+        p_config = Mapping(world_size=world_size, tp_size=world_size, rank=rank)
         return allgather(tensor, p_config, gather_dim=dim)
 
     def trtllm_allreduce(tensor, op, all_reduce_params=None):
         rank, world_size = get_rank_world_size()
         assert op == ReduceOp.SUM, "TRT-LLM all reduce only supports SUM op."
-        p_config = ParallelConfig(tensor_parallel_size=world_size, tensor_parallel_rank=rank)
+        p_config = Mapping(world_size=world_size, tp_size=world_size, rank=rank)
         torch_op = AllReduce(p_config)
         return torch_op(tensor, all_reduce_params=all_reduce_params)
 
diff --git a/tensorrt_llm/_torch/distributed.py b/tensorrt_llm/_torch/distributed.py
index 1044b2aa4cfe..a85f1e36e581 100644
--- a/tensorrt_llm/_torch/distributed.py
+++ b/tensorrt_llm/_torch/distributed.py
@@ -1,8 +1,6 @@
 import atexit
-import enum
 import os
 import threading
-from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -17,37 +15,6 @@
 _thread_local = threading.local()
 
 
-class TensorParallelMode(str, enum.Enum):
-    COLUMN = 'column'
-    ROW = 'row'
-
-    @classmethod
-    def split_dim(cls, mode):
-        return 1 if mode == cls.ROW else 0
-
-
-@dataclass(kw_only=True)
-class ParallelConfig:
-    tensor_parallel_size: int = 1
-    tensor_parallel_rank: int = 0
-    gpus_per_node: int = 8
-    tensor_parallel_mode: Optional[TensorParallelMode] = None
-    gather_output: bool = False
-    # pipeline parallel parameter in case we have multiple parallel groups
-    # default to TP-only mode if not specified for backward compatibility
-    # TODO Remove redundant fields. Keep only parallel_rank, tp_size, pp_size in constructor
-    # and infer tp_rank, pp_rank, etc. automatically.
-    pipeline_parallel_size: int = 1
-    parallel_rank: Optional[int] = None
-
-    def __post_init__(self):
-        self.parallel_size = self.tensor_parallel_size * self.pipeline_parallel_size
-        if self.pipeline_parallel_size > 1:
-            assert self.parallel_rank is not None, "parallel_rank must be specified for PP mode"
-        else:
-            self.parallel_rank = self.tensor_parallel_rank
-
-
 def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor:
     if not hasattr(_thread_local, 'allreduce_workspaces'):
         _thread_local.allreduce_workspaces = {}
@@ -77,7 +44,7 @@ def get_deepseek_allreduce_workspace(mapping: Mapping) -> torch.LongTensor:
 def allreduce(
     input: torch.Tensor,
     workspace: Optional[torch.LongTensor],
-    parallel_config: ParallelConfig,
+    mapping: Mapping,
     strategy: AllReduceStrategy = AllReduceStrategy.AUTO,
     config: AllReduceConfig = AllReduceConfig(0),
     all_reduce_params: Optional[AllReduceParams] = None
@@ -98,7 +65,7 @@ def allreduce(
 
     Args:
         input (Tensor): The input tensor.
-        parallel_config (ParallelConfig):  The parallel config.
+        mapping (Mapping):  The parallel mapping.
         strategy (AllReduceStrategy): NCCL delegates all-reduce to NCCL while ONESHOT and TWOSHOT are custom latency-optimal algorithms.
             AUTO chooses amongst the three based on a message-size heuristic.
         config (AllReduceConfig): The config for custom allreduce kernels.
@@ -106,19 +73,10 @@ def allreduce(
     Returns:
         The reduced tensor and an optional intermediate tensor if fused.
     '''
-    if parallel_config.tensor_parallel_size == 1 or (
-            all_reduce_params is not None
-            and all_reduce_params.enable_allreduce == False):
+    if mapping.tp_size == 1 or (all_reduce_params is not None and
+                                all_reduce_params.enable_allreduce == False):
         return input
 
-    mapping = Mapping(
-        world_size=parallel_config.parallel_size,
-        tp_size=parallel_config.tensor_parallel_size,
-        pp_size=parallel_config.pipeline_parallel_size,
-        rank=parallel_config.parallel_rank,
-        gpus_per_node=parallel_config.gpus_per_node,
-    )
-
     if all_reduce_params is None:
         all_reduce_params = AllReduceParams()
     is_fused = all_reduce_params.fusion_op == AllReduceFusionOp.RESIDUAL_RMS_NORM or \
@@ -164,7 +122,7 @@ def userbuffers_allreduce_finalize(
 
 
 def allgather(input: torch.Tensor,
-              parallel_config: ParallelConfig,
+              mapping: Mapping,
               gather_dim: int = -1) -> torch.Tensor:
     '''
     Add an operation that performs a collective all-gather.
@@ -184,22 +142,14 @@ def allgather(input: torch.Tensor,
 
     Args:
         input (Tensor): The input tensor.
-        parallel_config (ParallelConfig):  The parallel config.
+        mapping (Mapping):  The parallel mapping.
         gather_dim (int): Gather along given dimension. By default -1.
     Returns:
         The gathered tensor.
     '''
-    if parallel_config.tensor_parallel_size == 1:
+    if mapping.tp_size == 1:
         return input
 
-    mapping = Mapping(
-        world_size=parallel_config.parallel_size,
-        tp_size=parallel_config.tensor_parallel_size,
-        pp_size=parallel_config.pipeline_parallel_size,
-        rank=parallel_config.parallel_rank,
-        gpus_per_node=parallel_config.gpus_per_node,
-    )
-
     output = torch.ops.trtllm.allgather(
         input,
         mapping.tp_group,
@@ -211,26 +161,17 @@ def allgather(input: torch.Tensor,
     output = torch.movedim(output, 0, gather_dim)
     input_shape = input.size()
     output = output.reshape(input_shape[:gather_dim] +
-                            (parallel_config.tensor_parallel_size *
-                             input_shape[gather_dim], ) +
+                            (mapping.tp_size * input_shape[gather_dim], ) +
                             input_shape[gather_dim + 1:])
     return output
 
 
 def reducescatter(input: torch.Tensor,
-                  parallel_config: ParallelConfig,
+                  mapping: Mapping,
                   scatter_dim: int = -1) -> torch.Tensor:
-    if parallel_config.tensor_parallel_size == 1:
+    if mapping.tp_size == 1:
         return input
 
-    mapping = Mapping(
-        world_size=parallel_config.parallel_size,
-        tp_size=parallel_config.tensor_parallel_size,
-        pp_size=parallel_config.pipeline_parallel_size,
-        rank=parallel_config.parallel_rank,
-        gpus_per_node=parallel_config.gpus_per_node,
-    )
-
     output = torch.ops.trtllm.reducescatter(
         input,
         mapping.tp_group,
@@ -242,8 +183,7 @@ def reducescatter(input: torch.Tensor,
     output = torch.movedim(output, 0, scatter_dim)
     input_shape = input.size()
     output = output.reshape(input_shape[:scatter_dim] +
-                            (input_shape[scatter_dim] //
-                             parallel_config.tensor_parallel_size, ) +
+                            (input_shape[scatter_dim] // mapping.tp_size, ) +
                             input_shape[scatter_dim + 1:])
     return output
 
@@ -251,25 +191,14 @@ def reducescatter(input: torch.Tensor,
 class AllReduce(nn.Module):
 
     def __init__(self,
-                 parallel_config: ParallelConfig,
+                 mapping: Mapping,
                  strategy: AllReduceStrategy = AllReduceStrategy.AUTO):
         super().__init__()
 
-        self.parallel_config = parallel_config
-        self.tp_size = self.parallel_config.tensor_parallel_size
-        self.rank = self.parallel_config.parallel_rank
-        self.gpus_per_node = self.parallel_config.gpus_per_node
-
+        self.mapping = mapping
         self.workspace = None
         self.strategy = strategy
-        if self.tp_size > 1:
-            mapping = Mapping(
-                world_size=self.parallel_config.parallel_size,
-                tp_size=self.tp_size,
-                pp_size=self.parallel_config.pipeline_parallel_size,
-                rank=self.rank,
-                gpus_per_node=self.gpus_per_node,
-            )
+        if self.mapping.tp_size > 1:
             if self.strategy != AllReduceStrategy.UB:
                 self.workspace = get_allreduce_workspace(mapping)
 
@@ -281,7 +210,7 @@ def forward(
     ) -> torch.Tensor:
         output = allreduce(input,
                            self.workspace,
-                           self.parallel_config,
+                           self.mapping,
                            all_reduce_params=all_reduce_params,
                            strategy=self.strategy)
         return output
@@ -289,22 +218,11 @@ def forward(
 
 class DeepseekAllReduce(nn.Module):
 
-    def __init__(self, parallel_config: ParallelConfig):
+    def __init__(self, mapping: Mapping):
         super().__init__()
-        self.parallel_config = parallel_config
-        self.tp_size = self.parallel_config.tensor_parallel_size
-        self.tp_rank = self.parallel_config.tensor_parallel_rank
-        self.gpus_per_node = self.parallel_config.gpus_per_node
-        self.rank = self.parallel_config.parallel_rank
+        self.mapping = mapping
         self.workspace = None
-        if self.tp_size > 1:
-            mapping = Mapping(
-                world_size=self.parallel_config.parallel_size,
-                tp_size=self.tp_size,
-                pp_size=self.parallel_config.pipeline_parallel_size,
-                rank=self.rank,
-                gpus_per_node=self.gpus_per_node,
-            )
+        if self.mapping.tp_size > 1:
             self.workspace = get_deepseek_allreduce_workspace(mapping)
 
     def forward(
@@ -331,8 +249,8 @@ def forward(
             input=hidden_states,
             workspace=self.workspace,
             reduce_fusion_inputs=reduce_fusion_inputs,
-            rank=self.parallel_config.tensor_parallel_rank,
-            nranks=self.parallel_config.tensor_parallel_size,
+            rank=self.mapping.tp_rank,
+            nranks=self.mapping.tp_size,
             eps=eps,
             fusion_op=fusion_op,
         )
diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
index 4cacddaa9f3e..1d1d664b9b87 100644
--- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py
+++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
@@ -15,7 +15,7 @@
 from ..attention_backend import AttentionMetadata
 from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams
 from ..distributed import (AllReduce, AllReduceFusionOp, AllReduceParams,
-                           DeepseekAllReduce, ParallelConfig, allgather)
+                           DeepseekAllReduce, allgather)
 from ..model_config import ModelConfig
 from ..models.modeling_utils import MissingLayer, ModelConfig, support_pp
 from ..modules.attention import MLA
@@ -306,13 +306,8 @@ def __init__(self,
             overridden_tp_size=shared_tp_size,
             is_expert=True)
 
-        self.parallel_config = ParallelConfig(
-            tensor_parallel_rank=model_config.mapping.tp_rank,
-            tensor_parallel_size=model_config.mapping.tp_size,
-            gpus_per_node=model_config.mapping.gpus_per_node,
-            pipeline_parallel_size=model_config.mapping.pp_size,
-            parallel_rank=model_config.mapping.rank)
-        self.all_reduce = AllReduce(self.parallel_config)
+        self.mapping = model_config.mapping
+        self.all_reduce = AllReduce(self.mapping)
         self.aux_stream = aux_stream_dict[AuxStreamType.MoeShared]
         self.event_dict = {
             key: torch.cuda.Event()
@@ -321,14 +316,14 @@ def __init__(self,
 
     def compute_routed_output(self, hidden_states, hidden_states_fp4,
                               all_rank_num_tokens, min_latency_mode):
-        if self.use_dp and self.parallel_config.tensor_parallel_size > 1:
+        if self.use_dp and self.mapping.tp_size > 1:
             max_num_token = max(all_rank_num_tokens)
             hidden_states = torch.nn.functional.pad(
                 hidden_states,
                 (0, 0, 0, max_num_token - hidden_states.shape[0]))
             if disable_fp4_allgather():
                 hidden_states = allgather(hidden_states,
-                                          self.parallel_config,
+                                          self.mapping,
                                           gather_dim=0)
         router_logits = self.gate(hidden_states)
 
@@ -383,7 +378,7 @@ def forward(
         assert shared_output.size() == routed_output.size(
         ), f'unmatched tensor shape'
         final_hidden_states = shared_output + routed_output
-        if not self.use_dp and self.parallel_config.tensor_parallel_size > 1:
+        if not self.use_dp and self.mapping.tp_size > 1:
             final_hidden_states = self.all_reduce(
                 final_hidden_states, all_reduce_params=final_all_reduce_params)
 
@@ -470,14 +465,9 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig],
         self.post_attention_layernorm = RMSNorm(hidden_size=config.hidden_size,
                                                 eps=config.rms_norm_eps,
                                                 dtype=config.torch_dtype)
-        self.parallel_config = ParallelConfig(
-            tensor_parallel_rank=model_config.mapping.tp_rank,
-            tensor_parallel_size=model_config.mapping.tp_size,
-            gpus_per_node=model_config.mapping.gpus_per_node,
-            pipeline_parallel_size=model_config.mapping.pp_size,
-            parallel_rank=model_config.mapping.rank)
+        self.mapping = model_config.mapping
         self.layer_idx = layer_idx
-        self.all_reduce = AllReduce(self.parallel_config)
+        self.all_reduce = AllReduce(self.mapping)
         self.next_layer_layernorm: RMSNorm = None
 
         self.deepseek_allreduce_disabled = os.environ.get(
@@ -486,7 +476,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig],
             self.deepseek_allreduce_disabled = True
 
         if not self.deepseek_allreduce_disabled:
-            self.deepseek_allreduce = DeepseekAllReduce(self.parallel_config)
+            self.deepseek_allreduce = DeepseekAllReduce(self.mapping)
 
     def forward(
         self,
@@ -511,9 +501,9 @@ def forward(
             hidden_states=hidden_states,
             attn_metadata=attn_metadata,
             all_reduce_params=AllReduceParams(enable_allreduce=not (
-                self.fusion_config.PRE_MOE_FUSION or self.fusion_config.
-                PRE_MLP_FUSION or self.parallel_config.tensor_parallel_size == 1
-                or self.enable_attention_dp)),
+                self.fusion_config.PRE_MOE_FUSION
+                or self.fusion_config.PRE_MLP_FUSION
+                or self.mapping.tp_size == 1 or self.enable_attention_dp)),
             **kwargs,
         )
 
@@ -586,9 +576,9 @@ def forward(
                 hidden_states_fp4,
                 all_rank_num_tokens=attn_metadata.all_rank_num_tokens,
                 final_all_reduce_params=AllReduceParams(enable_allreduce=not (
-                    self.fusion_config.POST_MOE_FUSION or self.fusion_config.
-                    POST_MLP_FUSION or self.parallel_config.tensor_parallel_size
-                    == 1 or self.enable_attention_dp)),
+                    self.fusion_config.POST_MOE_FUSION
+                    or self.fusion_config.POST_MLP_FUSION
+                    or self.mapping.tp_size == 1 or self.enable_attention_dp)),
                 min_latency_mode=min_latency_mode,
             )
         else:
@@ -596,9 +586,9 @@ def forward(
                 hidden_states,
                 all_rank_num_tokens=attn_metadata.all_rank_num_tokens,
                 final_all_reduce_params=AllReduceParams(enable_allreduce=not (
-                    self.fusion_config.POST_MOE_FUSION or self.fusion_config.
-                    POST_MLP_FUSION or self.parallel_config.tensor_parallel_size
-                    == 1 or self.enable_attention_dp)),
+                    self.fusion_config.POST_MOE_FUSION
+                    or self.fusion_config.POST_MLP_FUSION
+                    or self.mapping.tp_size == 1 or self.enable_attention_dp)),
                 min_latency_mode=min_latency_mode,
             )
 
@@ -729,8 +719,8 @@ def forward(
             hidden_states=hidden_states,
             attn_metadata=attn_metadata,
             all_reduce_params=AllReduceParams(enable_allreduce=not (
-                self.fusion_config.PRE_MOE_FUSION or self.parallel_config.
-                tensor_parallel_size == 1 or self.enable_attention_dp)),
+                self.fusion_config.PRE_MOE_FUSION or self.mapping.tp_size == 1
+                or self.enable_attention_dp)),
             **kwargs,
         )
 
@@ -762,8 +752,8 @@ def forward(
             hidden_states,
             all_rank_num_tokens=spec_metadata.all_rank_num_tokens,
             final_all_reduce_params=AllReduceParams(enable_allreduce=not (
-                self.fusion_config.POST_MOE_FUSION or self.parallel_config.
-                tensor_parallel_size == 1 or self.enable_attention_dp)),
+                self.fusion_config.POST_MOE_FUSION or self.mapping.tp_size == 1
+                or self.enable_attention_dp)),
         )
 
         if self.fusion_config.POST_MOE_FUSION:
diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py
index d02f92549032..adb6b9c549d9 100644
--- a/tensorrt_llm/_torch/models/modeling_llama.py
+++ b/tensorrt_llm/_torch/models/modeling_llama.py
@@ -6,8 +6,7 @@
 from transformers import Llama4Config, LlamaConfig
 
 from tensorrt_llm._torch.distributed import (AllReduce, AllReduceFusionOp,
-                                             AllReduceParams, DeepseekAllReduce,
-                                             ParallelConfig, TensorParallelMode)
+                                             AllReduceParams, DeepseekAllReduce)
 from tensorrt_llm.functional import PositionEmbeddingType
 
 from ..attention_backend import AttentionMetadata
@@ -19,7 +18,8 @@
 from ..modules.fused_moe import (FusedMoE, Llama4RenormalizeMoeRoutingMethod,
                                  MoEWeightLoadingMode)
 from ..modules.gated_mlp import GatedMLP
-from ..modules.linear import Linear, WeightMode, WeightsLoadingConfig
+from ..modules.linear import (Linear, TensorParallelMode, WeightMode,
+                              WeightsLoadingConfig)
 from ..modules.rms_norm import RMSNorm
 from ..modules.rotary_embedding import RotaryEmbedding
 from ..speculative import Eagle3SpecMetadata, SpecMetadata
@@ -132,11 +132,8 @@ def __init__(
                              dtype=config.torch_dtype,
                              quant_config=None)
 
-        self.parallel_config = ParallelConfig(
-            tensor_parallel_rank=model_config.mapping.tp_rank,
-            tensor_parallel_size=model_config.mapping.tp_size,
-            gpus_per_node=model_config.mapping.gpus_per_node)
-        self.all_reduce = AllReduce(self.parallel_config)
+        self.mapping = model_config.mapping
+        self.all_reduce = AllReduce(self.mapping)
         self.moe_event = [torch.cuda.Event(), torch.cuda.Event()]
         self.aux_stream = aux_stream
 
@@ -177,7 +174,7 @@ def forward(
         assert shared_output.size() == routed_output.size(
         ), f'unmatched tensor shape'
         final_hidden_states = shared_output + routed_output
-        if self.parallel_config.tensor_parallel_size > 1:
+        if self.mapping.tp_size > 1:
             final_hidden_states = self.all_reduce(
                 final_hidden_states, all_reduce_params=final_all_reduce_params)
 
@@ -268,14 +265,11 @@ def __init__(
                                                 eps=config.rms_norm_eps,
                                                 dtype=config.torch_dtype)
 
-        self.parallel_config = ParallelConfig(
-            tensor_parallel_rank=model_config.mapping.tp_rank,
-            tensor_parallel_size=model_config.mapping.tp_size,
-            gpus_per_node=model_config.mapping.gpus_per_node)
-        self.all_reduce = AllReduce(self.parallel_config)
+        self.mapping = model_config.mapping
+        self.all_reduce = AllReduce(self.mapping)
         self.next_layer_layernorm: RMSNorm = None
 
-        self.deepseek_allreduce = DeepseekAllReduce(self.parallel_config)
+        self.deepseek_allreduce = DeepseekAllReduce(self.mapping)
 
     def forward(
         self,
@@ -307,9 +301,9 @@ def forward(
             position_ids=position_ids,
             hidden_states=hidden_states,
             attn_metadata=attn_metadata,
-            all_reduce_params=AllReduceParams(enable_allreduce=not (
-                self.fusion_config.PRE_MOE_FUSION
-                or self.parallel_config.tensor_parallel_size == 1)),
+            all_reduce_params=AllReduceParams(
+                enable_allreduce=not (self.fusion_config.PRE_MOE_FUSION
+                                      or self.mapping.tp_size == 1)),
             **kwargs,
         )
 
@@ -332,9 +326,8 @@ def forward(
             hidden_states,
             all_rank_num_tokens=attn_metadata.all_rank_num_tokens,
             final_all_reduce_params=AllReduceParams(enable_allreduce=not (
-                self.fusion_config.POST_MOE_FUSION
-                or self.fusion_config.POST_MLP_FUSION
-                or self.parallel_config.tensor_parallel_size == 1)),
+                self.fusion_config.POST_MOE_FUSION or self.fusion_config.
+                POST_MLP_FUSION or self.mapping.tp_size == 1)),
             min_latency_mode=min_latency_mode,
         )
         if spec_metadata is not None:
@@ -388,8 +381,6 @@ def __init__(
         config = model_config.pretrained_config
 
         tp_size = model_config.mapping.tp_size
-        tp_rank = model_config.mapping.tp_rank
-        gpus_per_node = model_config.mapping.gpus_per_node
 
         # Override the QKV projection. The number of input features
         # is twice as big for EAGLE3 draft models.
@@ -398,13 +389,8 @@ def __init__(
             tp_size * self.q_size + 2 * tp_size * self.kv_size,
             bias=config.attention_bias,
             dtype=config.torch_dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_size=tp_size,
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gpus_per_node=gpus_per_node,
-                pipeline_parallel_size=model_config.mapping.pp_size,
-                parallel_rank=model_config.mapping.rank),
+            mapping=model_config.mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
             weights_loading_config=WeightsLoadingConfig(
                 weight_mode=WeightMode.FUSED_QKV_LINEAR),
             quant_config=model_config.get_quant_config(),
@@ -486,15 +472,9 @@ def __init__(self, model_config: ModelConfig[LlamaConfig]):
             config.vocab_size,
             config.hidden_size,
             dtype=config.torch_dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=model_config.mapping.tp_rank,
-                tensor_parallel_size=model_config.mapping.tp_size,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                pipeline_parallel_size=model_config.mapping.pp_size,
-                parallel_rank=model_config.mapping.rank,
-                gather_output=True,
-                gpus_per_node=model_config.mapping.gpus_per_node,
-            ),
+            mapping=model_config.mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
+            gather_output=True,
         )
         self.layers = nn.ModuleList([
             LlamaDecoderLayer(
diff --git a/tensorrt_llm/_torch/models/modeling_mixtral.py b/tensorrt_llm/_torch/models/modeling_mixtral.py
index 451f698ae4ac..8d2beed94f7b 100644
--- a/tensorrt_llm/_torch/models/modeling_mixtral.py
+++ b/tensorrt_llm/_torch/models/modeling_mixtral.py
@@ -8,7 +8,6 @@
 
 from ..attention_backend import AttentionMetadata
 from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams
-from ..distributed import ParallelConfig
 from ..model_config import ModelConfig
 from ..models.modeling_utils import ModelConfig
 from ..modules.attention import Attention
@@ -133,11 +132,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig],
         self.post_attention_layernorm = RMSNorm(hidden_size=config.hidden_size,
                                                 eps=config.rms_norm_eps,
                                                 dtype=config.torch_dtype)
-        # TODO: add pipeline parallel config
-        self.parallel_config = ParallelConfig(
-            tensor_parallel_rank=model_config.mapping.tp_rank,
-            tensor_parallel_size=model_config.mapping.tp_size,
-            gpus_per_node=model_config.mapping.gpus_per_node)
+        self.mapping = model_config.mapping
         self.layer_idx = layer_idx
 
     def forward(
diff --git a/tensorrt_llm/_torch/models/modeling_mllama.py b/tensorrt_llm/_torch/models/modeling_mllama.py
index e389108b0eb8..af487a077be9 100644
--- a/tensorrt_llm/_torch/models/modeling_mllama.py
+++ b/tensorrt_llm/_torch/models/modeling_mllama.py
@@ -26,11 +26,11 @@
 
 from ...logger import logger
 from ..attention_backend.interface import AttentionMetadata
-from ..distributed import ParallelConfig, TensorParallelMode
 from ..model_config import ModelConfig
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding, LMHead
 from ..modules.gated_mlp import GatedMLP
+from ..modules.linear import TensorParallelMode
 from ..modules.logits_procesor import LogitsProcessor
 from .modeling_llama import LlamaAttention
 from .modeling_utils import duplicate_kv_weight, register_auto_model
@@ -224,13 +224,9 @@ def __init__(
             text_config.vocab_size,
             text_config.hidden_size,
             dtype=text_config.torch_dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=config.mapping.tp_rank,
-                tensor_parallel_size=config.mapping.tp_size,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gather_output=True,
-                gpus_per_node=config.mapping.gpus_per_node,
-            ),
+            mapping=config.mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
+            gather_output=True,
         )
         # use embedding weights in lm_head if tie word embedding is enabled
         if text_config.tie_word_embeddings:
diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_nas.py b/tensorrt_llm/_torch/models/modeling_nemotron_nas.py
index 435dd32c4abd..6d02765d8c52 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron_nas.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron_nas.py
@@ -4,8 +4,6 @@
 from torch import nn
 from transformers import PretrainedConfig
 
-from tensorrt_llm._torch.distributed import ParallelConfig, TensorParallelMode
-from tensorrt_llm._torch.modules.linear import Linear
 from tensorrt_llm.functional import PositionEmbeddingType
 
 from ..attention_backend import AttentionMetadata
@@ -15,6 +13,7 @@
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.gated_mlp import GatedMLP
+from ..modules.linear import Linear, TensorParallelMode
 from ..modules.rms_norm import RMSNorm
 from ..modules.rotary_embedding import RotaryEmbedding
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM,
@@ -55,15 +54,9 @@ def _create_linear_from_configs(model_config: ModelConfig[PretrainedConfig],
         config.hidden_size,
         bias=False,
         dtype=config.torch_dtype,
-        parallel_config=ParallelConfig(
-            tensor_parallel_rank=model_config.mapping.tp_rank,
-            tensor_parallel_size=model_config.mapping.tp_size,
-            tensor_parallel_mode=TensorParallelMode.COLUMN,
-            pipeline_parallel_size=model_config.mapping.pp_size,
-            parallel_rank=model_config.mapping.rank,
-            gather_output=True,
-            gpus_per_node=model_config.mapping.gpus_per_node,
-        ),
+        mapping=model_config.mapping,
+        tensor_parallel_mode=TensorParallelMode.COLUMN,
+        gather_output=True,
         quant_config=model_config.get_quant_config(),
         skip_create_weights=model_config.skip_create_weights,
     )
diff --git a/tensorrt_llm/_torch/models/modeling_qwen.py b/tensorrt_llm/_torch/models/modeling_qwen.py
index 09f7569a01ca..558bd06b03a3 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen.py
@@ -8,13 +8,12 @@
 
 from ..attention_backend import AttentionMetadata
 from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams
-from ..distributed import ParallelConfig, TensorParallelMode
 from ..model_config import ModelConfig
 from ..modules.attention import Attention
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.gated_mlp import GatedMLP
-from ..modules.linear import Linear
+from ..modules.linear import Linear, TensorParallelMode
 from ..modules.rms_norm import RMSNorm
 from ..modules.rotary_embedding import RotaryEmbedding
 from ..pipeline_interface import PipelineInterface
@@ -147,13 +146,9 @@ def __init__(self, model_config: ModelConfig[Qwen2Config]):
             config.pretrained_config.vocab_size,
             config.pretrained_config.hidden_size,
             dtype=config.pretrained_config.torch_dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=config.mapping.tp_rank,
-                tensor_parallel_size=config.mapping.tp_size,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gather_output=True,
-                gpus_per_node=config.mapping.gpus_per_node,
-            ),
+            mapping=config.mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
+            gather_output=True,
         )
         self.layers = nn.ModuleList([
             QwenDecoderLayer(
diff --git a/tensorrt_llm/_torch/models/modeling_qwen_moe.py b/tensorrt_llm/_torch/models/modeling_qwen_moe.py
index f047ba2a183a..9cdecd9ceb73 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen_moe.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen_moe.py
@@ -10,14 +10,13 @@
 
 from ..attention_backend import AttentionMetadata
 from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams
-from ..distributed import ParallelConfig, TensorParallelMode
 from ..model_config import ModelConfig
 from ..modules.attention import Attention
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.fused_moe import DefaultMoeRoutingMethod, FusedMoE
 from ..modules.gated_mlp import GatedMLP
-from ..modules.linear import Linear
+from ..modules.linear import Linear, TensorParallelMode
 from ..modules.rms_norm import RMSNorm
 from ..modules.rotary_embedding import RotaryEmbedding
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM,
@@ -175,11 +174,6 @@ def __init__(self, model_config: ModelConfig[Qwen2MoeConfig],
         self.post_attention_layernorm = RMSNorm(hidden_size=config.hidden_size,
                                                 eps=config.rms_norm_eps,
                                                 dtype=config.torch_dtype)
-        # TODO: add pipeline parallel config
-        self.parallel_config = ParallelConfig(
-            tensor_parallel_rank=model_config.mapping.tp_rank,
-            tensor_parallel_size=model_config.mapping.tp_size,
-            gpus_per_node=model_config.mapping.gpus_per_node)
         self.layer_idx = layer_idx
 
     def forward(
@@ -224,13 +218,9 @@ def __init__(self, model_config: ModelConfig[Qwen2MoeConfig]):
             config.pretrained_config.vocab_size,
             config.pretrained_config.hidden_size,
             dtype=config.pretrained_config.torch_dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=config.mapping.tp_rank,
-                tensor_parallel_size=config.mapping.tp_size,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gather_output=True,
-                gpus_per_node=config.mapping.gpus_per_node,
-            ),
+            mapping=config.mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
+            gather_output=True,
         )
         self.layers = nn.ModuleList([
             QwenMoeDecoderLayer(
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index ac2577ee94fd..57a5698a603b 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_utils.py
@@ -11,14 +11,15 @@
 from torch.utils._pytree import tree_any_only
 from tqdm import tqdm
 
+from tensorrt_llm.mapping import Mapping
+
 from ...logger import logger
 from ..attention_backend import AttentionMetadata
-from ..distributed import ParallelConfig, TensorParallelMode
 from ..model_config import ModelConfig, TConfig
 from ..modules.attention import Attention
 from ..modules.embedding import Embedding, LMHead
 from ..modules.fused_moe import FusedMoE
-from ..modules.linear import Linear, WeightMode
+from ..modules.linear import Linear, TensorParallelMode, WeightMode
 from ..modules.logits_procesor import LogitsProcessor
 from ..modules.rms_norm import RMSNorm
 from ..pipeline_interface import PipelineInterface
@@ -356,14 +357,13 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig],
                     vocab_size,
                     hidden_size,
                     dtype=config.pretrained_config.torch_dtype,
-                    parallel_config=ParallelConfig(
-                        tensor_parallel_rank=0,
-                        tensor_parallel_size=1,
-                        tensor_parallel_mode=None,
-                        gather_output=False,
-                        pipeline_parallel_size=config.mapping.pp_size,
-                        parallel_rank=config.mapping.rank,
+                    mapping=Mapping(
+                        world_size=1,
+                        tp_size=1,
+                        rank=0,
                     ),
+                    tensor_parallel_mode=None,
+                    gather_output=False,
                 )
             else:
                 # TODO(zhenhuanc): Currently lm_head Linear will not accept QuantConfig
@@ -372,15 +372,9 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig],
                     vocab_size,
                     hidden_size,
                     dtype=config.pretrained_config.torch_dtype,
-                    parallel_config=ParallelConfig(
-                        tensor_parallel_size=config.mapping.tp_size,
-                        tensor_parallel_rank=config.mapping.tp_rank,
-                        tensor_parallel_mode=TensorParallelMode.COLUMN,
-                        gather_output=True,
-                        gpus_per_node=config.mapping.gpus_per_node,
-                        pipeline_parallel_size=config.mapping.pp_size,
-                        parallel_rank=config.mapping.rank,
-                    ),
+                    mapping=config.mapping,
+                    tensor_parallel_mode=TensorParallelMode.COLUMN,
+                    gather_output=True,
                 )
 
             # use embedding weights in lm_head if tie word embedding is enabled
diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index 4983a9d27366..b333ca452c6a 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -4,15 +4,17 @@
 import torch
 from torch import nn
 
+from tensorrt_llm.mapping import Mapping
+
 from ..attention_backend import (AttentionInputType, AttentionMetadata,
                                  TrtllmAttention)
 from ..attention_backend.interface import (PositionalEmbeddingParams,
                                            PredefinedAttentionMask)
 from ..attention_backend.utils import create_attention
-from ..distributed import AllReduceParams, ParallelConfig, TensorParallelMode
+from ..distributed import AllReduceParams
 from ..model_config import ModelConfig
 from ..peft.lora.layer import LoraLayer, LoraModuleType
-from .linear import Linear, WeightMode, WeightsLoadingConfig
+from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
 from .rms_norm import RMSNorm
 from .rotary_embedding import RotaryEmbedding
 
@@ -80,12 +82,18 @@ def __init__(
         # tensor parallel
         config = config or ModelConfig()
         tp_size = config.mapping.tp_size
-        tp_rank = config.mapping.tp_rank
-        gpus_per_node = config.mapping.gpus_per_node
+        pp_size = config.mapping.pp_size
         if config.mapping.enable_attention_dp:
             tp_size = 1
-            tp_rank = 0
 
+        mapping = Mapping(
+            world_size=tp_size * pp_size,
+            tp_size=tp_size,
+            pp_size=pp_size,
+            rank=config.mapping.rank,
+            gpus_per_node=config.mapping.gpus_per_node,
+            enable_attention_dp=config.mapping.enable_attention_dp,
+        )
         assert self.num_heads % tp_size == 0
         self.num_heads = self.num_heads // tp_size
         self.num_key_value_heads = (self.num_key_value_heads + tp_size -
@@ -103,13 +111,8 @@ def __init__(
             tp_size * self.q_size + 2 * tp_size * self.kv_size,
             bias=bias,
             dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_size=tp_size,
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gpus_per_node=gpus_per_node,
-                pipeline_parallel_size=config.mapping.pp_size,
-                parallel_rank=config.mapping.rank),
+            mapping=mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
             weights_loading_config=WeightsLoadingConfig(
                 weight_mode=WeightMode.FUSED_QKV_LINEAR),
             quant_config=config.get_quant_config(),
@@ -120,13 +123,8 @@ def __init__(
             self.hidden_size,
             bias=self.dense_bias,
             dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_size=tp_size,
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_mode=TensorParallelMode.ROW,
-                gpus_per_node=gpus_per_node,
-                pipeline_parallel_size=config.mapping.pp_size,
-                parallel_rank=config.mapping.rank),
+            mapping=mapping,
+            tensor_parallel_mode=TensorParallelMode.ROW,
             quant_config=config.get_quant_config(),
             skip_create_weights=config.skip_create_weights,
         )
@@ -310,27 +308,17 @@ def __init__(
         # tensor parallel
         config = config or ModelConfig()
         tp_size = config.mapping.tp_size
-        tp_rank = config.mapping.tp_rank
-        gpus_per_node = config.mapping.gpus_per_node
+        pp_size = config.mapping.pp_size
         if config.mapping.enable_attention_dp:
             tp_size = 1
-            tp_rank = 0
 
-        row_parallel_config = ParallelConfig(
-            tensor_parallel_rank=tp_rank,
-            tensor_parallel_size=tp_size,
-            tensor_parallel_mode=TensorParallelMode.ROW,
-            gpus_per_node=gpus_per_node,
-            pipeline_parallel_size=config.mapping.pp_size,
-            parallel_rank=config.mapping.rank,
-        )
-        col_parallel_config = ParallelConfig(
-            tensor_parallel_rank=tp_rank,
-            tensor_parallel_size=tp_size,
-            tensor_parallel_mode=TensorParallelMode.COLUMN,
-            gpus_per_node=gpus_per_node,
-            pipeline_parallel_size=config.mapping.pp_size,
-            parallel_rank=config.mapping.rank,
+        mapping = Mapping(
+            world_size=tp_size * pp_size,
+            tp_size=tp_size,
+            pp_size=pp_size,
+            rank=config.mapping.rank,
+            gpus_per_node=config.mapping.gpus_per_node,
+            enable_attention_dp=config.mapping.enable_attention_dp,
         )
 
         assert self.num_heads % tp_size == 0
@@ -362,7 +350,8 @@ def __init__(
                 (self.qk_nope_head_dim + self.qk_rope_head_dim),
                 bias=bias,
                 dtype=dtype,
-                parallel_config=col_parallel_config,
+                mapping=mapping,
+                tensor_parallel_mode=TensorParallelMode.COLUMN,
                 quant_config=quant_config,
                 skip_create_weights=config.skip_create_weights)
         else:
@@ -381,7 +370,8 @@ def __init__(
                 (self.qk_nope_head_dim + self.qk_rope_head_dim),
                 bias=bias,
                 dtype=dtype,
-                parallel_config=col_parallel_config,
+                mapping=mapping,
+                tensor_parallel_mode=TensorParallelMode.COLUMN,
                 quant_config=quant_config,
                 skip_create_weights=config.skip_create_weights)
             self.q_b_proj = self.q_proj
@@ -400,7 +390,8 @@ def __init__(
                                 (self.qk_nope_head_dim + self.v_head_dim),
                                 bias=bias,
                                 dtype=dtype,
-                                parallel_config=col_parallel_config,
+                                mapping=mapping,
+                                tensor_parallel_mode=TensorParallelMode.COLUMN,
                                 quant_config=quant_config,
                                 skip_create_weights=config.skip_create_weights)
         # This parameter will view into self.kv_b_proj.weight after loading weights.
@@ -455,7 +446,8 @@ def __init__(
             self.hidden_size,
             bias=self.dense_bias,
             dtype=dtype,
-            parallel_config=row_parallel_config,
+            mapping=mapping,
+            tensor_parallel_mode=TensorParallelMode.ROW,
             quant_config=quant_config,
             skip_create_weights=config.skip_create_weights,
         )
diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py
index 49bde47e01a4..15cc50ab1fef 100644
--- a/tensorrt_llm/_torch/modules/embedding.py
+++ b/tensorrt_llm/_torch/modules/embedding.py
@@ -6,9 +6,10 @@
 from torch.nn.parameter import Parameter
 
 from tensorrt_llm.functional import AllReduceParams
+from tensorrt_llm.mapping import Mapping
 
-from ..distributed import ParallelConfig, TensorParallelMode, allgather
-from .linear import Linear
+from ..distributed import allgather
+from .linear import Linear, TensorParallelMode
 
 
 class LMHead(Linear):
@@ -18,7 +19,7 @@ class LMHead(Linear):
         num_embeddings (int): vocabulary size.
         embedding_dim (int): size of hidden state.
         dtype (Optional[torch.dtype]): type of the parameters.
-        parallel_config (Optional[ParallelConfig]): parallelism configuration.
+        mapping (Optional[Mapping]): parallelism configuration.
             If not provided, the embedding is not parallelized.
     """
 
@@ -27,17 +28,19 @@ def __init__(
         num_embeddings: int,
         embedding_dim: int,
         dtype: torch.dtype = None,
-        parallel_config: Optional[ParallelConfig] = None,
+        mapping: Optional[Mapping] = None,
+        tensor_parallel_mode: Optional[TensorParallelMode] = None,
+        gather_output: bool = False,
     ):
         local_in_features = embedding_dim
         local_out_features = num_embeddings
-        parallel_config = parallel_config or ParallelConfig()
-        tp_size = parallel_config.tensor_parallel_size
+        mapping = mapping or Mapping()
+        tp_size = mapping.tp_size
 
-        if parallel_config.tensor_parallel_mode == TensorParallelMode.ROW:
+        if tensor_parallel_mode == TensorParallelMode.ROW:
             local_in_features = math.ceil(embedding_dim / tp_size)
             self.padding_size = tp_size * local_in_features - embedding_dim
-        elif parallel_config.tensor_parallel_mode == TensorParallelMode.COLUMN:
+        elif tensor_parallel_mode == TensorParallelMode.COLUMN:
             local_out_features = math.ceil(num_embeddings / tp_size)
             self.padding_size = tp_size * local_out_features - num_embeddings
 
@@ -46,10 +49,12 @@ def __init__(
             local_out_features * tp_size,
             bias=False,
             dtype=dtype,
-            parallel_config=parallel_config,
+            mapping=mapping,
+            tensor_parallel_mode=tensor_parallel_mode,
+            gather_output=gather_output,
         )
 
-        if parallel_config.tensor_parallel_mode == TensorParallelMode.ROW:
+        if tensor_parallel_mode == TensorParallelMode.ROW:
             if self.tp_rank == self.tp_size - 1:
                 local_in_features -= self.padding_size
         self.in_features = local_in_features
@@ -61,7 +66,7 @@ def __init__(
 
     @property
     def vocab_size_padded(self) -> int:
-        if self.parallel_config.tensor_parallel_mode == TensorParallelMode.COLUMN:
+        if self.tp_mode == TensorParallelMode.COLUMN:
             return self.out_features * self.tp_size
         else:
             return self.out_features
@@ -73,8 +78,7 @@ def forward(
             all_reduce_params: Optional[AllReduceParams] = None
     ) -> torch.Tensor:
         output = super().forward(input, all_reduce_params=all_reduce_params)
-        if (self.tp_mode == TensorParallelMode.COLUMN
-                and self.parallel_config.gather_output
+        if (self.tp_mode == TensorParallelMode.COLUMN and self.gather_output
                 and self.padding_size > 0):
             output = output[..., :-self.padding_size]
 
@@ -82,7 +86,7 @@ def forward(
 
     def load_weights(self, weights: List[Dict]):
         original_weight = None
-        if self.parallel_config.tensor_parallel_mode == TensorParallelMode.COLUMN:
+        if self.tp_mode == TensorParallelMode.COLUMN:
             if self.tp_rank == self.tp_size - 1 and self.padding_size > 0:
                 original_weight = self.weight.data.zero_()
                 self.weight.data = self.weight[:-self.padding_size, :]
@@ -113,7 +117,7 @@ class Embedding(LMHead):
         num_embeddings (int): vocabulary size.
         embedding_dim (int): size of hidden state.
         dtype (Optional[torch.dtype]): type of the parameters.
-        parallel_config (Optional[ParallelConfig]): parallelism configuration.
+        mapping (Optional[Mapping]): parallelism configuration.
             If not provided, the embedding is not parallelized.
     """
 
@@ -122,13 +126,17 @@ def __init__(
         num_embeddings: int,
         embedding_dim: int,
         dtype: Optional[torch.dtype] = None,
-        parallel_config: Optional[ParallelConfig] = None,
+        mapping: Optional[Mapping] = None,
+        tensor_parallel_mode: Optional[TensorParallelMode] = None,
+        gather_output: bool = False,
     ):
         super().__init__(
             embedding_dim=embedding_dim,
             num_embeddings=num_embeddings,
             dtype=dtype,
-            parallel_config=parallel_config,
+            mapping=mapping,
+            tensor_parallel_mode=tensor_parallel_mode,
+            gather_output=gather_output,
         )
         if self.tp_size > 1:
             slice_width = math.ceil(num_embeddings / self.tp_size)
@@ -138,7 +146,7 @@ def __init__(
 
     def forward(self, input):
         if self.tp_size > 1:
-            if self.parallel_config.tensor_parallel_mode == TensorParallelMode.COLUMN:
+            if self.tp_mode == TensorParallelMode.COLUMN:
                 # Build the mask.
                 input, input_mask = get_masked_input_and_mask(
                     input,
@@ -149,15 +157,15 @@ def forward(self, input):
         output = F.embedding(input, self.weight)
         # Mask the output embedding.
         if self.tp_size > 1:
-            if self.parallel_config.tensor_parallel_mode == TensorParallelMode.COLUMN:
+            if self.tp_mode == TensorParallelMode.COLUMN:
                 output.masked_fill_(input_mask, 0)
                 # Reduce across all the model parallel GPUs.
                 output = self.all_reduce(output)
-            elif self.parallel_config.tensor_parallel_mode == TensorParallelMode.ROW:
-                if self.parallel_config.gather_output:
+            elif self.tp_mode == TensorParallelMode.ROW:
+                if self.gather_output:
                     if self.tp_rank == self.tp_size - 1 and self.padding_size > 0:
                         output = F.pad(output, (0, self.padding_size))
-                    output = allgather(output, self.parallel_config)
+                    output = allgather(output, self.mapping)
                     if self.padding_size > 0:
                         output = output[..., :-self.padding_size]
 
diff --git a/tensorrt_llm/_torch/modules/fused_moe.py b/tensorrt_llm/_torch/modules/fused_moe.py
index d5311523c21e..f6ef183e5731 100755
--- a/tensorrt_llm/_torch/modules/fused_moe.py
+++ b/tensorrt_llm/_torch/modules/fused_moe.py
@@ -10,7 +10,7 @@
 from ..model_config import ModelConfig
 from ..utils import (EventType, Fp4QuantizedTensor, disable_fp4_allgather,
                      reswizzle_sf)
-from .linear import ParallelConfig, TensorParallelMode, load_weight_shard
+from .linear import TensorParallelMode, load_weight_shard
 
 # The declarations aligns with moe_kernels.h
 # pack inputs into int64, e.g. 4 x bf16 input values
@@ -270,15 +270,10 @@ def __init__(
         self.use_dp = model_config.mapping.enable_attention_dp
 
         # All ranks participate in allreduce regardless of EP/TP combination
-        self.parallel_config = ParallelConfig(
-            tensor_parallel_rank=model_config.mapping.tp_rank,
-            tensor_parallel_size=model_config.mapping.tp_size,
-            gpus_per_node=model_config.mapping.gpus_per_node,
-            pipeline_parallel_size=model_config.mapping.pp_size,
-            parallel_rank=model_config.mapping.rank)
-        self.parallel_size = self.parallel_config.tensor_parallel_size
+        self.mapping = model_config.mapping
+        self.parallel_size = self.mapping.tp_size
 
-        self.all_reduce = AllReduce(self.parallel_config)
+        self.all_reduce = AllReduce(self.mapping)
 
         self.intermediate_size_per_partition = intermediate_size // self.tp_size
 
@@ -510,7 +505,7 @@ def all_gather(self, input_tensors):
 
         flatten_outputs = allgather(
             torch.cat(flatten_inputs),
-            self.parallel_config,
+            self.mapping,
             gather_dim=0,
         ).view(self.parallel_size, -1)
 
@@ -532,9 +527,7 @@ def reducescatter_or_allreduce(self, inputs):
         outputs = inputs
         if self.parallel_size > 1:
             if self.use_dp:
-                outputs = reducescatter(inputs,
-                                        self.parallel_config,
-                                        scatter_dim=0)
+                outputs = reducescatter(inputs, self.mapping, scatter_dim=0)
             elif self.reduce_results:
                 outputs = self.all_reduce(inputs)
         return outputs
@@ -595,7 +588,7 @@ def forward_chunk(
         ):
             x_sf, token_selected_experts, token_final_scales = self.all_gather(
                 [x_sf, token_selected_experts, token_final_scales])
-            x = allgather(x, self.parallel_config, gather_dim=0)
+            x = allgather(x, self.mapping, gather_dim=0)
             token_selected_experts = token_selected_experts.flatten(
                 0, 1).contiguous()
             token_final_scales = token_final_scales.flatten(0, 1).contiguous()
@@ -701,7 +694,7 @@ def forward(
             self.event_dict[EventType.MoeChunkingOverlap].wait()
             outputs = torch.cat(outputs_list)
         if self.use_dp:
-            rank = self.parallel_config.tensor_parallel_rank
+            rank = self.mapping.tp_rank
             outputs = outputs[:all_rank_num_tokens[rank]]
         return outputs
 
diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py
index 069882c2334a..73a8f1327069 100644
--- a/tensorrt_llm/_torch/modules/gated_mlp.py
+++ b/tensorrt_llm/_torch/modules/gated_mlp.py
@@ -5,13 +5,14 @@
 import torch.nn.functional as F
 from torch import nn
 
-from tensorrt_llm._torch.peft.lora.layer import LoraLayer, LoraModuleType
+from tensorrt_llm.mapping import Mapping
 
 from ..custom_ops import IS_FLASHINFER_AVAIABLE
-from ..distributed import AllReduceParams, ParallelConfig, TensorParallelMode
+from ..distributed import AllReduceParams
 from ..model_config import ModelConfig
+from ..peft.lora.layer import LoraLayer, LoraModuleType
 from ..utils import Fp4QuantizedTensor
-from .linear import Linear, WeightMode, WeightsLoadingConfig
+from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
 
 
 def swiglu(x):
@@ -44,30 +45,31 @@ def __init__(self,
         self.activation = activation
 
         config = config or ModelConfig()
+        self.mapping = config.mapping
         if overridden_tp_size is not None:
             assert config.mapping.tp_size % overridden_tp_size == 0
-            tp_rank = config.mapping.tp_rank % overridden_tp_size
             tp_size = overridden_tp_size
             # "Misuse" pp_size here to perform all-reduce within smaller groups
             pp_size = config.mapping.pp_size * config.mapping.tp_size // overridden_tp_size
         else:
-            tp_rank = config.mapping.tp_rank
             tp_size = config.mapping.tp_size
             pp_size = config.mapping.pp_size
-        gpus_per_node = config.mapping.gpus_per_node
+
+        mapping = Mapping(
+            world_size=tp_size * pp_size,
+            rank=self.mapping.rank,
+            gpus_per_node=self.mapping.gpus_per_node,
+            tp_size=tp_size,
+            pp_size=pp_size,
+        )
 
         self.gate_up_proj = Linear(
             self.hidden_size,
             self.intermediate_size * 2,
             bias=bias,
             dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_size=tp_size,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gpus_per_node=gpus_per_node,
-                pipeline_parallel_size=pp_size,
-                parallel_rank=config.mapping.rank),
+            mapping=mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
             weights_loading_config=WeightsLoadingConfig(
                 weight_mode=WeightMode.FUSED_GATE_UP_LINEAR),
             quant_config=config.get_quant_config(),
@@ -79,13 +81,8 @@ def __init__(self,
             self.hidden_size,
             bias=bias,
             dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_size=tp_size,
-                tensor_parallel_mode=TensorParallelMode.ROW,
-                gpus_per_node=gpus_per_node,
-                pipeline_parallel_size=pp_size,
-                parallel_rank=config.mapping.rank),
+            mapping=mapping,
+            tensor_parallel_mode=TensorParallelMode.ROW,
             quant_config=config.get_quant_config(),
             is_expert=is_expert,
             skip_create_weights=config.skip_create_weights,
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index 1747e94bee7b..b8df5586f628 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -10,9 +10,9 @@
 
 import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils
 from tensorrt_llm.functional import AllReduceFusionOp, AllReduceParams
+from tensorrt_llm.mapping import Mapping
 
 from ...models.modeling_utils import QuantConfig
-from ..distributed import ParallelConfig, TensorParallelMode
 from ..utils import Fp4QuantizedTensor
 
 E2M1_MAX = 6.0
@@ -33,6 +33,15 @@ class WeightsLoadingConfig:
     ignore_tensor_parallel: bool = False
 
 
+class TensorParallelMode(str, enum.Enum):
+    COLUMN = 'column'
+    ROW = 'row'
+
+    @classmethod
+    def split_dim(cls, mode):
+        return 1 if mode == cls.ROW else 0
+
+
 def load_weight_shard(
         weight,
         tensor_parallel_size: int = 1,
@@ -135,7 +144,9 @@ def __init__(
         out_features: int,
         bias: bool = True,
         dtype: torch.dtype = None,
-        parallel_config: Optional[ParallelConfig] = None,
+        mapping: Optional[Mapping] = None,
+        tensor_parallel_mode: Optional[TensorParallelMode] = None,
+        gather_output: bool = False,
         quant_config: Optional[QuantConfig] = None,
         weights_loading_config: Optional[WeightsLoadingConfig] = None,
         is_expert: bool = False,
@@ -147,38 +158,37 @@ def __init__(
         super().__init__()
         self.has_bias = bias
         self.dtype = dtype
-        self.parallel_config = parallel_config or ParallelConfig()
+        self.mapping = mapping or Mapping()
         # could be modified later
         self.quant_config = quant_config
         self.weights_loading_config = weights_loading_config or WeightsLoadingConfig(
         )
-        self.tp_size = self.parallel_config.tensor_parallel_size
-        self.tp_rank = self.parallel_config.tensor_parallel_rank
-        self.tp_mode = self.parallel_config.tensor_parallel_mode
+        self.tp_size = self.mapping.tp_size
+        self.tp_rank = self.mapping.tp_rank
+        self.tp_mode = tensor_parallel_mode
+        self.gather_output = gather_output
 
         local_in_features = in_features
         local_out_features = out_features
 
-        if self.parallel_config.tensor_parallel_mode == TensorParallelMode.ROW:
+        if self.tp_mode == TensorParallelMode.ROW:
             assert in_features % self.tp_size == 0, (
                 f'in_features {in_features} must be divisible by tp_size {self.tp_size}'
             )
             local_in_features = in_features // self.tp_size
-        elif self.parallel_config.tensor_parallel_mode == TensorParallelMode.COLUMN:
+        elif self.tp_mode == TensorParallelMode.COLUMN:
             assert out_features % self.tp_size == 0, (
                 f'out_features {out_features} must be divisible by tp_size {self.tp_size}'
             )
             local_out_features = out_features // self.tp_size
         else:
-            assert self.parallel_config.tensor_parallel_mode is None, (
-                'unsupported tensor parallel mode: {self.parallel_config.tensor_parallel_mode}'
-            )
+            assert self.tp_mode is None, (
+                'unsupported tensor parallel mode: {self.tp_mode}')
 
         self.in_features = local_in_features
         self.out_features = local_out_features
 
-        self.all_reduce = AllReduce(
-            self.parallel_config) if not is_expert else None
+        self.all_reduce = AllReduce(self.mapping) if not is_expert else None
         self._weights_created = False
         self.is_expert = is_expert
         self.use_custom_cublas_mm = use_custom_cublas_mm
@@ -385,8 +395,8 @@ def forward(
                 output = self.apply_linear(input, self.weight, bias)
         elif self.tp_mode == TensorParallelMode.COLUMN:
             output = self.apply_linear(input, self.weight, self.bias)
-            if self.parallel_config.gather_output:
-                output = allgather(output, self.parallel_config)
+            if self.gather_output:
+                output = allgather(output, self.mapping)
         else:
             output = self.apply_linear(input, self.weight, self.bias)
 
diff --git a/tensorrt_llm/_torch/modules/mamba.py b/tensorrt_llm/_torch/modules/mamba.py
index c0af7d9085a0..a2b44e9d2dd1 100644
--- a/tensorrt_llm/_torch/modules/mamba.py
+++ b/tensorrt_llm/_torch/modules/mamba.py
@@ -5,9 +5,8 @@
 from torch import nn
 
 from ..attention_backend import AttentionMetadata
-from ..distributed import ParallelConfig, TensorParallelMode
 from ..model_config import ModelConfig
-from .linear import Linear
+from .linear import Linear, TensorParallelMode
 
 
 @dataclass
@@ -166,9 +165,9 @@ def __init__(
         super().__init__()
 
         config = config or ModelConfig()
+        self.mapping = config.mapping
         tp_rank = config.mapping.tp_rank
         tp_size = config.mapping.tp_size
-        gpus_per_node = config.mapping.gpus_per_node
 
         d_inner = d_model * expand
         nheads = d_inner // head_dim
@@ -205,11 +204,8 @@ def __init__(
             d_in_proj,
             bias=bias,
             dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_size=tp_size,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gpus_per_node=gpus_per_node),
+            mapping=self.mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
             quant_config=config.get_quant_config(),
         )
 
@@ -219,11 +215,8 @@ def __init__(
             conv_dim,
             bias=conv_bias,
             dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_size=tp_size,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gpus_per_node=gpus_per_node),
+            mapping=self.mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
             quant_config=config.get_quant_config(),
             skip_create_weights=config.skip_create_weights,
         )
@@ -258,11 +251,8 @@ def __init__(
                                d_model,
                                bias=bias,
                                dtype=dtype,
-                               parallel_config=ParallelConfig(
-                                   tensor_parallel_rank=tp_rank,
-                                   tensor_parallel_size=tp_size,
-                                   tensor_parallel_mode=TensorParallelMode.ROW,
-                                   gpus_per_node=gpus_per_node),
+                               mapping=self.mapping,
+                               tensor_parallel_mode=TensorParallelMode.ROW,
                                quant_config=config.get_quant_config())
 
     def forward(
diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py
index 4e824f6251ca..96104848db8f 100644
--- a/tensorrt_llm/_torch/modules/mlp.py
+++ b/tensorrt_llm/_torch/modules/mlp.py
@@ -4,10 +4,9 @@
 import torch
 from torch import nn
 
-from ..distributed import ParallelConfig, TensorParallelMode
 from ..model_config import ModelConfig
 from ..peft.lora.layer import LoraLayer, LoraModuleType
-from .linear import Linear, WeightMode, WeightsLoadingConfig
+from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
 
 
 class MLP(nn.Module):
@@ -29,40 +28,25 @@ def __init__(self,
         self.activation = activation
 
         config = config or ModelConfig()
-        tp_rank = config.mapping.tp_rank
-        tp_size = config.mapping.tp_size
-        gpus_per_node = config.mapping.gpus_per_node
-        self.up_proj = Linear(
-            self.hidden_size,
-            self.intermediate_size,
-            bias=bias,
-            dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_size=tp_size,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gpus_per_node=gpus_per_node,
-                pipeline_parallel_size=config.mapping.pp_size,
-                parallel_rank=config.mapping.rank),
-            weights_loading_config=WeightsLoadingConfig(
-                weight_mode=WeightMode.VANILLA),
-            quant_config=config.get_quant_config(),
-            skip_create_weights=config.skip_create_weights)
+        self.up_proj = Linear(self.hidden_size,
+                              self.intermediate_size,
+                              bias=bias,
+                              dtype=dtype,
+                              mapping=config.mapping,
+                              tensor_parallel_mode=TensorParallelMode.COLUMN,
+                              weights_loading_config=WeightsLoadingConfig(
+                                  weight_mode=WeightMode.VANILLA),
+                              quant_config=config.get_quant_config(),
+                              skip_create_weights=config.skip_create_weights)
 
-        self.down_proj = Linear(
-            self.intermediate_size,
-            self.hidden_size,
-            bias=bias,
-            dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_size=tp_size,
-                tensor_parallel_mode=TensorParallelMode.ROW,
-                gpus_per_node=gpus_per_node,
-                pipeline_parallel_size=config.mapping.pp_size,
-                parallel_rank=config.mapping.rank),
-            quant_config=config.get_quant_config(),
-            skip_create_weights=config.skip_create_weights)
+        self.down_proj = Linear(self.intermediate_size,
+                                self.hidden_size,
+                                bias=bias,
+                                dtype=dtype,
+                                mapping=config.mapping,
+                                tensor_parallel_mode=TensorParallelMode.ROW,
+                                quant_config=config.get_quant_config(),
+                                skip_create_weights=config.skip_create_weights)
 
         self.up_lora = LoraLayer([LoraModuleType.MLP_H_TO_4H],
                                  [self.intermediate_size])
diff --git a/tensorrt_llm/mapping.py b/tensorrt_llm/mapping.py
index 1cce3f813e24..1a5b44f20e49 100644
--- a/tensorrt_llm/mapping.py
+++ b/tensorrt_llm/mapping.py
@@ -191,9 +191,9 @@ def __init__(
         self.attn_cp_size = attn_cp_size
         self.auto_parallel = auto_parallel
         self.world_size = world_size
+        self.enable_attention_dp = enable_attention_dp
         self.rank = rank
         self.gpus_per_node = gpus_per_node
-        self.enable_attention_dp = enable_attention_dp
         self.pp_groups = []
         self.cp_groups = []
         self.tp_groups = []
@@ -262,10 +262,13 @@ def rank(self):
 
     @rank.setter
     def rank(self, rank: int):
-        if not isinstance(rank, int) or rank < 0 or rank >= self.world_size:
-            raise ValueError(
-                f"Rank should be an integer between 0 and {self.world_size-1}, but got {rank}."
-            )
+        # TODO(qijun): skip check for enable_attention_dp temporarily, will support attention_dp_size
+        if not self.enable_attention_dp:
+            if not isinstance(rank,
+                              int) or rank < 0 and rank >= self.world_size:
+                raise ValueError(
+                    f"Rank should be an integer between 0 and {self.world_size-1}, but got {rank}."
+                )
         self._rank = rank
 
     @property
diff --git a/tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py b/tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py
index dc6e8b3c744e..2d97dcfbb940 100644
--- a/tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py
+++ b/tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py
@@ -28,9 +28,9 @@
 
 import tensorrt_llm
 from tensorrt_llm._torch.compilation.backend import Backend
-from tensorrt_llm._torch.distributed import ParallelConfig, TensorParallelMode
-from tensorrt_llm._torch.modules.linear import Linear
+from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
 from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+from tensorrt_llm.mapping import Mapping
 
 cloudpickle.register_pickle_by_value(sys.modules[__name__])
 MPI.pickle.__init__(
@@ -77,11 +77,10 @@ def row_linear_residual_norm_fusion_forward(
         out_features=hidden_size,
         bias=False,
         dtype=dtype,
-        parallel_config=ParallelConfig(
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_rank=tensor_parallel_rank,
-            tensor_parallel_mode=TensorParallelMode.ROW,
-        ),
+        tensor_parallel_mode=TensorParallelMode.ROW,
+        mapping=Mapping(world_size=tensor_parallel_size,
+                        tp_size=tensor_parallel_size,
+                        rank=tensor_parallel_rank),
     ).cuda()
     norm = RMSNorm(hidden_size=hidden_size, eps=eps, dtype=dtype).cuda()
 
diff --git a/tests/unittest/_torch/multi_gpu/test_deepseek_allreduce.py b/tests/unittest/_torch/multi_gpu/test_deepseek_allreduce.py
index c4c85af19503..589948895b14 100644
--- a/tests/unittest/_torch/multi_gpu/test_deepseek_allreduce.py
+++ b/tests/unittest/_torch/multi_gpu/test_deepseek_allreduce.py
@@ -26,10 +26,10 @@
 
 import tensorrt_llm
 from tensorrt_llm._torch.distributed import (AllReduce, AllReduceFusionOp,
-                                             AllReduceParams, DeepseekAllReduce,
-                                             ParallelConfig, TensorParallelMode)
-from tensorrt_llm._torch.modules.linear import Linear
+                                             AllReduceParams, DeepseekAllReduce)
+from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
 from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+from tensorrt_llm.mapping import Mapping
 
 cloudpickle.register_pickle_by_value(sys.modules[__name__])
 MPI.pickle.__init__(
@@ -87,16 +87,16 @@ def row_linear_residual_norm_fusion_forward(
 
     norm = RMSNorm(hidden_size=hidden_size, eps=eps, dtype=dtype).cuda()
 
-    allreduce = AllReduce(parallel_config=ParallelConfig(
-        tensor_parallel_size=tensor_parallel_size,
-        tensor_parallel_rank=tensor_parallel_rank,
-        tensor_parallel_mode=TensorParallelMode.ROW,
+    allreduce = AllReduce(mapping=Mapping(
+        world_size=tensor_parallel_size,
+        tp_size=tensor_parallel_size,
+        rank=tensor_parallel_rank,
     ), ).cuda()
 
-    deepseek_allreduce = DeepseekAllReduce(parallel_config=ParallelConfig(
-        tensor_parallel_size=tensor_parallel_size,
-        tensor_parallel_rank=tensor_parallel_rank,
-        tensor_parallel_mode=TensorParallelMode.ROW,
+    deepseek_allreduce = DeepseekAllReduce(mapping=Mapping(
+        world_size=tensor_parallel_size,
+        tp_size=tensor_parallel_size,
+        rank=tensor_parallel_rank,
     ), ).cuda()
 
     scale = torch.tensor(1.0, dtype=torch.float32).cuda()
@@ -268,10 +268,10 @@ def moe_residual_norm_fusion_forward(
     norm_weight = torch.randn((hidden_size, ), dtype=dtype, device="cuda")
 
     # Initialize DeepseekAllReduce and AllReduce
-    deepseek_allreduce = DeepseekAllReduce(parallel_config=ParallelConfig(
-        tensor_parallel_size=tensor_parallel_size,
-        tensor_parallel_rank=tensor_parallel_rank,
-        tensor_parallel_mode=TensorParallelMode.ROW,
+    deepseek_allreduce = DeepseekAllReduce(mapping=Mapping(
+        world_size=tensor_parallel_size,
+        tp_size=tensor_parallel_size,
+        rank=tensor_parallel_rank,
     )).cuda()
 
     # Initialize RMSNorm
@@ -283,11 +283,12 @@ def moe_residual_norm_fusion_forward(
         out_features=hidden_size,
         bias=False,
         dtype=dtype,
-        parallel_config=ParallelConfig(
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_rank=tensor_parallel_rank,
-            tensor_parallel_mode=TensorParallelMode.ROW,
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
         ),
+        tensor_parallel_mode=TensorParallelMode.ROW,
     ).cuda()
     l0.load_weights([dict(weight=l0_weight)])
     token_input_chunked = torch.chunk(token_input.clone(),
diff --git a/tests/unittest/_torch/multi_gpu/test_embedding.py b/tests/unittest/_torch/multi_gpu/test_embedding.py
index 3b0b50b3bc3f..68bc74412395 100644
--- a/tests/unittest/_torch/multi_gpu/test_embedding.py
+++ b/tests/unittest/_torch/multi_gpu/test_embedding.py
@@ -10,8 +10,9 @@
 from torch import nn
 
 import tensorrt_llm
-from tensorrt_llm._torch.distributed import ParallelConfig, TensorParallelMode
 from tensorrt_llm._torch.modules.embedding import Embedding, LMHead
+from tensorrt_llm._torch.modules.linear import TensorParallelMode
+from tensorrt_llm.mapping import Mapping
 
 cloudpickle.register_pickle_by_value(sys.modules[__name__])
 MPI.pickle.__init__(
@@ -44,10 +45,12 @@ def column_embedding_forward(x, vocab_size, hidden_size, dtype,
         num_embeddings=vocab_size,
         embedding_dim=hidden_size,
         dtype=dtype,
-        parallel_config=ParallelConfig(
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_rank=tensor_parallel_rank,
-            tensor_parallel_mode=TensorParallelMode.COLUMN),
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
+        ),
+        tensor_parallel_mode=TensorParallelMode.COLUMN,
     )
     embedding.load_weights([dict(weight=weight)])
     embedding.cuda()
@@ -79,11 +82,13 @@ def row_embedding_forward(x, vocab_size, hidden_size, dtype,
         num_embeddings=vocab_size,
         embedding_dim=hidden_size,
         dtype=dtype,
-        parallel_config=ParallelConfig(
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_rank=tensor_parallel_rank,
-            tensor_parallel_mode=TensorParallelMode.ROW,
-            gather_output=True),
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
+        ),
+        tensor_parallel_mode=TensorParallelMode.ROW,
+        gather_output=True,
     )
     embedding.load_weights([dict(weight=weight)])
     embedding.cuda()
@@ -115,11 +120,13 @@ def column_lm_head_forward(x, vocab_size, hidden_size, dtype,
         num_embeddings=vocab_size,
         embedding_dim=hidden_size,
         dtype=dtype,
-        parallel_config=ParallelConfig(
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_rank=tensor_parallel_rank,
-            tensor_parallel_mode=TensorParallelMode.COLUMN,
-            gather_output=True),
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
+        ),
+        tensor_parallel_mode=TensorParallelMode.COLUMN,
+        gather_output=True,
     )
     lm_head.load_weights([dict(weight=weight)])
     lm_head.cuda()
@@ -152,10 +159,12 @@ def row_lm_head_forward(x, vocab_size, hidden_size, dtype, tensor_parallel_size,
         num_embeddings=vocab_size,
         embedding_dim=hidden_size,
         dtype=dtype,
-        parallel_config=ParallelConfig(
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_rank=tensor_parallel_rank,
-            tensor_parallel_mode=TensorParallelMode.ROW),
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
+        ),
+        tensor_parallel_mode=TensorParallelMode.ROW,
     )
     lm_head.load_weights([dict(weight=weight)])
     lm_head.cuda()
diff --git a/tests/unittest/_torch/multi_gpu/test_linear.py b/tests/unittest/_torch/multi_gpu/test_linear.py
index b69c2fa8d01d..40cbf1205840 100644
--- a/tests/unittest/_torch/multi_gpu/test_linear.py
+++ b/tests/unittest/_torch/multi_gpu/test_linear.py
@@ -10,9 +10,9 @@
 from torch import nn
 
 import tensorrt_llm
-from tensorrt_llm._torch.distributed import ParallelConfig, TensorParallelMode
-from tensorrt_llm._torch.modules.linear import Linear
+from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
 from tensorrt_llm.functional import AllReduceFusionOp, AllReduceParams
+from tensorrt_llm.mapping import Mapping
 
 cloudpickle.register_pickle_by_value(sys.modules[__name__])
 MPI.pickle.__init__(
@@ -52,23 +52,27 @@ def mlp_forward(x, hidden_size, dtype, tensor_parallel_size,
         out_features=4 * hidden_size,
         bias=False,
         dtype=dtype,
-        parallel_config=ParallelConfig(
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_rank=tensor_parallel_rank,
-            tensor_parallel_mode=TensorParallelMode.COLUMN,
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
         ),
+        tensor_parallel_mode=TensorParallelMode.COLUMN,
     )
     l0.load_weights([dict(weight=weights[0])])
     l0.cuda()
-    l1 = Linear(in_features=4 * hidden_size,
-                out_features=hidden_size,
-                bias=False,
-                dtype=dtype,
-                parallel_config=ParallelConfig(
-                    tensor_parallel_size=tensor_parallel_size,
-                    tensor_parallel_rank=tensor_parallel_rank,
-                    tensor_parallel_mode=TensorParallelMode.ROW,
-                ))
+    l1 = Linear(
+        in_features=4 * hidden_size,
+        out_features=hidden_size,
+        bias=False,
+        dtype=dtype,
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
+        ),
+        tensor_parallel_mode=TensorParallelMode.ROW,
+    )
     l1.load_weights([dict(weight=weights[1])])
     l1.cuda()
 
@@ -102,15 +106,19 @@ def column_linear_forward(x, hidden_size, dtype, tensor_parallel_size,
                           tensor_parallel_rank, weights):
 
     x = x.cuda()
-    l0 = Linear(in_features=hidden_size,
-                out_features=hidden_size,
-                bias=False,
-                dtype=dtype,
-                parallel_config=ParallelConfig(
-                    tensor_parallel_size=tensor_parallel_size,
-                    tensor_parallel_rank=tensor_parallel_rank,
-                    tensor_parallel_mode=TensorParallelMode.COLUMN,
-                    gather_output=True))
+    l0 = Linear(
+        in_features=hidden_size,
+        out_features=hidden_size,
+        bias=False,
+        dtype=dtype,
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
+        ),
+        tensor_parallel_mode=TensorParallelMode.COLUMN,
+        gather_output=True,
+    )
     l0.load_weights([dict(weight=weights[0])])
     l0.cuda()
 
@@ -137,15 +145,18 @@ def row_linear_forward(x, hidden_size, dtype, tensor_parallel_size,
                        tensor_parallel_rank, weights):
 
     x = x.cuda()
-    l0 = Linear(in_features=hidden_size,
-                out_features=hidden_size,
-                bias=False,
-                dtype=dtype,
-                parallel_config=ParallelConfig(
-                    tensor_parallel_size=tensor_parallel_size,
-                    tensor_parallel_rank=tensor_parallel_rank,
-                    tensor_parallel_mode=TensorParallelMode.ROW,
-                ))
+    l0 = Linear(
+        in_features=hidden_size,
+        out_features=hidden_size,
+        bias=False,
+        dtype=dtype,
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
+        ),
+        tensor_parallel_mode=TensorParallelMode.ROW,
+    )
     l0.load_weights([dict(weight=weights[0])])
     l0.cuda()
 
@@ -183,15 +194,18 @@ def row_linear_norm_fusion_forward(x, hidden_size, dtype, tensor_parallel_size,
         eps=eps,
     )
 
-    l0 = Linear(in_features=hidden_size,
-                out_features=hidden_size,
-                bias=False,
-                dtype=dtype,
-                parallel_config=ParallelConfig(
-                    tensor_parallel_size=tensor_parallel_size,
-                    tensor_parallel_rank=tensor_parallel_rank,
-                    tensor_parallel_mode=TensorParallelMode.ROW,
-                ))
+    l0 = Linear(
+        in_features=hidden_size,
+        out_features=hidden_size,
+        bias=False,
+        dtype=dtype,
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
+        ),
+        tensor_parallel_mode=TensorParallelMode.ROW,
+    )
     l0.load_weights([dict(weight=weights[0])])
     l0.cuda()
 
diff --git a/tests/unittest/_torch/multi_gpu/test_user_buffers.py b/tests/unittest/_torch/multi_gpu/test_user_buffers.py
index be48acf5c973..c51f80a556a8 100644
--- a/tests/unittest/_torch/multi_gpu/test_user_buffers.py
+++ b/tests/unittest/_torch/multi_gpu/test_user_buffers.py
@@ -16,10 +16,10 @@
 from tensorrt_llm._torch.compilation.backend import Backend
 from tensorrt_llm._torch.distributed import (AllReduce, AllReduceFusionOp,
                                              AllReduceParams, AllReduceStrategy,
-                                             ParallelConfig, TensorParallelMode,
                                              userbuffers_allreduce_finalize)
-from tensorrt_llm._torch.modules.linear import Linear
+from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
 from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.models.modeling_utils import QuantConfig
 from tensorrt_llm.quantization import QuantAlgo
 
@@ -123,11 +123,12 @@ def run_single_rank_ar_rms_norm(tensor_parallel_size, a, b, c, gamma):
 
         ub0_tensor = create_userbuffers_tensor(c.size(), a.dtype)
         hidden = torch.matmul(a_local, b_local, out=ub0_tensor)
-        parallel_config = ParallelConfig(
-            tensor_parallel_rank=rank,
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_mode=TensorParallelMode.COLUMN)
-        ar = AllReduce(parallel_config, strategy=AllReduceStrategy.UB)
+        mapping = Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=rank,
+        )
+        ar = AllReduce(mapping, strategy=AllReduceStrategy.UB)
         ar_params = AllReduceParams(
             strategy=AllReduceStrategy.UB,
             fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
@@ -214,11 +215,12 @@ def run_single_rank_ar_rms_norm_fp8(tensor_parallel_size, a, b, c, gamma,
 
         ub0_tensor = create_userbuffers_tensor(c.size(), a.dtype)
         hidden = torch.matmul(a_local, b_local, out=ub0_tensor)
-        parallel_config = ParallelConfig(
-            tensor_parallel_rank=rank,
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_mode=TensorParallelMode.COLUMN)
-        ar = AllReduce(parallel_config, strategy=AllReduceStrategy.UB)
+        mapping = Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=rank,
+        )
+        ar = AllReduce(mapping, strategy=AllReduceStrategy.UB)
         ar_params = AllReduceParams(
             strategy=AllReduceStrategy.UB,
             fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8,
@@ -318,55 +320,45 @@ def __init__(self, tp_size, rank, hidden_size, dtype, eps, l0_weight,
         quant_config.layer_quant_mode
         self.rank = rank
         self.tp_size = tp_size
+        mapping = Mapping(
+            world_size=tp_size,
+            tp_size=tp_size,
+            rank=rank,
+        )
         self.l0 = Linear(in_features=hidden_size,
                          out_features=hidden_size,
                          bias=False,
                          dtype=dtype,
-                         parallel_config=ParallelConfig(
-                             tensor_parallel_size=tp_size,
-                             tensor_parallel_rank=rank,
-                             tensor_parallel_mode=TensorParallelMode.ROW,
-                         ),
+                         mapping=mapping,
+                         tensor_parallel_mode=TensorParallelMode.ROW,
                          quant_config=quant_config).cuda()
         self.l1 = Linear(in_features=hidden_size,
                          out_features=hidden_size,
                          bias=False,
                          dtype=dtype,
-                         parallel_config=ParallelConfig(
-                             tensor_parallel_size=tp_size,
-                             tensor_parallel_rank=rank,
-                             tensor_parallel_mode=TensorParallelMode.COLUMN,
-                         ),
+                         mapping=mapping,
+                         tensor_parallel_mode=TensorParallelMode.COLUMN,
                          quant_config=quant_config).cuda()
         self.l2 = Linear(in_features=hidden_size,
                          out_features=hidden_size,
                          bias=False,
                          dtype=dtype,
-                         parallel_config=ParallelConfig(
-                             tensor_parallel_size=tp_size,
-                             tensor_parallel_rank=rank,
-                             tensor_parallel_mode=TensorParallelMode.ROW,
-                         ),
+                         mapping=mapping,
+                         tensor_parallel_mode=TensorParallelMode.ROW,
                          quant_config=quant_config).cuda()
         self.l3 = Linear(in_features=hidden_size,
                          out_features=hidden_size,
                          bias=False,
                          dtype=dtype,
-                         parallel_config=ParallelConfig(
-                             tensor_parallel_size=tp_size,
-                             tensor_parallel_rank=rank,
-                             tensor_parallel_mode=TensorParallelMode.COLUMN,
-                         ),
+                         mapping=mapping,
+                         tensor_parallel_mode=TensorParallelMode.COLUMN,
                          quant_config=quant_config).cuda()
         self.l4 = Linear(in_features=hidden_size,
                          out_features=hidden_size,
                          bias=False,
                          dtype=dtype,
-                         parallel_config=ParallelConfig(
-                             tensor_parallel_size=tp_size,
-                             tensor_parallel_rank=rank,
-                             tensor_parallel_mode=TensorParallelMode.ROW,
-                         ),
+                         mapping=mapping,
+                         tensor_parallel_mode=TensorParallelMode.ROW,
                          quant_config=quant_config).cuda()
         self.norm0 = RMSNorm(hidden_size=hidden_size, eps=eps,
                              dtype=dtype).cuda()
@@ -608,11 +600,12 @@ def run_single_rank_ar_rms_norm_fp4(tensor_parallel_size, a, b, c, gamma):
 
         ub0_tensor = create_userbuffers_tensor(c.size(), a.dtype)
         hidden = torch.matmul(a_local, b_local, out=ub0_tensor)
-        parallel_config = ParallelConfig(
-            tensor_parallel_rank=rank,
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_mode=TensorParallelMode.COLUMN)
-        ar = AllReduce(parallel_config, strategy=AllReduceStrategy.UB)
+        mapping = Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=rank,
+        )
+        ar = AllReduce(mapping, strategy=AllReduceStrategy.UB)
         ar_params = AllReduceParams(
             strategy=AllReduceStrategy.UB,
             fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4,
@@ -694,18 +687,14 @@ def __init__(self, tp_size, rank, hidden_size, dtype, eps, norm0_gamma,
         self.rank = rank
         self.hidden_size = hidden_size
         self.dtype = dtype
-        self.ar_0 = AllReduce(
-            ParallelConfig(tensor_parallel_size=tp_size,
-                           tensor_parallel_rank=rank,
-                           tensor_parallel_mode=TensorParallelMode.ROW)).cuda()
-        self.ar_1 = AllReduce(
-            ParallelConfig(tensor_parallel_size=tp_size,
-                           tensor_parallel_rank=rank,
-                           tensor_parallel_mode=TensorParallelMode.ROW)).cuda()
-        self.ar_2 = AllReduce(
-            ParallelConfig(tensor_parallel_size=tp_size,
-                           tensor_parallel_rank=rank,
-                           tensor_parallel_mode=TensorParallelMode.ROW)).cuda()
+        mapping = Mapping(
+            world_size=tp_size,
+            tp_size=tp_size,
+            rank=rank,
+        )
+        self.ar_0 = AllReduce(mapping).cuda()
+        self.ar_1 = AllReduce(mapping).cuda()
+        self.ar_2 = AllReduce(mapping).cuda()
         self.norm0 = RMSNorm(hidden_size=hidden_size, eps=eps,
                              dtype=dtype).cuda()
         self.norm1 = RMSNorm(hidden_size=hidden_size, eps=eps,


Reviews:
user: login: hlu1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T06:40:34Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T13:29:59Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T13:30:09Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T13:30:26Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T13:30:39Z
user: login: SimengLiu-nv
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-09T15:36:01Z
user: login: hlu1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T17:16:00Z
user: login: hlu1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T17:24:03Z
user: login: amukkara
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T18:32:51Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T06:22:09Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T06:25:40Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T06:46:02Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T06:46:14Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T06:46:26Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T06:46:37Z
user: login: QiJune
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T06:46:49Z
user: login: amukkara
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T14:28:44Z
user: login: hlu1
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-11T22:10:19Z

Review comments:
diff_hunk: @@ -75,13 +74,12 @@ def __init__(self,
             self.hidden_size,
             bias=bias,
             dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_size=tp_size,
-                tensor_parallel_mode=TensorParallelMode.ROW,
-                gpus_per_node=gpus_per_node,
-                pipeline_parallel_size=pp_size,
-                parallel_rank=config.mapping.rank),
+            mapping=Mapping(world_size=tp_size * pp_size,
path: tensorrt_llm/_torch/modules/gated_mlp.py
user: login: hlu1
body: This mapping is the same as the one on line 60. You can create one Mapping object and use it in both places.
created_at: 2025-04-09T06:33:47Z
author_association: COLLABORATOR
diff_hunk: @@ -79,12 +81,15 @@ def __init__(
         # tensor parallel
         config = config or ModelConfig()
         tp_size = config.mapping.tp_size
-        tp_rank = config.mapping.tp_rank
+        pp_size = config.mapping.pp_size
+        rank = config.mapping.rank
         gpus_per_node = config.mapping.gpus_per_node
         if config.mapping.enable_attention_dp:
             tp_size = 1
-            tp_rank = 0
+            pp_size = 1
path: tensorrt_llm/_torch/modules/attention.py
user: login: hlu1
body: Why is `pp_size` set to 1 here? You can have pp > 1 with attn_dp.
created_at: 2025-04-09T06:36:44Z
author_association: COLLABORATOR
diff_hunk: @@ -119,13 +125,14 @@ def __init__(
             self.hidden_size,
             bias=self.dense_bias,
             dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_size=tp_size,
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_mode=TensorParallelMode.ROW,
+            mapping=Mapping(
path: tensorrt_llm/_torch/modules/attention.py
user: login: hlu1
body: Use the same Mapping object created on line 110
created_at: 2025-04-09T06:37:35Z
author_association: COLLABORATOR
diff_hunk: @@ -279,27 +286,20 @@ def __init__(
         # tensor parallel
         config = config or ModelConfig()
         tp_size = config.mapping.tp_size
-        tp_rank = config.mapping.tp_rank
+        pp_size = config.mapping.pp_size
+        rank = config.mapping.rank
         gpus_per_node = config.mapping.gpus_per_node
         if config.mapping.enable_attention_dp:
             tp_size = 1
-            tp_rank = 0
-
-        row_parallel_config = ParallelConfig(
-            tensor_parallel_rank=tp_rank,
-            tensor_parallel_size=tp_size,
-            tensor_parallel_mode=TensorParallelMode.ROW,
-            gpus_per_node=gpus_per_node,
-            pipeline_parallel_size=config.mapping.pp_size,
-            parallel_rank=config.mapping.rank,
-        )
-        col_parallel_config = ParallelConfig(
-            tensor_parallel_rank=tp_rank,
-            tensor_parallel_size=tp_size,
-            tensor_parallel_mode=TensorParallelMode.COLUMN,
+            pp_size = 1
path: tensorrt_llm/_torch/modules/attention.py
user: login: hlu1
body: Ditto
created_at: 2025-04-09T06:38:08Z
author_association: COLLABORATOR
diff_hunk: @@ -75,13 +74,12 @@ def __init__(self,
             self.hidden_size,
             bias=bias,
             dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_size=tp_size,
-                tensor_parallel_mode=TensorParallelMode.ROW,
-                gpus_per_node=gpus_per_node,
-                pipeline_parallel_size=pp_size,
-                parallel_rank=config.mapping.rank),
+            mapping=Mapping(world_size=tp_size * pp_size,
path: tensorrt_llm/_torch/modules/gated_mlp.py
user: login: QiJune
body: Done
created_at: 2025-04-09T13:29:59Z
author_association: COLLABORATOR
diff_hunk: @@ -119,13 +125,14 @@ def __init__(
             self.hidden_size,
             bias=self.dense_bias,
             dtype=dtype,
-            parallel_config=ParallelConfig(
-                tensor_parallel_size=tp_size,
-                tensor_parallel_rank=tp_rank,
-                tensor_parallel_mode=TensorParallelMode.ROW,
+            mapping=Mapping(
path: tensorrt_llm/_torch/modules/attention.py
user: login: QiJune
body: Done
created_at: 2025-04-09T13:30:09Z
author_association: COLLABORATOR
diff_hunk: @@ -79,12 +81,15 @@ def __init__(
         # tensor parallel
         config = config or ModelConfig()
         tp_size = config.mapping.tp_size
-        tp_rank = config.mapping.tp_rank
+        pp_size = config.mapping.pp_size
+        rank = config.mapping.rank
         gpus_per_node = config.mapping.gpus_per_node
         if config.mapping.enable_attention_dp:
             tp_size = 1
-            tp_rank = 0
+            pp_size = 1
path: tensorrt_llm/_torch/modules/attention.py
user: login: QiJune
body: Done, revert this change.
created_at: 2025-04-09T13:30:25Z
author_association: COLLABORATOR
diff_hunk: @@ -279,27 +286,20 @@ def __init__(
         # tensor parallel
         config = config or ModelConfig()
         tp_size = config.mapping.tp_size
-        tp_rank = config.mapping.tp_rank
+        pp_size = config.mapping.pp_size
+        rank = config.mapping.rank
         gpus_per_node = config.mapping.gpus_per_node
         if config.mapping.enable_attention_dp:
             tp_size = 1
-            tp_rank = 0
-
-        row_parallel_config = ParallelConfig(
-            tensor_parallel_rank=tp_rank,
-            tensor_parallel_size=tp_size,
-            tensor_parallel_mode=TensorParallelMode.ROW,
-            gpus_per_node=gpus_per_node,
-            pipeline_parallel_size=config.mapping.pp_size,
-            parallel_rank=config.mapping.rank,
-        )
-        col_parallel_config = ParallelConfig(
-            tensor_parallel_rank=tp_rank,
-            tensor_parallel_size=tp_size,
-            tensor_parallel_mode=TensorParallelMode.COLUMN,
+            pp_size = 1
path: tensorrt_llm/_torch/modules/attention.py
user: login: QiJune
body: Done, revert this change
created_at: 2025-04-09T13:30:39Z
author_association: COLLABORATOR
diff_hunk: @@ -262,10 +262,12 @@ def rank(self):
 
     @rank.setter
     def rank(self, rank: int):
-        if not isinstance(rank, int) or rank < 0 or rank >= self.world_size:
-            raise ValueError(
-                f"Rank should be an integer between 0 and {self.world_size-1}, but got {rank}."
-            )
+        if not isinstance(rank, int) or rank < 0:
+            # TODO(qijun): support attention_dp_size
+            if not self.enable_attention_dp and rank >= self.world_size:
+                raise ValueError(
+                    f"Rank should be an integer between 0 and {self.world_size-1}, but got {rank}."
+                )
path: tensorrt_llm/mapping.py
user: login: amukkara
body: this change seems erroneous to me. can you split the check for attention dp on/off cleanly:

```
if enable_attention_dp:
    if ..:
       raise()
else:
   if ..:
       raise()
```
created_at: 2025-04-09T14:48:48Z
author_association: COLLABORATOR
diff_hunk: @@ -242,34 +183,26 @@ def reducescatter(input: torch.Tensor,
     output = torch.movedim(output, 0, scatter_dim)
     input_shape = input.size()
     output = output.reshape(input_shape[:scatter_dim] +
-                            (input_shape[scatter_dim] //
-                             parallel_config.tensor_parallel_size, ) +
+                            (input_shape[scatter_dim] // mapping.tp_size, ) +
                             input_shape[scatter_dim + 1:])
     return output
 
 
 class AllReduce(nn.Module):
 
     def __init__(self,
-                 parallel_config: ParallelConfig,
+                 mapping: Mapping,
                  strategy: AllReduceStrategy = AllReduceStrategy.AUTO):
         super().__init__()
 
-        self.parallel_config = parallel_config
-        self.tp_size = self.parallel_config.tensor_parallel_size
-        self.rank = self.parallel_config.parallel_rank
-        self.gpus_per_node = self.parallel_config.gpus_per_node
+        self.mapping = mapping
+        self.tp_size = self.mapping.tp_size
+        self.rank = self.mapping.rank
+        self.gpus_per_node = self.mapping.gpus_per_node
path: tensorrt_llm/_torch/distributed.py
user: login: amukkara
body: why make a copy (for `tp_size`, `rank`, `gpus_per_node`) instead of directly accessing from `self.mapping`? 
created_at: 2025-04-09T14:52:22Z
author_association: COLLABORATOR
diff_hunk: @@ -281,30 +214,23 @@ def forward(
     ) -> torch.Tensor:
         output = allreduce(input,
                            self.workspace,
-                           self.parallel_config,
+                           self.mapping,
                            all_reduce_params=all_reduce_params,
                            strategy=self.strategy)
         return output
 
 
 class DeepseekAllReduce(nn.Module):
 
-    def __init__(self, parallel_config: ParallelConfig):
+    def __init__(self, mapping: Mapping):
         super().__init__()
-        self.parallel_config = parallel_config
-        self.tp_size = self.parallel_config.tensor_parallel_size
-        self.tp_rank = self.parallel_config.tensor_parallel_rank
-        self.gpus_per_node = self.parallel_config.gpus_per_node
-        self.rank = self.parallel_config.parallel_rank
+        self.mapping = mapping
+        self.tp_size = self.mapping.tp_size
+        self.tp_rank = self.mapping.tp_rank
+        self.gpus_per_node = self.mapping.gpus_per_node
+        self.rank = self.mapping.rank
path: tensorrt_llm/_torch/distributed.py
user: login: amukkara
body: same comment on making copies.
created_at: 2025-04-09T14:53:16Z
author_association: COLLABORATOR
diff_hunk: @@ -166,9 +165,10 @@ def __init__(
         super().__init__()
 
         config = config or ModelConfig()
+        self.mapping = config.mapping
         tp_rank = config.mapping.tp_rank
         tp_size = config.mapping.tp_size
-        gpus_per_node = config.mapping.gpus_per_node
+        config.mapping.gpus_per_node
path: tensorrt_llm/_torch/modules/mamba.py
user: login: amukkara
body: delete unused line.
created_at: 2025-04-09T15:04:22Z
author_association: COLLABORATOR
diff_hunk: @@ -191,9 +191,9 @@ def __init__(
         self.attn_cp_size = attn_cp_size
         self.auto_parallel = auto_parallel
         self.world_size = world_size
+        self.enable_attention_dp = enable_attention_dp
         self.rank = rank
         self.gpus_per_node = gpus_per_node
-        self.enable_attention_dp = enable_attention_dp
path: tensorrt_llm/mapping.py
user: login: amukkara
body: why move `enable_attention_dp` before `rank`?
created_at: 2025-04-09T15:05:34Z
author_association: COLLABORATOR
diff_hunk: @@ -134,7 +143,9 @@ def __init__(self,
                  out_features: int,
                  bias: bool = True,
                  dtype: torch.dtype = None,
-                 parallel_config: Optional[ParallelConfig] = None,
+                 mapping: Optional[Mapping] = None,
+                 gather_output: bool = False,
+                 tensor_parallel_mode: Optional[TensorParallelMode] = None,
path: tensorrt_llm/_torch/modules/linear.py
user: login: amukkara
body: since only `mapping` and `tensor_parallel_mode` are used in most cases, does it make sense to change the order
to `mapping`, `tensor_parallel_mode`, `gather_output`? We should keep class definition and instantiation order same whenever possible.

Please check the order is consistent across all definitions in `_torch/modules` (e.g `_torch/modules/embedding.py`) and their instantiations. 
created_at: 2025-04-09T15:13:56Z
author_association: COLLABORATOR
diff_hunk: @@ -283,10 +283,11 @@ def moe_residual_norm_fusion_forward(
         out_features=hidden_size,
         bias=False,
         dtype=dtype,
-        parallel_config=ParallelConfig(
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_rank=tensor_parallel_rank,
-            tensor_parallel_mode=TensorParallelMode.ROW,
+        tensor_parallel_mode=TensorParallelMode.ROW,
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
path: tests/unittest/_torch/multi_gpu/test_deepseek_allreduce.py
user: login: amukkara
body: nit: `tensor_parallel_mode` after `mapping`
created_at: 2025-04-09T15:16:56Z
author_association: COLLABORATOR
diff_hunk: @@ -279,27 +277,17 @@ def __init__(
         # tensor parallel
         config = config or ModelConfig()
         tp_size = config.mapping.tp_size
-        tp_rank = config.mapping.tp_rank
-        gpus_per_node = config.mapping.gpus_per_node
+        pp_size = config.mapping.pp_size
         if config.mapping.enable_attention_dp:
             tp_size = 1
-            tp_rank = 0
path: tensorrt_llm/_torch/modules/attention.py
user: login: hlu1
body: Is the rank in Mapping correct for attn_dp = True? Mapping doesn't modify rank based on enable_attention_dp.
created_at: 2025-04-09T17:16:00Z
author_association: COLLABORATOR
diff_hunk: @@ -279,27 +277,17 @@ def __init__(
         # tensor parallel
         config = config or ModelConfig()
         tp_size = config.mapping.tp_size
-        tp_rank = config.mapping.tp_rank
-        gpus_per_node = config.mapping.gpus_per_node
+        pp_size = config.mapping.pp_size
         if config.mapping.enable_attention_dp:
             tp_size = 1
-            tp_rank = 0
path: tensorrt_llm/_torch/modules/attention.py
user: login: hlu1
body: As a follow-up, should we move the attention_dp logics to inside Mapping? It currently accepts a `enable_attention_dp` flag but does nothing.

https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/mapping.py#L129
created_at: 2025-04-09T17:24:03Z
author_association: COLLABORATOR
diff_hunk: @@ -279,27 +277,17 @@ def __init__(
         # tensor parallel
         config = config or ModelConfig()
         tp_size = config.mapping.tp_size
-        tp_rank = config.mapping.tp_rank
-        gpus_per_node = config.mapping.gpus_per_node
+        pp_size = config.mapping.pp_size
         if config.mapping.enable_attention_dp:
             tp_size = 1
-            tp_rank = 0
path: tensorrt_llm/_torch/modules/attention.py
user: login: QiJune
body: In our current implementation, we only have a bool `enable_attention_dp` flag.  Actually, `enable_attention_dp=True` means `attention_tp_size=1`, and `attention_dp_size=tp_size`. Thus, `world_size` should be `pp_size * attention_tp_size(1) * attention_dp_size(tp_size)`, which keeps unchanged. In the future, we should remove `enable_attention_dp` field in Mapping, and replace it with  `attention_dp_size`  and `attention_tp_size` fields.


> Is the rank in Mapping correct for attn_dp = True? Mapping doesn't modify rank based on enable_attention_dp.

The rank is correct, but world_size is not correct now, since we don't have a `attention_dp_size` yet.


> should we move the attention_dp logics to inside Mapping?

No,  `attention_dp_size` and `attention_tp_size` still need to be decided first, and set to Mapping.

created_at: 2025-04-10T06:22:08Z
author_association: COLLABORATOR
diff_hunk: @@ -191,9 +191,9 @@ def __init__(
         self.attn_cp_size = attn_cp_size
         self.auto_parallel = auto_parallel
         self.world_size = world_size
+        self.enable_attention_dp = enable_attention_dp
         self.rank = rank
         self.gpus_per_node = gpus_per_node
-        self.enable_attention_dp = enable_attention_dp
path: tensorrt_llm/mapping.py
user: login: QiJune
body: Please see here: https://github.com/NVIDIA/TensorRT-LLM/pull/3370#discussion_r2036589711

We don't have an int `attention_dp_size` now, it's a bool `enable_attention_dp`. So, the `world_size` is not correct when `enable_attention_dp=True`. We have some check logic in

```
    @rank.setter
    def rank(self, rank: int):
         xxx
```

So, I move `enable_attention_dp` before `rank` to ensure `enable_attention_dp` can be passed into the check logic.

This will be refactored once we support `attention_dp_size` in the future.
created_at: 2025-04-10T06:25:40Z
author_association: COLLABORATOR
diff_hunk: @@ -262,10 +262,12 @@ def rank(self):
 
     @rank.setter
     def rank(self, rank: int):
-        if not isinstance(rank, int) or rank < 0 or rank >= self.world_size:
-            raise ValueError(
-                f"Rank should be an integer between 0 and {self.world_size-1}, but got {rank}."
-            )
+        if not isinstance(rank, int) or rank < 0:
+            # TODO(qijun): support attention_dp_size
+            if not self.enable_attention_dp and rank >= self.world_size:
+                raise ValueError(
+                    f"Rank should be an integer between 0 and {self.world_size-1}, but got {rank}."
+                )
path: tensorrt_llm/mapping.py
user: login: QiJune
body: Done
created_at: 2025-04-10T06:46:02Z
author_association: COLLABORATOR
diff_hunk: @@ -242,34 +183,26 @@ def reducescatter(input: torch.Tensor,
     output = torch.movedim(output, 0, scatter_dim)
     input_shape = input.size()
     output = output.reshape(input_shape[:scatter_dim] +
-                            (input_shape[scatter_dim] //
-                             parallel_config.tensor_parallel_size, ) +
+                            (input_shape[scatter_dim] // mapping.tp_size, ) +
                             input_shape[scatter_dim + 1:])
     return output
 
 
 class AllReduce(nn.Module):
 
     def __init__(self,
-                 parallel_config: ParallelConfig,
+                 mapping: Mapping,
                  strategy: AllReduceStrategy = AllReduceStrategy.AUTO):
         super().__init__()
 
-        self.parallel_config = parallel_config
-        self.tp_size = self.parallel_config.tensor_parallel_size
-        self.rank = self.parallel_config.parallel_rank
-        self.gpus_per_node = self.parallel_config.gpus_per_node
+        self.mapping = mapping
+        self.tp_size = self.mapping.tp_size
+        self.rank = self.mapping.rank
+        self.gpus_per_node = self.mapping.gpus_per_node
path: tensorrt_llm/_torch/distributed.py
user: login: QiJune
body: removed
created_at: 2025-04-10T06:46:14Z
author_association: COLLABORATOR
diff_hunk: @@ -166,9 +165,10 @@ def __init__(
         super().__init__()
 
         config = config or ModelConfig()
+        self.mapping = config.mapping
         tp_rank = config.mapping.tp_rank
         tp_size = config.mapping.tp_size
-        gpus_per_node = config.mapping.gpus_per_node
+        config.mapping.gpus_per_node
path: tensorrt_llm/_torch/modules/mamba.py
user: login: QiJune
body: Done
created_at: 2025-04-10T06:46:26Z
author_association: COLLABORATOR
diff_hunk: @@ -134,7 +143,9 @@ def __init__(self,
                  out_features: int,
                  bias: bool = True,
                  dtype: torch.dtype = None,
-                 parallel_config: Optional[ParallelConfig] = None,
+                 mapping: Optional[Mapping] = None,
+                 gather_output: bool = False,
+                 tensor_parallel_mode: Optional[TensorParallelMode] = None,
path: tensorrt_llm/_torch/modules/linear.py
user: login: QiJune
body: Done
created_at: 2025-04-10T06:46:36Z
author_association: COLLABORATOR
diff_hunk: @@ -283,10 +283,11 @@ def moe_residual_norm_fusion_forward(
         out_features=hidden_size,
         bias=False,
         dtype=dtype,
-        parallel_config=ParallelConfig(
-            tensor_parallel_size=tensor_parallel_size,
-            tensor_parallel_rank=tensor_parallel_rank,
-            tensor_parallel_mode=TensorParallelMode.ROW,
+        tensor_parallel_mode=TensorParallelMode.ROW,
+        mapping=Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
path: tests/unittest/_torch/multi_gpu/test_deepseek_allreduce.py
user: login: QiJune
body: Done
created_at: 2025-04-10T06:46:49Z
author_association: COLLABORATOR
diff_hunk: @@ -191,9 +191,9 @@ def __init__(
         self.attn_cp_size = attn_cp_size
         self.auto_parallel = auto_parallel
         self.world_size = world_size
+        self.enable_attention_dp = enable_attention_dp
         self.rank = rank
         self.gpus_per_node = gpus_per_node
-        self.enable_attention_dp = enable_attention_dp
path: tensorrt_llm/mapping.py
user: login: amukkara
body: makes sense. thanks.
created_at: 2025-04-10T14:28:44Z
author_association: COLLABORATOR

PR comments:
user: login: QiJune
created_at: 2025-04-08T11:25:56Z
body: /bot run --only-multi-gpu-test
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-08T11:31:21Z
body: [PR_Github #1453](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1453/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-08T13:18:36Z
body: [PR_Github #1453](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1453/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1087 `(Partly Tested)`](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1087/) completed with status: 'FAILURE'

author_association: COLLABORATOR
user: login: QiJune
created_at: 2025-04-08T13:21:56Z
body: /bot --help
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-04-08T13:22:07Z
body: ## GitHub Bot Help

`/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...`

Provide a user friendly way for developers to interact with a Jenkins server.

Run `/bot [-h|--help]` to print this help message.

See details below for each supported subcommand.

<details>

`run  [--disable-fail-fast --skip-test --stage-list "A10-1, xxx" --gpu-type "A30, H100_PCIe" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-[Post-Merge]-1, xxx"]`

Launch build/test pipelines. All previously running jobs will be killed.

`--disable-fail-fast ` *(OPTIONAL)* : Disable fail fast on build/tests/infra failures.

`--skip-test ` *(OPTIONAL)* : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does **NOT** update GitHub check status.

`--stage-list "A10-1, xxx"` *(OPTIONAL)* : Only run the specified test stages. Examples: "A10-1, xxx". Note: Does **NOT** update GitHub check status.

`--gpu-type "A30, H100_PCIe"` *(OPTIONAL)* : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does **NOT** update GitHub check status.

`--only-multi-gpu-test ` *(OPTIONAL)* : Only run the multi-GPU tests. Note: Does **NOT** update GitHub check status.

`--disable-multi-gpu-test ` *(OPTIONAL)* : Disable the multi-GPU tests. Note: Does **NOT** update GitHub check status.

`--add-multi-gpu-test ` *(OPTIONAL)* : Force run the multi-GPU tests. Will also run L0 pre-merge pipeline.

`--post-merge ` *(OPTIONAL)* : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

`--extra-stage "H100_PCIe-[Post-Merge]-1, xxx"` *(OPTIONAL)* : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-[Post-Merge]-1, xxx".

### kill

`kill  `

Kill all running builds associated with pull request.

### skip

`skip --comment COMMENT `

Skip testing for latest commit on pull request. `--comment "Reason for skipping build/test"` is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

### reuse-pipeline

`reuse-pipeline `

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

</details>
author_association: NONE
user: login: QiJune
created_at: 2025-04-08T13:22:27Z
body: /bot run --add-multi-gpu-test 
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-08T13:28:01Z
body: [PR_Github #1465](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1465/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-08T14:29:54Z
body: [PR_Github #1465](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1465/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1095](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1095/) completed with status: 'FAILURE'

author_association: COLLABORATOR
user: login: QiJune
created_at: 2025-04-09T00:31:01Z
body: /bot run --add-multi-gpu-test
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-09T00:36:22Z
body: [PR_Github #1513](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1513/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-09T01:27:37Z
body: [PR_Github #1513](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1513/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1130](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1130/) completed with status: 'FAILURE'

author_association: COLLABORATOR
user: login: QiJune
created_at: 2025-04-09T04:33:02Z
body: /bot run --add-multi-gpu-test
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-09T04:38:13Z
body: [PR_Github #1547](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1547/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-09T06:06:49Z
body: [PR_Github #1547](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1547/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1157](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1157/) completed with status: 'FAILURE'

author_association: COLLABORATOR
user: login: QiJune
created_at: 2025-04-09T06:10:46Z
body: /bot run --add-multi-gpu-test --disable-fail-fast
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-09T06:16:01Z
body: [PR_Github #1573](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1573/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-09T10:35:12Z
body: [PR_Github #1573](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1573/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1176](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1176/) completed with status: 'FAILURE'

author_association: COLLABORATOR
user: login: QiJune
created_at: 2025-04-09T13:33:33Z
body: /bot run --add-multi-gpu-test
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-09T13:38:46Z
body: [PR_Github #1620](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1620/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-09T19:03:22Z
body: [PR_Github #1620](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1620/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1213](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1213/) completed with status: 'FAILURE'

author_association: COLLABORATOR
user: login: QiJune
created_at: 2025-04-10T06:47:09Z
body: /bot run
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-10T06:52:27Z
body: [PR_Github #1721](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1721/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: QiJune
created_at: 2025-04-10T09:15:02Z
body: /bot run
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-10T09:17:29Z
body: [PR_Github #1721](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1721/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1285](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1285/) completed with status: 'FAILURE'

author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-10T09:20:13Z
body: [PR_Github #1751](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1751/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: QiJune
created_at: 2025-04-10T14:08:52Z
body: /bot run
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-10T14:14:10Z
body: [PR_Github #1787](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1787/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-10T14:16:41Z
body: [PR_Github #1751](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1751/) [ run ] completed with state `ABORTED`

author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-10T17:35:53Z
body: [PR_Github #1787](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1787/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1324](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1324/) completed with status: 'FAILURE'

author_association: COLLABORATOR
user: login: QiJune
created_at: 2025-04-11T05:21:25Z
body: /bot run
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-11T05:26:46Z
body: [PR_Github #1867](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1867/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-11T13:13:43Z
body: [PR_Github #1867](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/1867/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1384](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1384/) completed with status: 'SUCCESS'

author_association: COLLABORATOR