[https://nvbugs/6517844][fix] Fall back to DeepEP when NCCL-EP lacks shared memory - #16956
Conversation
nv-lschneider
changed the title
|
/bot run --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughNCCL EP eligibility now accounts for LL-combine dynamic shared-memory requirements and warp-group limits. Tests cover requirement calculation, forced-method validation, supported configurations, automatic DeepEP fallback, and A10 test-suite registration. ChangesNCCL EP shared-memory eligibility
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant CommunicationFactory
participant CUDADeviceProperties
participant NcclEP
participant DeepEP
CommunicationFactory->>CUDADeviceProperties: Query SM count and dynamic shared-memory capacity
CommunicationFactory->>CommunicationFactory: Calculate LL-combine SMEM requirement
alt Requirement fits device capacity
CommunicationFactory->>NcclEP: Select NCCL EP
else Requirement exceeds capacity
CommunicationFactory->>DeepEP: Select DeepEP
end
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.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py (1)
42-44: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
int | Nonehere.Optional[int]can be replaced with PEP 604 union syntax in this helper.🤖 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 `@tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py` around lines 42 - 44, Update the return annotation of _get_nccl_ep_ll_combine_smem_requirement to use the PEP 604 int | None syntax instead of Optional[int], leaving the helper’s behavior unchanged.Sources: Coding guidelines, Learnings
tests/unittest/_torch/modules/moe/test_communication_factory.py (1)
237-283: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a warp-group boundary test
test_forced_nccl_ep_validates_preconditionsalready covers the forcedNCCL_EPpath. Add a case for_get_nccl_ep_ll_combine_smem_requirement(...)=Nonewhennum_warp_groups > 32.Test coverage summary:
test_nccl_ep_ll_combine_smem_requirementandtest_auto_selection_skips_nccl_ep_when_ll_combine_exceeds_dynamic_smemcover the helper value and fallback path;test_forced_nccl_ep_validates_preconditionscovers the forced path. Coverage verdict: needs follow-up.🤖 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 `@tests/unittest/_torch/modules/moe/test_communication_factory.py` around lines 237 - 283, Add a unit test covering the warp-group boundary where _get_nccl_ep_ll_combine_smem_requirement returns None when num_warp_groups exceeds 32. Extend the existing NCCL_EP-related test coverage, particularly near test_forced_nccl_ep_validates_preconditions, and assert the expected forced-path validation behavior without changing the existing helper or fallback tests.Source: Path instructions
🤖 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.
Nitpick comments:
In
`@tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py`:
- Around line 42-44: Update the return annotation of
_get_nccl_ep_ll_combine_smem_requirement to use the PEP 604 int | None syntax
instead of Optional[int], leaving the helper’s behavior unchanged.
In `@tests/unittest/_torch/modules/moe/test_communication_factory.py`:
- Around line 237-283: Add a unit test covering the warp-group boundary where
_get_nccl_ep_ll_combine_smem_requirement returns None when num_warp_groups
exceeds 32. Extend the existing NCCL_EP-related test coverage, particularly near
test_forced_nccl_ep_validates_preconditions, and assert the expected forced-path
validation behavior without changing the existing helper or fallback tests.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 909818cc-f0b5-42b1-b3ca-fa3eb67a7ac5
📒 Files selected for processing (2)
tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.pytests/unittest/_torch/modules/moe/test_communication_factory.py
|
PR_Github #62220 [ run ] triggered by Bot. Commit: |
38996b2 to
aa181e2
Compare
|
/bot run --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" |
|
PR_Github #62230 [ run ] triggered by Bot. Commit: |
|
PR_Github #62220 [ run ] completed with state |
|
PR_Github #62230 [ run ] completed with state
|
|
/bot run --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" --disable-fail-fast |
|
PR_Github #62265 [ run ] triggered by Bot. Commit: |
|
/bot run --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" --disable-fail-fast |
|
PR_Github #62275 [ run ] triggered by Bot. Commit: |
|
PR_Github #62265 [ run ] completed with state |
|
PR_Github #62275 [ run ] completed with state
|
|
/bot run --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" --disable-fail-fast |
|
PR_Github #62323 [ run ] triggered by Bot. Commit: |
|
PR_Github #62323 [ run ] completed with state
|
|
PR_Github #62907 [ run ] triggered by Bot. Commit: |
|
PR_Github #62834 [ run ] completed with state |
|
PR_Github #62907 [ run ] completed with state
|
|
/bot run --disable-fail-fast --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" |
1 similar comment
|
/bot run --disable-fail-fast --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" |
|
PR_Github #63031 [ run ] triggered by Bot. Commit: |
|
PR_Github #63035 [ run ] triggered by Bot. Commit: |
|
PR_Github #63031 [ run ] completed with state |
|
PR_Github #63035 [ run ] completed with state
|
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
0ad30c3 to
e3040fc
Compare
|
/bot run --disable-fail-fast --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" |
|
PR_Github #63107 [ run ] triggered by Bot. Commit: |
|
PR_Github #63107 [ run ] completed with state
|
|
/bot run --disable-fail-fast --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" |
|
PR_Github #63158 [ run ] triggered by Bot. Commit: |
|
PR_Github #63158 [ run ] completed with state
|
|
/bot run --disable-fail-fast --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" |
|
PR_Github #63455 [ run ] triggered by Bot. Commit: |
|
PR_Github #63455 [ run ] completed with state
|
|
/bot run --disable-fail-fast --extra-stage "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" |
|
PR_Github #63474 [ run ] triggered by Bot. Commit: |
|
PR_Github #63474 [ run ] completed with state |
6 participants
Dev Engineer Review
ll_combinedynamic SMEM (bytes) from MoE slot count, hidden size, and GPU SM (multi-processor) count, returningNonewhen the derived warp-group configuration is invalid._get_nccl_ep_unavailable_reason(when CUDA is available) to:required_smemand compare it to the device’s maximum supported dynamic SMEM viashared_memory_per_block_optin(falling back toshared_memory_per_block).NCCL_EPselection validates the configuration, rejecting invalid/oversubscribed cases (including the “at most 32 expert warp groups” constraint) with aValueError.QA Engineer Review
Tests touched:
tests/unittest/_torch/modules/moe/test_communication_factory.py_get_nccl_ep_ll_combine_smem_requirementreturning the expected SMEM value for a known input._FakeDeepEPstub to detect DeepEP selection without instantiating the real DeepEP implementation.Test-list coverage mapping:
tests/integration/test_lists/test-db/l0_a10.yml: addedunittest/_torch/modules/moe/test_communication_factory.pyto the pre-merge A10 (ubuntu) PyTorch test list.Verdict: needs follow-up.
Description
Fixes NVBug 6517844 on RTX Pro 6000D.
NCCL-EP low-latency combine can require more opt-in dynamic shared memory than the device supports. Previously, this was detected only during NCCL-EP kernel setup, aborting the worker.
Preflight the NCCL-EP low-latency combine shared-memory requirement during communication-strategy selection. When unsupported, automatic selection skips NCCL-EP and falls back to DeepEP. A forced
NCCL_EPselection reports a clear validation error.The affected EP4 DeepSeek-V3-Lite test remains active in the RTX Pro 6000D L0 test list and has no waiver.
Test Coverage
PR Checklist