[TRTLLM-13968][fix] Enable MiniMax M3 piecewise CUDA graphs by liji-nv · Pull Request #15923 · NVIDIA/TensorRT-LLM

liji-nv · 2026-07-03T12:28:00Z

Wrap the MiniMax M3 metadata- and cache-dependent attention core in an inplace custom op so torch.compile can split it out of piecewise CUDA graphs. Keep QKV/index projections, QK normalization, RoPE, and the output projection visible to the compiled graph.

Write dense and sparse attention results into the custom-op output buffer. Preserve FP32 sparse GQA accumulation until the final copy/cast, and expose the output buffer through MiniMaxM3SparseRuntimeBackend.forward.

Register attention boundaries and mutation metadata through optional TRT-LLM op lookup, matching the latest GDN registration pattern from PR #15594. This avoids depending on model-specific custom ops being imported when compilation utilities initialize.

Track piecewise runners owned by the compile backend and reset their CUDA graphs, captured addresses, outputs, and warmup state when phase-1 KV-cache estimation is released. Phase 2 then recaptures against the final allocations instead of replaying stale graph pointers.

Add an 8-GPU MiniMax-M3-MXFP8 torch.compile E2E variant covering TP8/EP8, attention DP, TRTLLM MoE, padding CUDA graphs, multi-stream piecewise capture, and phase-2 recapture.

Summary by CodeRabbit

New Features
- Added support for writing attention results into preallocated output buffers, improving flexibility for MiniMax-M3 attention paths.
- Expanded compile-time piecewise CUDA graph handling so captured runners can be tracked and cleared later.
- Introduced broader support for MiniMax-M3 dense and sparse attention execution under compile-aware workflows.
Bug Fixes
- Improved cleanup of CUDA graph resources during model shutdown.
- Added coverage for a new MiniMax-M3 MXFP8 piecewise CUDA graph scenario.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-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.

coderabbitai · 2026-07-03T12:30:21Z

📝 Walkthrough

Walkthrough

Adds an optional preallocated output tensor parameter threaded through MiniMax-M3 sparse attention functions, runtime backend, and dense/sparse attention cores, with a compiled custom-op boundary for the dense path. Also adds tracking and cleanup of piecewise CUDA graph runners in the torch.compile backend, invoked from model engine cleanup.

Changes

MiniMax-M3 output tensor plumbing

Layer / File(s)	Summary
Sparse GQA masked output buffer support `tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/backend.py`	`_sparse_gqa_masked` and its decode/prefill entry points accept an optional `output` tensor and write results directly into it after FP32 accumulation, validating shape when provided.
Runtime backend output resolution and forwarding `tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/backend.py`	`MiniMaxM3SparseRuntimeBackend.forward_sparse` and `forward` accept/merge `output` from `AttentionForwardArgs` and kwargs, raise on conflicting values, and pass the resolved buffer into decode/prefill calls.
MiniMax-M3 attention core writes into output buffer `tensorrt_llm/_torch/models/modeling_minimaxm3.py`	Dense and sparse paths in `MiniMaxM3Attention` add `_forward_attention_core`/`_attention_core`/`_sparse_attention_core`, validate metadata earlier, and write results directly into a caller-provided output buffer instead of building intermediate tensors.
Compiled custom-op boundary and optional op lookup `tensorrt_llm/_torch/models/modeling_minimaxm3.py`, `tensorrt_llm/_torch/compilation/utils.py`	Adds `minimax_m3_attn_custom_op_inplace` registered via `torch.library.custom_op`, a `get_optional_trtllm_op` helper, and conditional `inplace_map` entries exposing available ops to the piecewise compiler.
MXFP8 piecewise CUDA graph test coverage `tests/integration/defs/accuracy/test_llm_api_pytorch.py`, `tests/integration/test_lists/qa/llm_function_core.txt`	Adds an integration test running MiniMax-M3-MXFP8 with piecewise CUDA graph and torch.compile settings, plus a QA test-list entry.

Piecewise CUDA graph runner tracking and cleanup

Layer / File(s)	Summary
PiecewiseRunner collection and clear_cuda_graphs `tensorrt_llm/_torch/compilation/piecewise_optimizer.py`	`PiecewiseInterpreter` collects created `PiecewiseRunner` instances and `piecewise_optimizer` returns them; `PiecewiseRunner.clear_cuda_graphs()` resets captured graphs and cached state.
Backend tracking and clear_piecewise_cuda_graphs `tensorrt_llm/_torch/compilation/backend.py`	`Backend` stores returned runners in a `WeakSet` and exposes `clear_piecewise_cuda_graphs()` to invoke `clear_cuda_graphs()` on each tracked runner.
Model engine invokes piecewise cleanup `tensorrt_llm/_torch/pyexecutor/model_engine.py`	`_release_cuda_graphs` calls `clear_piecewise_cuda_graphs()` on `_torch_compile_backend` before clearing other CUDA graph runners.

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

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant RuntimeBackend as MiniMaxM3SparseRuntimeBackend
  participant Decode as minimax_m3_sparse_decode/prefill
  participant Masked as _sparse_gqa_masked

  Caller->>RuntimeBackend: forward(forward_args, output, kwargs)
  RuntimeBackend->>RuntimeBackend: merge forward_args/kwargs, resolve output
  RuntimeBackend->>Decode: forward_sparse(output=resolved_output)
  Decode->>Masked: _sparse_gqa_masked(output=resolved_output)
  Masked-->>Decode: writes result into output tensor

sequenceDiagram
  participant ModelEngine as PyTorchModelEngine
  participant Backend
  participant Runner as PiecewiseRunner

  ModelEngine->>Backend: clear_piecewise_cuda_graphs()
  Backend->>Runner: clear_cuda_graphs() (for each tracked runner)
  Runner-->>Backend: resets captured graphs and cached state

Suggested reviewers: QiJune, fredricz-20070104, Tabrizian, jieli-matrix

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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 matches the main change and follows the required ticket-and-type format.
Description check	✅ Passed	The description explains the change and includes test context and checklist items, though the sections are not explicitly formatted to the template.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

liji-nv · 2026-07-03T12:34:19Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-03T12:39:43Z

PR_Github #57464 [ run ] triggered by Bot. Commit: c333655 Link to invocation

tensorrt-cicd · 2026-07-03T20:49:14Z

PR_Github #57464 [ run ] completed with state SUCCESS. Commit: c333655
/LLM/main/L0_MergeRequest_PR pipeline #46201 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

liji-nv · 2026-07-06T06:23:10Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-06T06:29:31Z

PR_Github #57706 [ run ] triggered by Bot. Commit: 3e5569f Link to invocation

nv-guomingz

LGTM

Wrap the MiniMax M3 metadata- and cache-dependent attention core in an inplace custom op so torch.compile can split it out of piecewise CUDA graphs. Keep QKV/index projections, QK normalization, RoPE, and the output projection visible to the compiled graph. Write dense and sparse attention results into the custom-op output buffer. Preserve FP32 sparse GQA accumulation until the final copy/cast, and expose the output buffer through MiniMaxM3SparseRuntimeBackend.forward. Register attention boundaries and mutation metadata through optional TRT-LLM op lookup, matching the latest GDN registration pattern from PR NVIDIA#15594. This avoids depending on model-specific custom ops being imported when compilation utilities initialize. Track piecewise runners owned by the compile backend and reset their CUDA graphs, captured addresses, outputs, and warmup state when phase-1 KV-cache estimation is released. Phase 2 then recaptures against the final allocations instead of replaying stale graph pointers. Add an 8-GPU MiniMax-M3-MXFP8 torch.compile E2E variant covering TP8/EP8, attention DP, TRTLLM MoE, padding CUDA graphs, multi-stream piecewise capture, and phase-2 recapture. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

Resetting the last piecewise CUDA graph drops the caching allocator private pool use count to zero, but the runner still retains the opaque pool ID. Reusing that ID during phase-2 recapture can hit the allocator use_count assertion when a live allocation keeps the released pool in the allocator table. After clearing all captures, allocate one fresh pool handle for each previous pool and update every runner in that group. This preserves pool sharing across partitions from the same compiled graph while avoiding reuse of released pool IDs. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

tensorrt-cicd · 2026-07-06T11:27:36Z

PR_Github #57755 [ run ] triggered by Bot. Commit: 5361295 Link to invocation

tensorrt-cicd · 2026-07-06T11:32:46Z

PR_Github #57706 [ run ] completed with state ABORTED. Commit: 3e5569f

Link to invocation

pcicotti · 2026-07-06T12:33:29Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-06T12:39:48Z

PR_Github #57763 [ run ] triggered by Bot. Commit: 5361295 Link to invocation

tensorrt-cicd · 2026-07-06T12:44:05Z

PR_Github #57755 [ run ] completed with state ABORTED. Commit: 5361295

Link to invocation

brb-nv

LGTM.

tensorrt-cicd · 2026-07-06T21:36:51Z

PR_Github #57763 [ run ] completed with state SUCCESS. Commit: 5361295
/LLM/main/L0_MergeRequest_PR pipeline #46466 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

brb-nv · 2026-07-06T21:39:44Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-06T21:47:07Z

PR_Github #57831 [ run ] triggered by Bot. Commit: 5361295 Link to invocation

tensorrt-cicd · 2026-07-06T23:43:41Z

PR_Github #57831 [ run ] completed with state SUCCESS. Commit: 5361295
/LLM/main/L0_MergeRequest_PR pipeline #46526 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

Link to invocation

longlee0622 · 2026-07-06T23:54:24Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-07T00:00:05Z

PR_Github #57859 [ run ] triggered by Bot. Commit: 5361295 Link to invocation

pcicotti · 2026-07-07T00:47:56Z

@NVIDIA/trt-llm-qa @NVIDIA/trt-llm-qa-function @NVIDIA/trt-llm-torch-graph-compiler
ping for owners review.
Thanks

tensorrt-cicd · 2026-07-07T01:43:43Z

PR_Github #57859 [ run ] completed with state SUCCESS. Commit: 5361295
/LLM/main/L0_MergeRequest_PR pipeline #46553 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

pcicotti · 2026-07-07T01:46:03Z

/bot run --disable-fail-fast

longlee0622 · 2026-07-07T01:50:31Z

Need to waive one AutoDeploy test first: #16023

tensorrt-cicd · 2026-07-07T01:52:51Z

PR_Github #57880 [ run ] triggered by Bot. Commit: 5361295 Link to invocation

longlee0622 · 2026-07-07T02:14:10Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-07T02:21:59Z

PR_Github #57890 [ run ] triggered by Bot. Commit: 5361295 Link to invocation

tensorrt-cicd · 2026-07-07T02:25:31Z

PR_Github #57880 [ run ] completed with state ABORTED. Commit: 5361295

Link to invocation

tensorrt-cicd · 2026-07-07T03:00:21Z

PR_Github #57890 [ run ] completed with state SUCCESS. Commit: 5361295
/LLM/main/L0_MergeRequest_PR pipeline #46585 completed with status: 'SUCCESS'

CI Report

Link to invocation

…5809) Squashed cherry-pick of NVIDIA/TensorRT-LLM PR NVIDIA#15809, combining: - [None][feat] Integrate M3 sparse attention kernels form MSA - [None][fix] Fix parallel config of Minimax M3 MXFP8 test - [TRTLLM-14019][feat] Support MSA kernel with enable cuda graph - [None][fix] Fix M3 attention semantic issue Conflicts with the already-integrated PR NVIDIA#15923 were resolved as follows: - requirements.txt: keep flashinfer-python==0.6.14 (main) and add the fmha_sm100 (MSA) dependency. - modeling_minimaxm3.py: graft the _msa_geometry publishing block onto the current _sparse_attention_core (which returns self.attn.forward(..., output=output, ...)). - llm_function_core.txt: keep the MiniMax M3 piecewise/nvfp4 entries. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

…5923) Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv requested review from a team as code owners July 3, 2026 12:28

liji-nv requested review from yechank-nvidia, yizhang-nv and yuxianq July 3, 2026 12:28

github-actions Bot assigned liji-nv Jul 3, 2026

liji-nv force-pushed the fix/minimax-m3-piecewise-cudagraph branch 2 times, most recently from 31496ff to c333655 Compare July 3, 2026 12:29

liji-nv force-pushed the fix/minimax-m3-piecewise-cudagraph branch 2 times, most recently from f7ee59f to 3e5569f Compare July 6, 2026 06:23

yuxianq reviewed Jul 6, 2026

View reviewed changes

Comment thread tests/integration/defs/accuracy/test_llm_api_pytorch.py

yuxianq reviewed Jul 6, 2026

View reviewed changes

Comment thread tests/integration/defs/accuracy/test_llm_api_pytorch.py

yuxianq reviewed Jul 6, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/compilation/piecewise_optimizer.py

yuxianq approved these changes Jul 6, 2026

View reviewed changes

nv-guomingz approved these changes Jul 6, 2026

View reviewed changes

liji-nv added 4 commits July 6, 2026 04:09

[None][test] Restore MiniMax M3 GSM8K coverage

f699c16

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

[None][test] Add MiniMax M3 piecewise graph coverage

2786771

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv force-pushed the fix/minimax-m3-piecewise-cudagraph branch from 68ebcbe to 2786771 Compare July 6, 2026 11:11

liji-nv changed the title ~~[None][fix] Enable MiniMax M3 piecewise CUDA graphs~~ [TRTLLM-13968][fix] Enable MiniMax M3 piecewise CUDA graphs Jul 6, 2026

brb-nv approved these changes Jul 6, 2026

View reviewed changes

hyukn approved these changes Jul 7, 2026

View reviewed changes

StanleySun639 approved these changes Jul 7, 2026

View reviewed changes

pcicotti enabled auto-merge (squash) July 7, 2026 01:46

pcicotti merged commit 211482b into NVIDIA:main Jul 7, 2026
11 checks passed

BrianLi23 pushed a commit to BrianLi23/TensorRT-LLM that referenced this pull request Jul 9, 2026

[TRTLLM-13968][fix] Enable MiniMax M3 piecewise CUDA graphs (NVIDIA#1…

53c4f1d

…5923) Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

pengbowang-nv mentioned this pull request Jul 20, 2026

[https://nvbugs/6305365][chore] Unwaive piecewise cudagraph related tests #16591

Merged

1 task

Uh oh!

Conversation

liji-nv commented Jul 3, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

liji-nv commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

liji-nv commented Jul 6, 2026

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

Uh oh!

nv-guomingz left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

pcicotti commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

brb-nv left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

brb-nv commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

longlee0622 commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

pcicotti commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

pcicotti commented Jul 7, 2026

Uh oh!

longlee0622 commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

longlee0622 commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

liji-nv commented Jul 3, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 3, 2026 •

edited

Loading