[None][perf] Fuse MiniMax-M3 MSA per-layer KV-cache writes into one kernel

[zheyuf]

Dev Engineer Review

• Added fused Triton cache writes for MiniMax-M3 MSA.
• The fused path writes K, V, and optional index-K in one launch per layer.
• Unsupported layouts use the legacy cache-write path.
• Metadata prevents duplicate K/V and index-K writes.
• The run_indexer API now accepts idx_k_prewritten.
• Main indexer KV dtype validation remains unchanged.
• No configuration or test-list files changed.
• Review focus: validate layout checks, cache strides, dtype casts, metadata reset and consumption, and fallback behavior.

QA Engineer Review

• Added test_fused_scatter_matches_reference.
• The test compares fused and reference K/V/index-K writes across cache dtypes, KV head counts, optional index-K writes, non-contiguous HND views, and strided inputs.
• The test is CUDA-gated.
• No corresponding tests/integration/test_lists/ coverage was identified.
• Verdict: needs follow-up.

Description

Port of #16755 (merged on Minimax-m3 side branch for AgentX submission) to main. Same change, re-applied on top of main's MSA backend.

Where this PR targets at:

MiniMax-M3 MSA per-layer KV-cache writes fusion (K, V, index-K):

Take a sparse layer as example:

1. input RMSNorm
2. QKV proj  (+ index-Q/K proj for the indexer)
3. ►► WRITE: scatter new K, V, index-K into the paged caches   ◄◄ our fix
4. indexer:  proxy-FMHA (reads index-K) → max-reduce → SelectBlocks (top-16 of 128-token blocks)
5. main sparse attention (reads K,V from selected blocks) → o_proj
6. post-attn RMSNorm
7. MoE FFN (router → top-4 experts + shared) + allreduce

The problem

Each MSA layer writes its new-token main K, main V, and (sparse layers) index-K. Writing of K, V and index-K takes four tiny kernels: division, remainder, index cast, index_put scatter. So that's 3*4=12 kernel launches per sparse layer. At 60 layers that is ~720 tiny kernels per decode step, all captured into decode CUDA graphs and re-executed on every step. It accumulates to large launch overhead.

The fix

This PR replaces them with one Triton launch per layer (~720 → 60) launches on the write path each iter).

Measured impact (4x B300, InferenceMAX-style serving benchmark)

trtllm-serve + benchmark_serving, NVFP4 + fp8 KV cache, MSA, decode CUDA graphs + overlap scheduler, no spec decode, random 8k/1k, identical seeded prompt sets, JIT/autotuner warmed.

metric	conc 64, TP4/EP4	conc 32, TP2/EP1
output throughput	1,817.1 → 1,942.5 tok/s (+6.9%)	1,082.1 → 1,171.7 tok/s (+8.3%)
median TPOT (≈ decode iteration time)	33.86 → 31.76 ms (−2.1 ms)	28.39 → 26.20 ms (−2.2 ms)
mean TTFT	995 → 954 ms (−4.1%)	715 → 690 ms (−3.5%)

Nsys trace for a layer (look at green parts)

Before fix: 4 tiny kernels-to-be-fused for index-K, then 3 kernels for indexer, then 8 tiny kernels-to-be-fused for K and V.

After fix: 1 fused Triton launch (contains previous 12 kernels for K, V and K-index), then 3 kernels for indexer. So it's fuse + reorder: compute the page/offset once, write K + V + index-K together, before the indexer runs.

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.

[zheyuf]

/bot help

[github-actions]

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental) --high-priority]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Supports wildcard * for pattern matching (e.g., "*PerfSanity*" matches all stages containing PerfSanity). Examples: "A10-PyTorch-1, xxx", "PerfSanity". The patterns "*", "*Post-Merge*", and "*PerfSanity*", including equivalent escaped or repeated-star forms and their use in comma-separated lists, require the ci: post-merge approved PR label. Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Requires the ci: full pre-merge approved label on the PR (ask a member of NVIDIA/trt-llm-ci-approvers). Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline. Requires the ci: full pre-merge approved label on the PR (ask a member of NVIDIA/trt-llm-ci-approvers).

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline. Requires the ci: post-merge approved PR label applied by an active member of NVIDIA/trt-llm-ci-approvers. The approval label remains in place when new commits are pushed.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Supports wildcard * for pattern matching. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx", --extra-stage "Post-Merge". The patterns "*", "*Post-Merge*", and "*PerfSanity*", including equivalent escaped or repeated-star forms and their use in comma-separated lists, require the ci: post-merge approved PR label.

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

--high-priority (OPTIONAL) : Run the pipeline with high priority. This option is restricted to authorized users only and will route the job to a high-priority queue.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[zheyuf]

/bot run --disable-fail-fast

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8ff510b3-3679-43c5-bc7d-2d0a0854b438

📥 Commits

Reviewing files that changed from the base of the PR and between 30316da and 8e75f98.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py
• tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
• tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_scatter.py
• tensorrt_llm/_torch/models/modeling_minimaxm3.py
• tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

Changes

The MiniMax-M3 MSA path now performs fused paged K/V and optional index-K cache writes. Metadata tracks prewritten layers so indexer and attention paths skip duplicate writes. CUDA-gated tests compare fused and reference writes across cache dtypes, head counts, and strided layouts.

MiniMax-M3 fused cache write flow

Layer / File(s)	Summary
Fused paged cache scatter tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_scatter.py	A Triton kernel writes K/V and optional index-K values in one launch. The wrapper validates layouts, strides, dimensions, devices, and empty inputs.
Prewritten-layer metadata and indexer integration tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py, tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py	Metadata records fused writes and clears stale markers. BF16 index-K and main K/V writes skip duplicate operations when caches are prewritten.
Model wiring and validation tensorrt_llm/_torch/models/modeling_minimaxm3.py, tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py	MSA layers call the shared cache-writing path before index selection. Tests compare fused and reference writes across dtypes, head counts, optional index-K writes, and strided layouts.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: 🔵 Low · up to 8e75f

The PR consolidates per-layer K/V/index-K cache writes into a fused operation, improving decode performance while preserving the established fallback for unsupported layouts. If the fused write fails after partially updating cache state, recovery and slot reuse are not fully established, so merge is appropriate with explicit owner follow-up on failure handling.

Sequence Diagram(s)

sequenceDiagram
  participant MiniMaxM3Model
  participant MiniMaxM3MsaSparseAttention
  participant fused_write_layer_caches
  participant PagedCaches
  participant run_indexer
  MiniMaxM3Model->>MiniMaxM3MsaSparseAttention: provide K/V and optional index-K
  MiniMaxM3MsaSparseAttention->>fused_write_layer_caches: write layer caches
  fused_write_layer_caches->>PagedCaches: store K/V and index-K
  MiniMaxM3MsaSparseAttention->>run_indexer: run with idx_k_prewritten
  run_indexer->>PagedCaches: skip duplicate BF16 index-K write

Loading

Suggested reviewers: bowenfu

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 84.62% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 13 functions across 5 files.
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 follows the required [None][type] format and clearly describes the main change: fusing MiniMax-M3 MSA per-layer KV-cache writes.
Description check	✅ Passed	The description clearly explains the problem, solution, performance impact, implementation flow, and checklist status. It omits the dedicated Test Coverage heading, but the overall description is suff…

Full details: Description check

Explanation

The description clearly explains the problem, solution, performance impact, implementation flow, and checklist status. It omits the dedicated Test Coverage heading, but the overall description is sufficiently complete and references the relevant validation context.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[tensorrt-cicd]

PR_Github #71089 [ run ] triggered by Bot. Commit: 8e75f98 Link to invocation

[tensorrt-cicd]

PR_Github #71089 [ run ] completed with state SUCCESS. Commit: 8e75f98
/LLM/main/L0_MergeRequest_PR pipeline #58239 completed with status: 'SUCCESS'

CI Report

Link to invocation