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

[nv-guomingz]

…m + RoPE + gate) into a single Triton kernel

Summary by CodeRabbit

• New Features

  • Added an optimized fused attention path combining QK normalization, rotary position processing, and output gating.
  • Added support for interleaved multi-dimensional rotary positioning in the fused path.
  • Added optional in-place output gating for improved performance and reduced memory use.

• Bug Fixes

  • Improved attention gating behavior while retaining compatibility with existing fallback processing.

• Tests

  • Added comprehensive coverage across supported data types, tensor layouts, rotary configurations, token counts, and in-place operations.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[nv-guomingz]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59613 [ run ] triggered by Bot. Commit: 43b9b5d Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

The attention implementation now supports an overridable gated-QKV preparation path and reusable output gating. QKNormRoPEAttention uses new Triton kernels for fused RMSNorm, RoPE, QKV packing, and sigmoid gating when runtime conditions allow it; Qwen3NextAttention enables the path. CUDA tests cover standard layouts, mRoPE, production compatibility, and fallback behavior.

Fused QK-Norm RoPE Gating

Layer / File(s)	Summary
Attention gate hooks tensorrt_llm/_torch/modules/attention.py	Adds _prepare_qkv_gate, centralizes output gating, and updates forward() to support prepared QKV/gate values while retaining the existing fallback path.
Fused Triton kernels tensorrt_llm/_torch/modules/fused_qk_norm_rope_gate.py	Adds fused QKV RMSNorm/RoPE/gate preprocessing and fused sigmoid multiplication with layout validation, empty-input handling, and optional in-place output.
QK attention integration tensorrt_llm/_torch/modules/qk_norm_attention.py, tensorrt_llm/_torch/models/modeling_qwen3_next.py	Adds fused-path eligibility checks and dispatch in QKNormRoPEAttention, derives rotary dimensions through a helper, and enables the feature for Qwen3NextAttention.
Fused kernel validation tests/unittest/_torch/modules/test_fused_qk_norm_rope_gate.py	Tests fused outputs against Python and production references across dtypes, token counts, layouts, mRoPE configurations, and output-gating fallbacks.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Qwen3NextAttention
  participant Attention
  participant QKNormRoPEAttention
  participant FusedKernels
  Qwen3NextAttention->>QKNormRoPEAttention: enable fused QK-norm RoPE gate
  QKNormRoPEAttention->>Attention: provide prepared QKV and gate
  Attention->>FusedKernels: normalize, apply RoPE, and pack QKV
  FusedKernels-->>Attention: return QKV and gate
  Attention->>FusedKernels: apply sigmoid gate to attention output

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#16142: Adds related Qwen Image fastpaths that select and test fused QK-norm, RoPE, and gating behavior.

Suggested reviewers: pengbowang-nv, yunruis

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description leaves the required Description and Test Coverage sections empty, so it doesn't explain the change or validation.	Fill in the Description with the problem and solution, and add concrete Test Coverage details for the new fused-kernel paths.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title matches the PR objective by naming the fused attention preprocessing optimization and the relevant ticket.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (2)

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

268-298: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add the missing Python annotations across the new fused path.

• tensorrt_llm/_torch/modules/qk_norm_attention.py#L268-L298: annotate both override signatures with precise tensor and optional tuple types.
• tests/unittest/_torch/modules/test_fused_qk_norm_rope_gate.py#L26-L294: annotate helper parameters/returns and add -> None to test functions.

As per coding guidelines, “Annotate every function” and prefer built-in generics and |.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/modules/qk_norm_attention.py` around lines 268 - 298, The
fused path in QKNormAttention and its tests lacks required Python annotations.
In tensorrt_llm/_torch/modules/qk_norm_attention.py lines 268-298, annotate
_prepare_qkv_gate and apply_output_gate with precise tensor and optional tuple
types; in tests/unittest/_torch/modules/test_fused_qk_norm_rope_gate.py lines
26-294, annotate all helper parameters and return values and add -> None to
every test function, using built-in generics and | for unions.

Source: Coding guidelines

tests/unittest/_torch/modules/test_fused_qk_norm_rope_gate.py (1)

210-285: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy lift

Add a production-path dispatch test.

Coverage is insufficient for the end-to-end objective: these tests call the kernels and base fallback directly, but never exercise QKNormRoPEAttention._prepare_qkv_gate or Qwen3Next enablement. Add test_qwen3_next_fused_qk_norm_rope_gate_dispatches here, asserting eligible inputs invoke both fused stages and an ineligible layout falls back.

As per path instructions, provide a concrete test filename and state whether coverage is sufficient.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modules/test_fused_qk_norm_rope_gate.py` around lines
210 - 285, Coverage is not sufficient because the production dispatch path is
untested. In tests/unittest/_torch/modules/test_fused_qk_norm_rope_gate.py, add
test_qwen3_next_fused_qk_norm_rope_gate_dispatches to exercise
QKNormRoPEAttention._prepare_qkv_gate with Qwen3Next enabled, asserting eligible
inputs invoke both fused stages and an ineligible layout uses the fallback path.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tensorrt_llm/_torch/modules/qk_norm_attention.py`:
- Around line 268-298: The fused path in QKNormAttention and its tests lacks
required Python annotations. In tensorrt_llm/_torch/modules/qk_norm_attention.py
lines 268-298, annotate _prepare_qkv_gate and apply_output_gate with precise
tensor and optional tuple types; in
tests/unittest/_torch/modules/test_fused_qk_norm_rope_gate.py lines 26-294,
annotate all helper parameters and return values and add -> None to every test
function, using built-in generics and | for unions.

In `@tests/unittest/_torch/modules/test_fused_qk_norm_rope_gate.py`:
- Around line 210-285: Coverage is not sufficient because the production
dispatch path is untested. In
tests/unittest/_torch/modules/test_fused_qk_norm_rope_gate.py, add
test_qwen3_next_fused_qk_norm_rope_gate_dispatches to exercise
QKNormRoPEAttention._prepare_qkv_gate with Qwen3Next enabled, asserting eligible
inputs invoke both fused stages and an ineligible layout uses the fallback path.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2da38cc4-97a7-412b-9fb9-5521c9ef493f

📥 Commits

Reviewing files that changed from the base of the PR and between 6013944 and 43b9b5d.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/models/modeling_qwen3_next.py
• tensorrt_llm/_torch/modules/attention.py
• tensorrt_llm/_torch/modules/fused_qk_norm_rope_gate.py
• tensorrt_llm/_torch/modules/qk_norm_attention.py
• tests/unittest/_torch/modules/test_fused_qk_norm_rope_gate.py

[tensorrt-cicd]

PR_Github #59613 [ run ] completed with state SUCCESS. Commit: 43b9b5d
/LLM/main/L0_MergeRequest_PR pipeline #48052 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation