From 68f89bf29176ebbfce3be9acf41eaae937a8c935 Mon Sep 17 00:00:00 2001 From: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com> Date: Tue, 9 Jun 2026 14:52:14 +0000 Subject: [PATCH] [TRTLLM-13349][perf] Fuse gemma RMSNorm into AllReduce for Qwen3-Next/Qwen3.5 TEP Enable eager AllReduce + RMSNorm fusion on the full-attention and GDN decoder layers in the tensor-parallel, non-attention-DP path. Defer each block reduction to a single fused AllReduce owner and keep MTP on its shared-head norm path. Precompute the Gemma weight offset in cache_derived_state so custom, NCCL, and symmetric-memory backends use one consistent norm contract across regular, GMS, staged, and reload loading paths. Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com> --- .../_torch/models/modeling_qwen3_next.py | 122 ++++++++++--- .../fla/fused_sigmoid_gating_recurrent.py | 18 ++ .../models/test_qwen3_next_eager_fusion.py | 167 ++++++++++++++++++ 3 files changed, 278 insertions(+), 29 deletions(-) create mode 100644 tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 62bea2f9f19d..6a8a2a9495c3 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -62,6 +62,58 @@ from .modeling_utils import DecoderModel, EagerFusionConfig, register_auto_model +def _fused_norm_weight(norm: RMSNorm) -> torch.Tensor: + """Weight to feed the fused AllReduce+RMSNorm op for ``norm``. + + Gemma RMSNorm scales by ``(1 + weight)`` (see RMSNorm.forward), but the + fused AR+RMSNorm kernels (and the NCCL / NCCL_SYMMETRIC fallbacks the + AUTO strategy may pick) only apply ``weight``. Baking the ``+1`` into the + weight makes EVERY allreduce backend produce the correct gemma result + without any backend-specific flag. + + For gemma norms the ``(1 + weight)`` tensor is precomputed once in + ``Qwen3NextForCausalLM.cache_derived_state`` and cached on the module as + ``_fused_norm_weight``. Computing it inline here would re-run a cast+add + elementwise kernel every forward inside the CUDA graph. The inline path + below is only a correctness fallback if the cache is absent. + """ + cached = getattr(norm, "_fused_norm_weight", None) + if cached is not None: + return cached + w = norm.weight + if getattr(norm, "use_gemma", False): + return (w.float() + 1.0).to(w.dtype) + return w + + +def _precompute_fused_norm_weights(module: nn.Module) -> None: + """Bake ``(1 + weight)`` once for every gemma RMSNorm under ``module``. + + Caches the result on each norm as ``_fused_norm_weight`` so the fused + AllReduce+RMSNorm path reads a ready tensor instead of recomputing the + cast+add every forward. Non-gemma norms are left untouched (the fused op + uses their ``weight`` directly). Must run after weights are loaded onto the + device; the cached tensor is a plain attribute, not a registered buffer, so + it stays out of the state dict. + + Norms whose ``weight`` was stripped are skipped: the layer-wise benchmark + runs ``remove_weights`` on unused layers (``skip_forward``), leaving a + ``use_gemma`` norm without a ``weight`` parameter; those layers never run + the fused path, so there is nothing to precompute. + """ + for norm in module.modules(): + if isinstance(norm, RMSNorm) and getattr(norm, "use_gemma", False): + w = getattr(norm, "weight", None) + if w is None: + continue + norm._fused_norm_weight = (w.float() + 1.0).to(w.dtype) + + +def _eager_fusion_enabled(enable_attention_dp: bool) -> bool: + return (os.environ.get("TRTLLM_QWEN3_EAGER_FUSION_DISABLED", "0") == "0" + and not enable_attention_dp) + + class Qwen3NextGate(nn.Module): def __init__( @@ -368,17 +420,15 @@ def __init__( self.next_layer_layernorm: RMSNorm = None self.fusion_config = EagerFusionConfig() - ### TODO: enable eager_fusion by default - self.enable_fusion = os.environ.get( - "TRTLLM_QWEN3_EAGER_FUSION_DISABLED", "1") == "0" - self.enable_fusion &= not self.enable_attention_dp + self.enable_fusion = _eager_fusion_enabled(self.enable_attention_dp) has_tp = self.mapping.has_tp() has_pp = self.mapping.has_pp() self.fusion_config.PRE_MOE_FUSION = self.enable_fusion and has_tp - self.fusion_config.POST_MOE_FUSION = self.fusion_config.PRE_MOE_FUSION and not has_pp and self.enable_attention_dp - self.disable_attn_allreduce = (self.mapping.tp_size == 1 + self.fusion_config.POST_MOE_FUSION = self.fusion_config.PRE_MOE_FUSION and not has_pp + self.disable_attn_allreduce = (self.fusion_config.PRE_MOE_FUSION + or self.mapping.tp_size == 1 or self.enable_attention_dp) self.moe_allreduce = MoEAllReduce(mapping=model_config.mapping) @@ -415,21 +465,18 @@ def forward( all_reduce_params=AllReduceParams( fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, residual=residual, - norm_weight=self.post_attention_layernorm.weight, + norm_weight=_fused_norm_weight( + self.post_attention_layernorm), eps=self.post_attention_layernorm.variance_epsilon, - enable_allreduce=not self.disable_attn_allreduce, )) else: # No fusion hidden_states, residual = self.post_attention_layernorm( hidden_states, residual) - # Note: this fusion pattern is only supported for TRTLLM-nvfp4 backend now - do_finalize = not (self.fusion_config.POST_MOE_FUSION - and hidden_states.shape[0] - <= self.moe_allreduce.max_token - and self.model_config.moe_backend == 'TRTLLM' - and self.mlp.experts.has_nvfp4) + # Qwen3NextSparseMoeBlock does not implement do_finalize=False. Defer + # only its final all-reduce so the decoder can fuse it with RMSNorm. + do_finalize = True hidden_states = self.mlp( hidden_states, @@ -448,7 +495,8 @@ def forward( all_reduce_params=AllReduceParams( fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, residual=residual, - norm_weight=self.next_layer_layernorm.weight, + norm_weight=_fused_norm_weight( + self.next_layer_layernorm), eps=self.next_layer_layernorm.variance_epsilon, )) else: @@ -534,17 +582,24 @@ def __init__(self, model_config: ModelConfig[Qwen3NextConfig], self.next_layer_layernorm: RMSNorm = None self.fusion_config = EagerFusionConfig() - self.enable_fusion = os.environ.get( - "TRTLLM_QWEN3_EAGER_FUSION_DISABLED", "0") == "0" - self.enable_fusion &= not self.enable_attention_dp + self.enable_fusion = _eager_fusion_enabled(self.enable_attention_dp) has_tp = self.mapping.has_tp() has_pp = self.mapping.has_pp() self.fusion_config.PRE_MOE_FUSION = self.enable_fusion and has_tp - self.fusion_config.POST_MOE_FUSION = self.fusion_config.PRE_MOE_FUSION and not has_pp and self.enable_attention_dp - self.disable_attn_allreduce = (self.mapping.tp_size == 1 + # POST_MOE_FUSION fuses the MoE-output all-reduce with the next layer's + # RMSNorm. It is a tensor-parallel (TEP) optimization: it is only valid + # when ranks share the same tokens (not attention_dp, where each rank holds + # different tokens and the MoE block does no cross-rank all-reduce). This + # mirrors the DeepSeek-V3 pattern (POST == PRE in the non-attention_dp path). + self.fusion_config.POST_MOE_FUSION = self.fusion_config.PRE_MOE_FUSION and not has_pp + # When PRE_MOE_FUSION is on, the attention all-reduce is deferred to the + # fused PRE all-reduce+RMSNorm, so disable the in-attention all-reduce to + # avoid reducing twice. + self.disable_attn_allreduce = (self.fusion_config.PRE_MOE_FUSION + or self.mapping.tp_size == 1 or self.enable_attention_dp) self.moe_allreduce = MoEAllReduce(mapping=model_config.mapping) @@ -577,13 +632,14 @@ def forward( **kwargs, ) - if self.fusion_config.PRE_MOE_FUSION and self.enable_attention_dp: + if self.fusion_config.PRE_MOE_FUSION: hidden_states, residual = self.allreduce( hidden_states, all_reduce_params=AllReduceParams( fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, residual=residual, - norm_weight=self.post_attention_layernorm.weight, + norm_weight=_fused_norm_weight( + self.post_attention_layernorm), eps=self.post_attention_layernorm.variance_epsilon, )) else: @@ -591,12 +647,12 @@ def forward( hidden_states, residual = self.post_attention_layernorm( hidden_states, residual) - # Note: this fusion pattern is only supported for TRTLLM-nvfp4 backend now - do_finalize = not (hidden_states.shape[0] - <= self.moe_allreduce.max_token - and self.fusion_config.POST_MOE_FUSION - and self.model_config.moe_backend == 'TRTLLM' - and self.mlp.experts.has_nvfp4) + # The fully-fused do_finalize=False MoE path (MoEAllReduce on the + # unfinalized expert output) is not implemented by Qwen3NextSparseMoeBlock + # (it raises NotImplementedError). Keep do_finalize=True so POST_MOE_FUSION + # still fuses the *finalized* MoE all-reduce with the next layer's RMSNorm + # via the do_finalize branch below, without hitting the unimplemented path. + do_finalize = True hidden_states = self.mlp( hidden_states, attn_metadata, @@ -614,7 +670,8 @@ def forward( all_reduce_params=AllReduceParams( fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, residual=residual, - norm_weight=self.next_layer_layernorm.weight, + norm_weight=_fused_norm_weight( + self.next_layer_layernorm), eps=self.next_layer_layernorm.variance_epsilon, )) else: @@ -780,6 +837,9 @@ def __init__(self, model_config: ModelConfig[Qwen3NextConfig], use_cute_dsl_blockscaling_mm=False, ) self.shared_head = Qwen3NextMTPHead(mtp_model_config) + # MTP applies shared_head.norm after the base decoder forward, so its + # MoE-output all-reduce cannot consume next_layer_layernorm. + self.fusion_config.POST_MOE_FUSION = False @staticmethod def _is_mtp_excluded_from_quant( @@ -1016,3 +1076,7 @@ def setup_aliases(self) -> None: else: layer.next_layer_layernorm = self.model.layers[ idx + 1].input_layernorm + + def cache_derived_state(self) -> None: + super().cache_derived_state() + _precompute_fused_norm_weights(self) diff --git a/tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py b/tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py index 961189d97346..a2ede4bb54ad 100644 --- a/tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py +++ b/tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py @@ -265,6 +265,24 @@ def _flashinfer_gdn_decode( HV = v.shape[2] V = v.shape[3] + # The FlashInfer CuTe-DSL kernel requires every input tensor's data pointer + # to be 32-byte aligned (enforced in build_memref_desc). ``a`` and ``b`` are + # per-head-scalar slices of the fused ``in_proj_ba`` output: ``b`` starts at + # offset 0 (aligned) but ``a`` starts ``num_v_heads_per_tp`` bf16 elements in, + # so when ``num_v_heads_per_tp`` is not a multiple of 16 (e.g. Qwen3.6-35B-A3B + # TEP4: 32 v-heads / 4 = 8 -> 16-byte offset) the slice base is not 32-byte + # aligned and the kernel aborts. ``.contiguous()`` is NOT enough: at decode + # the token dim is 1, so the strided/offset slice already reports as + # contiguous (size-1 dims are ignored by is_contiguous) and ``.contiguous()`` + # is a no-op that keeps the misaligned pointer. Clone into fresh (allocator- + # aligned) storage instead, and only when misaligned so the common aligned + # case (e.g. Qwen3.5-397B TEP4: 64 / 4 = 16 -> 32-byte offset) stays zero-copy. + # q/k/v are sliced on 128-element head boundaries (>=256 B), always aligned. + if a.data_ptr() % 32 != 0: + a = a.clone(memory_format=torch.contiguous_format) + if b.data_ptr() % 32 != 0: + b = b.clone(memory_format=torch.contiguous_format) + # Reshape from packed varlen [1, N*T, ...] to batched [N, T, ...]. q_bat = q.view(N, T_per_seq, q.shape[2], q.shape[3]) k_bat = k.view(N, T_per_seq, k.shape[2], k.shape[3]) diff --git a/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py b/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py new file mode 100644 index 000000000000..e0a73d9c1f05 --- /dev/null +++ b/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import torch +from torch import nn + +from tensorrt_llm._torch.distributed import AllReduceFusionOp +from tensorrt_llm._torch.models.modeling_qwen3_next import ( + Qwen3NextForCausalLM, + Qwen3NextLinearDecoderLayer, + _eager_fusion_enabled, +) +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + +def _new_causal_lm() -> Qwen3NextForCausalLM: + model = Qwen3NextForCausalLM.__new__(Qwen3NextForCausalLM) + nn.Module.__init__(model) + return model + + +@torch.no_grad() +def test_setup_aliases_does_not_read_meta_weights() -> None: + model = _new_causal_lm() + model.model_config = SimpleNamespace(pretrained_config=SimpleNamespace(num_hidden_layers=2)) + model.model = nn.Module() + model.model.layers = nn.ModuleList([nn.Module(), nn.Module()]) + for layer in model.model.layers: + layer.input_layernorm = RMSNorm( + hidden_size=4, + eps=1e-6, + dtype=torch.bfloat16, + device=torch.device("meta"), + use_gemma=True, + ) + layer.next_layer_layernorm = None + model.model.norm = RMSNorm( + hidden_size=4, eps=1e-6, dtype=torch.bfloat16, device=torch.device("meta"), use_gemma=True + ) + + model.setup_aliases() + + assert model.model.layers[0].next_layer_layernorm is model.model.layers[1].input_layernorm + assert model.model.layers[1].next_layer_layernorm is model.model.norm + assert not hasattr(model.model.norm, "_fused_norm_weight") + + +@torch.no_grad() +def test_cache_derived_state_refreshes_gemma_norm_weight() -> None: + model = _new_causal_lm() + model.gemma_norm = RMSNorm(hidden_size=4, eps=1e-6, dtype=torch.bfloat16, use_gemma=True) + model.standard_norm = RMSNorm(hidden_size=4, eps=1e-6, dtype=torch.bfloat16) + model.gemma_norm.weight.copy_(torch.tensor([-0.5, 0.0, 0.5, 1.0], dtype=torch.bfloat16)) + + model.cache_derived_state() + + expected = (model.gemma_norm.weight.float() + 1.0).to(torch.bfloat16) + # Exact: cache_derived_state bakes (1+weight) with the same fp32-add-then- + # cast recomputed here, so the result must be bitwise-identical. + torch.testing.assert_close(model.gemma_norm._fused_norm_weight, expected, atol=0.0, rtol=0.0) + assert not hasattr(model.standard_norm, "_fused_norm_weight") + + model.gemma_norm.weight.add_(1.0) + model.cache_derived_state() + expected = (model.gemma_norm.weight.float() + 1.0).to(torch.bfloat16) + # Exact: cache_derived_state bakes (1+weight) with the same fp32-add-then- + # cast recomputed here, so the result must be bitwise-identical. + torch.testing.assert_close(model.gemma_norm._fused_norm_weight, expected, atol=0.0, rtol=0.0) + + +@torch.no_grad() +def test_eager_fusion_is_enabled_for_gdn_by_default(monkeypatch) -> None: + monkeypatch.delenv("TRTLLM_QWEN3_EAGER_FUSION_DISABLED", raising=False) + assert _eager_fusion_enabled(enable_attention_dp=False) + assert not _eager_fusion_enabled(enable_attention_dp=True) + + monkeypatch.setenv("TRTLLM_QWEN3_EAGER_FUSION_DISABLED", "1") + assert not _eager_fusion_enabled(enable_attention_dp=False) + + +@torch.no_grad() +def test_gdn_fusion_has_single_allreduce_owner() -> None: + hidden_states = torch.randn(2, 4, dtype=torch.bfloat16) + residual = torch.randn_like(hidden_states) + # Use real RMSNorm modules (lightweight) rather than mocks: their + # weight / use_gemma / variance_epsilon are exactly what the fused-norm + # path reads. The linear_attn/allreduce/mlp below stay mocks because the + # test asserts on *how they are called* (call_count / call_args). + post_attention_norm = RMSNorm(hidden_size=4, eps=1e-6, dtype=torch.bfloat16, use_gemma=True) + next_layer_norm = RMSNorm(hidden_size=4, eps=1e-6, dtype=torch.bfloat16, use_gemma=True) + post_attention_norm.weight.copy_(torch.tensor([-0.5, 0.0, 0.5, 1.0], dtype=torch.bfloat16)) + next_layer_norm.weight.copy_(torch.tensor([0.0, 0.25, 0.5, 0.75], dtype=torch.bfloat16)) + linear_attn = MagicMock(side_effect=lambda hidden_states, *args, **kwargs: hidden_states) + allreduce = MagicMock( + side_effect=lambda hidden_states, *, all_reduce_params: ( + hidden_states, + all_reduce_params.residual, + ) + ) + mlp = MagicMock(side_effect=lambda hidden_states, *args, **kwargs: hidden_states) + layer = SimpleNamespace( + layer_idx=0, + input_layernorm=MagicMock(), + linear_attn=linear_attn, + post_attention_layernorm=post_attention_norm, + next_layer_layernorm=next_layer_norm, + fusion_config=SimpleNamespace(PRE_MOE_FUSION=True, POST_MOE_FUSION=True), + disable_attn_allreduce=True, + allreduce=allreduce, + mlp=mlp, + mapping=SimpleNamespace(tp_size=2), + moe_allreduce=MagicMock(), + ) + + Qwen3NextLinearDecoderLayer.forward( + layer, + position_ids=torch.arange(2), + hidden_states=hidden_states, + attn_metadata=SimpleNamespace(), + residual=residual, + ) + + internal_ar_params = linear_attn.call_args.kwargs["all_reduce_params"] + assert not internal_ar_params.enable_allreduce + + # Two module-level allreduces at the layer boundary: pre-MoE and post-MoE. + # The GDN linear_attn's own allreduce is disabled (single-owner, asserted + # just above), so it does not add a third. + assert allreduce.call_count == 2 + pre_ar_params = allreduce.call_args_list[0].kwargs["all_reduce_params"] + post_ar_params = allreduce.call_args_list[1].kwargs["all_reduce_params"] + assert pre_ar_params.enable_allreduce + assert post_ar_params.enable_allreduce + assert pre_ar_params.fusion_op == AllReduceFusionOp.RESIDUAL_RMS_NORM + assert post_ar_params.fusion_op == AllReduceFusionOp.RESIDUAL_RMS_NORM + # Exact: norm_weight is the same baked (1+weight) recomputed here. + torch.testing.assert_close( + pre_ar_params.norm_weight, + (post_attention_norm.weight.float() + 1.0).to(torch.bfloat16), + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + post_ar_params.norm_weight, + (next_layer_norm.weight.float() + 1.0).to(torch.bfloat16), + atol=0.0, + rtol=0.0, + ) + + mlp_ar_params = mlp.call_args.kwargs["all_reduce_params"] + assert not mlp_ar_params.enable_allreduce + assert mlp.call_args.kwargs["do_finalize"]