[None][perf] Fuse Qwen3.5 gated QKV preprocessing - #16024
Closed
mingyangHao wants to merge 3 commits into
Closed
Conversation
Add an opt-in QKV gate-tail layout to remove intermediate layout copies, and optionally fold Gemma QK RMSNorm plus RoPE into the KV-cache preprocessing kernel. Keep unsupported configurations on guarded fallback paths. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
mingyangHao
force-pushed
the
user/mingyangh/qwen35-gate-tail-fused-preprocess
branch
from
July 7, 2026 02:19
c3dc37f to
2bdf761
Compare
Contributor
📝 WalkthroughWalkthroughAdds fused Q/K RMS normalization plus RoPE preprocessing and strided (non-packed) QKV token layout support across C++ attention kernels (XQAParams, EnqueueParams, QKVPreprocessingParams, unfused attention kernels, xqaDispatcher), torch/nanobind bindings, and Python attention backends, including a gate-tail QKV weight layout permutation. ChangesFused QK-Norm/RoPE + Strided QKV Attention
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant QKNormRoPEAttention
participant TrtllmAttention
participant RunnerRun as thop RunnerRun
participant UnfusedKernel as QKV Preprocessing Kernel
QKNormRoPEAttention->>TrtllmAttention: configure_qk_norm_preprocessing(q_norm_weight, k_norm_weight, eps)
TrtllmAttention->>RunnerRun: run(q, q_norm_weight, k_norm_weight, qk_norm_eps)
RunnerRun->>UnfusedKernel: invokeQKVPreprocessing(input_token_stride, q_norm_weight, k_norm_weight)
UnfusedKernel-->>RunnerRun: normalized+rotated Q/K written to KV cache
RunnerRun-->>TrtllmAttention: attention output
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h (1)
42-45: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUndocumented new public fields.
q_norm_weight,k_norm_weight,qk_norm_eps,qk_norm_use_gemmahave no comments, unlikeqkv_token_strideabove them which is well documented. Given these gate a numerically sensitive fused-kernel path (head_size==256, eps>0 requirement enforced elsewhere), a short comment on semantics/units and the "TRTLLM-Gen only" constraint would help future readers avoid misuse.As per coding guidelines, "Follow Doxygen rules for documenting class interfaces and function prototypes... use
//!for C++ style comments and//!<for member comments."🤖 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 `@cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h` around lines 42 - 45, The new public members in xqaParams are undocumented, so add Doxygen-style member comments for q_norm_weight, k_norm_weight, qk_norm_eps, and qk_norm_use_gemma, matching the style used for qkv_token_stride. Describe their semantics/units, note that they are only for the TRTLLM-Gen fused-kernel path, and mention any relevant constraints such as the numerically sensitive head_size==256/eps>0 usage so future callers do not misuse these fields.Source: Coding guidelines
🤖 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.
Inline comments:
In `@cpp/tensorrt_llm/common/attentionOp.cpp`:
- Around line 1770-1773: The packed-QKV token stride formula is duplicated
across the context FMHA, unfused context MHA fallback, and masked MHA generation
fallback checks in attentionOp.cpp. Extract the shared `(mNumAttnHeads + 2 *
mNumAttnKVHeads) * getHeadSize()` logic into a helper with a clear name, then
update each strided `attention_input` safety check to call that helper so future
head-size changes only need one update.
In
`@cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h`:
- Around line 762-773: The RMS normalization in the q/k accumulation block is
using warpReduceSum in a way that can mix values from multiple tokens when
VECS_PER_HEAD is less than a warp, so adjust the reduction in
unfusedAttentionKernels_2_template.h to be token-scoped instead of warp-scoped.
Update the qSumSquares/kSumSquares path near the qRmsRcp and kRmsRcp computation
so the sum only covers lanes belonging to the same token, or gate this code path
behind a 32-lane-per-token layout constraint in the relevant unfused attention
kernel mapping.
In `@tensorrt_llm/_torch/modules/qk_norm_attention.py`:
- Around line 306-314: The caching added in _adjust_position_ids_for_spec_dec()
should be removed because it attaches `_tllm_flat_int32` directly to
`position_ids`, which can become stale when CUDA graph replay reuses the same
buffer. Update the logic in qk_norm_attention.py so the flattened int32
`position_ids_arg` is computed fresh on each call instead of memoized on the
tensor object, while keeping the rest of the position-id adjustment flow
unchanged.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h`:
- Around line 42-45: The new public members in xqaParams are undocumented, so
add Doxygen-style member comments for q_norm_weight, k_norm_weight, qk_norm_eps,
and qk_norm_use_gemma, matching the style used for qkv_token_stride. Describe
their semantics/units, note that they are only for the TRTLLM-Gen fused-kernel
path, and mention any relevant constraints such as the numerically sensitive
head_size==256/eps>0 usage so future callers do not misuse these fields.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 61590bc4-faed-44b0-bbae-d59f5a386e7b
📒 Files selected for processing (14)
cpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/common/attentionOp.hcpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.hcpp/tensorrt_llm/kernels/unfusedAttentionKernels.hcpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.hcpp/tensorrt_llm/kernels/xqaDispatcher.cppcpp/tensorrt_llm/nanobind/thop/bindings.cppcpp/tensorrt_llm/thop/attentionOp.cppcpp/tensorrt_llm/thop/attentionOp.htensorrt_llm/_torch/attention_backend/fmha/fallback.pytensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/qk_norm_attention.py
Comment on lines
+1770
to
+1773
| TLLM_CHECK_WITH_INFO( | ||
| params.attention_input_token_stride == (mNumAttnHeads + 2 * mNumAttnKVHeads) * getHeadSize() | ||
| || (fmha_reads_preprocessed_buffers && mCpSize == 1 && !mIsMLAEnabled && !useSageAttnSeparateQkv), | ||
| "Strided attention_input is not supported by this context attention configuration."); |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Extract the duplicated packed-QKV-stride formula into a shared helper.
(mNumAttnHeads + 2 * mNumAttnKVHeads) * getHeadSize() is repeated identically at three separate call sites (context FMHA path, unfused context MHA fallback, masked MHA generation fallback), each used as the safety boundary for rejecting unsupported strided attention_input. Any future change to head-dim handling (e.g. MLA/vision head sizing) risks being updated in only some of these copies, silently reopening the exact hazard the checks exist to prevent.
♻️ Proposed refactor: shared helper for packed QKV token stride
+ // Packed (non-strided) token stride of attention_input, in elements: Q + 2*KV heads * head size.
+ [[nodiscard]] int32_t getPackedAttentionInputTokenStride() const
+ {
+ return (mNumAttnHeads + 2 * mNumAttnKVHeads) * getHeadSize();
+ }Then reuse it at each check site, e.g.:
- params.attention_input_token_stride == (mNumAttnHeads + 2 * mNumAttnKVHeads) * getHeadSize()
+ params.attention_input_token_stride == getPackedAttentionInputTokenStride()Also applies to: 2115-2117, 2621-2624
🤖 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 `@cpp/tensorrt_llm/common/attentionOp.cpp` around lines 1770 - 1773, The
packed-QKV token stride formula is duplicated across the context FMHA, unfused
context MHA fallback, and masked MHA generation fallback checks in
attentionOp.cpp. Extract the shared `(mNumAttnHeads + 2 * mNumAttnKVHeads) *
getHeadSize()` logic into a helper with a clear name, then update each strided
`attention_input` safety check to call that helper so future head-size changes
only need one update.
Comment on lines
+762
to
+773
| float qSumSquares = 0.0f; | ||
| float kSumSquares = 0.0f; | ||
| #pragma unroll | ||
| for (int i = 0; i < kElementsPerVector; ++i) | ||
| { | ||
| qSumSquares += qValues[i] * qValues[i]; | ||
| kSumSquares += kValues[i] * kValues[i]; | ||
| } | ||
| qSumSquares = tensorrt_llm::common::warpReduceSum(qSumSquares); | ||
| kSumSquares = tensorrt_llm::common::warpReduceSum(kSumSquares); | ||
| float const qRmsRcp = rsqrtf(qSumSquares / static_cast<float>(headSize) + eps); | ||
| float const kRmsRcp = rsqrtf(kSumSquares / static_cast<float>(headSize) + eps); |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect how applyBiasRopeUpdateKVCacheV2 maps threads/warps to head dims and its launch config.
fd unfusedAttentionKernels_2_template.h --exec sed -n '800,1030p' {}
rg -nP -C4 'applyBiasRopeUpdateKVCacheV2\s*<<<' cpp/tensorrt_llm/kernels
rg -nP -C3 'head_dim_idx|rotated_head_dim_offset|first_half|ELTS_PER_VEC' cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.hRepository: NVIDIA/TensorRT-LLM
Length of output: 24541
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the V2 kernel declaration and launch/configuration sites.
sed -n '730,840p' cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h
printf '\n--- launch/search ---\n'
rg -n -C4 'applyBiasRopeUpdateKVCacheV2|BLOCK_SIZE|warpReduceSum|size_per_head == 256|head_size == 256|head_size.*256' cpp/tensorrt_llm/kernels/unfusedAttentionKernels
printf '\n--- warpReduceSum definition ---\n'
rg -n -C3 'warpReduceSum' cpp/tensorrt_llm/kernelsRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1419,1455p' cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h
printf '\n--- head-size dispatch callers ---\n'
rg -n -C3 'kernelV2DispatchHeadSize|invoke.*V2|Dh == 128|Dh == 256|size_per_head' cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.hRepository: NVIDIA/TensorRT-LLM
Length of output: 15562
Fix the RMS reduction to stay within one token’s vector group. The V2 mapping assigns VECS_PER_HEAD threads per token, so for most supported head sizes a warp spans multiple tokens. warpReduceSum then mixes unrelated tokens and produces the wrong normalization factor. Use a token-scoped reduction (or restrict this path to a 32-lane-per-token layout).
🤖 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
`@cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h`
around lines 762 - 773, The RMS normalization in the q/k accumulation block is
using warpReduceSum in a way that can mix values from multiple tokens when
VECS_PER_HEAD is less than a warp, so adjust the reduction in
unfusedAttentionKernels_2_template.h to be token-scoped instead of warp-scoped.
Update the qSumSquares/kSumSquares path near the qRmsRcp and kRmsRcp computation
so the sum only covers lanes belonging to the same token, or gate this code path
behind a 32-lane-per-token layout constraint in the relevant unfused attention
kernel mapping.
Comment on lines
+306
to
314
| # position_ids is identical for every layer; cache the flattened int32 | ||
| # copy on the tensor object so the cast runs once per forward instead | ||
| # of once per layer (fresh position_ids objects are created each step). | ||
| position_ids_arg = getattr(position_ids, "_tllm_flat_int32", None) | ||
| if position_ids_arg is None: | ||
| position_ids_arg = position_ids.reshape(-1).contiguous().to( | ||
| torch.int32) | ||
| position_ids._tllm_flat_int32 = position_ids_arg | ||
| mrope_section1, mrope_section2 = 0, 0 |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Look for in-place position_ids mutation and static/reused position_ids buffers.
rg -nP -C3 '_adjust_position_ids_for_spec_dec|position_ids\.view\(-1\)\[' tensorrt_llm/_torch/modules/attention.py
rg -nP -C3 'position_ids' tensorrt_llm/_torch/pyexecutor/*.py | rg -nP -C3 'static|graph|buffer|copy_|zero_'Repository: NVIDIA/TensorRT-LLM
Length of output: 26758
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the exact call flow around the reuse/mutation sites.
sed -n '1150,1185p' tensorrt_llm/_torch/modules/attention.py
printf '\n---\n'
sed -n '4880,5005p' tensorrt_llm/_torch/pyexecutor/model_engine.py
printf '\n---\n'
sed -n '4300,4375p' tensorrt_llm/_torch/pyexecutor/model_engine.py
printf '\n---\n'
sed -n '450,490p' tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
printf '\n---\n'
sed -n '1018,1095p' tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 17659
Avoid tensor-attached caching for position_ids
tensorrt_llm/_torch/modules/qk_norm_attention.py:306-314 — _adjust_position_ids_for_spec_dec() mutates position_ids in place, and CUDA-graph replay reuses static position_ids buffers. Memoizing _tllm_flat_int32 on the tensor object can return stale values on a later forward; compute the flattened int32 view per call instead of storing it on position_ids.
🤖 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 306 - 314, The
caching added in _adjust_position_ids_for_spec_dec() should be removed because
it attaches `_tllm_flat_int32` directly to `position_ids`, which can become
stale when CUDA graph replay reuses the same buffer. Update the logic in
qk_norm_attention.py so the flattened int32 `position_ids_arg` is computed fresh
on each call instead of memoized on the tensor object, while keeping the rest of
the position-id adjustment flow unchanged.
Remove the environment-variable gates so supported output-gated Qwen3.5 attention layers use the gate-tail layout and fused QK norm/RoPE preprocessing by default. Unsupported configurations retain the guarded fallback. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57903 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57903 [ run ] completed with state
|
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Collaborator
Author
|
/bot run |
1 task
Collaborator
Author
|
See #16322 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Background
Qwen3.5 full-attention layers with an output gate currently produce the fused projection in a per-head interleaved
layout:
[q0 | g0 | q1 | g1 | ... | K | V]
Before FMHA, this requires Q/gate de-interleaving copies, Q/K/V concatenation, and a standalone QK RMSNorm + RoPE
kernel. In the profiled decode path, the sequence between the QKV GEMM and FMHA contained nine kernels, including
four data-movement kernels totaling approximately 12.5 µs per full-attention layer per step.
This PR introduces two opt-in optimization stages for that path.
Summary
Phase 1: Gate-tail QKV layout
At weight-load time, the fused projection rows are permuted into:
[Q | K | V | G]
This makes Q, K, V, and the output gate zero-copy views of the GEMM result.
The implementation:
Removes the Q and gate de-interleaving copies.
Removes the cat(Q, K, V) operation.
Caches the flattened int32 position IDs instead of casting them once per layer.
Propagates the QKV input token stride through THOP, AttentionOp, XQA dispatch, and the shared QKV preprocessing
kernel.
Uses the existing QK norm/RoPE kernel with the gate tail represented as pass-through V heads.
The per-layer kernel sequence is reduced from:
QKV GEMM
→ Q copy
→ gate copy
→ cat(Q, K, V)
→ position IDs cast
→ fusedQKNormRope
→ QKV preprocessing
→ Fill
→ FMHA
to:
QKV GEMM
→ fusedQKNormRope
→ QKV preprocessing
→ Fill
→ FMHA
Phase 2: Fuse QK RMSNorm and RoPE into preprocessing
The second optimization folds Gemma Q/K RMSNorm and GPT-NeoX RoPE into applyBiasRopeUpdateKVCacheV2.
The fused preprocessing kernel now performs:
This removes the standalone fusedQKNormRopeKernel and reduces the target sequence to:
QKV GEMM
→ applyBiasRopeUpdateKVCacheV2(QK norm + RoPE + KV append)
→ Fill
→ FMHA
The PR also fixes text-only requests for mRoPE-configured models. These requests do not provide a per-token mRoPE
buffer and previously selected the V1 preprocessing kernel, where fused QK normalization is unavailable. With this
change, text-only mRoPE uses the V2 scalar-position cos/sin cache, while configurations that genuinely require V1
explicitly reject fused QK normalization instead of silently skipping it.
Feature flags
Both optimizations are disabled by default:
Phase 1
TRTLLM_ATTN_GATE_TAIL_LAYOUT=1
Phase 2; requires Phase 1
TRTLLM_ATTN_GATE_TAIL_FUSED_PREPROCESS=1
Unsupported configurations preserve the existing path or fail with an explicit check where fallback would be
unsafe.
Phase 2 is currently restricted to:
LoRA with the load-time gate-tail permutation is explicitly rejected.
Performance
Measured on one NVIDIA B200 with Qwen3.5-35B-A3B, TP1, BF16 model weights, FP8 KV cache, 512 requests, concurrency
512, ISL 128, OSL 64, and 32 warmup requests:
A separate historical Phase 1 measurement showed a +0.61% output-throughput improvement over the original layout.
Since that measurement used a different build/workload, it is not presented as a formal combined result.
A short non-CUDA-graph Nsight Systems run confirmed:
The short trace validates the execution path and is not a replacement for the concurrency-512 benchmark.
Correctness and validation
Successfully rebuilt the SM100 FAST_BUILD targets:
Focused Q/K numerical comparison against the standalone implementation:
Qwen3.5-35B-A3B Phase 1 repeatability was 3/4 token-exact prompts because of MoE nondeterminism.
Phase 1 versus Phase 2 reached the same 3/4 repeatability floor.
Every Phase 2 output matched an observed valid Phase 1 repeat.
All pre-commit hooks pass for the changed files.
DCO, Python syntax, formatting, and diff checks pass.
Compatibility and risk
and epsilon.
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.