[https://nvbugs/6529792][fix] Avoid GPT-OSS V2 cache estimation OOM

[jiaganc]

Dev Engineer Review

• Updated TestGPTOSS::test_w4_1gpu for the GPT-OSS 120B single-H100 accuracy case.
• Reduced max_batch_size from 720 to 512.
• Set max_seq_len to 12,416 tokens, including the GSM8K limits and a 128-token safety margin.
• Preserved V2 KV-cache, CUDA graph, and CUTLASS coverage.
• Pre-commit, Python compilation, and git diff --check passed.
• The full H100 accuracy test was not run locally.

QA Engineer Review

• Modified TestGPTOSS::test_w4_1gpu.
• Existing test-list coverage includes the H100 V2 KV-cache, CUDA graph, and CUTLASS configuration.
• No test-list files changed.
• The full H100 accuracy test remains pending.
• Verdict: needs follow-up.

Description

Fix the GPT-OSS 120B single-H100 V2 KV-cache accuracy test startup OOM tracked by https://nvbugs/6529792.

The test inherited the model's 131K-token sequence limit and configured a maximum batch size of 720. During engine initialization, temporary KV-cache estimation and executor resources exhausted H100 memory before the test reached profiling.

This change:

• caps max_seq_len at the GSM8K workload's configured maximum input plus output length plus a 128-token safety margin (12,416 tokens)
• reduces max_batch_size from 720 to 512

The test continues to cover the V2 KV-cache, CUDA graph, and CUTLASS paths while sizing initialization resources for its actual workload.

Test Coverage

• pre-commit run --files tests/integration/defs/accuracy/test_llm_api_pytorch.py
• PYTHONPYCACHEPREFIX=/tmp/trtllm_gptoss_pr_pycache python3 -m py_compile tests/integration/defs/accuracy/test_llm_api_pytorch.py
• git diff --check

The full H100 accuracy case was not run locally because it requires an H100 and the GPT-OSS 120B model weights. CI should rerun TestGPTOSS::test_w4_1gpu with v2_kv_cache=True, CUDA graph enabled, and the CUTLASS MoE backend.

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]

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: a5506c48-c7e5-4abc-9431-2302257d5543

📥 Commits

Reviewing files that changed from the base of the PR and between 1068940 and 6e193d6.

📒 Files selected for processing (1)

• tests/integration/defs/accuracy/test_llm_api_pytorch.py

---

Walkthrough

The GPT-OSS single-GPU accuracy test reduces max_batch_size from 720 to 512 and sets max_seq_len to the GSM8K input and output limits plus 128 tokens.

Changes

GPT-OSS accuracy configuration

Layer / File(s)	Summary
Configure single-GPU test limits tests/integration/defs/accuracy/test_llm_api_pytorch.py	test_w4_1gpu limits batches to 512 and sets an explicit maximum sequence length based on GSM8K limits plus 128 tokens.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Suggested reviewers: bowenfu, jadotu

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the GPT-OSS V2 cache estimation OOM fix and follows the required NVBugs and fix format.
Description check	✅ Passed	The description explains the problem, solution, affected test paths, validation steps, and the limitation that the full H100 test was not run locally.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[jiaganc]

/bot run --only-qa-verify test accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v2_kv_cache-True-True-cutlass-auto]

[tensorrt-cicd]

PR_Github #65195 [ run ] triggered by Bot. Commit: 6e193d6 Link to invocation

[tensorrt-cicd]

PR_Github #65195 [ run ] completed with state SUCCESS. Commit: 6e193d6
LLM_FUNCTION_AUTO_V2C #366 completed with status: 'SUCCESS'
QA verify test: accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v2_kv_cache-True-True-cutlass-auto] (NVBug 6529792, branch codex/bug-6529792-gptoss-oom, fork jiaganc, dry_run_close=true)

Link to invocation

[jiaganc]

/bot run

[jiaganc]

/bot run

[jiaganc]

/bot run

[yihwang-nv]

/bot run

[tensorrt-cicd]

PR_Github #65232 [ run ] triggered by Bot. Commit: 6e193d6 Link to invocation

[tensorrt-cicd]

PR_Github #65232 [ run ] completed with state FAILURE. Commit: 6e193d6
/LLM/main/L0_MergeRequest_PR pipeline #53014 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

[brnguyen2]

Approving — the comments below are optional touch-ups, not blockers.

Sizing checks out: the mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", 8192) above runs before LLM(...), so this is 4096+8192+128 = 12416, and accuracy_core truncates prompts to MAX_INPUT_LEN and caps generation at MAX_OUTPUT_LEN, so the bound can't be exceeded at runtime.

Two things:

• This is a test-side workaround. The engine-init allocation still scales with max_batch_size × max_seq_len independently of what the workload requests, so the same config outside this test hits the same wall. Is there a follow-up tracking that, or is the intent that the test config was simply unrealistic?
• Please confirm CI ran the failing parametrization (v2_kv_cache-True-True-cutlass-auto) on the hardware that originally reproduced it — the failure is memory-margin sensitive and won't show up on a card with more headroom.

[brnguyen2]

Add a one-line comment tying these two values to https://nvbugs/6529792 and noting that the +128 is headroom over the GSM8K limits. Without it, the next person tuning this test has no way to tell that raising max_seq_len back toward the model's 131k default reintroduces the init-time OOM.