diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index be6c20ec92e1..c49064bfe38c 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -831,6 +831,14 @@ def compute_routed_output(self, hidden_states, hidden_states_fp4, all_rank_num_tokens, do_finalize): # max-throughput use_dp_padding = False + # Add DP padding on SM120 for context comm performance + # TODO: Move this model-agonostic part to MoE + if self.use_dp and self.mapping.tp_size > 1 and get_sm_version() == 120: + use_dp_padding = True + hidden_states = torch.nn.functional.pad( + hidden_states, + (0, 0, 0, max(all_rank_num_tokens) - hidden_states.shape[0])) + router_logits = self.gate(hidden_states) routed_output = self.experts(