[chenfeiz0326]

Description

#20246 removed moe_a2a_steady_state_budget_for_capture from tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py. It now wraps the capture pass of the new _warmup_and_capture_cuda_graphs in moe_comm_serving_timeouts(). tests/unittest/_torch/compilation/test_phase_selective_forward.py::test_decoder_runner_warmup_bypasses_only_auxiliary_phases still monkeypatches the old name, so both parametrizations fail on any branch that contains #20246:

unittest/_torch/compilation/test_phase_selective_forward.py:266: in test_decoder_runner_warmup_bypasses_only_auxiliary_phases
    monkeypatch.setattr(
E   AttributeError: <module 'tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner' ...> has no attribute 'moe_a2a_steady_state_budget_for_capture'

It hit CPU-Generic-arm-1 in L0 #64229 (PR #20250). That test is in l0_cpu.yml, l0_h100.yml and l0_b200.yml.

This PR patches moe_comm_serving_timeouts instead, with the same contextlib.nullcontext stand-in. The test's purpose is unchanged: it checks which warmup phases run eagerly and which compiled, and the timeout context is just scaffolding.

Test Coverage

• tests/unittest/_torch/compilation/test_phase_selective_forward.py::test_decoder_runner_warmup_bypasses_only_auxiliary_phases[True] and [False].
• The other six tests in the file passed in the same CI run and don't touch this name.
• Not yet run locally. This PR's CI run is the check.

PR Checklist

• PR description clearly explains what and why.
• Follows the TensorRT-LLM coding guidelines.
• Test cases are provided.

🤖 Generated with Claude Code

Dev Engineer Review

The test now patches runner_module.moe_comm_serving_timeouts with contextlib.nullcontext, replacing the removed moe_a2a_steady_state_budget_for_capture. Its warmup-routing assertions are unchanged. Review finding counts are unavailable.

QA Engineer Review

tests/unittest/_torch/compilation/test_phase_selective_forward.py modifies test_decoder_runner_warmup_bypasses_only_auxiliary_phases. It still tests both enabled values and checks which warmup phases use eager or compiled forwarding. The supplied objectives report that both parametrizations previously failed with AttributeError, and the other six tests in the file passed in the same CI run. The PR was not run locally. No integration test-list match was reported. Coverage verdict: sufficient for this targeted regression.

Per-File QA Perspective

• tests/unittest/_torch/compilation/test_phase_selective_forward.py: The warmup test checks forwarding behavior across auxiliary and main warmup phases for both enabled values. No integration test-list entry was reported; this is a unit-test file.

[chenfeiz0326]

/bot run

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: eb6a2ffd-6b21-4cf6-ace0-f2d7d12f6c0a

📥 Commits

Reviewing files that changed from the base of the PR and between 15f060c and 08e5211.

📒 Files selected for processing: 1

• tests/unittest/_torch/compilation/test_phase_selective_forward.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The warmup test now patches moe_comm_serving_timeouts with contextlib.nullcontext instead of the removed capture-budget context manager. Its warmup assertions are unchanged.

Changes

Warmup test patch

Layer / File(s)	Summary
Update the warmup test patch target tests/unittest/_torch/compilation/test_phase_selective_forward.py	The test patches moe_comm_serving_timeouts instead of moe_a2a_steady_state_budget_for_capture. The warmup assertions remain unchanged.

Priority: ⬇️ Low

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

Change: Bug fix

Suggested reviewers: bowenfu

Merge Risk: ⚪ Minimal · up to 08e52

This test-only update is reported to pass in CI, with no remaining merge-blocking concern.

Pre-merge checks: 5

Passed checks: 5 passed

Check name	Status	Explanation
Description check		The description clearly explains the failure, the replacement patch, the unchanged test purpose, and the relevant test coverage. The required Description, Test Coverage, and PR Checklist sections are …
Title check		The title follows the required [None][type] format and clearly states that the warmup test now patches moe_comm_serving_timeouts.
Docstring Coverage		Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

[tensorrt-cicd]

PR_Github #77994 [ run ] triggered by Bot. Commit: 08e5211 Link to invocation

[chenfeiz0326]

CI result for this change: CPU-Generic-x86-1-cbts in L0 #64296 ran tests/unittest/_torch/compilation/test_phase_selective_forward.py, and all 8 cases passed. That includes test_decoder_runner_warmup_bypasses_only_auxiliary_phases[True] and [False], which fail with the AttributeError above on every branch without this change, for example #20250 (L0 #64229) and #20252 (L0 #64245).

The rest of the pipeline is still running.

[tensorrt-cicd]

PR_Github #77994 [ run ] completed with state SUCCESS. Commit: 08e5211
/LLM/main/L0_MergeRequest_PR pipeline #64296 completed with status: 'SUCCESS'

CI Report

Link to invocation

[chenfeiz0326]

@xxi-nv @AlessioNetti, CI is green on this one-line test fix (L0 #64296), and could one of you approve it? It needs a trt-llm-runtime-devs approval to merge.

Until it lands, every PR based on a main that includes #20246 fails test_phase_selective_forward.py on the CPU stage. #20250 and #20252 are both blocked on it.

[chenfeiz0326]

@QiJune, this needs a trt-llm-runtime-devs approval. It's a one-line test fix, with CI green in L0 #64296. It also unblocks #20250, which you're already reviewing.

[Funatiq]

Approving for @AlessioNetti

[github-actions]

Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.