[https://nvbugs/6171743][fix] Set PYTORCH_ALLOC_CONF=expandable_segments:True on MPI workers via `patch_mpi_

[tensorrt-cicd]

Summary

• Root cause: On RTX PRO 6000 Blackwell (96 GB, SM120), the DeepSeek-R1 NVFP4 throughput_tp8 config leaves only ~7 GiB free CUDA memory after kv_cache, so the MoE autotuner's ~14.3 GiB contiguous workspace request fails despite ~13 GiB held in PyTorch's fragmented cache.
• Fix: Set PYTORCH_ALLOC_CONF=expandable_segments:True on MPI workers via patch_mpi_pool_session_for_env so the allocator can satisfy the autotuner request from non-contiguous virtual segments; also lower free_gpu_memory_fraction from 0.70 to 0.5 for more headroom; remove the 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/6171743

Summary by CodeRabbit

• Tests
  • Enhanced test reliability by improving environment configuration handling for multi-GPU setups.
  • Adjusted memory allocation parameters to optimize resource usage in memory-constrained test scenarios.
  • Re-enabled a previously skipped test case with fixes to address underlying stability issues.

[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: 5292f179-6166-4964-baad-cf5e4f3e24bf

📥 Commits

Reviewing files that changed from the base of the PR and between e6ccf97 and 56e2162.

📒 Files selected for processing (2)

• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

---

📝 Walkthrough

Walkthrough

The PR updates the SM120 NVFP4 multi-GPU test to enable MPI environment variable propagation for PyTorch memory allocation expansion, adjusts the KV-cache memory fraction configuration, and removes the test's skip waiver.

Changes

SM120 NVFP4 Test Configuration

Layer / File(s)	Summary
SM120 NVFP4 test configuration and MPI environment patching tests/integration/defs/accuracy/test_llm_api_pytorch.py	Test function adds mocker fixture, patches MPI pool session environment to set PYTORCH_ALLOC_CONF=expandable_segments:True, and reduces KV cache free_gpu_memory_fraction from 0.70 to 0.5.
Test waiver removal tests/integration/test_lists/waives.txt	Skip waiver for TestDeepSeekR1::test_nvfp4_multi_gpus_sm120[throughput_tp8] is removed, enabling the test to run.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#14086: Modifies the same test waiver for test_nvfp4_multi_gpus_sm120[throughput_tp8].

Suggested reviewers

• xinhe-nv
• hyukn
• jieli-matrix

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The PR title is truncated and incomplete, ending abruptly at 'patch_mpi_', making it unclear and difficult to understand the full scope of the change.	Complete the PR title to clearly convey the full change: '[https://nvbugs/6171743][fix] Set PYTORCH_ALLOC_CONF=expandable_segments:True on MPI workers via patch_mpi_pool_session_for_env'.
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.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description check	✅ Passed	The PR description covers the root cause, fix, and test plan comprehensively. It clearly explains the issue, the solution implemented, and verification steps performed.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

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

[dominicshanshan]

/bot run

[tensorrt-cicd]

PR_Github #48483 [ run ] triggered by Bot. Commit: 56e2162 Link to invocation

[tensorrt-cicd]

PR_Github #48483 [ run ] completed with state SUCCESS. Commit: 56e2162
/LLM/main/L0_MergeRequest_PR pipeline #38279 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

[ziyixiong-nv]

/bot run

[tensorrt-cicd]

PR_Github #48555 [ run ] triggered by Bot. Commit: 6833844 Link to invocation

[tensorrt-cicd]

PR_Github #48555 [ run ] completed with state SUCCESS. Commit: 6833844
/LLM/main/L0_MergeRequest_PR pipeline #38345 completed with status: 'SUCCESS'

CI Report

Link to invocation

[dominicshanshan]

/bot run

[tensorrt-cicd]

PR_Github #48683 [ run ] triggered by Bot. Commit: 6833844 Link to invocation

[tensorrt-cicd]

PR_Github #48683 [ run ] completed with state SUCCESS. Commit: 6833844
/LLM/main/L0_MergeRequest_PR pipeline #38460 completed with status: 'SUCCESS'

CI Report

Link to invocation

[dominicshanshan]

/bot run

[tensorrt-cicd]

PR_Github #49625 [ run ] triggered by Bot. Commit: aeb0d85 Link to invocation

[tensorrt-cicd]

PR_Github #49625 [ run ] completed with state SUCCESS. Commit: aeb0d85
/LLM/main/L0_MergeRequest_PR pipeline #39245 completed with status: 'SUCCESS'

CI Report

Link to invocation