From 6df5c9aa75519c7d69e73f80b9a66d80a7165c88 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Tue, 13 Oct 2026 08:23:51 -0700 Subject: [PATCH 1/2] [https://nvbugs/6919428][fix] Keep DeepSeek-V4 TP decode all-reduces off NCCL DeepSeek-V4 TP decoder layers issued two NCCL all-reduces per layer that NCCL 2.31.2 (pulled in by the PyTorch 26.09 base image) no longer runs on its symmetric kernel for small, unregistered inputs. The 57 KB decode messages fell back to 1-channel RING/LL, which regressed GB300 DSv4 TEP8 gen_only device step time by ~18%. - Drop the pre-MoE fused all-reduce. The attention output projection (o_b_proj) already reduces across TP, so this all-reduce only scaled a replicated tensor by tp_size ahead of an RMSNorm that cancels the scale. DeepSeek-V4.1 disables it for the same reason. - Pass the model dtype to the MoE output AllReduce so AUTO can select MNNVL AllReduce on multi-node TP, as o_b_proj already does. Co-Authored-By: Claude Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- .../_torch/models/modeling_deepseekv4.py | 15 +- .../modeling/test_modeling_deepseekv4.py | 156 +++++++++++++++++- 2 files changed, 165 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index d30baee3523c..812b3955688d 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -1714,8 +1714,12 @@ def __init__( self.allreduce = None if not self.use_dp and self.mapping.tp_size > 1: + # The dtype lets AUTO route a multi-node TP group to MNNVL AllReduce; + # without it the reduction always takes the NCCL path. self.allreduce = AllReduce( - mapping=model_config.mapping, strategy=model_config.allreduce_strategy + mapping=model_config.mapping, + strategy=model_config.allreduce_strategy, + dtype=config.torch_dtype, ) self.aux_stream = aux_stream_dict[AuxStreamType.MoeShared] self.event_dict = {key: torch.cuda.Event() for key in [EventType.Main, EventType.MoeShared]} @@ -1962,8 +1966,6 @@ def __init__( ) self.fusion_config = EagerFusionConfig() - self.enable_fusion = os.environ.get("TRTLLM_DEEPSEEK_EAGER_FUSION_DISABLED", "0") == "0" - self.enable_fusion &= not self.enable_attention_dp self.hc_ffn = self._make_mhc() @@ -1984,8 +1986,11 @@ def __init__( ) self.moe_allreduce = MoEAllReduce(self.mapping) - has_tp = mapping.has_tp() - self.fusion_config.PRE_MOE_FUSION = self.enable_fusion and has_tp + # The attention output projection (o_b_proj) completes the TP reduction + # itself, so the tensor mHC hands to the MoE block is already replicated + # across TP ranks. A pre-MoE all-reduce would only scale it by tp_size + # ahead of an RMSNorm that cancels the scale, so none is issued. + self.fusion_config.PRE_MOE_FUSION = False # DeepSeek-V4 applies the next RMSNorm after mHC post_mapping and the # next layer's mHC pre_mapping. Fusing the post-MoE all-reduce with the # next RMSNorm would normalize the raw MoE output before mHC post_mapping, diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index a6f43f942671..56468a6d24e9 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -57,7 +57,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.utils import AuxStreamType, model_extra_attrs -from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType, RotaryScalingType from tensorrt_llm.llmapi.llm_args import ( DeepSeekV4SparseAttentionConfig, KvCacheConfig, @@ -777,6 +777,160 @@ def fake_decoder_layer_init(self, model_config, *_args, **_kwargs): assert hasattr(mtp_layer.h_proj, "weight_scale") +class _RecordingAllReduce: + """Stands in for ``AllReduce`` and records the constructor kwargs.""" + + instances: list["_RecordingAllReduce"] = [] + + def __init__(self, **kwargs: object) -> None: + self.kwargs = kwargs + self.calls = 0 + _RecordingAllReduce.instances.append(self) + + def __call__(self, *args: object, **kwargs: object) -> torch.Tensor: + self.calls += 1 + raise AssertionError("a replicated MoE input must not be all-reduced") + + +class _ReduceAttention(nn.Module): + def __init__(self, *args: object, reduce_output: bool, **kwargs: object) -> None: + super().__init__() + self.reduce_output = reduce_output + + +class _IdentityNorm(nn.Module): + def __init__(self, *, hidden_size: int, eps: float, dtype: torch.dtype) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype), requires_grad=False) + self.variance_epsilon = eps + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + return hidden * 2 + + +class _RecordingMoE(nn.Module): + def __init__(self, **kwargs: object) -> None: + super().__init__() + self.calls: list[tuple[torch.Tensor, bool]] = [] + + def forward( + self, hidden: torch.Tensor, hidden_fp4: None, *, final_all_reduce_params, **kwargs + ) -> torch.Tensor: + self.calls.append((hidden.clone(), final_all_reduce_params.enable_allreduce)) + return hidden + 1 + + +@pytest.mark.parametrize("eager_fusion", [True, False]) +@pytest.mark.parametrize( + "tp_size,attention_dp", [(1, False), (4, False), (4, True)], ids=["tp1", "tp4", "tp4-adp"] +) +def test_deepseek_v4_decoder_reduces_attention_output_exactly_once( + monkeypatch: pytest.MonkeyPatch, tp_size: int, attention_dp: bool, eager_fusion: bool +) -> None: + import tensorrt_llm._torch.models.modeling_deepseekv4 as v4 + + monkeypatch.setenv("TRTLLM_DEEPSEEK_EAGER_FUSION_DISABLED", "0" if eager_fusion else "1") + monkeypatch.setattr(v4, "can_access_peer", lambda mapping: False) + monkeypatch.setattr(v4, "RMSNorm", _IdentityNorm) + monkeypatch.setattr(v4, "DeepseekV4MoE", _RecordingMoE) + monkeypatch.setattr(v4, "MoEAllReduce", lambda mapping: None) + monkeypatch.setattr(v4, "_resolve_enable_fused_hc", lambda config: False) + monkeypatch.setattr(v4.DeepseekV4DecoderLayer, "attention_cls", _ReduceAttention) + monkeypatch.setattr(v4.DeepseekV4DecoderLayer, "_make_mhc", lambda self: nn.Identity()) + _RecordingAllReduce.instances = [] + monkeypatch.setattr(v4, "AllReduce", _RecordingAllReduce) + quant = SimpleNamespace( + layer_quant_mode=SimpleNamespace(has_nvfp4=lambda: False), quant_algo=None + ) + monkeypatch.setattr( + v4.DeepseekV4DecoderLayer, + "_get_decoder_layer_quant_config", + lambda self, model_config, layer_idx: quant, + ) + config = SimpleNamespace( + hidden_size=8, + moe_intermediate_size=16, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=1, + rms_norm_eps=1e-6, + torch_dtype=torch.bfloat16, + ) + mapping = Mapping(world_size=tp_size, tp_size=tp_size, rank=0, enable_attention_dp=attention_dp) + model_config = SimpleNamespace( + pretrained_config=config, mapping=mapping, allreduce_strategy=None, moe_backend="TRTLLM" + ) + + layer = DeepseekV4DecoderLayer(model_config, layer_idx=0, aux_stream_dict={}) + + # The attention output projection owns the TP reduction, so the pre-MoE + # all-reduce is never fused in and the attention all-reduce is never skipped. + requires_attention_reduction = tp_size > 1 and not attention_dp + assert layer.self_attn.reduce_output == requires_attention_reduction + assert layer.disable_attn_allreduce == (not requires_attention_reduction) + assert layer.fusion_config.PRE_MOE_FUSION is False + assert layer.fusion_config.POST_MOE_FUSION is False + + hidden = torch.arange(24, dtype=torch.bfloat16).reshape(3, 8) / 16 + metadata = SimpleNamespace(all_rank_num_tokens=[3] * tp_size) + output = layer.forward_MoE(hidden, metadata) + + torch.testing.assert_close(output, hidden * 2 + 1, rtol=0, atol=0) + ((moe_input, final_reduce),) = layer.mlp.calls + torch.testing.assert_close(moe_input, hidden * 2, rtol=0, atol=0) + assert final_reduce == (tp_size > 1) + assert all(instance.calls == 0 for instance in _RecordingAllReduce.instances) + + +def test_deepseek_v4_moe_allreduce_is_mnnvl_eligible(monkeypatch: pytest.MonkeyPatch) -> None: + import tensorrt_llm._torch.distributed as distributed + import tensorrt_llm._torch.models.modeling_deepseekv4 as v4 + + _RecordingAllReduce.instances = [] + monkeypatch.setattr(distributed, "AllReduce", _RecordingAllReduce) + monkeypatch.setattr( + v4, "DeepseekV4Gate", lambda *args, **kwargs: SimpleNamespace(routing_method=None) + ) + monkeypatch.setattr(v4, "create_moe", lambda **kwargs: nn.Identity()) + monkeypatch.setattr(v4, "GatedMLP", lambda **kwargs: nn.Identity()) + monkeypatch.setattr(torch.cuda, "Event", lambda: object()) + + mapping = Mapping(world_size=8, tp_size=8, rank=0, gpus_per_node=4) + config = SimpleNamespace( + n_group=1, + topk_group=1, + routed_scaling_factor=1.0, + n_hash_layers=0, + vocab_size=16, + torch_dtype=torch.bfloat16, + ) + model_config = SimpleNamespace( + pretrained_config=config, + mapping=mapping, + allreduce_strategy=AllReduceStrategy.AUTO, + moe_backend="TRTLLM", + quant_config=QuantConfig(), + use_cute_dsl_blockscaling_mm=False, + ) + moe = v4.DeepseekV4MoE( + num_experts=4, + top_k=1, + hidden_size=8, + intermediate_size=128, + shared_expert_intermediate_size=128, + aux_stream_dict={AuxStreamType.MoeShared: object()}, + dtype=torch.bfloat16, + model_config=model_config, + layer_idx=0, + ) + + # Without a dtype AllReduce never constructs MNNVLAllReduce, which leaves a + # multi-node TP group on NCCL for this reduction. + assert moe.allreduce is _RecordingAllReduce.instances[-1] + assert moe.allreduce.kwargs["dtype"] is torch.bfloat16 + assert moe.allreduce.kwargs["strategy"] is AllReduceStrategy.AUTO + + def test_deepseek_v4_routed_moe_quant_config_ignores_fp8_header(tmp_path): tensor_name = "layers.0.ffn.experts.0.w1.weight" shard_name = "model-00001-of-00001.safetensors" From 60e6d4feadfb69baccc7b0e24ddc8e60ad8e7f35 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Thu, 15 Oct 2026 02:22:46 -0700 Subject: [PATCH 2/2] [https://nvbugs/6919428][test] Cover the DeepSeek-V4 decoder attention reduction and pre-MoE RMSNorm input Address review feedback on the DeepSeek-V4 all-reduce fix: - Run DeepseekV4DecoderLayer.forward with a recording attention stub and check the runtime reduction request: attention reduces its output exactly once under TP, and the MoE input is not reduced again. - Feed forward_MoE through the real RMSNorm with a large epsilon on a TP4 layer, and check that the MoE input is the RMSNorm of the unscaled replicated input rather than of a tp_size-scaled copy. - Rename the forward_MoE test to what it checks, and share the layer construction between the decoder tests. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- .../modeling/test_modeling_deepseekv4.py | 130 +++++++++++++++--- 1 file changed, 113 insertions(+), 17 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 56468a6d24e9..76e4c77ae22f 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -39,6 +39,7 @@ from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention from tensorrt_llm._torch.attention.mla import MLA from tensorrt_llm._torch.configs.deepseekv4 import DeepseekV4Config +from tensorrt_llm._torch.distributed import AllReduceParams from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv4 import ( @@ -793,9 +794,30 @@ def __call__(self, *args: object, **kwargs: object) -> torch.Tensor: class _ReduceAttention(nn.Module): + """Records the reduction each forward requests; adds 0.25 to its input.""" + def __init__(self, *args: object, reduce_output: bool, **kwargs: object) -> None: super().__init__() self.reduce_output = reduce_output + self.calls: list[bool] = [] + + def forward( + self, *, hidden_states: torch.Tensor, all_reduce_params: AllReduceParams, **kwargs: object + ) -> torch.Tensor: + self.calls.append(all_reduce_params.enable_allreduce) + return hidden_states + 0.25 + + +class _FirstStreamMHC(nn.Module): + """Reads the first residual stream and adds a sublayer output to every stream.""" + + def pre_mapping(self, residual: torch.Tensor) -> tuple[None, None, torch.Tensor]: + return None, None, residual[:, 0] + + def post_mapping( + self, *, x: torch.Tensor, residual: torch.Tensor, **mixes: object + ) -> torch.Tensor: + return residual + x.unsqueeze(1) class _IdentityNorm(nn.Module): @@ -820,23 +842,27 @@ def forward( return hidden + 1 -@pytest.mark.parametrize("eager_fusion", [True, False]) -@pytest.mark.parametrize( - "tp_size,attention_dp", [(1, False), (4, False), (4, True)], ids=["tp1", "tp4", "tp4-adp"] -) -def test_deepseek_v4_decoder_reduces_attention_output_exactly_once( - monkeypatch: pytest.MonkeyPatch, tp_size: int, attention_dp: bool, eager_fusion: bool -) -> None: +def _build_v4_decoder_layer( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, + attention_dp: bool, + *, + eager_fusion: bool = True, + norm_cls: type[nn.Module] = _IdentityNorm, + torch_dtype: torch.dtype = torch.bfloat16, + rms_norm_eps: float = 1e-6, +) -> DeepseekV4DecoderLayer: + """Run the production constructor, replacing the hardware-bound components.""" import tensorrt_llm._torch.models.modeling_deepseekv4 as v4 monkeypatch.setenv("TRTLLM_DEEPSEEK_EAGER_FUSION_DISABLED", "0" if eager_fusion else "1") monkeypatch.setattr(v4, "can_access_peer", lambda mapping: False) - monkeypatch.setattr(v4, "RMSNorm", _IdentityNorm) + monkeypatch.setattr(v4, "RMSNorm", norm_cls) monkeypatch.setattr(v4, "DeepseekV4MoE", _RecordingMoE) monkeypatch.setattr(v4, "MoEAllReduce", lambda mapping: None) monkeypatch.setattr(v4, "_resolve_enable_fused_hc", lambda config: False) monkeypatch.setattr(v4.DeepseekV4DecoderLayer, "attention_cls", _ReduceAttention) - monkeypatch.setattr(v4.DeepseekV4DecoderLayer, "_make_mhc", lambda self: nn.Identity()) + monkeypatch.setattr(v4.DeepseekV4DecoderLayer, "_make_mhc", lambda self: _FirstStreamMHC()) _RecordingAllReduce.instances = [] monkeypatch.setattr(v4, "AllReduce", _RecordingAllReduce) quant = SimpleNamespace( @@ -853,21 +879,24 @@ def test_deepseek_v4_decoder_reduces_attention_output_exactly_once( n_routed_experts=4, n_shared_experts=1, num_experts_per_tok=1, - rms_norm_eps=1e-6, - torch_dtype=torch.bfloat16, + rms_norm_eps=rms_norm_eps, + torch_dtype=torch_dtype, ) mapping = Mapping(world_size=tp_size, tp_size=tp_size, rank=0, enable_attention_dp=attention_dp) model_config = SimpleNamespace( pretrained_config=config, mapping=mapping, allreduce_strategy=None, moe_backend="TRTLLM" ) + return DeepseekV4DecoderLayer(model_config, layer_idx=0, aux_stream_dict={}) - layer = DeepseekV4DecoderLayer(model_config, layer_idx=0, aux_stream_dict={}) - # The attention output projection owns the TP reduction, so the pre-MoE - # all-reduce is never fused in and the attention all-reduce is never skipped. - requires_attention_reduction = tp_size > 1 and not attention_dp - assert layer.self_attn.reduce_output == requires_attention_reduction - assert layer.disable_attn_allreduce == (not requires_attention_reduction) +@pytest.mark.parametrize("eager_fusion", [True, False]) +@pytest.mark.parametrize( + "tp_size,attention_dp", [(1, False), (4, False), (4, True)], ids=["tp1", "tp4", "tp4-adp"] +) +def test_deepseek_v4_forward_moe_does_not_reduce_replicated_input( + monkeypatch: pytest.MonkeyPatch, tp_size: int, attention_dp: bool, eager_fusion: bool +) -> None: + layer = _build_v4_decoder_layer(monkeypatch, tp_size, attention_dp, eager_fusion=eager_fusion) assert layer.fusion_config.PRE_MOE_FUSION is False assert layer.fusion_config.POST_MOE_FUSION is False @@ -882,6 +911,73 @@ def test_deepseek_v4_decoder_reduces_attention_output_exactly_once( assert all(instance.calls == 0 for instance in _RecordingAllReduce.instances) +def test_deepseek_v4_forward_moe_normalizes_unscaled_input( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + tp_size, eps = 4, 0.25 + # Float32 keeps the real RMSNorm on its reference torch path. The large + # epsilon makes the norm sensitive to input scale, so normalizing a + # tp_size-scaled copy of the replicated input would give a different result. + layer = _build_v4_decoder_layer( + monkeypatch, + tp_size, + attention_dp=False, + norm_cls=RMSNorm, + torch_dtype=torch.float32, + rms_norm_eps=eps, + ) + norm_weight = torch.linspace(0.5, 2.0, 8) + with torch.no_grad(): + layer.post_attention_layernorm.weight.copy_(norm_weight) + + def reference_rms_norm(values: torch.Tensor) -> torch.Tensor: + return norm_weight * (values * torch.rsqrt(values.pow(2).mean(-1, keepdim=True) + eps)) + + hidden = torch.arange(24, dtype=torch.float32).reshape(3, 8) / 16 + expected = reference_rms_norm(hidden) + assert not torch.allclose(reference_rms_norm(tp_size * hidden), expected, rtol=1e-2, atol=1e-2) + + layer.forward_MoE(hidden, SimpleNamespace(all_rank_num_tokens=[3] * tp_size)) + + ((moe_input, _),) = layer.mlp.calls + torch.testing.assert_close(moe_input, expected) + + +@pytest.mark.parametrize("eager_fusion", [True, False]) +@pytest.mark.parametrize( + "tp_size,attention_dp", [(1, False), (4, False), (4, True)], ids=["tp1", "tp4", "tp4-adp"] +) +def test_deepseek_v4_decoder_forward_requests_attention_reduction( + monkeypatch: pytest.MonkeyPatch, tp_size: int, attention_dp: bool, eager_fusion: bool +) -> None: + layer = _build_v4_decoder_layer(monkeypatch, tp_size, attention_dp, eager_fusion=eager_fusion) + # The attention output projection owns the TP reduction, so the decoder + # requests it exactly when the projection was built to reduce. + requires_attention_reduction = tp_size > 1 and not attention_dp + assert layer.self_attn.reduce_output == requires_attention_reduction + assert layer.disable_attn_allreduce == (not requires_attention_reduction) + + residual = torch.arange(96, dtype=torch.float32).reshape(3, 4, 8) / 16 + metadata = SimpleNamespace(all_rank_num_tokens=[3] * tp_size) + output = layer(torch.arange(3), residual, metadata) + + assert layer.self_attn.calls == [requires_attention_reduction] + # Each norm stand-in doubles its input, attention adds 0.25 and the MoE + # stand-in adds 1, so the values pin the data flow between the stand-ins. + attention_output = 2 * residual[:, 0] + 0.25 + after_attention = residual + attention_output.unsqueeze(1) + ((moe_input, final_reduce),) = layer.mlp.calls + torch.testing.assert_close(moe_input, 2 * after_attention[:, 0], rtol=0, atol=0) + assert final_reduce == (tp_size > 1) + assert not output.is_deferred + torch.testing.assert_close( + output.residual, after_attention + (moe_input + 1).unsqueeze(1), rtol=0, atol=0 + ) + assert all(instance.calls == 0 for instance in _RecordingAllReduce.instances) + + def test_deepseek_v4_moe_allreduce_is_mnnvl_eligible(monkeypatch: pytest.MonkeyPatch) -> None: import tensorrt_llm._torch.distributed as distributed import tensorrt_llm._torch.models.modeling_deepseekv4 as v4