[https://nvbugs/5580099][fix] Cherry pick IMA issue fix from release/1.1

[JunyiXu-nv]

Summary by CodeRabbit

Release Notes

• Improvements

  • Enhanced CUDA graph memory workspace management to optimize inference efficiency during model execution.

• Tests

  • Improved integration test memory usage diagnostics with clearer error messages and updated baseline calibration values.

Description

This PR is cherry-picking PR: #8685 to fix the issue on main branch. The commit was excluded from mass integration due to a failed test case. (test_e2e.py::test_ptp_quickstart_advanced_mtp)

The test case failed due to the introduced cuda graph workspace, which increases the memory usage by ~300MB. As the memory usage has been close to the limit before (about 200MB more than expected, which is within 300MB delta), this increasing makes the test fail. But if we delete the memory used by the non-cuda-graph workspace, there would be a cudaFree issue (test: accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales). Then both of the workspaces need to exist. That's why I increased the memory usage by 0.3 GB in the test case, which is ~0.3% of the total expected peak memory usage.

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)

• 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

/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)]

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. Examples: "A10-PyTorch-1, xxx". 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. 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.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

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

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

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.

[JunyiXu-nv]

/bot run

[tensorrt-cicd]

PR_Github #23943 [ run ] triggered by Bot. Commit: 736cb9d

[tensorrt-cicd]

PR_Github #23943 [ run ] completed with state SUCCESS. Commit: 736cb9d
/LLM/main/L0_MergeRequest_PR pipeline #18029 completed with status: 'FAILURE'

[JunyiXu-nv]

/bot run

[tensorrt-cicd]

PR_Github #23960 [ run ] triggered by Bot. Commit: 736cb9d

[tensorrt-cicd]

PR_Github #23960 [ run ] completed with state SUCCESS. Commit: 736cb9d
/LLM/main/L0_MergeRequest_PR pipeline #18043 completed with status: 'FAILURE'

[JunyiXu-nv]

/bot run

[tensorrt-cicd]

PR_Github #24431 [ run ] triggered by Bot. Commit: 1b5ee49

[tensorrt-cicd]

PR_Github #24431 [ run ] completed with state FAILURE. Commit: 1b5ee49
/LLM/main/L0_MergeRequest_PR pipeline #18436 completed with status: 'FAILURE'

[JunyiXu-nv]

/bot run

[tensorrt-cicd]

PR_Github #24441 [ run ] triggered by Bot. Commit: aee8df8

[tensorrt-cicd]

PR_Github #24441 [ run ] completed with state SUCCESS. Commit: aee8df8
/LLM/main/L0_MergeRequest_PR pipeline #18441 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes add CUDA graph workspace support to the TensorRT LLM attention backend metadata by introducing a new optional field with conditional initialization, and update integration test memory assertions with clearer intermediate variable calculations and adjusted expected memory thresholds.

Changes

Cohort / File(s)	Summary
Attention backend workspace enhancement tensorrt_llm/_torch/attention_backend/trtllm.py	Added cuda_graph_workspace: Optional[torch.Tensor] field to TrtllmAttentionMetadata. Initialized to empty CUDA tensor by default. Updated no-cache path to create cuda_graph_workspace if None. Modified forward planning to conditionally branch between metadata.workspace and metadata.cuda_graph_workspace based on CUDA graph mode.
Integration test memory assertion refactoring tests/integration/defs/test_e2e.py	Refactored _check_mem_usage memory assertion from direct inequality to explicit intermediate variable increased_peak_mem with more detailed error messaging. Updated test_ptp_quickstart_advanced_mtp expected memory threshold from [54.60, 0, 0, 0] to [54.90, 0, 0, 0].

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

• Focus on verifying the CUDA graph workspace branching logic is correctly conditioned in forward planning
• Confirm memory threshold increase (54.60 → 54.90) is justified by actual test execution results
• Validate that workspace initialization path handles None case consistently in both regular and no-cache branches

Pre-merge checks and finishing touches

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.
Description check	❓ Inconclusive	The PR description explains the issue and solution, but the Test Coverage section is empty.	Complete the Test Coverage section by listing the relevant tests that validate the changes (e.g., test_e2e.py::test_ptp_quickstart_advanced_mtp and accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales).

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly and specifically references the NVBugs issue being fixed and indicates it's cherry-picking a fix from the release branch to main.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (2)

tensorrt_llm/_torch/attention_backend/trtllm.py (2)

580-580: Consider adding field documentation.

While the field declaration is correct, consider adding a docstring or comment explaining the purpose of cuda_graph_workspace versus the regular workspace, especially since the distinction is important for avoiding CUDA free issues in FP8 MLA scenarios.

Example:

     workspace: Optional[torch.Tensor] = None
+    # Separate workspace for CUDA graph mode to avoid cudaFree issues with FP8 MLA
     cuda_graph_workspace: Optional[torch.Tensor] = None

---

1347-1349: Suggest clarifying the comment.

The workspace selection logic is correct, but the comment could be clearer about why separate workspaces are needed.

Consider rephrasing the comment:

-            # re-enable it, if pass None to it, fp8 mla will encounter invalid cuda free issue.
+            # Use separate workspace for CUDA graph mode; passing None causes cudaFree issues in FP8 MLA.
             workspace=metadata.workspace
             if not metadata.is_cuda_graph else metadata.cuda_graph_workspace,

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between d0ea417 and aee8df8.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/attention_backend/trtllm.py (3 hunks)
• tests/integration/defs/test_e2e.py (2 hunks)

🧰 Additional context used

🧠 Learnings (4)

📓 Common learnings

Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: tensorrt_llm/_torch/attention_backend/trtllm.py:259-262
Timestamp: 2025-08-14T15:43:23.107Z
Learning: In TensorRT-LLM's attention backend, tensor parameters in the plan() method are assigned directly without validation (dtype, device, contiguity checks). This maintains consistency across all tensor inputs and follows the pattern of trusting callers to provide correctly formatted tensors.

Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:53.813Z
Learning: In the TensorRT-LLM KV cache manager, SWA (Sliding Window Attention) combined with beam search is currently in a broken/non-functional state and is planned for future rework. During preparatory refactoring phases, code related to SWA+beam search may intentionally remain in a non-working state until the broader rework is completed.

Learnt from: galagam
Repo: NVIDIA/TensorRT-LLM PR: 6487
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py:1-12
Timestamp: 2025-08-06T13:58:07.506Z
Learning: In TensorRT-LLM, test files (files under tests/ directories) do not require NVIDIA copyright headers, unlike production source code files. Test files typically start directly with imports, docstrings, or code.

Learnt from: achartier
Repo: NVIDIA/TensorRT-LLM PR: 6763
File: tests/integration/defs/triton_server/conftest.py:16-22
Timestamp: 2025-08-11T20:09:24.389Z
Learning: In the TensorRT-LLM test infrastructure, the team prefers simple, direct solutions (like hard-coding directory traversal counts) over more complex but robust approaches when dealing with stable directory structures. They accept the maintenance cost of updating tests if the layout changes.

Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/thop/allreduceOp.cpp:352-446
Timestamp: 2025-09-23T15:12:38.312Z
Learning: In TensorRT-LLM NCCL device implementation, NCCL version 2.28+ requirements are handled at runtime in the nccl_device/config layer rather than with compile-time guards. This allows the allreduceOp to remain version-agnostic and delegates version compatibility validation to the appropriate lower-level components that can gracefully handle unsupported configurations.

Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/multimem.h:20-30
Timestamp: 2025-09-23T15:13:48.819Z
Learning: TRT-LLM targets modern CUDA toolkits that support FP8 datatypes, so cuda_fp8.h can be included unconditionally without version guards in TRT-LLM code.

Learnt from: pengbowang-nv
Repo: NVIDIA/TensorRT-LLM PR: 7192
File: tests/integration/test_lists/test-db/l0_dgx_b200.yml:56-72
Timestamp: 2025-08-26T09:49:04.956Z
Learning: In TensorRT-LLM test configuration files, the test scheduling system handles wildcard matching with special rules that prevent duplicate test execution even when the same tests appear in multiple yaml files with overlapping GPU wildcards (e.g., "*b200*" and "*gb200*").

Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7645
File: tests/integration/test_lists/qa/llm_function_core.txt:648-648
Timestamp: 2025-09-09T09:40:45.658Z
Learning: In TensorRT-LLM test lists, it's common and intentional for the same test to appear in multiple test list files when they serve different purposes (e.g., llm_function_core.txt for comprehensive core functionality testing and llm_function_core_sanity.txt for quick sanity checks). This duplication allows tests to be run in different testing contexts.

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6728
File: cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp:966-966
Timestamp: 2025-08-08T04:10:19.038Z
Learning: TensorRT plugins currently don't support padding functionality, and TensorRT is not getting new features (in maintenance mode). This means that duplicating parameters like mExpertHiddenSize in function calls, even with TODO comments, can be acceptable as pragmatic solutions within these constraints.

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp:399-417
Timestamp: 2025-08-21T21:48:35.135Z
Learning: CUTLASS extensions in TensorRT-LLM (located under cpp/tensorrt_llm/cutlass_extensions/) are designed to integrate with and extend functionality in the external CUTLASS repository. When analyzing these extensions, their consumers and functionality wiring may exist in the CUTLASS codebase rather than within TensorRT-LLM itself.

📚 Learning: 2025-08-15T06:46:53.813Z

Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:53.813Z
Learning: In the TensorRT-LLM KV cache manager, SWA (Sliding Window Attention) combined with beam search is currently in a broken/non-functional state and is planned for future rework. During preparatory refactoring phases, code related to SWA+beam search may intentionally remain in a non-working state until the broader rework is completed.

Applied to files:

• tensorrt_llm/_torch/attention_backend/trtllm.py

📚 Learning: 2025-08-14T15:43:23.107Z

Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: tensorrt_llm/_torch/attention_backend/trtllm.py:259-262
Timestamp: 2025-08-14T15:43:23.107Z
Learning: In TensorRT-LLM's attention backend, tensor parameters in the plan() method are assigned directly without validation (dtype, device, contiguity checks). This maintains consistency across all tensor inputs and follows the pattern of trusting callers to provide correctly formatted tensors.

Applied to files:

• tensorrt_llm/_torch/attention_backend/trtllm.py

📚 Learning: 2025-08-21T00:16:56.457Z

Learnt from: farshadghodsian
Repo: NVIDIA/TensorRT-LLM PR: 7101
File: docs/source/blogs/tech_blog/blog9_Deploying_GPT_OSS_on_TRTLLM.md:36-36
Timestamp: 2025-08-21T00:16:56.457Z
Learning: TensorRT-LLM container release tags in documentation should only reference published NGC container images. The README badge version may be ahead of the actual published container versions.

Applied to files:

• tensorrt_llm/_torch/attention_backend/trtllm.py

🔇 Additional comments (3)

tests/integration/defs/test_e2e.py (2)

151-156: LGTM! Clear refactoring improves debuggability.

The introduction of the intermediate increased_peak_mem variable makes the memory assertion more explicit and the error message more informative. This will help diagnose memory usage issues more effectively.

---

1958-1958: Memory threshold adjustment aligns with CUDA graph workspace addition.

The 0.3 GB increase corresponds to the new CUDA graph workspace mentioned in the PR description, representing approximately 0.5% of the total expected peak memory.

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

697-702: LGTM! Consistent initialization pattern.

The cuda_graph_workspace initialization follows the same pattern as the regular workspace initialization above it (lines 690-695), ensuring consistency in how both workspace tensors are managed.