[None][perf] Fold q/k/v quantization into qknorm_rope_fused kernel & remove contiguous - #17093
Conversation
60af8fe to
be54863
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe change adds an out-of-place BF16-to-FP8 fused QK normalization and RoPE kernel. It adds Torch operator registration and MiniMax-M3 integration with backend-aware tensor handling. Tests cover FP8 output, V conversion, input preservation, and BF16 reference agreement. ChangesFP8 QK normalization and RoPE
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant MiniMaxM3Attention
participant fused_qk_norm_rope_to_fp8
participant launchFusedQKNormRopeToFp8
participant fusedQKNormRopeKernel
participant AttentionBackend
MiniMaxM3Attention->>fused_qk_norm_rope_to_fp8: request FP8 Q/K/V processing
fused_qk_norm_rope_to_fp8->>launchFusedQKNormRopeToFp8: pass BF16 QKV and FP8 output
launchFusedQKNormRopeToFp8->>fusedQKNormRopeKernel: launch fused operation
fusedQKNormRopeKernel-->>fused_qk_norm_rope_to_fp8: write normalized Q/K and cast V
fused_qk_norm_rope_to_fp8-->>MiniMaxM3Attention: return FP8 Q/K/V
MiniMaxM3Attention->>AttentionBackend: pass backend-specific tensor views
AttentionBackend-->>MiniMaxM3Attention: return attention output in activation dtype
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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.
🧹 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
|
@brb-nv can we have a proper title and description? If it's not ready, please mark as Draft. THanks. |
be54863 to
2a0c68e
Compare
brb-nv
changed the title
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
|
/bot run --disable-fail-fast |
|
PR_Github #63577 [ run ] triggered by Bot. Commit: |
WeiHaocheng
left a comment
WeiHaocheng
left a comment
There was a problem hiding this comment.
Approve for the modeling parts.
|
PR_Github #63577 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63651 [ run ] triggered by Bot. Commit: |
yunruis
left a comment
yunruis
left a comment
There was a problem hiding this comment.
Reviewed the kernel, the op registration, and the MiniMax-M3 gating. The overall approach looks right, and I verified the parts I was most worried about:
- The two independent "is the KV cache FP8" decisions cannot diverge. The model side uses
quant_config.quant_mode.has_fp8_kv_cache()whilerun_msa_paged_gqausesk_paged.dtype == float8_e4m3fn, butTorchLlmArgs.sync_quant_config_with_kv_cache_config_dtypemapsKvCacheConfig(dtype="fp8")ontokv_cache_quant_algo = FP8, and_util.pyderives the cache dtype from the same predicate. - FP8 / strided
qcannot leak into the SDPA or Triton cores:_dispatch_attention_backend's first branch tests exactly the same condition as_msa_backend_active(). write_kv_slotshandles a strided FP8 source fine (values.to(cache.dtype)is a no-op, and the advanced-index assignment goes through TensorIterator).- The index branch correctly stays BF16, since
_torch_dtype_for_index_cache()always returns bfloat16. - Pinning the attention
outputdtype is a genuine bug fix; it previously inherited FP8 fromq.
Two things I would like to see addressed before merge, commented inline.
|
PR_Github #63651 [ run ] completed with state |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp (1)
121-135: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winValidate dimensions before shape arithmetic and CUDA dispatch.
Reject negative head counts and values outside the CUDA
intrange before computingtotal_heads * head_dimor allocatingout. Otherwise,(-1, 2, 0, 1)passes the shape check but produces invalid kernel offsets, while signed overflow can invalidate the shape check. Apply the same validation tofused_qk_norm_ropeand the Meta implementation.🤖 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 121 - 135, The fused QK norm RoPE implementations must validate all head counts and CUDA-dispatched dimensions before shape arithmetic or output allocation. In the shared validation used by fused_qk_norm_rope and its Meta implementation, reject negative num_heads_q, num_heads_k, and num_heads_v, plus any values exceeding the CUDA int range; also validate other dimensions cast to int as needed. Perform these checks before computing total_heads * head_dim, allocating out, or calling launchFusedQKNormRopeToFp8, while preserving valid-input behavior.
🧹 Nitpick comments (1)
cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h (1)
54-62: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument both new interfaces with Doxygen.
cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h#L54-L62: replace the//block with Doxygen documentation for the public kernel declaration.cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp#L92-L102: replace the//block with Doxygen documentation for the Torch operator interface.As per coding guidelines, “document new interfaces with Doxygen.”
🤖 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.h` around lines 54 - 62, Replace the existing line-comment blocks with Doxygen documentation for both new interfaces: document launchFusedQKNormRopeToFp8 in cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h at lines 54-62, and document the Torch operator interface in cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp at lines 92-102. Describe each interface’s purpose, inputs, outputs, and relevant layout/quantization behavior without changing implementation logic.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.
Outside diff comments:
In `@cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp`:
- Around line 121-135: The fused QK norm RoPE implementations must validate all
head counts and CUDA-dispatched dimensions before shape arithmetic or output
allocation. In the shared validation used by fused_qk_norm_rope and its Meta
implementation, reject negative num_heads_q, num_heads_k, and num_heads_v, plus
any values exceeding the CUDA int range; also validate other dimensions cast to
int as needed. Perform these checks before computing total_heads * head_dim,
allocating out, or calling launchFusedQKNormRopeToFp8, while preserving
valid-input behavior.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h`:
- Around line 54-62: Replace the existing line-comment blocks with Doxygen
documentation for both new interfaces: document launchFusedQKNormRopeToFp8 in
cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h at lines 54-62, and document
the Torch operator interface in cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp at
lines 92-102. Describe each interface’s purpose, inputs, outputs, and relevant
layout/quantization behavior without changing implementation logic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 20b9e2d5-4ffc-4de8-9424-f3e17b442d25
📒 Files selected for processing (4)
cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cucpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.hcpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpptests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py
- cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
|
/bot run --disable-fail-fast |
|
PR_Github #64159 [ run ] triggered by Bot. Commit: |
|
PR_Github #64355 [ run ] triggered by Bot. Commit: |
|
PR_Github #64355 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64363 [ run ] triggered by Bot. Commit: |
|
PR_Github #64363 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #64379 [ run ] triggered by Bot. Commit: |
|
PR_Github #64379 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64436 [ run ] triggered by Bot. Commit: |
|
PR_Github #64436 [ run ] completed with state
|
…remove contiguous (NVIDIA#16699) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Bound rotary_dim to 1..head_dim before dispatch; rotary_dim == 0 passed the evenness check and then divided by zero deriving the RoPE frequencies. This matches the predicate the sibling Triton path already gates on. Replace launchFusedQKNormRopeOut with an FP8-only launchFusedQKNormRopeToFp8. The BF16 out-of-place branch had no caller and no test, and left V uninitialized for any genuinely out-of-place caller. Guard the remaining process_v=false case on input and output aliasing so that combination cannot be reached. Tighten the FP8 test tolerance from rtol=0.2 (3x looser than E4M3 needs, so it would not catch swapped norm weights) to 0.07, assert V matches torch's E4M3 cast bit-exactly, and parametrize use_gemma so the Gemma norm MiniMax-M3 actually ships is covered. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Drop the .reshape() and dtype-guard changes in msa_sparse_gqa.py and msa_backend.py: .view() already handles the strided column-views the model now hands over, and .to() on an already-FP8 tensor is a no-op, so both files return to their original form. Cut the explanatory comments added across the kernel, op, model and test down to what is not already obvious from the code. Fold the duplicated input validation in fusedQKNormRopeOp.cpp into one helper shared by both operators, and reject negative head counts and out-of-int-range dimensions there before any shape arithmetic. Use static_cast rather than reinterpret_cast for the void* launcher conversions, tuple[...] instead of Tuple[...] in the new annotations, and rename the test's head-group constant to _FP8_NUM_HEADS_GROUPS. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
1cae558 to
31c4853
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp (1)
116-120: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse lower camel case for the new C++ callback names.
Keep the dispatcher schema name unchanged.
cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp#L116-L120: renamefused_qk_norm_rope_to_fp8tofusedQKNormRopeToFp8.cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp#L142-L150: renamefused_qk_norm_rope_to_fp8_metatofusedQKNormRopeToFp8Meta.As per coding guidelines, C++ functions use lowercase camelCase.
🤖 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 116 - 120, Rename the C++ callback function fused_qk_norm_rope_to_fp8 to fusedQKNormRopeToFp8 and fused_qk_norm_rope_to_fp8_meta to fusedQKNormRopeToFp8Meta in cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp at lines 116-120 and 142-150; keep the dispatcher schema name unchanged.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/thop/fusedQKNormRopeOp.cpp`:
- Around line 39-41: Update the validation around use_mrope and position_ids so
plain RoPE (use_mrope false) accepts only 1D IDs, while mRoPE accepts only the
validated 2D [3, num_tokens] layout. Replace the current independent checks near
the fused QK norm RoPE operation with mutually exclusive validation that rejects
conflicting dimensions before kernel execution.
- Around line 47-50: Update the input validation around the CHECK_INPUT calls in
the fused QKNorm/RoPE launcher to verify that position_ids, q_weight, and
k_weight are on the same CUDA device as qkv before passing their raw pointers to
the kernel. Preserve the existing dtype checks and reject any mixed-device
inputs before launch.
In `@tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py`:
- Line 23: Update the function parameter use_gemma to include the bool type
annotation while retaining its default value of False, consistent with the
requirement that every function parameter is annotated.
- Around line 379-413: Extend test_fused_qk_norm_rope_to_fp8 with a focused
mRoPE case that passes use_mrope=True, supplies int32 position_ids shaped [3,
num_tokens], and uses nonzero mrope_section1 and mrope_section2 values. Add
assertions against the appropriate FP8 reference output, while keeping qkv
contiguous and avoiding any strided-QKV scenario.
---
Nitpick comments:
In `@cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp`:
- Around line 116-120: Rename the C++ callback function
fused_qk_norm_rope_to_fp8 to fusedQKNormRopeToFp8 and
fused_qk_norm_rope_to_fp8_meta to fusedQKNormRopeToFp8Meta in
cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp at lines 116-120 and 142-150; keep
the dispatcher schema name unchanged.
🪄 Autofix
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: ea9a99c1-7d2d-481d-af2a-990b71f0c8a3
📒 Files selected for processing (5)
cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cucpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.hcpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpptensorrt_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 (3)
- cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h
- tensorrt_llm/_torch/models/modeling_minimaxm3.py
- cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
|
PR_Github #64490 [ run ] triggered by Bot. Commit: |
|
PR_Github #64490 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64629 [ run ] triggered by Bot. Commit: |
|
PR_Github #64629 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64694 [ run ] triggered by Bot. Commit: |
|
PR_Github #64694 [ run ] completed with state |
5 participants
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.Dev Engineer Review
launchFusedQKNormRopeToFp8for out-of-place BF16-to-FP8 E4M3 QKV processing.CODING_GUIDELINES.mdcompliance, API documentation, validation conventions, and backend-specific layout behavior.QA Engineer Review
tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py.torch_ref_rms_norm_ropewith optional Gemma RMSNorm support.tests/integration/test_lists/coverage was identified in the provided changes.