[https://nvbugs/6159132][fix] Differentiate the two paths via extra_acc_spec="tp_attn" when attention_dp=False

[tensorrt-cicd]

Summary

• Root cause: The GSM8K reference accuracy (93.75) for MiniMax-M2 FP8_BLOCK_SCALES was set based on the attention_dp=True path, but the attention_dp=False path uses the fused minimax_allreduce_rms_qk kernel whose numerics produce ~90.49, below the derived threshold of 90.547.
• Fix: Differentiate the two paths via extra_acc_spec="tp_attn" when attention_dp=False, add a dedicated reference entry of 92.0 in gsm8k.yaml (threshold 88.797), and remove the existing waiver.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6159132

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

CUDA graph cleanup now uses a shared safe-reset helper. Piecewise, decoder, and encoder capture paths restore mutable state, defer metadata publication, and reset partially created graphs when failures occur. New tests cover cleanup continuation, state restoration, and capture failure handling.

Changes

CUDA Graph Capture Cleanup

Layer / File(s)	Summary
Safe graph reset utility tensorrt_llm/_torch/utils.py	Adds safe_reset_cuda_graph, which logs RuntimeError failures instead of propagating them.
Piecewise capture cleanup tensorrt_llm/_torch/compilation/piecewise_optimizer.py, tests/unittest/_torch/compilation/test_piecewise_cuda_graph_cleanup.py	Restores stream metadata, resets failed captures, defers output-address publication, and continues cleanup after reset errors.
Decoder and encoder capture transactions tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py, tests/unittest/_torch/executor/test_cuda_graph_cleanup.py	Restores speculative-decoding state, defers graph metadata commits, resets partial graphs, and validates decoder and encoder cleanup behavior.

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

Sequence Diagram(s)

sequenceDiagram
  participant CUDAGraphRunner
  participant ModelForward
  participant safe_reset_cuda_graph
  participant GraphMetadata
  CUDAGraphRunner->>ModelForward: Run warmup and capture forwards
  ModelForward-->>CUDAGraphRunner: Return outputs or raise an exception
  CUDAGraphRunner->>CUDAGraphRunner: Restore KV-length and stream state
  alt Capture or output processing fails
    CUDAGraphRunner->>safe_reset_cuda_graph: Reset the partial graph
    safe_reset_cuda_graph-->>CUDAGraphRunner: Log reset failure if needed
  else Capture succeeds
    CUDAGraphRunner->>GraphMetadata: Publish outputs and graph metadata
  end

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#14914: Modifies the same CUDA graph orphan cleanup and transactional capture/reset logic.
• NVIDIA/TensorRT-LLM#16706: Also modifies CUDA graph capture and cleanup in cuda_graph_runner.py, including encoder-decoder graph handling.

Suggested reviewers: qijune, bowenfu

🚥 Pre-merge checks | ✅ 2 | ❌ 3

❌ Failed checks (3 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	⚠️ Warning	The title describes a GSM8K attention-path fix, but the changes implement CUDA graph capture and teardown hardening.	Update the title to describe CUDA graph capture, cleanup, and failure-handling changes, and use the relevant bug ID if available.
Description check	⚠️ Warning	The description explains a MiniMax-M2 GSM8K accuracy fix, but the changes and tests address CUDA graph cleanup and capture failures.	Rewrite the description to explain the CUDA graph hardening changes, affected paths, failure behavior, and the added cleanup tests.

✅ Passed checks (2 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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[tensorrt-cicd]

[Repair Bot] Attempted automated rebase/maintenance, but could not complete it safely. Leaving the PR unchanged.

Reason:

reserve_gpu: No partition matched for b200 (need 4 GPUs)

The bot will try again on a later cycle. Manual rebase is also fine.

[coderabbitai]

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.

[2ez4bz]

/bot run

[tensorrt-cicd]

PR_Github #64800 [ run ] triggered by Bot. Commit: 4067f46 Link to invocation

[tensorrt-cicd]

PR_Github #64800 [ run ] completed with state FAILURE. Commit: 4067f46
/LLM/main/L0_MergeRequest_PR pipeline #52643 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

[2ez4bz]

/bot run

[tensorrt-cicd]

PR_Github #64933 [ run ] triggered by Bot. Commit: 4067f46 Link to invocation

[tensorrt-cicd]

PR_Github #64933 [ run ] completed with state FAILURE. Commit: 4067f46
/LLM/main/L0_MergeRequest_PR pipeline #52774 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

[2ez4bz]

/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.

[2ez4bz]

/bot run --add-multi-gpu-test

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unittest/_torch/compilation/test_piecewise_cuda_graph_cleanup.py (1)

73-100: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Verify the current_stream call count assumption.

side_effect=[capture_stream, restored_stream] provides exactly two values. The capture path calls torch.cuda.current_stream() twice today. If a future change adds a third call inside the patched region, the test fails with StopIteration rather than a clear assertion. Consider side_effect=itertools.chain([capture_stream], itertools.repeat(restored_stream)) to keep the failure message meaningful.

🤖 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/compilation/test_piecewise_cuda_graph_cleanup.py`
around lines 73 - 100, Update
test_capture_failure_resets_graph_before_entry_commit so the patched
torch.cuda.current_stream side effect supplies capture_stream first and
restored_stream for all subsequent calls, using an unbounded iterator instead of
a fixed two-value list. Preserve the existing assertions and capture behavior
while preventing unrelated StopIteration failures if additional current_stream
calls are introduced.

🤖 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 `@tests/unittest/_torch/executor/test_cuda_graph_cleanup.py`:
- Around line 183-209: Update
test_encoder_rejects_nested_output_without_orphaning_graph so capture reaches
nested-output validation: stub _stage_inputs or provide its complete required
inputs, including input_ids, seq_lens, and attention metadata containing
_seq_lens_cuda and _seq_lens. Patch torch.cuda.current_stream for CPU-only
execution, preserving the existing cleanup assertions.

---

Nitpick comments:
In `@tests/unittest/_torch/compilation/test_piecewise_cuda_graph_cleanup.py`:
- Around line 73-100: Update
test_capture_failure_resets_graph_before_entry_commit so the patched
torch.cuda.current_stream side effect supplies capture_stream first and
restored_stream for all subsequent calls, using an unbounded iterator instead of
a fixed two-value list. Preserve the existing assertions and capture behavior
while preventing unrelated StopIteration failures if additional current_stream
calls are introduced.

🪄 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: 44e555e1-38d9-43b1-ab02-7cfc2d5d850e

📥 Commits

Reviewing files that changed from the base of the PR and between 4067f46 and 12ee846.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/compilation/piecewise_optimizer.py
• tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
• tensorrt_llm/_torch/utils.py
• tests/unittest/_torch/compilation/test_piecewise_cuda_graph_cleanup.py
• tests/unittest/_torch/executor/test_cuda_graph_cleanup.py

[tensorrt-cicd]

PR_Github #64967 [ run ] triggered by Bot. Commit: 12ee846 Link to invocation

[2ez4bz]

/bot run --add-multi-gpu-test

[github-actions]

⚠️ Bot command ignored: The /bot command must appear at the very beginning of the comment (no leading blank lines or spaces). Please post a new comment with /bot as the first character.

[2ez4bz]

/bot run --add-multi-gpu-test

[tensorrt-cicd]

PR_Github #64973 [ run ] triggered by Bot. Commit: 772a4bb Link to invocation

[tensorrt-cicd]

PR_Github #64967 [ run ] completed with state ABORTED. Commit: 12ee846

Link to invocation

[tensorrt-cicd]

PR_Github #64973 [ run ] completed with state FAILURE. Commit: 772a4bb
/LLM/main/L0_MergeRequest_PR pipeline #52807 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

[zhaoyangwang-nvidia]

The PR title/description describe a MiniMax-M2 GSM8K accuracy threshold fix, but the diff is CUDA graph capture/teardown hardening (nvbugs/6248648). Please update the title and body to match before merge.

Dismissing, as I accidentally force pushed completely unrelated changes.

[2ez4bz]

/bot run

[tensorrt-cicd]

PR_Github #65080 [ run ] triggered by Bot. Commit: be462ad Link to invocation

[tensorrt-cicd]

PR_Github #65080 [ run ] completed with state SUCCESS. Commit: be462ad
/LLM/main/L0_MergeRequest_PR pipeline #52883 completed with status: 'SUCCESS'

CI Report

Link to invocation