Repository navigation
[https://nvbugs/6482566][fix] Make MoE all-to-all completion-flag timeout phase-aware - #17376
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughThe change adds configurable steady-state and warmup timeout budgets for MoE all-to-all operations. It adds atomic warmup-state control, CUDA graph integration, DeepSeek-V4 indexer-Q prewarming, and unit tests. ChangesMoE all-to-all timeout and warmup control
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant PyTorchModelEngine
participant moe_a2a_set_warmup
participant moeA2AGetTimeoutCycles
participant DispatchCombineKernels
PyTorchModelEngine->>moe_a2a_set_warmup: update warmup state
DispatchCombineKernels->>moeA2AGetTimeoutCycles: resolve timeout budget
moeA2AGetTimeoutCycles->>DispatchCombineKernels: return timeout_cycles
DispatchCombineKernels->>DispatchCombineKernels: poll completion flags
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h (1)
166-173: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse Doxygen for the public timeout resolver.
Line 173 declares a new public interface. Replace the free-form comment with Doxygen that documents
is_warmupand the returned cycle budget.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h` around lines 166 - 173, Replace the free-form comment above moeA2AGetTimeoutCycles with a Doxygen comment documenting the is_warmup parameter and the returned timeout budget in clock64() cycles, while preserving the existing behavioral context and environment-variable references.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 303-315: Import Iterator and update
_moe_a2a_steady_state_budget_for_capture with an Iterator[None] return
annotation, preserving its existing context-manager behavior and cleanup logic.
- Around line 1440-1453: The prewarm block around
indexer._project_and_quantize_q currently suppresses all exceptions; narrow it
to catch only the CuTe JIT-specific failure type used for compilation misses.
Let RuntimeError and other configuration or execution exceptions propagate,
while preserving the existing warning for the expected prewarm miss.
In `@tests/unittest/_torch/misc/test_moe_a2a_warmup_timeout.py`:
- Around line 37-45: Add a test case alongside test_missing_op_is_tolerated that
mocks model_engine.torch.ops.trtllm.moe_a2a_set_warmup to raise RuntimeError,
then calls model_engine._set_moe_a2a_warmup(True) and verifies it does not
raise, covering the incompatible dispatcher binding path.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h`:
- Around line 166-173: Replace the free-form comment above
moeA2AGetTimeoutCycles with a Doxygen comment documenting the is_warmup
parameter and the returned timeout budget in clock64() cycles, while preserving
the existing behavioral context and environment-variable references.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 407ab46b-7e6b-4d2d-803c-f97b65553e27
📒 Files selected for processing (5)
cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cucpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.hcpp/tensorrt_llm/thop/moeAlltoAllOp.cpptensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/misc/test_moe_a2a_warmup_timeout.py
|
/bot run --disable-fail-fast |
…eout phase-aware The MoE all-to-all dispatch/combine kernels wait for every peer's completion flag under a hardcoded ~300 s device-side deadline and execute `trap;` when it expires, poisoning the CUDA context of every waiting rank. No collective separates a rank's first-touch JIT/autotune work from its dispatch launch, so that device-side budget is in effect a deadline on the slowest peer's *host-side* progress through its first forward. During startup warmup a rank can legitimately spend minutes in one-time work (CuTe DSL compilation, autotuning, module loading). When it does, its peers trap, and a survivable slow start becomes a fatal crash for the whole worker. Changes: * The completion-flag budget becomes a runtime value carried in the kernel params instead of a compile-time constant, resolved by `moeA2AGetTimeoutCycles()`. Steady state is unchanged at 300 s; warmup uses 1800 s, matching the order of magnitude of torch.distributed's default collective timeout. Both are overridable via `TRTLLM_MOE_A2A_TIMEOUT_SEC` and `TRTLLM_MOE_A2A_WARMUP_TIMEOUT_SEC`. * A new `trtllm::moe_a2a_set_warmup` op lets the engine select the phase. It is driven from the `is_warmup` setter, the single choke point every warmup transition passes through, so the elevated budget cannot latch on into steady state. Missing-op builds degrade to the previous behaviour. * CUDA-graph capture is wrapped so graphs bake in the steady-state budget; the timeout is a launch argument and would otherwise be frozen at warmup's value for the life of the process. * The DSv4 indexer-Q CuTe DSL kernels are pre-compiled before the first collective-bearing forward, removing the largest single contributor to the skew. This is a partial mitigation: other first-touch compiles sit on the all-to-all path itself, which is why the runtime budget is the primary fix. Validated with fault injection on GB200 and GB300: delaying one rank by 600 s inside the first warmup forward reproduces the reported signature (N-1 peers timing out in dispatch, all naming the delayed rank) under the old 300 s budget, and starts cleanly under the new warmup budget. Signed-off-by: Wang, Xiao <24860335+xwang233@users.noreply.github.com>
0cd9e53 to
82db459
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
|
PR_Github #64423 [ run ] triggered by Bot. Commit: |
longlee0622
added
the
Release Blocker
|
PR_Github #64423 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64504 [ run ] triggered by Bot. Commit: |
|
PR_Github #64504 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64558 [ run ] triggered by Bot. Commit: |
|
PR_Github #64558 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64598 [ run ] triggered by Bot. Commit: |
pcastonguay
added
ci: post-merge approved
|
/bot run --disable-fail-fast |
|
Another rebase to pick up the pre-merge CI fix. |
|
PR_Github #64647 [ run ] triggered by Bot. Commit: |
|
PR_Github #64598 [ run ] completed with state |
|
PR_Github #64647 [ run ] completed with state
|
|
/bot run --disable-fail-fast --reuse-test |
|
PR_Github #64712 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #64744 [ run ] triggered by Bot. Commit: |
|
PR_Github #64712 [ run ] completed with state
|
|
PR_Github #64744 [ run ] completed with state |
…main _prewarm_cute_dsl_indexer_q (NVIDIA#17376) imports DeepseekV4Indexer from attention_backend.sparse.deepseek_v4.deepseek_v4, a module the attention consolidation (NVIDIA#17968) moved to attention.backends.sparse.deepseek_v4.indexer. The ImportError is caught, so the phase has returned without compiling anything since then, on every model. Removing it changes no behaviour. The kernels it targeted are still compiled during warmup: the indexer-Q runner (CuteDSLIndexerQBlackwellRunner) goes through the autotuner, which compiles its tactics in-process on every startup. JIT prefetch does not build them. The rest of NVIDIA#17376 (the phase-aware MoE all-to-all timeout) is unchanged. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
Description
The MoE all-to-all dispatch/combine kernels wait for every peer's completion flag under a hardcoded ~300 s device-side deadline and execute
trap;when it expires, which poisons the CUDA context of every waiting rank.No collective separates a rank's first-touch JIT/autotune work from its dispatch launch, so that device-side budget is in effect a deadline on the slowest peer's host-side progress through its first forward. During startup warmup a rank can legitimately spend minutes in one-time work (CuTe DSL compilation, autotuning, module loading). When one does, its peers trap — a survivable slow start becomes a fatal crash for the whole worker.
Reported symptom: N-1 ranks print
dispatch: ---Rank X timed out waiting for completion flag from rank Yand fail with
CUDA_ERROR_LAUNCH_FAILED, while the blamed rank shows no error at all.Fixes nvbugs/6482566.
Changes
moeA2AGetTimeoutCycles(). Steady state is unchanged at 300 s; warmup uses 1800 s, the same order of magnitude astorch.distributed's default collective timeout. Overridable viaTRTLLM_MOE_A2A_TIMEOUT_SEC/TRTLLM_MOE_A2A_WARMUP_TIMEOUT_SEC.trtllm::moe_a2a_set_warmupop, driven from theis_warmupsetter — the single choke point every warmup transition passes through — so the elevated budget cannot latch on into steady state. Builds without the op degrade to the previous behaviour.Test Coverage
tests/unittest/_torch/misc/test_moe_a2a_warmup_timeout.py(new, CPU-only):is_warmupsetter switches the budget in both directions, so it cannot latch on after warmup.End-to-end validation used fault injection on GB200 and GB300: delaying one rank by 600 s inside the first warmup forward reproduces the reported signature exactly (N-1 peers timing out in dispatch, all naming the delayed rank) under the 300 s budget, and starts cleanly under the new warmup budget.
Notes for reviewers
trap;with a recoverable error path is a separate, larger piece of work: after a timeout the receive buffers are incomplete and the flag generation would need invalidating across ranks. Not attempted here.PR Checklist
[JIRA/NVBugs ID][type] SummaryDev Engineer Review
TRTLLM_MOE_A2A_TIMEOUT_SECandTRTLLM_MOE_A2A_WARMUP_TIMEOUT_SEC.moe_a2a_set_warmup(bool)with compatibility handling for unavailable native bindings.QA Engineer Review
tests/unittest/_torch/misc/test_moe_a2a_warmup_timeout.py.test_set_warmup_forwards_value_to_optest_missing_op_is_toleratedtest_capture_context_selects_steady_state_then_restorestest_is_warmup_setter_switches_budget_both_waystests/integration/test_lists/,test-db/, orqa/.