[https://nvbugs/6517844][fix] Fall back to DeepEP when NCCL-EP lacks shared memory

tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py

@@ -38,6 +38,34 @@
 from .nvlink_two_sided import NVLinkTwoSided
 from .nvlink_two_sided_flashinfer import NVLinkTwoSidedFlashinfer
 
+# Temporary NCCL-EP v0.1.0 limitation. The LL combine kernel derives its
+# warp-group count and dynamic-SMEM requirement here:
+# https://github.com/NVIDIA/nccl/blob/nccl-ep-v0.1.0/contrib/nccl_ep/device/low_latency.cu#L1990-L2025
+# Its v0.1.0 group initialization limits LL execution to 14 warp groups:
+# https://github.com/NVIDIA/nccl/blob/nccl-ep-v0.1.0/contrib/nccl_ep/nccl_ep.cc#L1302-L1314
+# TODO: Remove this compatibility check after upgrading to NCCL-EP v0.2,
+# which removes the v0.1.0 LL-combine launch limitation.
+_NCCL_EP_V0_1_LL_MAX_WARP_GROUPS = 14
+
+
+def _get_nccl_ep_ll_combine_smem_requirement(
+    num_slots: int, hidden_size: int, num_device_sms: int
+) -> int | None:
+    """Return the NCCL-EP LL combine dynamic-SMEM requirement in bytes."""
+    num_warp_groups = (num_slots + num_device_sms - 1) // num_device_sms
+    if num_warp_groups > _NCCL_EP_V0_1_LL_MAX_WARP_GROUPS:
+        return None
+    num_warps_per_group = 32 // num_warp_groups
+
+    num_warps = num_warp_groups * num_warps_per_group
+    num_meta_bytes = hidden_size // 128 * 4
+    num_send_tma_bytes = 32 * 16 * 4 + 16
+    smem_send_size = num_warps * (3 * num_send_tma_bytes + num_meta_bytes)
+
+    num_recv_tma_bytes = 16 + hidden_size * 2
+    smem_recv_size = 2 * (3 * num_recv_tma_bytes + hidden_size * 2 + 3 * num_meta_bytes * 3)
+    return max(smem_send_size, smem_recv_size)
+
 
 class CommunicationFactory:
     """
@@ -413,4 +441,26 @@ def _get_nccl_ep_unavailable_reason(
             )
         if top_k <= 0 or top_k > num_slots:
             return f"NcclEP requires 0 < top_k <= num_slots, got {top_k=}, {num_slots=}."
+        if torch.cuda.is_available():
+            device_properties = torch.cuda.get_device_properties(torch.cuda.current_device())
+            required_smem = _get_nccl_ep_ll_combine_smem_requirement(
+                num_slots, hidden_size, device_properties.multi_processor_count
+            )
+            max_dynamic_smem = getattr(
+                device_properties,
+                "shared_memory_per_block_optin",
+                device_properties.shared_memory_per_block,
+            )
+            if required_smem is None:
+                return (
+                    "NcclEP low-latency combine requires at most "
+                    f"{_NCCL_EP_V0_1_LL_MAX_WARP_GROUPS} expert warp groups, got "
+                    f"{num_slots=} and {device_properties.multi_processor_count=}."
+                )
+            if required_smem > max_dynamic_smem:
+                return (
+                    "NcclEP low-latency combine requires "
+                    f"{required_smem} bytes of dynamic shared memory, but the current device "
+                    f"supports only {max_dynamic_smem} bytes."
+                )
         return None

tests/integration/test_lists/test-db/l0_a10.yml

@@ -47,6 +47,7 @@ l0_a10:
   - unittest/_torch/modules/dwdp/test_dwdp_manager.py
   - unittest/_torch/modules/dwdp/test_dwdp_mapping.py
   - unittest/_torch/modules/dwdp/test_dwdp_peer_ranges.py
+  - unittest/_torch/modules/moe/test_communication_factory.py
   # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no
   # test list either).
   - unittest/_torch/models/checkpoints

tests/unittest/_torch/modules/moe/test_communication_factory.py

@@ -91,6 +91,11 @@ def __init__(
         self.top_k = top_k
 
 
+class _FakeDeepEP:
+    def __init__(self, *args: object, **kwargs: object) -> None:
+        pass
+
+
 @pytest.mark.parametrize(
     ("act_dtype", "moe_max_num_tokens", "match"),
     [
@@ -119,10 +124,60 @@ def test_forced_nccl_ep_validates_preconditions(
         )
 
 
+def test_forced_nccl_ep_rejects_more_than_14_warp_groups(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model_config = _make_model_config()
+    assert (
+        communication_factory._get_nccl_ep_ll_combine_smem_requirement(
+            num_slots=15,
+            hidden_size=4096,
+            num_device_sms=1,
+        )
+        is None
+    )
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda _: SimpleNamespace(
+            multi_processor_count=1,
+            shared_memory_per_block_optin=102400,
+            shared_memory_per_block=102400,
+        ),
+    )
+
+    with pytest.raises(ValueError, match="at most 14 expert warp groups"):
+        communication_factory.CommunicationFactory._create_forced_method(
+            "NCCL_EP",
+            model_config,
+            num_experts=16,
+            num_slots=15,
+            top_k=8,
+            expert_size_per_partition=8,
+            payload_in_workspace=False,
+            alltoall_result_do_sum=True,
+            use_flashinfer=False,
+            hidden_size=4096,
+        )
+
+
 def test_forced_nccl_ep_allows_missing_moe_max_num_tokens(
     monkeypatch: pytest.MonkeyPatch,
 ):
     model_config = _make_model_config(torch.bfloat16, None)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda _: SimpleNamespace(
+            multi_processor_count=72,
+            shared_memory_per_block_optin=232448,
+            shared_memory_per_block=102400,
+        ),
+    )
     monkeypatch.setattr(communication_factory, "NcclEP", _FakeNcclEP)
 
     strategy = communication_factory.CommunicationFactory._create_forced_method(
@@ -151,6 +206,17 @@ def test_auto_selection_uses_nccl_ep_with_missing_moe_max_num_tokens(
     monkeypatch.setattr(communication_factory, "NVLinkOneSided", _strategy_unavailable)
     monkeypatch.setattr(communication_factory, "NVLinkTwoSided", _strategy_unavailable)
     monkeypatch.setenv("TRTLLM_CAN_USE_DEEP_EP", "0")
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda _: SimpleNamespace(
+            multi_processor_count=72,
+            shared_memory_per_block_optin=232448,
+            shared_memory_per_block=102400,
+        ),
+    )
     monkeypatch.setattr(communication_factory, "NcclEP", _FakeNcclEP)
 
     strategy = communication_factory.CommunicationFactory.create_strategy(
@@ -229,6 +295,94 @@ def test_auto_selection_skips_nccl_ep_for_quantized_moe(
     assert isinstance(strategy, AllGatherReduceScatter)
 
 
+def test_nccl_ep_ll_combine_smem_requirement() -> None:
+    assert (
+        communication_factory._get_nccl_ep_ll_combine_smem_requirement(
+            num_slots=72,
+            hidden_size=2560,
+            num_device_sms=72,
+        )
+        == 200704
+    )
+
+
+@pytest.mark.parametrize(
+    "device_properties",
+    [
+        SimpleNamespace(
+            multi_processor_count=72,
+            shared_memory_per_block_optin=200704,
+            shared_memory_per_block=102400,
+        ),
+        SimpleNamespace(
+            multi_processor_count=72,
+            shared_memory_per_block=200704,
+        ),
+    ],
+    ids=["optin_shared_memory", "legacy_shared_memory"],
+)
+def test_forced_nccl_ep_accepts_supported_ll_combine_dynamic_smem(
+    monkeypatch: pytest.MonkeyPatch,
+    device_properties: SimpleNamespace,
+) -> None:
+    model_config = _make_model_config()
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _: device_properties)
+    monkeypatch.setattr(communication_factory, "NcclEP", _FakeNcclEP)
+
+    strategy = communication_factory.CommunicationFactory._create_forced_method(
+        "NCCL_EP",
+        model_config,
+        num_experts=72,
+        num_slots=72,
+        top_k=6,
+        expert_size_per_partition=18,
+        payload_in_workspace=False,
+        alltoall_result_do_sum=True,
+        use_flashinfer=False,
+        hidden_size=2560,
+    )
+
+    assert isinstance(strategy, _FakeNcclEP)
+
+
+def test_auto_selection_skips_nccl_ep_when_ll_combine_exceeds_dynamic_smem(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model_config = _make_model_config()
+    monkeypatch.setattr(communication_factory, "NVLinkOneSided", _strategy_unavailable)
+    monkeypatch.setattr(communication_factory, "NVLinkTwoSided", _strategy_unavailable)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda _: SimpleNamespace(
+            multi_processor_count=72,
+            shared_memory_per_block_optin=102400,
+            shared_memory_per_block=102400,
+        ),
+    )
+    monkeypatch.setattr(
+        communication_factory,
+        "NcclEP",
+        lambda *args, **kwargs: pytest.fail("NcclEP should not be constructed"),
+    )
+    monkeypatch.setattr(communication_factory, "DeepEP", _FakeDeepEP)
+
+    strategy = communication_factory.CommunicationFactory.create_strategy(
+        model_config,
+        num_experts=72,
+        num_slots=72,
+        top_k=6,
+        expert_size_per_partition=18,
+        hidden_size=2560,
+    )
+
+    assert isinstance(strategy, _FakeDeepEP)
+
+
 def test_auto_selection_falls_back_when_nccl_probe_runtime_fails(
     monkeypatch: pytest.MonkeyPatch,
 ):