Skip to content

[None][perf] Use FP8 MiniMax-M3 MSA indexer QK - #17318

Open
peihu-nv wants to merge 4 commits into
NVIDIA:mainfrom
peihu-nv:peihengh/m3-fp8-indexer-main
Open

[None][perf] Use FP8 MiniMax-M3 MSA indexer QK#17318
peihu-nv wants to merge 4 commits into
NVIDIA:mainfrom
peihu-nv:peihengh/m3-fp8-indexer-main

Conversation

@peihu-nv

@peihu-nv peihu-nv commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Added an opt-in fused CUDA FP8 indexer path for MiniMax-M3.
  • Added RMSNorm, partial NeoX RoPE, BF16 rounding, FP8 E4M3 conversion, and direct strided index-K cache insertion.
  • Preserved BF16 as the default.
  • Added indexer_kv_dtype configuration with bf16 and fp8 values.
  • Restricted FP8 to MSA with index values disabled.
  • Updated MSA and KV-cache handling for optional live index-K handoff.
  • Added CUDA operator support and tests for accuracy, cache writes, invalid cache slots, and CUDA graph replay.

Dev Engineer Review

  • Configuration values are consistent across llm_args.py, common.py, and cache_manager.py.
  • Validation covers unsupported dtype combinations and invalid index-K handoff states.
  • The fused path avoids separate index-K output and duplicate cache writes.
  • The BF16 path remains the default.
  • The CUDA kernel validates dimensions and cache locations before writing.
  • Main-branch build and focused integration validation remain follow-up items.

QA Engineer Review

Added test functions:

  • test_msa_fp8_indexer_config_is_explicit_and_lowered
  • test_msa_fp8_cache_converts_live_index_query_before_scoring
  • test_minimax_m3_fp8_indexer_matches_bf16_then_cast
  • test_minimax_m3_fp8_indexer_skips_invalid_cache_slots
  • test_minimax_m3_fp8_indexer_cuda_graph_replay_updates_outputs

No files under tests/integration/test_lists/test-db/ or tests/integration/test_lists/qa/ were modified. The new unit tests are not covered by those lists.

Verdict: needs follow-up.

Description

Ports #16742 from feat/m3_with_msa to main on 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 separate 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 index-K insertion in one kernel.

The path is controlled by the prototype indexer_kv_dtype option and keeps bf16 as the default. FP8 is restricted to the MSA implementation with the index-value branch disabled. The existing general MSA max-score FMHA and fused block selector remain unchanged; score accumulation remains FP32.

This main port adapts only the original focused optimization to the current V2 side-cache lifecycle and optional live index-K handoff. It does not include later packed-projection or mixed/decode feature-branch work.

The original matched GB200 disaggregated A/B from #16742 measured:

Metric BF16 indexer FP8 indexer Change
GEN/decode median TPOT 22.657 ms 22.296 ms -1.592%
GEN/decode p90 TPOT 22.769 ms 22.449 ms -1.407%
E2E total throughput 24,000.465 token/s 24,362.696 token/s +1.509%

Test Coverage

  • Changed-file pre-commit checks: passed for all 13 files.
  • Python byte-compilation for all changed Python files: passed with Python 3.12.
  • CodeRabbit's three actionable comments were addressed in 6904b7b57f; all review threads are resolved.
  • GitHub pre-commit, DCO, PR title, checklist, API compatibility, and base-freshness checks passed on the current head.
  • Full GPU CI (blossom-ci): pending.
  • 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-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.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
@peihu-nv peihu-nv added the api-compatible Accepted LLM API contract change that is backwards-compatible label Aug 5, 2026
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
@peihu-nv
peihu-nv marked this pull request as ready for review August 5, 2026 23:05
@peihu-nv
peihu-nv requested review from a team as code owners August 5, 2026 23:05
@coderabbitai

coderabbitai Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Walkthrough

The pull request adds MiniMax-M3 FP8 indexer configuration, a fused CUDA Q/K normalization and RoPE kernel, Torch integration, sparse-attention wiring, and CUDA tests for cache writes and graph replay.

Changes

MiniMax-M3 FP8 indexer

Layer / File(s) Summary
FP8 indexer configuration
tensorrt_llm/llmapi/llm_args.py, tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/..., tensorrt_llm/usage/llm_args_golden_manifest.json
Adds indexer_kv_dtype with bf16 and fp8 modes. Validates FP8 settings and selects torch.float8_e4m3fn cache storage.
CUDA kernel and launcher
cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.*
Adds BF16 RMSNorm, partial NeoX RoPE, FP8 E4M3 conversion, query output, and paged key-cache writes with invalid-slot checks.
Torch operator registration
cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp, cpp/tensorrt_llm/thop/CMakeLists.txt, tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
Adds validation, schema registration, fake shape inference, CUDA dispatch, and build integration.
Sparse attention integration
tensorrt_llm/_torch/models/modeling_minimaxm3.py, tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
Uses the fused FP8 path, writes index-K directly to the cache, and permits missing live index-K in the FP8 path.
FP8 indexer validation
tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py, tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py
Tests configuration validation, FP8 conversion, strided caches, invalid cache slots, multiple token counts, and CUDA graph replay.

Estimated code review effort: 4 (Complex) | ~60 minutes

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 for block selection
Loading

Possibly related PRs

Suggested reviewers: qijune, lowsfer, vallis-neria

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 27.27% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
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.
Title check ✅ Passed The title clearly identifies the FP8 MiniMax-M3 MSA indexer optimization and follows the repository's ticket and type format.
Description check ✅ Passed The description explains the motivation, implementation, scope, performance results, test coverage, and checklist status.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 value

Collapse the nested condition.

The two nested if statements 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 value

Use 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 | 🔵 Trivial

Test coverage summary (tests/** path instructions).

  1. Changed test functions in this new module:
    • Added: test_minimax_m3_fp8_indexer_matches_bf16_then_cast (parametrized over num_tokens in 1, 16, 129).
    • Added: test_minimax_m3_fp8_indexer_cuda_graph_replay_updates_outputs.
    • Helpers added: _assert_fp8_close, _reference, _strided_cache, _run.
  2. Test-list registration: this module is new, so it is not listed under tests/integration/test_lists/test-db/ or tests/integration/test_lists/qa/. Add it to the appropriate test-db/ list for CI execution.
  3. 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 == 0 early return in the operator, operator validation failures (wrong cache dtype, wrong cache rank, mismatched head_dim, outCacheLoc shorter than num_tokens), and the head_dim != 128 / rotary_dim != 64 launcher checks. Add validation cases with pytest.raises so the TORCH_CHECK and TLLM_CHECK_WITH_INFO guards 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 win

Document and check the vectorized-access alignment assumption.

The kernel loads uint2 from qk and stores uint32_t into k_cache. These accesses require 8-byte and 4-byte alignment. The operator validates indexKCache.stride(3) == 1 and stride(2) == headDim, but it does not validate that stride(0) is a multiple of four elements, and it does not validate the storage offset of qk. Add the missing checks in minimaxM3Fp8IndexerOp.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 win

Complete the FP8 cache coverage.

  • Added tests: test_msa_fp8_indexer_config_is_explicit_and_lowered and test_msa_fp8_cache_converts_live_index_query_before_scoring.
  • Neither test is listed under tests/integration/test_lists/test-db/ or tests/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 assert ValueError.
  • The __new__ setup does not cover MiniMaxM3MsaSparseAttention.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

📥 Commits

Reviewing files that changed from the base of the PR and between e5e3821 and 459bd9d.

📒 Files selected for processing (13)
  • cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cu
  • cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.h
  • cpp/tensorrt_llm/thop/CMakeLists.txt
  • cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp
  • tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
  • tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py
  • tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
  • tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
  • tensorrt_llm/_torch/models/modeling_minimaxm3.py
  • tensorrt_llm/llmapi/llm_args.py
  • tensorrt_llm/usage/llm_args_golden_manifest.json
  • tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
  • tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py

Comment thread cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cu
Comment thread tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py Outdated
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
@peihu-nv

peihu-nv commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64202 [ run ] triggered by Bot. Commit: cf6ae3b Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64202 [ run ] completed with state FAILURE. Commit: cf6ae3b
/LLM/main/L0_MergeRequest_PR pipeline #52110 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@peihu-nv

peihu-nv commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64374 [ run ] triggered by Bot. Commit: cf6ae3b Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64374 [ run ] completed with state SUCCESS. Commit: cf6ae3b
/LLM/main/L0_MergeRequest_PR pipeline #52266 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@peihu-nv

peihu-nv commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64450 [ run ] triggered by Bot. Commit: cf6ae3b Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64450 [ run ] completed with state FAILURE. Commit: cf6ae3b
/LLM/main/L0_MergeRequest_PR pipeline #52326 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@peihu-nv

peihu-nv commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64486 [ run ] triggered by Bot. Commit: cf6ae3b Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64486 [ run ] completed with state FAILURE. Commit: cf6ae3b
/LLM/main/L0_MergeRequest_PR pipeline #52361 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

# convention (layers 0..2 dense, 3..N-1 sparse,
# disable_index_value=True, sparse_index_dim=128).
sparse_attn_config = kwargs.get("sparse_attn_config")
sparse_attn_config = kwargs.get("sparse_attn_config") or kwargs.get(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sparse_attn_config is never passed by any caller — grep -rn "sparse_attn_config=" tensorrt_llm/ returns nothing, and every construction site (_util.py:2151/2264/2373/2419) passes sparse_attention_config=, with manager_extra_kwargs carrying only enable_stats/is_disagg. So this line is really a latent-bug fix: until now the cache manager always fell back to the hardcoded sparse_index_dim=128 at line 171 while the backend received the user's value via to_sparse_params(), so the two disagreed on the index width for any non-default sparse_index_dim under implementation="triton". Could the dead name be dropped rather than OR-ed in, and this fix be called out in the description (or split out) since it is independent of the FP8 work?

# 1/8; 2**-9 is its subnormal bin width.
def _assert_fp8_close(actual, expected):
torch.testing.assert_close(actual.float(), expected.float(), rtol=0.125, atol=2**-9)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rtol=0.125 is E4M3's relative bin width: enumerating all 126 finite positive E4M3 values, 125 of 125 adjacent pairs satisfy this assertion. So the test cannot detect a result that lands one bin away — which is exactly the failure mode the kernel guards against ("A reassociated pair sum can move a final BF16 value across an FP8 bin"). Reassociating the sum of squares, dropping the BF16 intermediate rounding, or swapping __log2f/exp2f for powf would all keep this green. Consider asserting byte equality with a small allowed mismatch rate instead, e.g. (actual.view(torch.uint8) == expected.view(torch.uint8)).float().mean() > 0.999.

"MiniMax M3 indexer_kv_dtype must be 'bf16' or 'fp8', got "
f"{self.indexer_kv_dtype!r}."
)
if self.indexer_kv_dtype == "fp8" and (

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Neither disable_index_value_layer_ids nor sparse_layer_ids is passed by any caller, so line 178 always sets the former to list(sparse_layer_ids) and this set difference is always empty — the raise is unreachable. llm_args.py:760 already rejects the same combination at a reachable point. Suggest dropping this, or making it an assert with a note on the invariant it guards.

# The fused production path arrives here with E4M3 Q and an already
# populated cache. The explicit conversion is retained for standalone
# callers that supply BF16 Q/K to an E4M3-configured backend.
if idx_k_cache.dtype == torch.float8_e4m3fn:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This conversion looks unreachable in production: the only non-test caller is modeling_minimaxm3.py:1357, and when indexer_kv_dtype == "fp8" (which llm_args.py:756 restricts to the MSA implementation) _fused_fp8_index_qk_norm_rope either returns a tensor or raises, so _index_norm_rope always yields (fp8_q, None) and idx_k is never non-None against an E4M3 cache. Only test_minimax_m3_msa_backend.py:232 exercises it. Could it be dropped, or the actual standalone caller named in the comment? Separately, the elif above states the negation of the invariant and the nested if could be a single and.

@brnguyen2 brnguyen2 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving — the comments below are optional touch-ups, not blockers.

Solid port. I checked the kernel against fusedQKNormRopeKernel: same __log2f/exp2f/__sincosf frequency scheme, the shfl-xor NeoX pairing and signs are correct for the 4-elems-per-lane layout, and the RMSNorm accumulation order matches, so the "bitwise aligned" claim holds. The aux-stream ordering also checks out: the fused cache write runs inside the maybe_execute_in_parallel closure and the event join orders it before the main-stream indexer read.

A few non-blocking items:

  • Tracking ticket: this is a nontrivial perf feature with measured wins; it should carry a JIRA/TRTLLM ticket rather than [None].
  • Docs: docs/source/developer-guide/telemetry.md enumerates the manifest fields (it has a row for sparse_attention_config.indexer_k_dtype); the new indexer_kv_dtype row is missing. Also note the new manifest field needs telemetry/privacy CODEOWNER sign-off per AGENTS.md.
  • Description correction: the QA section claims the new unit tests aren't covered by any test list, but both are picked up by existing directory-level entries: l0_h100.yml includes unittest/_torch/thop/parallel_hw_agnostic (SM90, so the FP8 skip doesn't fire) and l0_b200.yml/l0_b300.yml include unittest/_torch/attention. Worth fixing the description so nobody adds redundant entries.
  • The remaining unverified surface is fmha_sm100 consuming FP8 Q against the strided FP8 cache on all plan shapes, which only feature-branch runs and the pending blossom CI cover. Fine to land on green CI.

else
{
int const slot = out_cache_loc[token_idx];
if (slot < 0)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The slot < 0 / page >= num_pages guards silently skip the write, but the BF16 path this fuses away has the opposite behavior: write_kv_slots in minimax_m3/common.py computes page = out_long // tokens_per_block, so a -1 slot becomes cache[-1, :, 127], which Python negative indexing wraps to the last page (a real write into another request's cache line). If negative/out-of-range slots can actually occur (this PR's test treats them as expected input), the BF16 path has a latent corruption bug and deserves the same mask; if they can't, it would help to say in a comment which caller produces them. Worth clarifying which of the two is intended so the two paths don't diverge.

self.sparse_num_index_heads,
self.sparse_index_dim,
rotary_dim,
self.index_q_norm.variance_epsilon,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The kernel takes a single eps and applies it to both norms, but this passes only index_q_norm.variance_epsilon; index_k_norm's eps is silently assumed equal. Today both come from the same config so this holds, but if they ever diverge the K side would be normalized with the wrong epsilon and nothing would flag it. Cheap fix: assert self.index_q_norm.variance_epsilon == self.index_k_norm.variance_epsilon next to the other precondition checks (they already fail loudly).

# convention (layers 0..2 dense, 3..N-1 sparse,
# disable_index_value=True, sparse_index_dim=128).
sparse_attn_config = kwargs.get("sparse_attn_config")
sparse_attn_config = kwargs.get("sparse_attn_config") or kwargs.get(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Worth noting the blast radius of this fallback: the executor constructs cache managers with sparse_attention_config= (see _util.py), so before this change the sparse_attn_config lookup was always None on the production path and the whole resolution branch below (including sparse_index_dim from config) was dead. This one-liner makes it live for the first time, i.e. it doesn't just plumb indexer_kv_dtype, it also starts honoring sparse_index_dim from user config where the 128 fallback previously always won. That looks intended and correct, but it's a behavior change for BF16 users too, so it deserves a mention in the PR description.

@VALLIS-NERIA VALLIS-NERIA left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the KV cache manager portion only. The V2-managed FP8 index-K buffer sizing and dtype selection, page-lifecycle integration, and disaggregated role mapping look consistent. Approving from the KV cache manager ownership scope.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

api-compatible Accepted LLM API contract change that is backwards-compatible

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants