[None][perf] Fuse Qwen3.5 gated QKV preprocessing for qwen3.5 v2 ckpt#16322
[None][perf] Fuse Qwen3.5 gated QKV preprocessing for qwen3.5 v2 ckpt#16322nv-guomingz wants to merge 7 commits into
Conversation
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThe change propagates QK normalization and token-stride parameters through attention APIs and preprocessing kernels, adds BF16 fused QK normalization with GPT-NeoX RoPE, enables backend selection for strided fused QKV layouts, and introduces gate-tail QKV projection support with tests. ChangesAttention parameter and preprocessing pipeline
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant QKNormRoPEAttention
participant TrtllmAttention
participant thop_attention
participant QKVPreprocessing
QKNormRoPEAttention->>TrtllmAttention: configure QK normalization preprocessing
TrtllmAttention->>thop_attention: pass Q/K weights and normalization settings
thop_attention->>QKVPreprocessing: enqueue strided QKV preprocessing
QKVPreprocessing->>TrtllmAttention: return normalized and RoPE-transformed Q/K
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
tests/unittest/_torch/attention/test_gate_tail_layout.py (1)
23-88: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winTest coverage for
_maybe_permute_gate_tail_layoutis insufficient; runtimeforward()gate-tail path is entirely untested.As per path instructions, coverage assessment for
tests/**: this suite only exercises the load-time permutation helper, and even there leaves gaps:
- No test sets
attention.qkv_proj = nn.Linear(..., bias=True), so the bias-permutation branch inattention.py(if self.qkv_proj.bias is not None: ...) is never exercised. Suggest adding e.g.test_gate_tail_layout_permutes_bias.- The guard chain in
_maybe_permute_gate_tail_layout(attn_backend == "TRTLLM",support_fused_qkv,fuse_qk_norm_rope,not skip_rope,cp_size == 1) is only exercised through thesupported/quant_kindparameters; none of these individual flags are flipped toFalsein isolation, so a regression that drops one of these checks wouldn't be caught. Suggest parametrizing_make_gate_tail_attentionover each guard flag independently.Attention.forward()'s gate-tail branch (attention.pylines 1049-1066 and 1086-1090) — the q/k/v/gate extraction and q re-slicing that this permutation exists to enable — has no test at all in this file. A numerical-equivalence test (e.g.test_gate_tail_layout_forward_matches_unpermuted_reference) comparingforward()output/gate values between the permuted-active and interleaved-inactive paths on identical inputs would give real confidence the permutation preserves inference correctness, which the current unit tests can't demonstrate.Coverage is currently insufficient for the runtime-critical path; the load-time permutation itself is reasonably well covered.
🤖 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/attention/test_gate_tail_layout.py` around lines 23 - 88, The tests need broader coverage for gate-tail layout behavior. Extend _make_gate_tail_attention and add focused tests that exercise qkv_proj bias permutation, each guard in _maybe_permute_gate_tail_layout independently disabled, and Attention.forward’s active gate-tail branch. Include a numerical comparison between permuted-active and unpermuted-reference attention paths, validating equivalent outputs and gate values while preserving the existing permutation tests.Source: Path instructions
tensorrt_llm/_torch/modules/qk_norm_attention.py (1)
304-312: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winCaching via a monkey-patched tensor attribute is fragile under graph capture.
position_ids._tllm_flat_int32 = position_ids_argattaches an ad-hoc Python attribute to theposition_idstensor to memoize the flattened int32 cast across layers within a forward. This works for eager execution, but attribute-based side channels on tensors can behave unpredictably undertorch.compile/dynamo tracing or CUDA-graph capture (attribute reads/writes on captured/traced tensors aren't guaranteed to be preserved or may trigger recompilation), especially since this module already threadsdisable_on_compileelsewhere for compile-awareness.Consider caching by
id(position_ids)in an external dict (cleared per forward) instead of mutating the tensor object itself.🤖 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 304 - 312, Replace the monkey-patched _tllm_flat_int32 attribute access in the position_ids handling with an external cache keyed by id(position_ids), initialized and cleared per forward so it is shared across layers without mutating tensors. Preserve the existing reshape, contiguous, int32 conversion, and reuse behavior.
🤖 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 2127-2129: Update the validation guards for the unfused context
MHA and masked-MHA generation fallback paths near the existing
attention_input_token_stride checks to reject non-null q_norm_weight or
k_norm_weight before dispatch. Preserve the current stride validation, and
ensure both fallback paths cannot silently proceed when QK normalization
parameters are provided.
- Around line 250-254: Update enqueueGeneration where xqaParams.qkv is replaced
by the packed Ulysses workspace for mCpSize > 1: reject non-packed inputs or
reset qkv_token_stride to the packed workspace’s row stride so it remains
consistent with xqaParams.qkv. Preserve the existing stride for non-Ulysses
generation.
In
`@cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h`:
- Around line 993-1006: Update invokeQKVPreprocessing around qk_norm_applied so
a non-null params.q_norm_weight with T other than __nv_bfloat16 fails fast
instead of entering the plain GPT-NeoX rotation fallback. Preserve the existing
BF16 applyQkNormRopeGptNeox path and only allow the fallback when QK-norm is not
requested.
---
Nitpick comments:
In `@tensorrt_llm/_torch/modules/qk_norm_attention.py`:
- Around line 304-312: Replace the monkey-patched _tllm_flat_int32 attribute
access in the position_ids handling with an external cache keyed by
id(position_ids), initialized and cleared per forward so it is shared across
layers without mutating tensors. Preserve the existing reshape, contiguous,
int32 conversion, and reuse behavior.
In `@tests/unittest/_torch/attention/test_gate_tail_layout.py`:
- Around line 23-88: The tests need broader coverage for gate-tail layout
behavior. Extend _make_gate_tail_attention and add focused tests that exercise
qkv_proj bias permutation, each guard in _maybe_permute_gate_tail_layout
independently disabled, and Attention.forward’s active gate-tail branch. Include
a numerical comparison between permuted-active and unpermuted-reference
attention paths, validating equivalent outputs and gate values while preserving
the existing permutation tests.
🪄 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: 6bc2544d-98b1-4bd7-be9e-0db484d4d291
📒 Files selected for processing (16)
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/interface.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/qk_norm_attention.pytests/unittest/_torch/attention/test_gate_tail_layout.py
|
PR_Github #58961 [ run ] triggered by Bot. Commit: |
|
PR_Github #58961 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58995 [ run ] triggered by Bot. Commit: |
|
PR_Github #58995 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59047 [ run ] triggered by Bot. Commit: |
|
PR_Github #59047 [ run ] completed with state
|
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>
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>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
5b3d188 to
bf9342a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59140 [ run ] triggered by Bot. Commit: |
|
PR_Github #59140 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59225 [ run ] triggered by Bot. Commit: |
|
PR_Github #59225 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59309 [ run ] triggered by Bot. Commit: |
|
PR_Github #59309 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #59371 [ run ] triggered by Bot. Commit: |
|
PR_Github #59371 [ run ] completed with state
|
Guard the strided (gate-tail) QKV / fused QK-norm preprocessing so unsupported combinations raise instead of silently corrupting output: - Reject strided attention_input under Ulysses/CP (packed-layout preprocess). - Reject fused QK-norm for non-bf16 inputs (bf16-only kernel). - Reject fused QK-norm weights on the unfused context / masked-MHA fallbacks. Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
…ycle The gate-tail permutation ran only in post_load_weights with a permanent flag, giving wrong output on the GMS read-only and RLHF reload paths. Split it across the staged hooks: transform_weights() does the one-shot permute (guarded by _weights_transformed), cache_derived_state() recomputes the layout state and QK-norm backend config, and pre_reload_weights() resets the guard so reloads re-permute. Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
Gate-tail has no LoRA-aware path and asserts on non-empty lora_params, so every LoRA request failed. Gate the layout on the absence of a LoRA config. Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
f69fded to
47a0e5d
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59395 [ run ] triggered by Bot. Commit: |
|
PR_Github #59395 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59473 [ run ] triggered by Bot. Commit: |
|
PR_Github #59473 [ run ] completed with state
|
3 participants
It's a successor version of #16024
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-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.Summary by CodeRabbit