[TRTLLM-13349][perf] Fuse gemma RMSNorm into AllReduce for Qwen3-Next/Qwen3.5…

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,26 +632,27 @@ 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:
             # 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 (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)

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])