[None][perf] Use FP8 MiniMax-M3 MSA indexer QK - #17318
Conversation
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
peihu-nv
added
the
api-compatible
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
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 pull request adds MiniMax-M3 FP8 indexer configuration, a fused CUDA Q/K normalization and RoPE kernel, Torch integration, sparse-attention wiring, telemetry updates, and CUDA tests. ChangesMiniMax-M3 FP8 indexer
Estimated code review effort: 4 (Complex) | ~60 minutes Merge Risk: 🟡 Moderate · up to This PR adds an opt-in FP8 fused MiniMax-M3 indexer path and changes index-cache configuration. At the current head, invalid sparse_index_dim values can cause cache-shape/runtime failures, while some valid main-cache configurations may select the wrong BF16 index-cache storage. Merge readiness is moderate until these configuration and dtype-handling risks are fixed or explicitly accepted. Sequence Diagram(s)sequenceDiagram
participant MiniMaxM3Model
participant TorchOperator
participant MinimaxM3Fp8IndexerKernel
participant PagedKeyCache
participant MiniMaxM3MsaSparseAttention
MiniMaxM3Model->>TorchOperator: Submit BF16 index-Q/K and cache metadata
TorchOperator->>MinimaxM3Fp8IndexerKernel: Validate and launch fused operation
MinimaxM3Fp8IndexerKernel->>PagedKeyCache: Store FP8 index-K
MinimaxM3Fp8IndexerKernel->>MiniMaxM3Model: Return FP8 index-Q
MiniMaxM3Model->>MiniMaxM3MsaSparseAttention: Provide FP8 index-Q without index-K
MiniMaxM3MsaSparseAttention->>PagedKeyCache: Read cached index-K
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: 3
🧹 Nitpick comments (5)
tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py (1)
847-849: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCollapse the nested condition.
The two nested
ifstatements test independent conditions and can be one statement.♻️ Proposed simplification
- if idx_k_cache.dtype == torch.float8_e4m3fn: - if idx_q_view.dtype != torch.float8_e4m3fn: - idx_q_view = idx_q_view.to(torch.float8_e4m3fn) + if idx_k_cache.dtype == torch.float8_e4m3fn and idx_q_view.dtype != torch.float8_e4m3fn: + idx_q_view = idx_q_view.to(torch.float8_e4m3fn)🤖 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/attention_backend/sparse/minimax_m3/msa_backend.py` around lines 847 - 849, Update the dtype conversion logic around idx_k_cache and idx_q_view to combine the two independent conditions into a single conditional, while preserving the existing conversion to torch.float8_e4m3fn only when both conditions are satisfied.cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.h (1)
29-37: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse Doxygen comments for the new public launcher.
The repository C++ guidelines require
//!and//!<Doxygen comments to document new interfaces. This block uses plain//comments. Convert the block to//!and document the cache-layout parameters (page_stride,token_stride,page_size), which are not self-explanatory from the signature.As per coding guidelines: "Use C++ comments, not C comments except special inline cases; use
//for single-line comments,//!and//!<for Doxygen comments, and 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/minimaxM3Fp8IndexerKernel.h` around lines 29 - 37, Convert the comment above launchMinimaxM3Fp8IndexerQKNormRope to Doxygen syntax using //! and document the cache-layout parameters page_stride, token_stride, and page_size, including their roles in the paged E4M3 cache layout.Source: Coding guidelines
tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py (1)
12-70: 📐 Maintainability & Code Quality | 🔵 TrivialTest coverage summary (
tests/**path instructions).
- Changed test functions in this new module:
- Added:
test_minimax_m3_fp8_indexer_matches_bf16_then_cast(parametrized overnum_tokensin 1, 16, 129).- Added:
test_minimax_m3_fp8_indexer_cuda_graph_replay_updates_outputs.- Helpers added:
_assert_fp8_close,_reference,_strided_cache,_run.- Test-list registration: this module is new, so it is not listed under
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/. Add it to the appropriatetest-db/list for CI execution.- Coverage verdict: insufficient.
- Covered: numerical equivalence against the BF16 fused kernel followed by an E4M3 cast, strided HND cache writes, page-boundary token counts, and CUDA-graph replay.
- Not covered: the
numTokens == 0early return in the operator, operator validation failures (wrong cache dtype, wrong cache rank, mismatchedhead_dim,outCacheLocshorter thannum_tokens), and thehead_dim != 128/rotary_dim != 64launcher checks. Add validation cases withpytest.raisesso theTORCH_CHECKandTLLM_CHECK_WITH_INFOguards stay enforced.Run the tests with
pytest tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py.As per path instructions: the summary must list changed test functions, state test-list registration, and give a coverage verdict.
Do you want me to generate the validation test cases?
Also applies to: 100-134
🤖 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_minimax_m3_fp8_indexer.py` around lines 12 - 70, Expand coverage around the existing _run and test_minimax_m3_fp8_indexer_* helpers by adding pytest cases for zero tokens, wrong cache dtype or rank, mismatched head_dim, undersized slots, and unsupported head_dim or rotary_dim values, asserting each raises the expected validation error. Register this new test module in the appropriate test-db list so CI executes it, while preserving the existing numerical and CUDA-graph tests.Source: Path instructions
cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cu (1)
66-69: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDocument and check the vectorized-access alignment assumption.
The kernel loads
uint2fromqkand storesuint32_tintok_cache. These accesses require 8-byte and 4-byte alignment. The operator validatesindexKCache.stride(3) == 1andstride(2) == headDim, but it does not validate thatstride(0)is a multiple of four elements, and it does not validate the storage offset ofqk. Add the missing checks inminimaxM3Fp8IndexerOp.cpp, or state the alignment contract in the launcher comment.Also applies to: 144-147
🤖 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/minimaxM3Fp8IndexerKernel.cu` around lines 66 - 69, Update the validation in minimaxM3Fp8IndexerOp.cpp for the vectorized accesses used by the kernel: require indexKCache.stride(0) to be a multiple of four elements and validate qk’s storage offset is aligned for the uint2 load. If qk alignment cannot be checked there, document the required alignment contract in the launcher comment near the uint2 and uint32_t accesses.tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py (1)
182-198: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winComplete the FP8 cache coverage.
- Added tests:
test_msa_fp8_indexer_config_is_explicit_and_loweredandtest_msa_fp8_cache_converts_live_index_query_before_scoring.- Neither test is listed under
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/.- Add a real FP8 cache assertion for BF16
idx_k; the fake writer only captures the input.- Add
run_indexer(bf16_q, None, metadata_with_bf16_cache)coverage and assertValueError.- The
__new__setup does not coverMiniMaxM3MsaSparseAttention.indexer_kv_dtype.- Coverage verdict: insufficient.
- Run
pytest tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.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 `@tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py` around lines 182 - 198, Complete the FP8 coverage in the added tests: register test_msa_fp8_indexer_config_is_explicit_and_lowered and test_msa_fp8_cache_converts_live_index_query_before_scoring in both relevant integration test lists, make the BF16 idx_k case perform a real FP8 cache assertion rather than only capturing fake-writer input, and add run_indexer(bf16_q, None, metadata_with_bf16_cache) coverage asserting ValueError. In test_msa_fp8_cache_converts_live_index_query_before_scoring, initialize MiniMaxM3MsaSparseAttention.indexer_kv_dtype in the __new__ setup so the test exercises the actual dtype path.Source: Path instructions
🤖 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/minimaxM3Fp8IndexerKernel.cu`:
- Around line 139-147: Update the cache-output path around the direct store to
accept the page count from indexKCache.size(0), then guard invalid slots before
computing the output pointer or writing packed_output: return when slot is
negative or the derived page is at least page_count. Preserve the existing
address calculation for valid slots.
In `@tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py`:
- Line 220: Rename the captured “idx_k” entry used by FakeIndexer.select_blocks
and the assertions to identify it as the index-K cache, not the live tensor.
Change self.cache initialization to an intentionally strided cache view, then
keep the dtype assertion and update the stride assertion to verify the cache’s
expected strided layout rather than relying on contiguous storage.
In
`@tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py`:
- Around line 73-74: Add CUDA availability skip guards to both tests in
test_minimax_m3_fp8_indexer.py, including the test parametrized by num_tokens
and the other CUDA-dependent test. Mirror the sibling module’s
torch.cuda.is_available() guard, and add a compute-capability check where needed
to skip environments without E4M3 FP8 hardware support.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cu`:
- Around line 66-69: Update the validation in minimaxM3Fp8IndexerOp.cpp for the
vectorized accesses used by the kernel: require indexKCache.stride(0) to be a
multiple of four elements and validate qk’s storage offset is aligned for the
uint2 load. If qk alignment cannot be checked there, document the required
alignment contract in the launcher comment near the uint2 and uint32_t accesses.
In `@cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.h`:
- Around line 29-37: Convert the comment above
launchMinimaxM3Fp8IndexerQKNormRope to Doxygen syntax using //! and document the
cache-layout parameters page_stride, token_stride, and page_size, including
their roles in the paged E4M3 cache layout.
In `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py`:
- Around line 847-849: Update the dtype conversion logic around idx_k_cache and
idx_q_view to combine the two independent conditions into a single conditional,
while preserving the existing conversion to torch.float8_e4m3fn only when both
conditions are satisfied.
In `@tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py`:
- Around line 182-198: Complete the FP8 coverage in the added tests: register
test_msa_fp8_indexer_config_is_explicit_and_lowered and
test_msa_fp8_cache_converts_live_index_query_before_scoring in both relevant
integration test lists, make the BF16 idx_k case perform a real FP8 cache
assertion rather than only capturing fake-writer input, and add
run_indexer(bf16_q, None, metadata_with_bf16_cache) coverage asserting
ValueError. In test_msa_fp8_cache_converts_live_index_query_before_scoring,
initialize MiniMaxM3MsaSparseAttention.indexer_kv_dtype in the __new__ setup so
the test exercises the actual dtype path.
In
`@tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py`:
- Around line 12-70: Expand coverage around the existing _run and
test_minimax_m3_fp8_indexer_* helpers by adding pytest cases for zero tokens,
wrong cache dtype or rank, mismatched head_dim, undersized slots, and
unsupported head_dim or rotary_dim values, asserting each raises the expected
validation error. Register this new test module in the appropriate test-db list
so CI executes it, while preserving the existing numerical and CUDA-graph tests.
🪄 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: 3892a2a2-9cd8-40e3-b57b-0b5dc1320694
📒 Files selected for processing (13)
cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cucpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpptensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
/bot run |
|
PR_Github #64202 [ run ] triggered by Bot. Commit: |
|
PR_Github #64202 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64374 [ run ] triggered by Bot. Commit: |
|
PR_Github #64374 [ run ] completed with state
|
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #68084 [ run ] triggered by Bot. Commit: |
|
PR_Github #68084 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #68181 [ run ] triggered by Bot. Commit: |
|
PR_Github #68182 [ run ] triggered by Bot. Commit: |
|
PR_Github #68181 [ run ] completed with state |
|
PR_Github #68182 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68520 [ run ] triggered by Bot. Commit: |
|
PR_Github #68520 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68540 [ run ] triggered by Bot. Commit: |
|
PR_Github #68540 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68578 [ run ] triggered by Bot. Commit: |
|
PR_Github #68578 [ run ] completed with state |
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #68597 [ run ] triggered by Bot. Commit: |
|
PR_Github #68597 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68604 [ run ] triggered by Bot. Commit: |
|
PR_Github #68604 [ run ] completed with state |
10 participants
Dev Engineer Review
indexer_kv_dtypeconfiguration and correctedsparse_attention_configpropagation.sparse_index_dimvalues.QA Engineer Review
tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.pytests/unittest/_torch/models/test_minimax_m3.pytests/integration/test_lists/test-db/l0_cpu.yml.Description
Ports #16742 from
feat/m3_with_msato currentmainon top of the revised MiniMax-M3 MSA and KV-cache-manager architecture.MiniMax-M3 MSA currently produces normalized/RoPE index Q/K in BF16, converts the indexer tensors separately, and launches a paged index-K cache write in every sparse layer. This PR adds an opt-in fused CUDA path that performs Gemma RMSNorm, NeoX partial RoPE, BF16 rounding, raw E4M3 index-Q output, and direct strided-HND E4M3 index-K insertion in one kernel. This removes the separate cast/scatter work from the decode graph while preserving FP32 score accumulation.
The path is controlled by the prototype
indexer_kv_dtypeoption.bf16remains the default;fp8is restricted to the MSA implementation with the index-value branch disabled. The fused producer hands MSA only states that occur in production: exact E4M3 index-Q with an already-populated E4M3 index-K cache, or exact BF16 index-Q with a live BF16 index-K tensor for the default path. Unsupported FP16/FP32 indexer handoffs now fail beforefmha_sm100.The index cache dtype is configured independently of the main KV-cache dtype. The default indexer therefore keeps an exact BF16 index cache even when the main cache uses another supported dtype.
This also fixes an independent production configuration bug: the executor passes
sparse_attention_config, but the cache manager previously read the unused namesparse_attn_config. The cache layout now honors user-specifiedsparse_index_dimvalues for BF16 as well as the new indexer dtype. The telemetry manifest and generated reference document the new option.The original matched GB200 disaggregated A/B from #16742 measured:
Test Coverage
fmha_sm100proxy scoring with both BF16 and E4M3 index-Q against strided-HND and packed index-K caches, with exact score parity.l0_cpu.yml;l0_h100.ymlcollects the parallel-hardware-agnostic CUDA module; existing H100/B200/B300 entries collect the attention module; andl0_dgx_b200.ymlexplicitly schedulesTestMiniMaxM3::test_nvfp4[use_msa=True]with the FP8 indexer.fmha_sm100execution, and B200 end-to-end accuracy case. The original feature-branch implementation passed GB200 native parity, CUDA-graph, MSA integration, accuracy, serving A/B, and Nsys validation as documented in [None][perf] Use FP8 MiniMax-M3 MSA indexer QK #16742.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.