[None][perf] Fold q/k/v quantization into qknorm_rope_fused kernel & remove contiguous - #17093
Open
brb-nv wants to merge 1 commit into
Open
[None][perf] Fold q/k/v quantization into qknorm_rope_fused kernel & remove contiguous#17093brb-nv wants to merge 1 commit into
brb-nv wants to merge 1 commit into
Conversation
brb-nv
force-pushed
the
user/brb/fold-qkv-quant-qknorm-rope-main
branch
from
August 1, 2026 20:14
60af8fe to
be54863
Compare
brb-nv
marked this pull request as ready for review
August 1, 2026 20:16
Contributor
WalkthroughChangesThe fused QK normalization and RoPE kernel now supports out-of-place BF16 or FP8 E4M3 output, including optional V conversion. A Torch operator exposes the FP8 path. MiniMax-M3 attention selects it for supported FP8 KV-cache configurations and preserves backend-specific tensor layouts. Tests cover parameterized FP8 behavior. FP8 fused kernel and public API
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant MiniMaxM3Attention
participant fused_qk_norm_rope_to_fp8
participant fusedQKNormRopeKernel
participant FP8KVCache
MiniMaxM3Attention->>fused_qk_norm_rope_to_fp8: request FP8 fused QKV output
fused_qk_norm_rope_to_fp8->>fusedQKNormRopeKernel: validate inputs and launch kernel
fusedQKNormRopeKernel->>FP8KVCache: write FP8 E4M3 Q, K, and V
FP8KVCache-->>MiniMaxM3Attention: return FP8 QKV tensors
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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.
🧹 Nitpick comments (4)
tensorrt_llm/_torch/models/modeling_minimaxm3.py (1)
991-1023: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse built-in generic types in the new return annotations.
Replace
Tuple[...]withtuple[...]in both helper signatures. The project guidelines prefer built-in generic types.Proposed change
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... - def _split_index_qk(self, fused_idx: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def _split_index_qk(self, fused_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:🤖 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/models/modeling_minimaxm3.py` around lines 991 - 1023, Update the return annotations of _split_main_qkv and _split_index_qk to use the built-in tuple[...] generic instead of Tuple[...], preserving the existing tensor element types and method behavior.Source: Coding guidelines
tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py (1)
347-351: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse a private UPPER_SNAKE_CASE constant.
fp8_num_heads_groupsis a module-level non-public constant. Rename it to_FP8_NUM_HEADS_GROUPS. Prefer a tuple to prevent mutation.As per coding guidelines, “use … UPPER_SNAKE_CASE for constants” and “Prefix non-public names with
_.”🤖 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/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py` around lines 347 - 351, Rename the module-level constant fp8_num_heads_groups to _FP8_NUM_HEADS_GROUPS and change its collection type from list to tuple, updating all references accordingly while preserving the existing head-group values.Source: Coding guidelines
cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp (1)
92-138: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract shared input validation to avoid duplicated checks.
The validation block in
fused_qk_norm_rope_to_fp8(dim checks, position_ids shape, weight shape,CHECK_INPUTcalls,total_heads * head_dimcheck) duplicates the block infused_qk_norm_rope(Lines 57-77) almost verbatim. Extract a shared private helper that both functions call, so a future validation fix does not need to land in two places.♻️ Proposed refactor sketch
namespace { int64_t validateFusedQKNormRopeInputs(torch::Tensor const& qkv, torch::Tensor const& position_ids, torch::Tensor const& q_weight, torch::Tensor const& k_weight, int64_t num_heads_q, int64_t num_heads_k, int64_t num_heads_v, int64_t head_dim, bool use_mrope) { TORCH_CHECK(qkv.dim() == 2, "QKV tensor must be 2D: [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]"); TORCH_CHECK(position_ids.dim() == 1 || (position_ids.dim() == 2 && position_ids.size(0) == 3), "Position IDs must be 1D [num_tokens] (plain RoPE) or 2D [3, num_tokens] (mRoPE)"); TORCH_CHECK(!use_mrope || position_ids.dim() == 2, "use_mrope requires 2D [3, num_tokens] position_ids"); TORCH_CHECK(q_weight.dim() == 1, "Query weights must be 1D: [head_dim]"); TORCH_CHECK(k_weight.dim() == 1, "Key weights must be 1D: [head_dim]"); TORCH_CHECK(q_weight.size(0) == head_dim, "Query weights size must match head dimension"); TORCH_CHECK(k_weight.size(0) == head_dim, "Key weights size must match head dimension"); CHECK_INPUT(qkv, torch::kBFloat16); CHECK_INPUT(position_ids, torch::kInt32); CHECK_INPUT(q_weight, torch::kBFloat16); CHECK_INPUT(k_weight, torch::kBFloat16); int64_t num_tokens = qkv.size(0); TORCH_CHECK(position_ids.size(-1) == num_tokens, "Number of tokens in position_ids must match QKV"); int64_t total_heads = num_heads_q + num_heads_k + num_heads_v; TORCH_CHECK( qkv.size(1) == total_heads * head_dim, "QKV tensor size must match total number of heads and head dimension"); return num_tokens; } } // namespaceBoth
fused_qk_norm_ropeandfused_qk_norm_rope_to_fp8would call this helper instead of repeating the checks.🤖 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/thop/fusedQKNormRopeOp.cpp` around lines 92 - 138, Extract the duplicated validation from fused_qk_norm_rope and fused_qk_norm_rope_to_fp8 into a shared private validateFusedQKNormRopeInputs helper. Move all dimension, shape, dtype, token-count, and total-head checks into that helper, have both functions call it, and reuse its returned token count while preserving the existing validation behavior and messages.cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu (1)
435-472: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd coverage for the BF16 out-of-place path or remove it.
The only in-tree caller passes
out_fp8=trueandprocess_v=true. No repository call site exercisesout_fp8=false, process_v=true; add a BF16 out-of-place operation and test, or remove this unused branch.🤖 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/fusedQKNormRopeKernel.cu` around lines 435 - 472, The launchFusedQKNormRopeOut branch for out_fp8=false and process_v=true lacks repository coverage. Add a BF16 out-of-place caller and test that exercises this combination, or remove the unsupported unused branch while preserving the existing FP8 path and other valid 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.
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu`:
- Around line 435-472: The launchFusedQKNormRopeOut branch for out_fp8=false and
process_v=true lacks repository coverage. Add a BF16 out-of-place caller and
test that exercises this combination, or remove the unsupported unused branch
while preserving the existing FP8 path and other valid behavior.
In `@cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp`:
- Around line 92-138: Extract the duplicated validation from fused_qk_norm_rope
and fused_qk_norm_rope_to_fp8 into a shared private
validateFusedQKNormRopeInputs helper. Move all dimension, shape, dtype,
token-count, and total-head checks into that helper, have both functions call
it, and reuse its returned token count while preserving the existing validation
behavior and messages.
In `@tensorrt_llm/_torch/models/modeling_minimaxm3.py`:
- Around line 991-1023: Update the return annotations of _split_main_qkv and
_split_index_qk to use the built-in tuple[...] generic instead of Tuple[...],
preserving the existing tensor element types and method behavior.
In `@tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py`:
- Around line 347-351: Rename the module-level constant fp8_num_heads_groups to
_FP8_NUM_HEADS_GROUPS and change its collection type from list to tuple,
updating all references accordingly while preserving the existing head-group
values.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8b8c8f48-b47b-4c79-b3ea-31c92b6b08de
📒 Files selected for processing (7)
cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cucpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.hcpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpptensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py
Collaborator
|
@brb-nv can we have a proper title and description? If it's not ready, please mark as Draft. THanks. |
…remove contiguous (NVIDIA#16699) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
force-pushed
the
user/brb/fold-qkv-quant-qknorm-rope-main
branch
from
August 3, 2026 23:07
be54863 to
2a0c68e
Compare
brb-nv
changed the title
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu (2)
143-145: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd braces to the changed control-flow bodies.
Use braces for the token bounds check and the
defaultswitch case.Proposed fix
- if (tokenIdx >= num_tokens) - return; + if (tokenIdx >= num_tokens) + { + return; + } ... - default: TLLM_THROW("Unsupported head dimension for fusedQKNormRope: %d", head_dim); + default: + { + TLLM_THROW("Unsupported head dimension for fusedQKNormRope: %d", head_dim); + }As per coding guidelines, “use Allman braces” and “braced control-flow bodies.”
Also applies to: 431-432
🤖 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/fusedQKNormRopeKernel.cu` around lines 143 - 145, Update the token bounds check near the warp token handling to use Allman-style braces around its early-return body, and apply the same braced format to the switch statement’s default case. Leave the existing control-flow behavior unchanged.Source: Coding guidelines
442-446: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
static_castforvoid*conversions.These conversions start from
void*orvoid const*. Usestatic_castfor them. Keepreinterpret_castonly where representation reinterpretation is required.Proposed fix
- launchFusedQKNormRopeImpl<__nv_bfloat16>(reinterpret_cast<__nv_bfloat16 const*>(qkv), - reinterpret_cast<__nv_bfloat16*>(qkv), /*process_v=*/false, num_tokens, num_heads_q, num_heads_k, num_heads_v, + launchFusedQKNormRopeImpl<__nv_bfloat16>(static_cast<__nv_bfloat16 const*>(qkv), + static_cast<__nv_bfloat16*>(qkv), /*process_v=*/false, num_tokens, num_heads_q, num_heads_k, num_heads_v, ... - auto const* in = reinterpret_cast<__nv_bfloat16 const*>(qkv_in); + auto const* in = static_cast<__nv_bfloat16 const*>(qkv_in);As per coding guidelines, “use
static_castfromvoid*.”Also applies to: 455-469
🤖 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/fusedQKNormRopeKernel.cu` around lines 442 - 446, Update the QKV pointer conversions in the launchFusedQKNormRopeImpl calls around the shown code and the corresponding lines at 455–469: replace reinterpret_cast conversions from void* or void const* with static_cast, while preserving reinterpret_cast only for conversions that require representation reinterpretation.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/kernels/fusedQKNormRopeKernel.cu`:
- Around line 371-377: Validate rotary_dim in launchFusedQKNormRopeImpl before
calculating launch dimensions or dispatching the kernel, rejecting values less
than 1 or greater than head_dim while preserving the existing evenness
validation. Ensure invalid values cannot reach RoPE frequency calculation or
kernel launch.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu`:
- Around line 143-145: Update the token bounds check near the warp token
handling to use Allman-style braces around its early-return body, and apply the
same braced format to the switch statement’s default case. Leave the existing
control-flow behavior unchanged.
- Around line 442-446: Update the QKV pointer conversions in the
launchFusedQKNormRopeImpl calls around the shown code and the corresponding
lines at 455–469: replace reinterpret_cast conversions from void* or void const*
with static_cast, while preserving reinterpret_cast only for conversions that
require representation reinterpretation.
🪄 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: 63421f10-2851-4fde-b98a-74282933f119
📒 Files selected for processing (7)
cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cucpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.hcpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpptensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py
🚧 Files skipped from review as they are similar to previous changes (6)
- cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h
- tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
- tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py
- tensorrt_llm/_torch/models/modeling_minimaxm3.py
- cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp
Comment on lines
+371
to
+377
| template <typename OutT> | ||
| static void launchFusedQKNormRopeImpl(__nv_bfloat16 const* qkv_in, OutT* qkv_out, bool const process_v, | ||
| int const num_tokens, int const num_heads_q, int const num_heads_k, int const num_heads_v, int const head_dim, | ||
| int const rotary_dim, float const eps, __nv_bfloat16 const* q_weight, __nv_bfloat16 const* k_weight, | ||
| float const base, bool const interleave, int const* position_ids, float factor, float low, float high, | ||
| float attention_factor, cudaStream_t stream, bool is_qk_norm, bool use_gemma, bool use_mrope, int mrope_section1, | ||
| int mrope_section2) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Validate rotary_dim before dispatch.
The FP8 operator checks tensor shapes but does not constrain rotary_dim. A value of 0 passes the existing evenness check and then causes division by zero during RoPE frequency calculation. A value greater than head_dim also violates the kernel pairing assumptions. Reject values outside 1..head_dim before calculating launch dimensions.
Proposed fix
- TLLM_CHECK_WITH_INFO(rotary_dim % 2 == 0, "rotary_dim must be even");
+ TLLM_CHECK_WITH_INFO(
+ rotary_dim > 0 && rotary_dim <= head_dim && rotary_dim % 2 == 0,
+ "rotary_dim must be positive, no greater than head_dim, and even");🤖 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/fusedQKNormRopeKernel.cu` around lines 371 - 377,
Validate rotary_dim in launchFusedQKNormRopeImpl before calculating launch
dimensions or dispatching the kernel, rejecting values less than 1 or greater
than head_dim while preserving the existing evenness validation. Ensure invalid
values cannot reach RoPE frequency calculation or kernel launch.
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63577 [ run ] triggered by Bot. Commit: |
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
3 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
This MR does the following:
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.Overview
fused_qk_norm_rope_to_fp8with CUDA and Meta implementations.viewwithreshapefor strided tensor handling.Dev Engineer Review
launchFusedQKNormRopeOutprovides out-of-place QKV processing with optional V conversion.reshapechanges support strided views and avoid unnecessary copies.CODING_GUIDELINES.md.QA Engineer Review
tests/integration/test_lists/entries changed.