[TRTLLM-14352][perf] Fuse Qwen3.5/3.6 attention preprocessing (QK-nor…

tensorrt_llm/_torch/models/modeling_qwen3_next.py

@@ -1,5 +1,7 @@
 # Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
 # Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/configs/qwen3_next.py
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
 # coding=utf-8
 # Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
 #
@@ -489,6 +491,7 @@ def __init__(self, model_config: ModelConfig[Qwen3NextConfig],
                          fuse_qk_norm_rope=fuse_qk_norm_rope,
                          attn_output_gate=True,
                          use_gemma_rms_norm=True)
+        self._fuse_qk_norm_rope_gate = True
 
 
 class Qwen3NextFullAttentionDecoderLayer(DecoderLayer):

tensorrt_llm/_torch/modules/attention.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import math
 import weakref
 from typing import List, Optional, Tuple, Union
@@ -696,6 +699,20 @@ def split_qkv(self, q, k=None, v=None):
             q, k, v = q.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         return q, k, v
 
+    def _prepare_qkv_gate(
+        self, qkv: torch.Tensor, position_ids: Optional[torch.Tensor]
+    ) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor],
+                        Optional[torch.Tensor], torch.Tensor]]:
+        """Optionally prepare a gated QKV projection before the generic split path."""
+        return None
+
+    def apply_output_gate(self, attention_output: torch.Tensor,
+                          gate: torch.Tensor) -> torch.Tensor:
+        """Apply the attention output gate."""
+        if gate.shape != attention_output.shape:
+            gate = gate.reshape(attention_output.shape)
+        return attention_output * torch.sigmoid(gate)
+
     def convert_qkv(self, q, k, v):
         if k is None and v is None and not self.support_fused_qkv:
             q, k, v = self.split_qkv(q)
@@ -983,18 +1000,6 @@ def forward(
             if qkv_lora is not None:
                 qkv = qkv + qkv_lora
 
-        if self.attn_output_gate:
-            q_gate, k, v = qkv.split(
-                [self.q_size * 2, self.kv_size, self.kv_size], dim=-1)
-            orig_shape = q_gate.shape[:-1]
-            # Single line: view -> chunk -> reshape both q and gate
-            q, gate = [
-                t.reshape(*orig_shape, -1) for t in torch.chunk(
-                    q_gate.view(*orig_shape, self.num_heads, -1), 2, dim=-1)
-            ]
-        else:
-            q, k, v = qkv, None, None
-
         # For dynamic tree spec decoding with Python RoPE, adjust position_ids
         # to use tree offsets (same as C++ kernel: past_seq_len + offset).
         if (not self.rope_fusion
@@ -1008,7 +1013,26 @@ def forward(
             position_ids = self._adjust_position_ids_for_spec_dec(
                 position_ids, attn_metadata)
 
-        q, k, v = self.apply_rope(q, k, v, position_ids)
+        qkv_gate_prepared = False
+        if self.attn_output_gate:
+            prepared = self._prepare_qkv_gate(qkv, position_ids)
+            if prepared is not None:
+                q, k, v, gate = prepared
+                qkv_gate_prepared = True
+            else:
+                q_gate, k, v = qkv.split(
+                    [self.q_size * 2, self.kv_size, self.kv_size], dim=-1)
+                orig_shape = q_gate.shape[:-1]
+                # Single line: view -> chunk -> reshape both q and gate
+                q, gate = [
+                    t.reshape(*orig_shape, -1) for t in torch.chunk(
+                        q_gate.view(*orig_shape, self.num_heads, -1), 2, dim=-1)
+                ]
+        else:
+            q, k, v = qkv, None, None
+
+        if not qkv_gate_prepared:
+            q, k, v = self.apply_rope(q, k, v, position_ids)
         q, k, v = self.convert_qkv(q, k, v)
 
         if attention_sinks is not None:
@@ -1034,8 +1058,7 @@ def forward(
         )
 
         if self.attn_output_gate:
-            gate = torch.sigmoid(gate)
-            attn_output = attn_output * gate
+            attn_output = self.apply_output_gate(attn_output, gate)
 
         attn_output = _helix_cp_output_projection(self.o_proj, attn_output,
                                                   attn_metadata,