tensorrt_llm/_torch/modules/qk_norm_attention.py (1)

1-112: Please add comprehensive unit tests for the QKNormRoPEAttention module.

No existing tests cover this new module—let’s ensure it’s thoroughly exercised:

• Create a new test file (e.g. tests/unittest/_torch/test_qk_norm_attention.py)
• Cover both execution paths:
– fuse_qk_norm_rope=True (fused QK-norm+RoPE)
– fuse_qk_norm_rope=False (separate QK-norm then RoPE)
• Test a variety of input shapes and configurations:
– Single sequence vs. batch sequences
– Varying num_heads and num_key_value_heads
• Edge cases:
– Empty tensors (zero-length sequence)
– Invalid or out-of-range position_ids (should error or handle gracefully)
• Verify attention internals:
– Q and K outputs are correctly normalized (e.g. compare against standalone RMSNorm)
– RoPE is applied correctly (match baseline implementation)
• Parallel execution path via maybe_execute_in_parallel (ensure both streams complete)

Adding these tests will catch regressions in both fused and non-fused code paths and validate the QK normalization logic.

tensorrt_llm/_torch/modules/qk_norm_attention.py (2)

65-66: Consider adding CUDA availability check.

The code directly creates CUDA streams and events without checking if CUDA is available. This could cause issues in CPU-only environments or during testing.

Consider adding a CUDA availability check:

-        self.aux_stream = torch.cuda.Stream()
-        self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
+        if torch.cuda.is_available():
+            self.aux_stream = torch.cuda.Stream()
+            self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
+        else:
+            self.aux_stream = None
+            self.ln_events = None

And update the apply_qk_norm method to handle the None case appropriately.

---

109-111: Improve assertion error message.

The assertion provides a clear constraint but could be more informative about why this requirement exists.

-        assert k is None and v is None, "The input should be a concatenated qkv tensor to apply_qk_norm_rope"
+        assert k is None and v is None, (
+            "When fuse_qk_norm_rope=True, the input must be a concatenated qkv tensor "
+            "(k and v should be None) for the fused operator to work correctly"
+        )

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

tensorrt_llm/_torch/models/modeling_qwen3_moe.py (1)

1-1: Add NVIDIA copyright header (2025) at file top.

Required by coding guidelines for all source files.

Apply this diff:

+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
 import os

tensorrt_llm/_torch/models/modeling_qwen3.py (1)

1-1: Add NVIDIA copyright header (2025) at file top.

Required by coding guidelines.

Apply this diff:

+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
 from typing import Optional, Tuple

tensorrt_llm/_torch/modules/qk_norm_attention.py (2)

66-73: Use configured RMSNorm epsilon instead of hardcoded 1e-6.

Keep normalization consistent with model config by using the configured eps.

Apply this diff:

-        self.q_norm = RMSNorm(hidden_size=self.head_dim,
-                              eps=1e-6,
-                              dtype=pretrained_config.torch_dtype,
-                              has_weights=True)
-        self.k_norm = RMSNorm(hidden_size=self.head_dim,
-                              eps=1e-6,
-                              dtype=pretrained_config.torch_dtype,
-                              has_weights=True)
+        self.q_norm = RMSNorm(hidden_size=self.head_dim,
+                              eps=pretrained_config.rms_norm_eps,
+                              dtype=pretrained_config.torch_dtype,
+                              has_weights=True)
+        self.k_norm = RMSNorm(hidden_size=self.head_dim,
+                              eps=pretrained_config.rms_norm_eps,
+                              dtype=pretrained_config.torch_dtype,
+                              has_weights=True)

---

97-104: Wrap fused operator call with availability check and clear error handling.

Improve robustness and debuggability if the custom op is missing or fails.

Apply this diff:

-    def apply_qk_norm_rope(self, qkv, position_ids):
-        torch.ops.trtllm.fused_qk_norm_rope(
-            qkv, self.num_heads, self.num_key_value_heads,
-            self.num_key_value_heads, self.head_dim,
-            self.q_norm.variance_epsilon, self.q_norm.weight,
-            self.k_norm.weight, self.pos_embd_params.rope.theta,
-            self.pos_embd_params.is_neox, position_ids.view(-1))
-        return qkv, None, None
+    def apply_qk_norm_rope(self, qkv, position_ids):
+        op = getattr(torch.ops.trtllm, "fused_qk_norm_rope", None)
+        if op is None:
+            raise RuntimeError(
+                "torch.ops.trtllm.fused_qk_norm_rope is not available in this build."
+            )
+        try:
+            op(qkv,
+               self.num_heads,
+               self.num_key_value_heads,
+               self.num_key_value_heads,  # verify this argument; see follow-up comment
+               self.head_dim,
+               self.q_norm.variance_epsilon,
+               self.q_norm.weight,
+               self.k_norm.weight,
+               self.pos_embd_params.rope.theta,
+               self.pos_embd_params.is_neox,
+               position_ids.view(-1))
+        except Exception as e:
+            raise RuntimeError(f"fused_qk_norm_rope failed: {e}") from e
+        return qkv, None, None

tensorrt_llm/_torch/modules/qk_norm_attention.py (1)

77-86: Add short docstrings for internal helpers.

Concise docstrings aid maintainability and clarify expected shapes/semantics.

Apply this diff:

     def apply_qk_norm(self, q, k):
+        """Apply RMSNorm to Q and K in parallel using an auxiliary CUDA stream.
+
+        Args:
+            q (torch.Tensor): Query tensor shaped (..., q_size, head_dim) flattened internally.
+            k (torch.Tensor): Key tensor shaped (..., kv_size, head_dim) flattened internally.
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: Normalized Q and K with original leading dims restored.
+        """
@@
     def apply_qk_norm_rope(self, qkv, position_ids):
+        """Apply fused Q/K RMSNorm and RoPE via custom TRTLLM operator.
+
+        Args:
+            qkv (torch.Tensor): Concatenated QKV projection tensor before splitting.
+            position_ids (torch.Tensor): Positions shaped (batch*seq,) or broadcastable.
+        Returns:
+            Tuple[torch.Tensor, None, None]: Modified qkv and placeholders to match base interface.
+        """

Also applies to: 97-104

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

tensorrt_llm/_torch/models/modeling_qwen3.py (1)

56-66: Fix forward() return annotation to match actual return type

forward() currently annotates a single torch.Tensor but returns a tuple (hidden_states, residual). Update the annotation to prevent type-checking confusion and reader errors.

Apply this diff:

-    ) -> torch.Tensor:
+    ) -> Tuple[torch.Tensor, torch.Tensor]:

tensorrt_llm/_torch/models/modeling_qwen3.py (2)

22-29: Good fix: init no longer has an invalid return annotation

Previous feedback about init returning a Tuple is addressed; the signature now correctly returns None implicitly.

---

30-40: Correct removal of unsupported pos_embd_params when constructing QKNormRoPEAttention

The constructor no longer passes pos_embd_params, aligning with QKNormRoPEAttention’s API. This avoids a TypeError at instantiation.

tensorrt_llm/_torch/modules/qk_norm_attention.py (2)

81-88: Use model-configured RMSNorm epsilon instead of hardcoded 1e-6

Match other layers’ eps and avoid silent divergence across components.

Apply this diff:

-        self.q_norm = RMSNorm(hidden_size=self.head_dim,
-                              eps=1e-6,
+        self.q_norm = RMSNorm(hidden_size=self.head_dim,
+                              eps=pretrained_config.rms_norm_eps,
                               dtype=pretrained_config.torch_dtype,
                               has_weights=True)
-        self.k_norm = RMSNorm(hidden_size=self.head_dim,
-                              eps=1e-6,
+        self.k_norm = RMSNorm(hidden_size=self.head_dim,
+                              eps=pretrained_config.rms_norm_eps,
                               dtype=pretrained_config.torch_dtype,
                               has_weights=True)

---

112-120: Wrap fused operator call with availability check and clear error handling

Improve debuggability when the custom op is missing or fails.

Apply this diff:

-        torch.ops.trtllm.fused_qk_norm_rope(
-            qkv, self.num_heads, self.num_key_value_heads,
-            self.num_key_value_heads, self.head_dim,
-            self.q_norm.variance_epsilon, self.q_norm.weight,
-            self.k_norm.weight, self.pos_embd_params.rope.theta,
-            self.pos_embd_params.is_neox, position_ids.view(-1))
+        if not hasattr(torch.ops, "trtllm") or not hasattr(torch.ops.trtllm, "fused_qk_norm_rope"):
+            raise RuntimeError("torch.ops.trtllm.fused_qk_norm_rope is not available. "
+                               "Ensure the TRT-LLM extension is built and loaded.")
+        try:
+            torch.ops.trtllm.fused_qk_norm_rope(
+                qkv, self.num_heads, self.num_key_value_heads,
+                self.num_key_value_heads, self.head_dim,
+                self.q_norm.variance_epsilon, self.q_norm.weight,
+                self.k_norm.weight, self.pos_embd_params.rope.theta,
+                self.pos_embd_params.is_neox, position_ids.view(-1))
+        except (RuntimeError, AttributeError) as e:
+            raise RuntimeError(f"Failed to apply fused_qk_norm_rope: {e}") from e

If helpful, I can add a lightweight unit test stub that conditionally skips when the fused op is unavailable and validates the non-fused path equivalence on small inputs.

tensorrt_llm/_torch/models/modeling_qwen3.py (1)

130-134: Minor readability: simplify the XOR check for inputs selection

The XOR condition is correct but less common in codebases. A more idiomatic form is equality on None-ness.

Apply this diff:

-        if (input_ids is None) ^ (inputs_embeds is not None):
+        if (input_ids is None) == (inputs_embeds is None):
             raise ValueError(
                 "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
             )

tensorrt_llm/_torch/modules/qk_norm_attention.py (1)

28-33: Docstring could specify expected tensor shapes and execution paths

Consider briefly documenting the expected input shapes for q, k, v in both fused and non-fused modes and how position_ids should be shaped. This reduces ambiguity for future maintainers.

I can add a short Google-style docstring with shapes and examples if you’d like.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-12T10:28:57.288Z
Learning: Applies to **/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py} : Prepend NVIDIA copyright header (current year) to all source files

Learnt from: galagam
PR: NVIDIA/TensorRT-LLM#6487
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py:1-12
Timestamp: 2025-08-06T13:58:07.506Z
Learning: In TensorRT-LLM, test files (files under tests/ directories) do not require NVIDIA copyright headers, unlike production source code files. Test files typically start directly with imports, docstrings, or code.