[None][perf] Optimize MiniMax-M3 MSA block selection - #17236
Conversation
|
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:
WalkthroughChangesThe PR adds a CUDA Minimax-M3 block selector and exposes it through a registered Torch operator. The MSA backend now uses the fused operator. CUDA-gated tests cover selection behavior, invalid inputs, strided tensors, and graph replay. Minimax M3 selector
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant MSAUtils
participant TorchOperator
participant invokeMinimaxM3SelectBlocks
participant CUDAKernel
MSAUtils->>TorchOperator: pass scores and valid-block counts
TorchOperator->>TorchOperator: validate inputs and allocate int32 output
TorchOperator->>invokeMinimaxM3SelectBlocks: pass tensor pointers and strides
invokeMinimaxM3SelectBlocks->>CUDAKernel: launch on current CUDA stream
CUDAKernel-->>TorchOperator: sorted selections with -1 padding
TorchOperator-->>MSAUtils: return [total_queries, num_kv_heads, 16]
Possibly related PRs
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.
Actionable comments posted: 1
🧹 Nitpick comments (6)
cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h (1)
30-32: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new public launcher with a Doxygen comment.
The contract of this entry point is not obvious from the signature. Callers must know that strides are element counts, that
outputis[totalQueries, numKvHeads, 16]int32, that IDs are ascending, and that unused slots hold-1. Add a Doxygen block that states these facts, including the fixed top-k of 16 and the accepted range ofnValidBlocks.As per coding guidelines: "use Doxygen comments for new interfaces".
🤖 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/minimaxM3SelectBlocks.h` around lines 30 - 32, Add a Doxygen comment immediately above invokeMinimaxM3SelectBlocks documenting that all stride parameters are element counts, output has shape [totalQueries, numKvHeads, 16] with int32 entries, the fixed top-k is 16, selected block IDs are ascending, unused output slots contain -1, and each nValidBlocks value must be within the supported block-count range.Source: Coding guidelines
cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp (1)
32-33: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueShare the top-k and block-index limits with the kernel.
kRequiredTopKrepeatskTopKincpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.culine 36, andkMaxBlockIndexrepeats the index sentinel ofTopKRedType. A later change to the kernel would leave these checks silently wrong. Expose both constants fromcpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.hand use them here.🤖 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/minimaxM3SelectBlocksOp.cpp` around lines 32 - 33, Move the shared top-k and maximum block-index constants to the declarations in minimaxM3SelectBlocks.h, then update the kernel implementation and the validation logic in MinimaxM3SelectBlocksOp to reference those header symbols instead of local duplicates. Remove the duplicated kRequiredTopK and kMaxBlockIndex definitions while preserving their current values and checks.cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu (2)
36-41: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winDocument the Python source of truth for the matching forced-score constants.
kInitScoreandkLocalScorematch_INIT_SCOREand_LOCAL_SCOREintensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py.🤖 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/minimaxM3SelectBlocks.cu` around lines 36 - 41, Document that the CUDA constants kInitScore and kLocalScore must remain synchronized with the Python source-of-truth constants _INIT_SCORE and _LOCAL_SCORE in common.py, using concise comments next to their definitions.
159-171: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winCheck the CUDA launch status. Add
sync_check_cuda_error(stream);immediately after the kernel launch so errors surface at this call.🤖 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/minimaxM3SelectBlocks.cu` around lines 159 - 171, Update invokeMinimaxM3SelectBlocks by calling sync_check_cuda_error(stream) immediately after the minimaxM3SelectBlocksKernel launch, preserving the existing launch parameters and zero-row early return.tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py (2)
232-256: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd negative-path coverage for the selector contract.
Add
pytest.raisescases fortopk != 16, non-float32scores, and ann_valid_blockslength different fromtotal_q.Test coverage summary: six positive-path tests were added; none were modified or removed. All six are covered by existing directory-level CI entries. No QA-list entries cover these unit tests. Verdict: insufficient.
🤖 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/sparse/test_minimax_m3_msa_selector.py` around lines 232 - 256, Extend test_fused_selector_supports_strided_scores_and_cuda_validity with pytest.raises cases covering selector contract violations: topk values other than 16, scores with a non-float32 dtype, and n_valid_blocks whose length differs from total_q. Keep the existing valid strided CUDA comparison unchanged and assert the expected exception for each invalid input.Source: Path instructions
158-178: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse an explicit tie-case expectation.
torch.topkdoes not guarantee tie ordering. For 32 or 64 valid blocks, sorting its output does not make the selected set deterministic. The kernel selects lower block IDs first. Compare against explicit IDs instead of_reference_select_blocks.Test coverage:
test_fused_selector_matches_reference_equal_score_tiescovers three fill values and valid-block counts 15, 32, and 64. No matchingtest-db/orqa/entry exists.🤖 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/sparse/test_minimax_m3_msa_selector.py` around lines 158 - 178, Update test_fused_selector_matches_reference_equal_score_ties to stop using _reference_select_blocks for equal-score cases and assert the explicit expected block IDs selected in ascending order, matching the kernel’s lower-ID-first tie behavior. Preserve coverage for all fill_value and n_valid_blocks parameter combinations.
🤖 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/minimaxM3SelectBlocksOp.cpp`:
- Around line 64-70: Add a c10::cuda::CUDAGuard initialized with scores.device()
at the start of the operation, before torch::empty and the current-stream lookup
in the invokeMinimaxM3SelectBlocks flow, so allocation and kernel launch use the
scores CUDA device.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu`:
- Around line 36-41: Document that the CUDA constants kInitScore and kLocalScore
must remain synchronized with the Python source-of-truth constants _INIT_SCORE
and _LOCAL_SCORE in common.py, using concise comments next to their definitions.
- Around line 159-171: Update invokeMinimaxM3SelectBlocks by calling
sync_check_cuda_error(stream) immediately after the minimaxM3SelectBlocksKernel
launch, preserving the existing launch parameters and zero-row early return.
In `@cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h`:
- Around line 30-32: Add a Doxygen comment immediately above
invokeMinimaxM3SelectBlocks documenting that all stride parameters are element
counts, output has shape [totalQueries, numKvHeads, 16] with int32 entries, the
fixed top-k is 16, selected block IDs are ascending, unused output slots contain
-1, and each nValidBlocks value must be within the supported block-count range.
In `@cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp`:
- Around line 32-33: Move the shared top-k and maximum block-index constants to
the declarations in minimaxM3SelectBlocks.h, then update the kernel
implementation and the validation logic in MinimaxM3SelectBlocksOp to reference
those header symbols instead of local duplicates. Remove the duplicated
kRequiredTopK and kMaxBlockIndex definitions while preserving their current
values and checks.
In `@tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py`:
- Around line 232-256: Extend
test_fused_selector_supports_strided_scores_and_cuda_validity with pytest.raises
cases covering selector contract violations: topk values other than 16, scores
with a non-float32 dtype, and n_valid_blocks whose length differs from total_q.
Keep the existing valid strided CUDA comparison unchanged and assert the
expected exception for each invalid input.
- Around line 158-178: Update
test_fused_selector_matches_reference_equal_score_ties to stop using
_reference_select_blocks for equal-score cases and assert the explicit expected
block IDs selected in ascending order, matching the kernel’s lower-ID-first tie
behavior. Preserve coverage for all fill_value and n_valid_blocks parameter
combinations.
🪄 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: c3216ff4-5301-4747-ad80-166de425a8b2
📒 Files selected for processing (7)
cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cucpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpptensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.pytensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py
07a722d to
a6ae542
Compare
There was a problem hiding this comment.
♻️ Duplicate comments (1)
cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp (1)
64-70:⚠️ Potential issue | 🟠 MajorGuard the CUDA device before allocation and launch.
This remains the unresolved issue from the previous review. If
scoresis on a non-current CUDA device, the code obtains a stream forscores.get_device()without ac10::cuda::CUDAGuard. The kernel launch can then use a stream incompatible with the current device.Add
c10::cuda::CUDAGuard const deviceGuard{scores.device()};beforetorch::empty, and include<c10/cuda/CUDAGuard.h>directly.#!/usr/bin/env bash set -euo pipefail rg -n -C6 \ 'CUDAGuard|torch::empty|getCurrentCUDAStream|invokeMinimaxM3SelectBlocks' \ cpp/tensorrt_llm/thop🤖 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/minimaxM3SelectBlocksOp.cpp` around lines 64 - 70, Add the direct <c10/cuda/CUDAGuard.h> include and create a c10::cuda::CUDAGuard from scores.device() before the torch::empty allocation in the MinimaxM3 block-selection function, so allocation, stream retrieval, and invokeMinimaxM3SelectBlocks execute on scores’ CUDA device.
🧹 Nitpick comments (2)
cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp (2)
64-64: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeclare
outputasconst.Line 64 initializes
outputonce and never rebinds it. Useauto const output.As per coding guidelines, declare unmodified variables
const.🤖 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/minimaxM3SelectBlocksOp.cpp` at line 64, Update the output declaration in the minimax M3 select-blocks operation to use const qualification, changing the existing auto declaration for output while preserving its initialization and subsequent use.Source: Coding guidelines
20-22: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the direct
<cstdint>include.This file uses
int64_tandint32_ton Lines 29-69. Add<cstdint>instead of relying on transitive ATen or Torch includes.As per coding guidelines, use
<cstdint>for fixed-width integer types.🤖 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/minimaxM3SelectBlocksOp.cpp` around lines 20 - 22, Add a direct <cstdint> include alongside the standard library headers in minimaxM3SelectBlocksOp.cpp so the int64_t and int32_t usages are provided explicitly rather than through transitive ATen or Torch includes.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.
Duplicate comments:
In `@cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp`:
- Around line 64-70: Add the direct <c10/cuda/CUDAGuard.h> include and create a
c10::cuda::CUDAGuard from scores.device() before the torch::empty allocation in
the MinimaxM3 block-selection function, so allocation, stream retrieval, and
invokeMinimaxM3SelectBlocks execute on scores’ CUDA device.
---
Nitpick comments:
In `@cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp`:
- Line 64: Update the output declaration in the minimax M3 select-blocks
operation to use const qualification, changing the existing auto declaration for
output while preserving its initialization and subsequent use.
- Around line 20-22: Add a direct <cstdint> include alongside the standard
library headers in minimaxM3SelectBlocksOp.cpp so the int64_t and int32_t usages
are provided explicitly rather than through transitive ATen or Torch includes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a9fc03bf-1f6b-4a28-b416-81d544ab7631
📒 Files selected for processing (7)
cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cucpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpptensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.pytensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py
🚧 Files skipped from review as they are similar to previous changes (6)
- cpp/tensorrt_llm/thop/CMakeLists.txt
- tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
- tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py
- cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h
- cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
a6ae542 to
fc8db26
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (3)
cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp (3)
27-28: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse lower-camel-case for the extension namespace.
torch_extuses snake_case. Rename it totorchExtand update the closing comment and registration reference. Confirm that no other translation unit usestensorrt_llm::torch_ext.Proposed namespace rename
-namespace torch_ext +namespace torchExt ... -} // namespace torch_ext +} // namespace torchExt ... - &tensorrt_llm::torch_ext::minimaxM3SelectBlocks); + &tensorrt_llm::torchExt::minimaxM3SelectBlocks);As per coding guidelines, local namespaces use lower camel case.
Also applies to: 76-76, 89-89
🤖 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/minimaxM3SelectBlocksOp.cpp` around lines 27 - 28, Rename the extension namespace from torch_ext to torchExt throughout the affected translation unit, including its closing comment and registration reference. Update all matching qualified references in the file, and verify no other translation unit still references tensorrt_llm::torch_ext.Source: Coding guidelines
18-23: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winInclude
<cstdint>directly.This file uses
int64_tandint32_t. Add<cstdint>instead of relying on transitive includes.Proposed include
`#include` <ATen/cuda/CUDAContext.h> `#include` <c10/cuda/CUDAGuard.h> +#include <cstdint> `#include` <limits>As per coding guidelines, use
<cstdint>for fixed-width integer types.🤖 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/minimaxM3SelectBlocksOp.cpp` around lines 18 - 23, Update the include list in minimaxM3SelectBlocksOp.cpp to add the standard <cstdint> header directly, alongside the existing includes, so the int64_t and int32_t usages do not depend on transitive includes.Source: Coding guidelines
66-66: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeclare the output handle const.
outputis not reassigned after construction. Useauto const outputwhile allowing the CUDA kernel to modify the tensor storage.Proposed const qualification
- auto output = torch::empty({scores.size(2), scores.size(0), topK}, scores.options().dtype(torch::kInt32)); + auto const output + = torch::empty({scores.size(2), scores.size(0), topK}, scores.options().dtype(torch::kInt32));Verify that
data_ptr<int32_t>()remains available on a consttorch::Tensorwith the repository's PyTorch version.As per coding guidelines, declare unmodified variables const and use east-const style.
🤖 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/minimaxM3SelectBlocksOp.cpp` at line 66, Update the output declaration in the top-K block allocation flow to use east-const style with a const handle, while preserving mutable tensor storage for the CUDA kernel. Verify that the existing data_ptr<int32_t>() usage remains valid with the repository’s PyTorch version.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.
Nitpick comments:
In `@cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp`:
- Around line 27-28: Rename the extension namespace from torch_ext to torchExt
throughout the affected translation unit, including its closing comment and
registration reference. Update all matching qualified references in the file,
and verify no other translation unit still references tensorrt_llm::torch_ext.
- Around line 18-23: Update the include list in minimaxM3SelectBlocksOp.cpp to
add the standard <cstdint> header directly, alongside the existing includes, so
the int64_t and int32_t usages do not depend on transitive includes.
- Line 66: Update the output declaration in the top-K block allocation flow to
use east-const style with a const handle, while preserving mutable tensor
storage for the CUDA kernel. Verify that the existing data_ptr<int32_t>() usage
remains valid with the repository’s PyTorch version.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 70fa1e99-eeaf-4146-9e43-a8ff768e3eae
📒 Files selected for processing (7)
cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cucpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpptensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.pytensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py
🚧 Files skipped from review as they are similar to previous changes (6)
- cpp/tensorrt_llm/thop/CMakeLists.txt
- tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py
- cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h
- tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py
- cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu
|
/bot run |
|
PR_Github #63838 [ run ] triggered by Bot. Commit: |
|
PR_Github #63838 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63861 [ run ] triggered by Bot. Commit: |
|
PR_Github #63861 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63944 [ run ] triggered by Bot. Commit: |
|
PR_Github #63944 [ run ] completed with state |
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
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. |
There was a problem hiding this comment.
🧹 Nitpick comments (3)
cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp (3)
66-66: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake the Tensor handle const.
outputis not rebound after initialization. Declare it asauto const outputto enforce that invariant. Verify that the selected Torch headers supportdata_ptr<int32_t>()on a consttorch::Tensor.As per coding guidelines, “declare unmodified variables as const.”
🤖 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/minimaxM3SelectBlocksOp.cpp` at line 66, Update the output declaration in the relevant selection operation to use a const Tensor handle, preserving its existing allocation and subsequent use. Confirm the selected Torch headers support data_ptr<int32_t>() through the const handle and adjust only if required for compilation.Source: Coding guidelines
1-16: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse C++ line comments for the license header.
Lines 1-16 use a C-style block comment. Replace it with
//lines while retaining the NVIDIA copyright and SPDX text.As per coding guidelines, “Use C++ comments, not C comments except special inline cases.”
🤖 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/minimaxM3SelectBlocksOp.cpp` around lines 1 - 16, Replace the block comment at the top of minimaxM3SelectBlocksOp.cpp with C++ // line comments, preserving the complete NVIDIA copyright, SPDX license identifier, and license text unchanged.Source: Coding guidelines
30-31: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the operator contract with Doxygen.
minimaxM3SelectBlocksis registered as a Torch operator. Add Doxygen that specifies input shapes, dtypes, device requirements, fixedtopK, ascendingint32output IDs, and-1padding.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/thop/minimaxM3SelectBlocksOp.cpp` around lines 30 - 31, Add a Doxygen comment directly above minimaxM3SelectBlocks documenting its Torch operator contract: describe the shapes and dtypes of scores and nValidBlocks, required device placement, the fixed topK behavior, and that it returns ascending int32 block IDs with -1 padding.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.
Nitpick comments:
In `@cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp`:
- Line 66: Update the output declaration in the relevant selection operation to
use a const Tensor handle, preserving its existing allocation and subsequent
use. Confirm the selected Torch headers support data_ptr<int32_t>() through the
const handle and adjust only if required for compilation.
- Around line 1-16: Replace the block comment at the top of
minimaxM3SelectBlocksOp.cpp with C++ // line comments, preserving the complete
NVIDIA copyright, SPDX license identifier, and license text unchanged.
- Around line 30-31: Add a Doxygen comment directly above minimaxM3SelectBlocks
documenting its Torch operator contract: describe the shapes and dtypes of
scores and nValidBlocks, required device placement, the fixed topK behavior, and
that it returns ascending int32 block IDs with -1 padding.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 80ae80f2-7d14-475b-b972-5ffdf63d6024
📒 Files selected for processing (7)
cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cucpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpptensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.pytensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py
🚧 Files skipped from review as they are similar to previous changes (6)
- cpp/tensorrt_llm/thop/CMakeLists.txt
- tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
- cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu
- cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py
- tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py
|
/bot run |
|
PR_Github #64097 [ run ] triggered by Bot. Commit: |
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
/bot run |
|
PR_Github #64099 [ run ] triggered by Bot. Commit: |
|
PR_Github #64097 [ run ] completed with state |
|
PR_Github #64099 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64105 [ run ] triggered by Bot. Commit: |
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
peihu-nv
changed the title
|
PR_Github #64105 [ run ] completed with state
|
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
/bot run |
|
PR_Github #64196 [ run ] triggered by Bot. Commit: |
|
PR_Github #64196 [ run ] completed with state
|
2 participants
Description
This PR combines the MiniMax-M3 MSA selector implementation with its context/prefill optimization. It introduces one CUDA selector and then specializes its hot CTX path in two ways:
int32block IDs with-1padding.These are kept together because the CTX optimization directly specializes the selector introduced by this PR and shares its CUDA kernel, Torch custom op, Python registration/routing, and tests.
This main-branch port consolidates the implementations originally merged into
feat/m3_with_msain #16681 and #16818.Performance
Fused selector
Matched serving A/B results from the original main-port validation:
Optimized CTX selector and zero-copy handoff
Matched 8K context validation from the feature-branch optimization:
Test Coverage
fmha_sm100's actual K2Q conversion and verifies that it retains the same backing storage without an allocation or copy.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.