[https://nvbugs/6020038][feat] Add NCCL-EP v0.1 MoE communication support

requirements.txt

@@ -29,6 +29,8 @@ nvidia-modelopt[torch]~=0.37.0
 # https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-05.html#rel-26-05 uses 2.30.4
 # torch 2.11.0+cu130 depends on nvidia-nccl-cu13==2.28.9
 nvidia-nccl-cu13>=2.28.9,<=2.30.4
+# NcclEP uses nccl4py's nccl.ep package without changing the NCCL wheel constraint.
+nccl4py>=0.3.1,<0.4
 nvidia-cuda-nvrtc
 transformers==5.5.4
 prometheus_client

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

@@ -34,6 +34,7 @@
 from .communication_factory import CommunicationFactory
 from .deep_ep import DeepEP
 from .deep_ep_low_latency import DeepEPLowLatency
+from .nccl_ep import NcclEP
 from .nvlink_one_sided import NVLinkOneSided
 from .nvlink_two_sided import NVLinkTwoSided
 
@@ -46,6 +47,7 @@
     "NVLinkOneSided",
     "DeepEP",
     "DeepEPLowLatency",
+    "NcclEP",
     # Factory
     "CommunicationFactory",
 ]

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

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -33,6 +33,7 @@
 from .base import Communication
 from .deep_ep import DeepEP
 from .deep_ep_low_latency import DeepEPLowLatency
+from .nccl_ep import NcclEP
 from .nvlink_one_sided import NVLinkOneSided
 from .nvlink_two_sided import NVLinkTwoSided
 from .nvlink_two_sided_flashinfer import NVLinkTwoSidedFlashinfer
@@ -68,6 +69,7 @@ def create_strategy(
         2. Auto-selection (tries in order):
            - NVLinkOneSided (highest priority for throughput)
            - NVLinkTwoSided (high priority for latency)
+           - NcclEP (if nccl-ep is available)
            - DeepEP (if enabled via TRTLLM_CAN_USE_DEEP_EP)
            - DeepEPLowLatency (if enabled via TRTLLM_CAN_USE_DEEP_EP)
            - AllGather + ReduceScatter (fallback, always works)
@@ -130,7 +132,7 @@ def create_strategy(
             )
 
         # Auto-selection: Try strategies in priority order using try-catch
-        # Priority: NVLinkOneSided > NVLinkTwoSided > DeepEP > DeepEPLowLatency > AllGather
+        # Priority: NVLinkOneSided > NVLinkTwoSided > NcclEP > DeepEP > DeepEPLowLatency > AllGather
 
         try:
             enable_eplb = model_config.moe_load_balancer is not None
@@ -188,6 +190,34 @@ def create_strategy(
         except Exception as e:
             logger.info(f"NVLinkTwoSided not available: {e}")
 
+        # Try NCCL EP (rank-major LL). Falls through to DeepEP/AllGather if
+        # prerequisites are not met or libnccl_ep.so is not available.
+        nccl_ep_unavailable_reason = CommunicationFactory._get_nccl_ep_unavailable_reason(
+            act_dtype,
+            quant_config,
+            num_slots,
+            hidden_size,
+            max_num_tokens,
+            moe_max_num_tokens,
+            top_k,
+        )
+        if nccl_ep_unavailable_reason is None:
+            try:
+                strategy = NcclEP(
+                    mapping,
+                    num_slots,
+                    hidden_size,
+                    max_num_tokens,
+                    moe_max_num_tokens,
+                    top_k=top_k,
+                )
+                logger.info("Selected communication strategy: NcclEP")
+                return strategy
+            except RuntimeError as e:
+                logger.debug(f"NcclEP not available: {e}")
+        else:
+            logger.debug(f"NcclEP not available: {nccl_ep_unavailable_reason}")
+
         # Try DeepEP (if enabled and weight dtype is bfloat16)
         if os.environ.get("TRTLLM_CAN_USE_DEEP_EP", "1") == "1" and act_dtype == torch.bfloat16:
             try:
@@ -331,7 +361,56 @@ def _create_forced_method(
                 use_low_precision_combine,
                 moe_max_num_tokens,
             )
+        elif method == "NCCL_EP":
+            nccl_ep_unavailable_reason = CommunicationFactory._get_nccl_ep_unavailable_reason(
+                act_dtype,
+                quant_config,
+                num_slots,
+                hidden_size,
+                max_num_tokens,
+                moe_max_num_tokens,
+                top_k,
+            )
+            if nccl_ep_unavailable_reason is not None:
+                raise ValueError(nccl_ep_unavailable_reason)
+            return NcclEP(
+                mapping,
+                num_slots,
+                hidden_size,
+                max_num_tokens,
+                moe_max_num_tokens,
+                top_k=top_k,
+            )
         elif method == "ALLGATHER":
             return AllGatherReduceScatter(mapping)
         else:
             raise ValueError(f"Unknown communication method: {method}")
+
+    @staticmethod
+    def _get_nccl_ep_unavailable_reason(
+        act_dtype: torch.dtype,
+        quant_config,
+        num_slots: int,
+        hidden_size: int,
+        max_num_tokens: int,
+        moe_max_num_tokens: Optional[int],
+        top_k: int,
+    ) -> Optional[str]:
+        if act_dtype != torch.bfloat16:
+            return f"NcclEP requires act_dtype=torch.bfloat16, got {act_dtype}."
+        if quant_config is not None:
+            quant_mode = getattr(quant_config, "layer_quant_mode", None)
+            if quant_mode is not None and quant_mode.has_any_quant(exclude_kv_cache=True):
+                return "NcclEP v0.1 does not support quantized MoE communication."
+        if num_slots <= 0 or hidden_size <= 0 or max_num_tokens <= 0:
+            return (
+                "NcclEP requires positive num_slots, hidden_size, and max_num_tokens, got "
+                f"{num_slots=}, {hidden_size=}, {max_num_tokens=}."
+            )
+        if moe_max_num_tokens is not None and moe_max_num_tokens <= 0:
+            return (
+                f"NcclEP requires moe_max_num_tokens > 0 when provided, got {moe_max_num_tokens}."
+            )
+        if top_k <= 0 or top_k > num_slots:
+            return f"NcclEP requires 0 < top_k <= num_slots, got {top_k=}, {num_slots=}."
+        return None