User/brb/fold qkv quant qknorm rope main

[brb-nv]

Description

Must follow #16906.

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.

Overview

• Folded Q/K/V quantization into the fused QK normalization and RoPE kernel.
• Added out-of-place BF16 and FP8 E4M3 output support.
• Added fused_qk_norm_rope_to_fp8 with CUDA and Meta implementations.
• Updated MiniMax-M3 attention to support FP8 KV caches.
• Replaced view with reshape for strided tensor handling.
• Removed the need for the previous contiguous operation.
• Added parameterized FP8 correctness tests.

Dev Engineer Review

• The kernel now supports templated BF16 and FP8 E4M3 output storage.
• launchFusedQKNormRopeOut provides out-of-place QKV processing with optional V conversion.
• The operator validates BF16 inputs and preserves the original input tensor.
• The MiniMax-M3 paths select FP8 or BF16 behavior based on the configured backend and KV-cache dtype.
• The reshape changes support strided views and avoid unnecessary copies.
• The new API and operator require consistency checks against repository API and coding guidelines.
• No configuration or test-list files changed.
• Review risk is high because the kernel launch sizing, head indexing, output layout, and FP8 conversion paths changed together.

QA Engineer Review

• Added parameterized coverage for the out-of-place FP8 fused QK normalization and RoPE path.
• Coverage includes multiple sequence lengths, GQA head configurations, rotary modes, and partial rotary dimensions.
• The test verifies FP8 E4M3 output shape and dtype, unchanged BF16 input, and dequantized agreement with the BF16 reference.
• No tests/integration/test_lists/ entries changed, so CI or manual QA coverage is not established from the available changes.
• Verdict: needs follow-up.

[coderabbitai]

Walkthrough

Changes

The 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

Layer / File(s)	Summary
Templated fused kernel and public API cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu, cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h	The kernel reads separate BF16 input and output buffers, stores BF16 or FP8 E4M3 values, optionally processes V heads, and updates head indexing and launch sizing. A public out-of-place launcher is added.
FP8 Torch operator and Meta registration cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp	The new operator validates BF16 QKV inputs, allocates FP8 output, launches the CUDA kernel, and registers CUDA and Meta implementations.
MiniMax-M3 FP8 attention integration tensorrt_llm/_torch/models/modeling_minimaxm3.py, tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py, tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py	MiniMax-M3 detects FP8 KV-cache use, selects FP8 fused output when supported, preserves strided MSA views, uses contiguous tensors for other backends, and allocates attention outputs with the configured activation dtype.
FP8 output validation tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py	Parameterized tests verify FP8 E4M3 shape and dtype, unchanged BF16 input, out-of-place behavior, and agreement with the BF16 reference after dequantization.

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

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#16906: Extends the fused QK normalization and RoPE path used by modeling_minimaxm3.py.
• NVIDIA/TensorRT-LLM#17047: Overlaps with MiniMax-M3 FP8 attention performance configuration.
• NVIDIA/TensorRT-LLM#17091: Also changes MiniMax-M3 attention forward paths.

Suggested reviewers: qijune

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description references a prior pull request but does not explain the problem or solution, and it leaves the Test Coverage section empty despite adding tests.	Describe the problem and solution, list the relevant tests and results, and complete the checklist items that apply to this API and kernel change.
Docstring Coverage	⚠️ Warning	Docstring coverage is 73.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title identifies the main QKV quantization and QK norm/RoPE fusion change, although it is fragmented and omits the required ticket and type format.
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.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (4)

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

991-1023: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Use built-in generic types in the new return annotations.

Replace Tuple[...] with tuple[...] 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 value

Use a private UPPER_SNAKE_CASE constant.

fp8_num_heads_groups is 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 win

Extract 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_INPUT calls, total_heads * head_dim check) duplicates the block in fused_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;
}
} // namespace

Both fused_qk_norm_rope and fused_qk_norm_rope_to_fp8 would 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 value

Add coverage for the BF16 out-of-place path or remove it.

The only in-tree caller passes out_fp8=true and process_v=true. No repository call site exercises out_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

📥 Commits

Reviewing files that changed from the base of the PR and between fdf7bd5 and be54863.

📒 Files selected for processing (7)

• cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
• cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h
• cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp
• tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py
• tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
• tensorrt_llm/_torch/models/modeling_minimaxm3.py
• tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py

[pcastonguay]

@brb-nv can we have a proper title and description? If it's not ready, please mark as Draft. THanks.