[None][fix] Merge back #19593 from release/1.3: Make MoE communication timeouts follow the warmup phase

[xxi-nv]

Dev Engineer Review

• MoECommTimeoutGuard centralizes warmup and serving timeout budgets across NVLinkOneSided, DeepGEMM MegaMoE, DeepEP, and NCCL-EP. It validates values, resolves deprecated aliases, and rejects conflicting values.
• The CUDA-graph capture pass uses serving timeouts. The warmup-only pass remains on warmup timeouts. Verify that each capture path applies serving timeouts and restores the prior state.
• Backend timeout APIs use different units and reset behavior. The proxies convert units as needed and use each backend’s native default when the serving budget is unset. Verify timeout defaults and proxy cleanup across backend lifecycles.
• CMake now applies dependency patches in order. Verify patch application with the supported dependency revisions.
• The DeepGEMM patch adds runtime barrier timeout control and restores DG_JIT_BARRIER_TIMEOUT_SECONDS as the JIT-time default. DeepEP and NCCL-EP patches add runtime timeout APIs. Verify compatibility with the pinned dependencies and CUDA graph replay behavior.
• The change removes runner-level MoE A2A timeout helpers and adds a timeout getter operator. Verify callers use the process-wide guard and that external users do not depend on the removed APIs.
• No current review findings or test execution results were supplied. Review severity counts and pass/fail status are unavailable.

QA Engineer Review

The change adds tests for budget resolution, aliases, validation, guard transitions, backend proxies, engine warmup transitions, and CUDA-graph capture behavior. It also adds communication-factory and DeepGEMM import compatibility tests. The B200 test list includes the timeout-guard test and the DeepGEMM import test. GPU-dependent proxy tests may skip when runtime requirements are unavailable. Test execution results were not supplied. Coverage verdict: needs follow-up, pending CI results for GPU-dependent behavior and dependency patch builds.

Per-File QA Perspective

• 3rdparty/CMakeLists.txt: Verify ordered patch application and the failure paths for inconsistent patch state.
• 3rdparty/fetch_content.json: Verify the DeepEP, NCCL-EP, and DeepGEMM patch assignments against the pinned dependency sources.
• 3rdparty/patches/deep_ep_runtime_timeout.patch: Verify setter bounds, default restoration, and timeout behavior during kernel launches and graph replay.
• 3rdparty/patches/deepgemm_configurable_barrier_timeout.patch: Verify JIT defaults, runtime overrides, and timeout reporting on supported kernels.
• 3rdparty/patches/nccl_ep_group_timeout.patch: Verify group timeout conversion, reset behavior, and invalid-value handling.
• cpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp: Verify that zero restores the native default and the getter reports the current timeout.
• docs/source/developer-guide/overview.md: Verify the documented warmup and serving defaults match runtime behavior.
• tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md: Verify the documented aliases, transitions, and proxy ownership match implementation.
• tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py: Verify invalid timeout values fail before backend construction and unsupported NCCL-EP versions are rejected clearly.
• tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py: Verify the proxy applies and resets NVLinkOneSided timeouts. The previous module-level timeout API was removed.
• tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py: Verify proxy registration and confirm which DeepEP kernel modes use the process-wide timeout.
• tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py: Verify DeepGEMM proxy registration and runtime timeout updates.
• tensorrt_llm/_torch/moe/fused_moe/moe_comm_timeout_guard.py: Verify budget parsing, validation, proxy lifecycle, capture-time transition rejection, and serving-timeout restoration.
• tensorrt_llm/_torch/moe/fused_moe/nccl_ep_utils.py: Verify NCCL-EP capability detection, timeout conversion, registration, and cleanup.
• tensorrt_llm/_torch/moe/fused_moe/quantization.py: Verify DeepGEMM builds without the new setter fail with the intended compatibility error.
• tensorrt_llm/_torch/pyexecutor/engine/runners/common.py: Verify no remaining call paths require the removed timeout helpers.
• tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py: Verify only serving graph capture uses serving timeouts and warmup remains on the warmup budget.
• tensorrt_llm/_torch/pyexecutor/model_engine.py: Verify each engine warmup-state transition updates the communication guard.
• tests/integration/test_lists/test-db/l0_b200.yml: Adds the timeout-guard and DeepGEMM import-rejection tests to the B200 test list.
• tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py: Covers engine warmup transitions and graph warmup/capture timeout selection. The supplied test-list change does not show this test in the B200 list.
• tests/unittest/_torch/moe/multi_gpu/test_nvlink_one_sided.py: Verify MPI workers use the canonical timeout variables and retain the intended values.
• tests/unittest/_torch/moe/test_communication_factory.py: Covers early validation and NCCL-EP timeout capability detection. The supplied test-list change does not show this test in the B200 list.
• tests/unittest/_torch/moe/test_moe_backend.py: Covers rejection of DeepGEMM without the runtime setter. The B200 list includes the associated import-rejection test.
• tests/unittest/_torch/moe/test_moe_comm_timeout_guard.py: Covers budget resolution, guard transitions, and backend proxies. GPU-dependent tests can skip when runtime requirements are unavailable. The supplied test-list change does not show this test in the B200 list.

Description

This PR is a merge-back of #19593 from release/1.3 (merged there as f5eaf58) to main. It brings the same feature to main. The differences listed below exist because main changed the same areas after release/1.3 branched.

What #19593 does

Several MoE communication kernels wait for their peers inside the kernel and trap once a peer is late by more than a device-side timeout. The trap is a sticky CUDA error, so every waiting rank dies. During warmup, a healthy rank can arrive minutes late while it pays for one-time work such as kernel JIT, autotuning or module loading. On GB300, single ranks reached their first DeepGEMM MegaMoE launch 55-73 s late, and the other ranks trapped with DeepGEMM grid sync timeout (60s).

A process-wide MoECommTimeoutGuard (tensorrt_llm/_torch/moe/fused_moe/moe_comm_timeout_guard.py) relaxes the timeouts of every MoE communication backend while the engine warms up, and restores them for serving:

• TRTLLM_MOE_COMM_WARMUP_TIMEOUT_SEC sets the warmup budget for every backend (default 1800 s).
• TRTLLM_MOE_COMM_TIMEOUT_SEC sets the serving budget. When it is unset, each backend keeps its native timeout.
• The PyTorchModelEngine.is_warmup setter switches the guard. The CUDA-graph pass that captures runs with the serving timeouts, because a captured graph keeps the timeouts that were current at capture and serving replays it.
• Each backend registers a proxy when it is constructed:

Backend	Timeout set through	Native serving timeout
NVLinkOneSided	torch.ops.trtllm.moe_a2a_set_timeout	300 s
DeepGEMM MegaMoE	deep_gemm.set_barrier_timeout_seconds (3rdparty/patches/deepgemm_configurable_barrier_timeout.patch)	60 s, or DG_JIT_BARRIER_TIMEOUT_SECONDS
DeepEP normal kernels	deep_ep.set_timeout_seconds (3rdparty/patches/deep_ep_runtime_timeout.patch)	~100 s
NcclEP	nccl.ep.Group.set_timeout_ns (3rdparty/patches/nccl_ep_group_timeout.patch)	NCCL_EP_TIMEOUT_MS, else GroupConfig.timeout_ns, else ~100 s

The design discussion and review are in #19593.

Differences from #19593

NVLinkOneSided timeouts (#19610). On main, #19610 had already given NVLinkOneSided its own warmup budget, read from TRTLLM_NVLINK_ONE_SIDED_A2A_[WARMUP_]TIMEOUT_SEC and applied through NVLinkOneSided.set_timeout. The guard replaces it:

• _NVLinkOneSidedTimeoutProxy replaces NVLinkOneSided.set_timeout and get_timeout_seconds. The helpers in engine/runners/common.py that drove them are removed.
• moe_a2a_set_timeout now accepts 0, which restores the native 300 s budget. A new moe_a2a_get_timeout op reads the budget back for tests.
• Two older pairs of names become deprecated aliases of TRTLLM_MOE_COMM_[WARMUP_]TIMEOUT_SEC: TRTLLM_NVLINK_ONE_SIDED_A2A_[WARMUP_]TIMEOUT_SEC from [None][refactor] overhaul NVLink one-sided all-to-all #19610, and TRTLLM_MOE_A2A_[WARMUP_]TIMEOUT_SEC, which main used before [None][refactor] overhaul NVLink one-sided all-to-all #19610 and release/1.3 still accepts. Each alias logs a deprecation warning and now applies to every backend. Conflicting values raise.
• The NVLinkOneSided default budgets do not change: 1800 s during warmup, 300 s while serving.
• Behavior change: an invalid value now raises a ValueError that names the variable. CommunicationFactory.create_strategy raises it before trying any backend. [None][refactor] overhaul NVLink one-sided all-to-all #19610 logged a warning and used the default instead. The accepted range is still 1..86400 s.

CUDA-graph capture (#19896). Capture moved from the model engine to the decoder runner, so the capture-pass switch now lives in a new DecoderRunner._warmup_and_capture_cuda_graphs. #19610 deleted the engine-level test that #19593 rewrote (tests/unittest/_torch/misc/test_moe_a2a_warmup_timeout.py). It is re-added as tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py.

DeepGEMM patch (#19727). #19727 bumped DeepGEMM from 8b1392b9 to b64107f2 and dropped deepgemm_configurable_barrier_timeout.patch from #18179. As a result, the in-kernel barrier timeout on main is a fixed 60 s, and DG_JIT_BARRIER_TIMEOUT_SECONDS, which docs/source/developer-guide/overview.md documents, has no effect. The patch is re-authored for b64107f2:

• It restores DG_JIT_BARRIER_TIMEOUT_SECONDS as the JIT-time default.
• It adds the runtime set_barrier_timeout_seconds, which is passed as a kernel argument to the SM100 BF16, FP8×FP4 and FP4×FP4 MegaMoE kernels and to the SM90 FP8 MegaMoE kernel. Changing it does not recompile.
• overview.md gains a paragraph on how the TRT-LLM runtime overrides the default.

Several patches per dependency. 3rdparty/CMakeLists.txt takes #19593's support for a list of patches in patch_file, applied and verified in order. It builds on #19957's single-argument patch command. DeepEP now carries two patches.

Test Coverage

• New tests/unittest/_torch/moe/test_moe_comm_timeout_guard.py:
  • TestBudgetResolution and TestGuard are cpu_only and run in l0_cpu.yml through unittest/_torch/moe. They cover budget resolution, both deprecated alias families (including conflicts and invalid values), validation, warmup transitions, weak proxy registration, and the capture check.
  • TestBackendProxies is added to l0_b200.yml. Each proxy round-trips through its native setter.
• New tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py is cpu_only and runs in l0_cpu.yml through unittest/_torch/executor. It checks that the engine switches the timeouts both ways, and that only the capturing CUDA-graph pass uses the serving timeouts.
• tests/unittest/_torch/moe/test_communication_factory.py: an invalid timeout fails before any strategy is constructed, and NcclEP is rejected with a reason when Group.set_timeout_ns is missing.
• tests/unittest/_torch/moe/test_moe_backend.py::test_import_deep_gemm_rejects_mega_moe_without_barrier_timeout is new and added to l0_b200.yml.
• tests/unittest/_torch/moe/multi_gpu/test_nvlink_one_sided.py now sets the new variable names.

Checked locally without a GPU:

• Each patch applies through the CMake patch step to its pinned upstream tarball, and re-running the step is idempotent.
• The patched DeepGEMM MegaMoE kernels compile with the CUDA 13.4 toolchain of the CI image, and the compiled kernel parameter layout matches the host-side launch arguments.

The GPU tests run in CI.

Notes for reviewers

• As noted in [https://nvbugs/6713426][fix] Make MoE communication timeouts follow the warmup phase #19593, the nccl-extensions change should also go upstream to NVIDIA/nccl-extensions. It is carried as a patch until the pin moves.
• If a local build tree reports Patch state for <dep> is inconsistent, remove _deps/<dep>-src and reconfigure. CI builds from scratch.

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: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 92d39bd7-8d91-4d38-b345-494d142332e5

📥 Commits

Reviewing files that changed from the base of the PR and between a5e1d84 and d9ece37.

📒 Files selected for processing: 3

• 3rdparty/CMakeLists.txt
• tensorrt_llm/_torch/moe/fused_moe/quantization.py
• tests/integration/test_lists/test-db/l0_b200.yml

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

This change adds runtime timeout controls for MoE communication backends. A shared guard resolves warmup and serving budgets, applies them through backend timeout proxies, and switches timeout state during engine warmup and CUDA-graph capture.

Changes

MoE communication timeout control

Layer / File(s)	Summary
Vendor timeout APIs and patch setup 3rdparty/CMakeLists.txt, 3rdparty/fetch_content.json, 3rdparty/patches/*	Patch configuration accepts ordered patch paths. DeepEP, DeepGEMM, and NCCL EP receive timeout APIs and use configured values in their communication waits.
Shared timeout budgets and guard tensorrt_llm/_torch/moe/fused_moe/moe_comm_timeout_guard.py, tests/unittest/_torch/moe/test_moe_comm_timeout_guard.py, tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md, docs/source/developer-guide/overview.md	The guard resolves and validates warmup and serving budgets, tracks backend proxies, and changes timeouts with guard state. Tests and documentation cover the guard and budget behavior.
Backend timeout adapters and capability checks tensorrt_llm/_torch/moe/fused_moe/communication/*, tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py, tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py, tensorrt_llm/_torch/moe/fused_moe/nccl_ep_utils.py, tensorrt_llm/_torch/moe/fused_moe/quantization.py, cpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp, tests/unittest/_torch/moe/*	Backend proxies forward timeout updates to native APIs. The communication factory validates timeout configuration and checks NCCL EP group-timeout support. Tests cover proxy updates, capability checks, and the all-to-all timeout accessor.
Engine warmup and graph-capture transitions tensorrt_llm/_torch/pyexecutor/engine/runners/common.py, tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py, tensorrt_llm/_torch/pyexecutor/model_engine.py, tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py, tests/integration/test_lists/test-db/l0_b200.yml	The model engine updates guard state when warmup changes. The graph runner uses warmup timeouts for its warmup pass and serving timeouts for capture. Tests cover the state transitions and capture sequence.

Priority: ➖ Normal

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

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant PyTorchModelEngine
  participant MoECommTimeoutGuard
  participant TimeoutProxies
  participant CudaGraphRunner
  PyTorchModelEngine->>MoECommTimeoutGuard: set_moe_comm_warmup updates timeout state
  MoECommTimeoutGuard->>TimeoutProxies: apply warmup timeout
  CudaGraphRunner->>MoECommTimeoutGuard: enter serving_timeouts during capture
  MoECommTimeoutGuard->>TimeoutProxies: apply serving timeout
  CudaGraphRunner->>MoECommTimeoutGuard: exit serving_timeouts
  MoECommTimeoutGuard->>TimeoutProxies: restore prior timeout state

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#19593: Implements the same MoE communication timeout change; this pull request is described as a merge-back of that change into main.

Suggested reviewers: juney-nvidia

Merge Risk: ⚪ Minimal · up to d9ece

No actionable issue is established for these changes; they are mergeable after normal checks.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 29.73% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 111 functions across 15 files. (2 skipped…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the merge-back and the main change: MoE communication timeouts now follow the warmup phase.
Description check		The description explains the problem, solution, backend coverage, dependency changes, test coverage, local validation, reviewer notes, and checklist status. It is complete and aligned with the reposit…
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.

Full details: Docstring Coverage

Explanation

Docstring coverage is 29.73% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 111 functions across 15 files. (2 skipped: 2 unsupported.)

• Fix all pre-merge checks with AI

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

[xxi-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #77625 [ run ] triggered by Bot. Commit: a5e1d84 Link to invocation

[QiJune]

LGTM

[Barry-Delaney]

Reviewed a5e1d84 for code correctness and shared-path regressions. Traced warmup/serving transitions, the separate graph warmup and capture passes, proxy lifetime and native-default restoration, timeout units, and DeepGEMM host/kernel argument propagation. All 39 focused CPU tests passed against the actual guard and extracted engine/runner methods with torch/logger dependencies stubbed. The DeepEP, NCCL-EP and DeepGEMM patches apply to their exact pinned sources and pass reverse-applicability checks; existing DeepEP CMake enables separate device compilation. All 15 changed Python sources parse, removed-helper callers are gone, and diff --check passes using the PR attributes. No blocking code findings. Native CUDA/GPU/MPI execution was not run locally.

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

[xxi-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #77674 [ run ] triggered by Bot. Commit: d9ece37 Link to invocation

[tensorrt-cicd]

PR_Github #77625 [ run ] completed with state ABORTED. Commit: a5e1d84

Link to invocation

[Barry-Delaney]

Re-reviewed d9ece37 after the rebase. Range-diff against the previously approved a5e1d84 shows the MoE timeout implementation and dependency patches are unchanged; the substantive adaptation is the CMake integration with main's patch_dependency.py. Configured the actual 3rdparty CMake declarations in an isolated project and verified fresh application, repeated application and reverse checks for DeepEP (including the concatenated patch), NCCL-EP and DeepGEMM against their pinned sources. No new blocking code findings. The earlier 39 isolated CPU checks remain applicable to unchanged runtime code; native GPU/MPI execution was not run locally.

[tensorrt-cicd]

PR_Github #77674 [ run ] completed with state SUCCESS. Commit: d9ece37
/LLM/main/L0_MergeRequest_PR pipeline #64034 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

[xxi-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #77777 [ run ] triggered by Bot. Commit: d9ece37 Link to invocation

[tensorrt-cicd]

PR_Github #77777 [ run ] completed with state SUCCESS. Commit: d9ece37
/LLM/main/L0_MergeRequest_PR pipeline #64115 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

[xxi-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #77792 [ run ] triggered by Bot. Commit: d9ece37 Link to invocation

[tensorrt-cicd]

PR_Github #77792 [ run ] completed with state SUCCESS. Commit: d9ece37
/LLM/main/L0_MergeRequest_PR pipeline #64130 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

[xxi-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #77829 [ run ] triggered by Bot. Commit: d9ece37 Link to invocation

[tensorrt-cicd]

PR_Github #77829 [ run ] completed with state SUCCESS. Commit: d9ece37
/LLM/main/L0_MergeRequest_PR pipeline #64160 completed with status: 'SUCCESS'

CI Report

Link to invocation