From e8d753d88097fc7a306b19a5887c5723aa922225 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 28 Jul 2026 09:37:08 -0500 Subject: [PATCH 1/5] [6517844][fix] fall back to DeepEP when NCCL-EP lacks shared memory Signed-off-by: Ludwig Schneider --- .../communication/communication_factory.py | 40 ++++++++ .../modules/moe/test_communication_factory.py | 91 +++++++++++++++++++ 2 files changed, 131 insertions(+) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py index a514291bdca5..6785b65479f0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -39,6 +39,25 @@ from .nvlink_two_sided_flashinfer import NVLinkTwoSidedFlashinfer +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 + num_warps_per_group = 32 // num_warp_groups + if num_warps_per_group == 0: + return None + + 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: """ Factory for creating MoE communication methods @@ -413,4 +432,25 @@ 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 32 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 diff --git a/tests/unittest/_torch/modules/moe/test_communication_factory.py b/tests/unittest/_torch/modules/moe/test_communication_factory.py index 841f0e9157cc..613c0d7017eb 100644 --- a/tests/unittest/_torch/modules/moe/test_communication_factory.py +++ b/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,6 +124,45 @@ def test_forced_nccl_ep_validates_preconditions( ) +def test_forced_nccl_ep_rejects_more_than_32_warp_groups( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_config = _make_model_config() + assert ( + communication_factory._get_nccl_ep_ll_combine_smem_requirement( + num_slots=33, + 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 32 expert warp groups"): + communication_factory.CommunicationFactory._create_forced_method( + "NCCL_EP", + model_config, + num_experts=34, + num_slots=33, + top_k=8, + expert_size_per_partition=17, + 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, ): @@ -229,6 +273,53 @@ 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 + ) + + +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, ): From 3ee6061ec25f6797010182e7924a42aa5c35ebbd Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Wed, 29 Jul 2026 08:15:30 -0500 Subject: [PATCH 2/5] [6517844][test] cover NCCL-EP shared-memory device variants Signed-off-by: Ludwig Schneider --- .../modules/moe/test_communication_factory.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/unittest/_torch/modules/moe/test_communication_factory.py b/tests/unittest/_torch/modules/moe/test_communication_factory.py index 613c0d7017eb..c461a9aa8445 100644 --- a/tests/unittest/_torch/modules/moe/test_communication_factory.py +++ b/tests/unittest/_torch/modules/moe/test_communication_factory.py @@ -284,6 +284,47 @@ def test_nccl_ep_ll_combine_smem_requirement() -> None: ) +@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: From 1b63396d5aca4ef2f90f8b8c61333472ffb1441c Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Wed, 29 Jul 2026 08:26:12 -0500 Subject: [PATCH 3/5] [6517844][test] register NCCL-EP factory coverage in CI Signed-off-by: Ludwig Schneider --- tests/integration/test_lists/test-db/l0_a10.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 6c4b63f424d5..b1982a32cd67 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/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 From b3c194e280cd7346d1780c01f47a966a46223d87 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Thu, 30 Jul 2026 12:33:25 -0500 Subject: [PATCH 4/5] [6517844][test] mock NCCL-EP device requirements Signed-off-by: Ludwig Schneider --- .../modules/moe/test_communication_factory.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/unittest/_torch/modules/moe/test_communication_factory.py b/tests/unittest/_torch/modules/moe/test_communication_factory.py index c461a9aa8445..0f8e98316781 100644 --- a/tests/unittest/_torch/modules/moe/test_communication_factory.py +++ b/tests/unittest/_torch/modules/moe/test_communication_factory.py @@ -167,6 +167,17 @@ 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( @@ -195,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( From e3040fcd62efffc019a7b69f9ae5aaff892bba6a Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Thu, 30 Jul 2026 13:42:45 -0500 Subject: [PATCH 5/5] [6517844][fix] document NCCL-EP LL compatibility limit Signed-off-by: Ludwig Schneider --- .../communication/communication_factory.py | 16 +++++++++++++--- .../modules/moe/test_communication_factory.py | 12 ++++++------ 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py index 6785b65479f0..92b0425384b4 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -38,15 +38,24 @@ 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 - num_warps_per_group = 32 // num_warp_groups - if num_warps_per_group == 0: + 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 @@ -444,7 +453,8 @@ def _get_nccl_ep_unavailable_reason( ) if required_smem is None: return ( - "NcclEP low-latency combine requires at most 32 expert warp groups, got " + "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: diff --git a/tests/unittest/_torch/modules/moe/test_communication_factory.py b/tests/unittest/_torch/modules/moe/test_communication_factory.py index 0f8e98316781..3919ea233af4 100644 --- a/tests/unittest/_torch/modules/moe/test_communication_factory.py +++ b/tests/unittest/_torch/modules/moe/test_communication_factory.py @@ -124,13 +124,13 @@ def test_forced_nccl_ep_validates_preconditions( ) -def test_forced_nccl_ep_rejects_more_than_32_warp_groups( +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=33, + num_slots=15, hidden_size=4096, num_device_sms=1, ) @@ -148,14 +148,14 @@ def test_forced_nccl_ep_rejects_more_than_32_warp_groups( ), ) - with pytest.raises(ValueError, match="at most 32 expert warp groups"): + with pytest.raises(ValueError, match="at most 14 expert warp groups"): communication_factory.CommunicationFactory._create_forced_method( "NCCL_EP", model_config, - num_experts=34, - num_slots=33, + num_experts=16, + num_slots=15, top_k=8, - expert_size_per_partition=17, + expert_size_per_partition=8, payload_in_workspace=False, alltoall_result_do_sum=True, use_flashinfer=False,