diff --git a/docs/advanced_features/expert_parallelism.md b/docs/advanced_features/expert_parallelism.md
index 73bab333c..5c052114b 100644
--- a/docs/advanced_features/expert_parallelism.md
+++ b/docs/advanced_features/expert_parallelism.md
@@ -15,13 +15,14 @@ SGLang's EP integrates diverse, highly efficient backends for different use case
| **`none` (default)** | Disables all-to-all for EP. Uses All-Reduce or All-Gather for token dispatch. | Hybrid EP and TP setups. |
| `deepep` | DeepEP, a communication library for efficient token shuffling in MoE models. | Large-scale EP deployments. |
| `mooncake` | An extension of DeepEP for elastic inference, leveraging RDMA for high-performance data transfers. | Elastic EP serving. |
+| `nixl` | [NIXL-EP](https://github.com/ai-dynamo/nixl/tree/main/examples/device/ep), an elastic EP communication library built on NVIDIA's [NIXL](https://github.com/ai-dynamo/nixl) framework with native RDMA and NVLink support. | Elastic EP serving with fault tolerance and dynamic scaling. |
| `mori` | MORI-EP, AMD's native all-to-all communication implementation optimized for ROCm. | AMD GPU deployments. |
| `flashinfer` | Flashinfer implementation of all-to-all. | Large-scale EP deployments. |
| `ascend_fuseep` | Ascend NPU native fused all-to-all communication. | Ascend NPU deployments. |
-DeepEP and Mooncake backends support two modes for token dispatch: `normal` mode (optimized for prefill workloads with high throughput) and `low_latency` mode (optimized for decode workloads with low latency and CUDA Graph compatibility). MORI backend only supports `normal` mode now. Users are recommended to set `--deepep-mode auto` to enable automatic dispatch mode switching during runtime. Setting `--deepep-mode normal` or `--deepep-mode low_latency` is useful for debugging or development purposes.
+DeepEP and Mooncake backends support two modes for token dispatch: `normal` mode (optimized for prefill workloads with high throughput) and `low_latency` mode (optimized for decode workloads with low latency and CUDA Graph compatibility). MORI backend only supports `normal` mode now. NIXL-EP currently operates in low-latency mode with CUDA Graph support. Users are recommended to set `--deepep-mode auto` to enable automatic dispatch mode switching during runtime. Setting `--deepep-mode normal` or `--deepep-mode low_latency` is useful for debugging or development purposes.
-Currently, DeepEP, Mooncake, `ascend_fuseep` and MORI only support cases where `ep_size = tp_size`. For hybrid EP and TP (i.e., `ep_size < tp_size`), only the `none` backend (All-Reduce or All-Gather-based dispatching) is supported.
+Currently, DeepEP, Mooncake, NIXL-EP, `ascend_fuseep` and MORI only support cases where `ep_size = tp_size`. For hybrid EP and TP (i.e., `ep_size < tp_size`), only the `none` backend (All-Reduce or All-Gather-based dispatching) is supported.
### Backends for MoE Computation
diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md
index 61fc9059b..1317f04f5 100644
--- a/docs/advanced_features/server_arguments.md
+++ b/docs/advanced_features/server_arguments.md
@@ -311,7 +311,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
| Argument | Description | Defaults | Options |
| --- | --- | --- | --- |
| `--expert-parallel-size`
`--ep-size`
`--ep` | The expert parallelism size. | `1` | Type: int |
-| `--moe-a2a-backend` | Select the backend for all-to-all communication for expert parallelism. | `none` | `none`, `deepep`, `mooncake`, `mori`, `ascend_fuseep`|
+| `--moe-a2a-backend` | Select the backend for all-to-all communication for expert parallelism. | `none` | `none`, `deepep`, `mooncake`, `mori`, `nixl`, `ascend_fuseep`|
| `--moe-runner-backend` | Choose the runner backend for MoE. | `auto` | `auto`, `deep_gemm`, `triton`, `triton_kernel`, `flashinfer_trtllm`, `flashinfer_trtllm_routed`, `flashinfer_cutlass`, `flashinfer_mxfp4`, `flashinfer_cutedsl`, `cutlass` |
| `--flashinfer-mxfp4-moe-precision` | Choose the computation precision of flashinfer mxfp4 moe | `default` | `default`, `bf16` |
| `--enable-flashinfer-allreduce-fusion` | Enable FlashInfer allreduce fusion with Residual RMSNorm. | `False` | bool flag (set to enable) |
diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py
index c0d1a4923..b66c39551 100644
--- a/python/sglang/srt/batch_overlap/two_batch_overlap.py
+++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py
@@ -31,6 +31,7 @@ from sglang.srt.layers.moe.token_dispatcher import (
DeepEPDispatcher,
MooncakeEPDispatcher,
MoriEPDispatcher,
+ NixlEPDispatcher,
)
from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher
from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -1036,6 +1037,10 @@ class MaybeTboDeepEPDispatcher(BaseDispatcher):
self._inners = [
MoriEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
]
+ elif get_moe_a2a_backend().is_nixl():
+ self._inners = [
+ NixlEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
+ ]
def _execute(self, name, tbo_subbatch_index: Optional[int] = None, **kwargs):
return getattr(self._inners[tbo_subbatch_index or 0], name)(**kwargs)
diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py
index 9df3668d0..a2b6a84cc 100644
--- a/python/sglang/srt/distributed/parallel_state.py
+++ b/python/sglang/srt/distributed/parallel_state.py
@@ -42,6 +42,7 @@ from torch.distributed import Backend, ProcessGroup
from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
+from sglang.srt.distributed.utils import set_global_tcp_store
from sglang.srt.environ import envs
from sglang.srt.utils import (
get_bool_env_var,
@@ -1611,6 +1612,61 @@ def get_default_distributed_backend(device: str) -> str:
return _DEVICE_TO_DISTRIBUTED_BACKEND.get(device, "gloo")
+def _create_global_tcp_store(rank: int, world_size: int) -> None:
+ """Create a global TCPStore for coordination across ranks.
+
+ This function creates a TCPStore that all ranks can use for coordination
+ (e.g., for NIXL buffer setup).
+ """
+ from torch.distributed import TCPStore
+
+ master_ip = os.environ.get("MASTER_ADDR")
+
+ if not master_ip:
+ logger.warning(
+ "Could not determine master IP for global TCPStore. "
+ "Broadcasting from rank 0 to all ranks."
+ )
+
+ base_store_port = envs.SGLANG_TCP_STORE_PORT.get()
+
+ # Rank 0 gets its local IP and broadcasts it to all ranks
+ # Use broadcast_object_list which works with any backend (handles CPU/GPU automatically)
+ if not master_ip:
+ if rank == 0:
+ master_ip = get_local_ip_auto()
+ ip_list = [master_ip]
+ else:
+ ip_list = [None]
+
+ torch.distributed.broadcast_object_list(ip_list, src=0)
+ master_ip = ip_list[0]
+
+ try:
+ tcp_store = TCPStore(
+ host_name=master_ip,
+ port=base_store_port,
+ world_size=world_size,
+ is_master=(rank == 0),
+ )
+ set_global_tcp_store(tcp_store)
+ logger.info(
+ "Created global TCPStore at %s:%d (rank=%d, world_size=%d)",
+ master_ip,
+ base_store_port,
+ rank,
+ world_size,
+ )
+ except Exception as e:
+ logger.warning(
+ "Failed to create global TCPStore at %s:%d: %s. "
+ "Components requiring TCPStore (like NIXL) may not work.",
+ master_ip,
+ base_store_port,
+ e,
+ )
+
+
def init_distributed_environment(
world_size: int = -1,
rank: int = -1,
@@ -1618,6 +1674,7 @@ def init_distributed_environment(
local_rank: int = -1,
backend: str = "nccl",
timeout: Optional[int] = None,
+ moe_a2a_backend: Optional[str] = None,
):
logger.debug(
"world_size=%d rank=%d local_rank=%d " "distributed_init_method=%s backend=%s",
@@ -1660,6 +1717,10 @@ def init_distributed_environment(
pg_options=pg_options,
)
+ # Create a global TCPStore for coordination (used by NIXL)
+ if moe_a2a_backend == "nixl":
+ _create_global_tcp_store(rank, world_size)
+
# set the local rank
# local_rank is not available in torch ProcessGroup,
# see https://github.com/pytorch/pytorch/issues/122816
diff --git a/python/sglang/srt/distributed/utils.py b/python/sglang/srt/distributed/utils.py
index 2fde4fc92..c0090c086 100644
--- a/python/sglang/srt/distributed/utils.py
+++ b/python/sglang/srt/distributed/utils.py
@@ -17,6 +17,42 @@ from torch.distributed import TCPStore
logger = logging.getLogger(__name__)
+# Global TCPStore that is created during distributed initialization
+# This is the single shared store that all components should use
+_global_tcp_store: Optional[TCPStore] = None
+
+
+def set_global_tcp_store(store: TCPStore) -> None:
+ """Set the global TCPStore instance.
+
+ This should be called during distributed initialization to make
+ the store available to all components that need it.
+ """
+ global _global_tcp_store
+ _global_tcp_store = store
+ logger.info("Global TCPStore has been set")
+
+
+def get_global_tcp_store() -> Optional[TCPStore]:
+ """Get the existing global TCPStore.
+
+ This function provides access to the shared TCPStore instance that was
+ created during distributed initialization. All components (like NIXL buffers)
+ should use this same store for coordination.
+
+ Returns:
+ The global TCPStore instance, or None if not initialized yet.
+ """
+ global _global_tcp_store
+
+ if _global_tcp_store is None:
+ logger.warning(
+ "Global TCPStore not found. Make sure init_distributed_environment "
+ "was called with a tcp:// init method."
+ )
+
+ return _global_tcp_store
+
def ensure_divisibility(numerator, denominator):
"""Ensure that numerator is divisible by the denominator."""
diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index fd828bf77..48a9ea122 100644
--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -271,6 +271,7 @@ class Envs:
# Override the distributed init method used by torch.distributed.init_process_group.
# Set to "env://" to use an externally-created TCPStore via MASTER_ADDR/MASTER_PORT.
SGLANG_DISTRIBUTED_INIT_METHOD_OVERRIDE = EnvStr(None)
+ SGLANG_TCP_STORE_PORT = EnvInt(29600)
# Tool Calling
SGLANG_FORWARD_UNKNOWN_TOOLS = EnvBool(False)
@@ -378,6 +379,10 @@ class Envs:
SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32)
SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False)
+ # NIXL-EP
+ SGLANG_NIXL_EP_BF16_DISPATCH = EnvBool(False)
+ SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
+
# NSA Backend
SGLANG_NSA_FUSE_TOPK = EnvBool(True)
SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBool(True)
diff --git a/python/sglang/srt/eplb/eplb_algorithms/__init__.py b/python/sglang/srt/eplb/eplb_algorithms/__init__.py
index b09a14175..d5e9c4460 100644
--- a/python/sglang/srt/eplb/eplb_algorithms/__init__.py
+++ b/python/sglang/srt/eplb/eplb_algorithms/__init__.py
@@ -13,6 +13,7 @@ class EplbAlgorithm(Enum):
deepseek_vec = auto()
deepseek_vec_hierarchical = auto()
elasticity_aware = auto()
+ elasticity_aware_hierarchical = auto()
# TODO may have more algorithm later
@@ -47,14 +48,19 @@ def rebalance_experts(
enable_hierarchical=algorithm == EplbAlgorithm.deepseek_vec_hierarchical,
)
- if algorithm == EplbAlgorithm.elasticity_aware:
+ if algorithm in [
+ EplbAlgorithm.elasticity_aware,
+ EplbAlgorithm.elasticity_aware_hierarchical,
+ ]:
return elasticity_aware.rebalance_experts(
weight=tokens_per_expert.sum(dim=0),
num_replicas=num_physical_experts,
num_groups=num_groups,
num_nodes=num_nodes,
num_gpus=num_physical_experts // num_local_physical_experts,
- enable_hierarchical=False,
+ enable_hierarchical=(
+ algorithm == EplbAlgorithm.elasticity_aware_hierarchical
+ ),
active_ranks=(
ElasticEPStateManager.instance().active_ranks
if ElasticEPStateManager.instance() is not None
diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py
index 43e8bcce0..d507dad12 100644
--- a/python/sglang/srt/layers/moe/ep_moe/layer.py
+++ b/python/sglang/srt/layers/moe/ep_moe/layer.py
@@ -747,7 +747,11 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]):
# [TODO] kk, temporary solution
if get_moe_a2a_backend().is_mori():
return MoriEPMoE
- if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
+ if (
+ get_moe_a2a_backend().is_deepep()
+ or get_moe_a2a_backend().is_mooncake()
+ or get_moe_a2a_backend().is_nixl()
+ ):
return DeepEPMoE
if get_moe_a2a_backend().is_ascend_fuseep():
return NpuFuseEPMoE
diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py
index 144b05092..a54e9ea7d 100644
--- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py
@@ -95,7 +95,12 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
a2a_backend = get_moe_a2a_backend()
if a2a_backend.is_none():
return StandardDispatcher(moe_runner_config)
- elif a2a_backend.is_deepep() or a2a_backend.is_mooncake() or a2a_backend.is_mori():
+ elif (
+ a2a_backend.is_deepep()
+ or a2a_backend.is_mooncake()
+ or a2a_backend.is_mori()
+ or a2a_backend.is_nixl()
+ ):
return MaybeTboDeepEPDispatcher(
group=(
get_tp_group().device_group
diff --git a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py
index dd40a8d98..cb6909660 100644
--- a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py
+++ b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py
@@ -33,6 +33,11 @@ from sglang.srt.layers.moe.token_dispatcher.moriep import (
MoriEPNormalCombineInput,
MoriEPNormalDispatchOutput,
)
+from sglang.srt.layers.moe.token_dispatcher.nixl import (
+ NixlEPCombineInput,
+ NixlEPDispatcher,
+ NixlEPDispatchOutput,
+)
from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardCombineInput,
StandardDispatcher,
@@ -58,6 +63,9 @@ __all__ = [
"MoriEPLLDispatchOutput",
"MoriEPLLCombineInput",
"MoriEPDispatcher",
+ "NixlEPCombineInput",
+ "NixlEPDispatchOutput",
+ "NixlEPDispatcher",
"StandardDispatcher",
"StandardDispatchOutput",
"StandardCombineInput",
diff --git a/python/sglang/srt/layers/moe/token_dispatcher/nixl.py b/python/sglang/srt/layers/moe/token_dispatcher/nixl.py
new file mode 100644
index 000000000..e1977f362
--- /dev/null
+++ b/python/sglang/srt/layers/moe/token_dispatcher/nixl.py
@@ -0,0 +1,465 @@
+from __future__ import annotations
+
+import logging
+from enum import Enum, auto
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+
+from sglang.srt.distributed.utils import get_global_tcp_store
+from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
+from sglang.srt.environ import envs
+from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
+from sglang.srt.layers import deep_gemm_wrapper
+from sglang.srt.layers.dp_attention import get_is_extend_in_batch
+from sglang.srt.layers.moe.token_dispatcher.base import (
+ BaseDispatcher,
+ CombineInput,
+ DispatchOutput,
+)
+from sglang.srt.layers.moe.token_dispatcher.deepep import (
+ DeepEPLLCombineInput,
+ DeepEPLLDispatchOutput,
+)
+from sglang.srt.layers.moe.topk import TopKOutput
+from sglang.srt.layers.moe.utils import DeepEPMode
+
+try:
+ from nixl_ep import Buffer
+
+ use_nixl = True
+except ImportError:
+ use_nixl = False
+
+logger = logging.getLogger(__name__)
+
+NixlEPDispatchOutput = DeepEPLLDispatchOutput
+NixlEPCombineInput = DeepEPLLCombineInput
+
+
+class NixlEPBuffer:
+ _buffer = None
+ _hidden_size: Optional[int] = None
+ _num_max_dispatch_tokens_per_rank: Optional[int] = None
+ _num_experts: Optional[int] = None
+ _num_local_experts: Optional[int] = None
+
+ @classmethod
+ def get_nixl_buffer(
+ cls,
+ group: dist.ProcessGroup,
+ hidden_size: int,
+ deepep_mode: DeepEPMode,
+ num_max_dispatch_tokens_per_rank: int = -1,
+ num_experts: int = -1,
+ num_local_experts: int = -1,
+ ):
+ if cls._buffer is not None:
+ return cls._buffer
+
+ cls._hidden_size = hidden_size
+ cls._num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank
+ cls._num_experts = num_experts
+ cls._num_local_experts = num_local_experts
+
+ num_rdma_bytes = 0
+ if deepep_mode.enable_normal():
+ raise NotImplementedError("Normal mode is not supported for Nixl EP yet.")
+ if deepep_mode.enable_low_latency():
+ assert num_max_dispatch_tokens_per_rank != -1
+ assert num_experts != -1 and num_experts % group.size() == 0
+ num_rdma_bytes = Buffer.get_rdma_size_hint(
+ num_max_dispatch_tokens_per_rank,
+ hidden_size,
+ group.size(),
+ num_experts,
+ )
+
+ rank = dist.get_rank(group)
+ world_size = dist.get_world_size(group)
+
+ # Get the global TCPStore for coordination
+ tcp_store = get_global_tcp_store()
+ if tcp_store is None:
+ raise RuntimeError(
+ "Global TCPStore is not initialized. "
+ "Make sure init_distributed_environment was called before using NIXL EP."
+ )
+
+ logger.info(
+ f"Using NIXL EP (world_size={world_size}, rank={rank}, "
+ f"num_experts={cls._num_experts}, num_experts_per_rank={cls._num_local_experts}) "
+ )
+
+ cls._buffer = Buffer(
+ rank=rank,
+ tcp_store_group=tcp_store,
+ )
+
+ cls._buffer.update_memory_buffers(
+ num_ranks=world_size,
+ num_experts_per_rank=cls._num_local_experts,
+ num_rdma_bytes=num_rdma_bytes,
+ )
+ all_ranks = list(range(world_size))
+ cls._buffer.connect_ranks(all_ranks)
+
+ return cls._buffer
+
+ @classmethod
+ def clean_buffer(cls):
+ cls._buffer.clean_buffer(
+ cls._num_max_dispatch_tokens_per_rank,
+ cls._hidden_size,
+ cls._num_experts,
+ )
+
+
+class _NixlEPDispatcherImplBase:
+ def __init__(
+ self,
+ group: torch.distributed.ProcessGroup,
+ router_topk: int,
+ permute_fusion: bool,
+ num_experts: int,
+ num_local_experts: int,
+ hidden_size: int,
+ params_dtype: torch.dtype,
+ deepep_mode: DeepEPMode,
+ ):
+ if not use_nixl:
+ raise ImportError(
+ "NixlEP is not installed. Please install NixlEP package from "
+ "https://github.com/ai-dynamo/nixl."
+ )
+
+ self.group = group
+ self.router_topk = router_topk
+ self.permute_fusion = permute_fusion
+ self.num_experts = num_experts
+ self.num_local_experts = num_local_experts
+ self.hidden_size = hidden_size
+ self.params_dtype = params_dtype
+ self.deepep_mode = deepep_mode
+
+ self.num_max_dispatch_tokens_per_rank = (
+ envs.SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get()
+ )
+ # NixlEP internode_ll dispatch uses FINISHED_SUM_TAG=1024
+ # and the logic requires num-tokens-sent-from-one-rank-to-another-rank less than it
+ assert self.num_max_dispatch_tokens_per_rank <= 1024
+ elastic_state = ElasticEPStateManager.instance()
+ self.active_ranks = (
+ elastic_state.active_ranks if elastic_state is not None else None
+ )
+ self._mask_buffer = (
+ torch.zeros_like(self.active_ranks)
+ if self.active_ranks is not None
+ else None
+ )
+
+ self.handle = None
+ self.quant_config = None
+ self.overlap_args = None
+ self.meta_overlap_args = None
+
+ def set_quant_config(self, quant_config: dict) -> None:
+ self.quant_config = quant_config
+
+ def set_overlap_args(self, combine_overlap_args, meta_overlap_args) -> None:
+ self.overlap_args = combine_overlap_args
+ self.meta_overlap_args = meta_overlap_args
+
+ def dispatch_a(
+ self,
+ hidden_states: torch.Tensor,
+ topk_output: TopKOutput,
+ ):
+ raise NotImplementedError
+
+ def dispatch_b(self, *args, **kwargs):
+ raise NotImplementedError
+
+ def combine_a(
+ self,
+ hidden_states: torch.Tensor,
+ topk_ids: torch.Tensor,
+ topk_weights: torch.Tensor,
+ ):
+ raise NotImplementedError
+
+ def combine_b(self, *args, **kwargs):
+ raise NotImplementedError
+
+ def _get_buffer(self):
+ raise NotImplementedError
+
+
+class _NixlEPDispatcherImpl(_NixlEPDispatcherImplBase):
+ def __init__(self, return_recv_hook: bool, **kwargs):
+ super().__init__(**kwargs)
+
+ """
+ num_max_dispatch_tokens_per_rank: the actual batch size in the decoding engine should be less than 256
+ https://github.com/ai-dynamo/nixl
+ """
+ self.return_recv_hook = return_recv_hook
+ self.device_module = torch.get_device_module()
+
+ def dispatch_a(
+ self,
+ hidden_states: torch.Tensor,
+ topk_output: TopKOutput,
+ ):
+ buffer = self._get_buffer()
+ topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
+ topk_ids = topk_ids.to(torch.int64)
+ expected_m = (
+ hidden_states.shape[0] * buffer.group_size * topk_ids.shape[1]
+ + self.num_experts
+ ) // self.num_experts
+ hidden_states, masked_m, event, hook = self._dispatch_core(
+ hidden_states,
+ topk_ids,
+ )
+ return (
+ hidden_states,
+ topk_ids,
+ topk_weights,
+ masked_m,
+ expected_m,
+ event,
+ hook,
+ )
+
+ def dispatch_b(
+ self,
+ hidden_states,
+ topk_ids,
+ topk_weights,
+ masked_m,
+ expected_m,
+ event,
+ hook,
+ ):
+ hook() if self.return_recv_hook else event.current_stream_wait()
+
+ get_global_expert_distribution_recorder().on_deepep_dispatch_low_latency(
+ masked_m
+ )
+
+ if isinstance(hidden_states, tuple):
+ hidden_states, hidden_states_scale = hidden_states
+ else:
+ hidden_states_scale = None
+
+ nixl_output = NixlEPDispatchOutput(
+ hidden_states,
+ hidden_states_scale,
+ topk_ids,
+ topk_weights,
+ masked_m,
+ expected_m,
+ )
+ return nixl_output
+
+ def _dispatch_core(
+ self,
+ hidden_states: torch.Tensor,
+ topk_idx: torch.Tensor,
+ ):
+ use_fp8 = not envs.SGLANG_NIXL_EP_BF16_DISPATCH.get()
+
+ buffer = self._get_buffer()
+ packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = (
+ buffer.dispatch(
+ hidden_states,
+ topk_idx,
+ self.num_max_dispatch_tokens_per_rank,
+ self.num_experts,
+ use_fp8=use_fp8,
+ async_finish=not self.return_recv_hook,
+ return_recv_hook=self.return_recv_hook,
+ round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+ and deep_gemm_wrapper.DEEPGEMM_BLACKWELL,
+ use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+ and deep_gemm_wrapper.DEEPGEMM_BLACKWELL,
+ )
+ )
+ return packed_recv_hidden, self.packed_recv_count, event, hook
+
+ def combine_a(
+ self,
+ hidden_states: torch.Tensor,
+ topk_ids: torch.Tensor,
+ topk_weights: torch.Tensor,
+ ):
+ hidden_states, event, hook = self._combine_core(
+ hidden_states,
+ topk_ids,
+ topk_weights,
+ )
+ return hidden_states, event, hook
+
+ def combine_b(self, hidden_states, event, hook):
+ hook() if self.return_recv_hook else event.current_stream_wait()
+ return hidden_states
+
+ def _combine_core(
+ self,
+ hidden_states: torch.Tensor,
+ topk_ids: torch.Tensor,
+ topk_weights: torch.Tensor,
+ ):
+ buffer = self._get_buffer()
+
+ combined_hidden_states, event, hook = buffer.combine(
+ x=hidden_states,
+ topk_idx=topk_ids,
+ topk_weights=topk_weights,
+ handle=self.handle,
+ async_finish=not self.return_recv_hook,
+ return_recv_hook=self.return_recv_hook,
+ )
+ if self._mask_buffer is not None:
+ buffer.query_mask_buffer(self._mask_buffer)
+ self.active_ranks.copy_(1 - self._mask_buffer)
+
+ self.packed_recv_count = self.handle = None
+ return combined_hidden_states, event, hook
+
+ def _get_buffer(self):
+ return NixlEPBuffer.get_nixl_buffer(
+ self.group,
+ self.hidden_size,
+ self.deepep_mode,
+ self.num_max_dispatch_tokens_per_rank,
+ self.num_experts,
+ self.num_local_experts,
+ )
+
+
+class _Stage(Enum):
+ INITIAL = auto()
+ AFTER_DISPATCH_A = auto()
+ AFTER_DISPATCH_B = auto()
+ AFTER_COMBINE_A = auto()
+
+
+class NixlEPDispatcher(BaseDispatcher):
+ def __init__(
+ self,
+ group: torch.distributed.ProcessGroup,
+ router_topk: int,
+ permute_fusion: bool = False,
+ num_experts: int = None,
+ num_local_experts: int = None,
+ hidden_size: int = None,
+ params_dtype: torch.dtype = None,
+ deepep_mode: DeepEPMode = DeepEPMode.LOW_LATENCY,
+ async_finish: bool = False,
+ return_recv_hook: bool = False,
+ ):
+ self.deepep_mode = deepep_mode
+
+ common_kwargs = dict(
+ group=group,
+ router_topk=router_topk,
+ permute_fusion=permute_fusion,
+ num_experts=num_experts,
+ num_local_experts=num_local_experts,
+ hidden_size=hidden_size,
+ params_dtype=params_dtype,
+ deepep_mode=deepep_mode,
+ )
+
+ if self.deepep_mode.enable_low_latency():
+ self._low_latency_dispatcher = _NixlEPDispatcherImpl(
+ return_recv_hook=return_recv_hook,
+ **common_kwargs,
+ )
+ if self.deepep_mode.enable_normal():
+ raise NotImplementedError("Normal mode is not supported for Nixl EP yet.")
+
+ self._stage = _Stage.INITIAL
+
+ def dispatch(
+ self,
+ hidden_states: torch.Tensor,
+ topk_output: TopKOutput,
+ ) -> DispatchOutput:
+ self.dispatch_a(hidden_states=hidden_states, topk_output=topk_output)
+ ret = self.dispatch_b()
+ return ret
+
+ def dispatch_a(
+ self,
+ hidden_states: torch.Tensor,
+ topk_output: TopKOutput,
+ ):
+ self._update_stage(_Stage.INITIAL, _Stage.AFTER_DISPATCH_A)
+ inner_state = self._get_impl().dispatch_a(
+ hidden_states=hidden_states,
+ topk_output=topk_output,
+ )
+ self._dispatch_intermediate_state = inner_state
+
+ def dispatch_b(self):
+ self._update_stage(_Stage.AFTER_DISPATCH_A, _Stage.AFTER_DISPATCH_B)
+ inner_state = self._dispatch_intermediate_state
+ del self._dispatch_intermediate_state
+ return self._get_impl().dispatch_b(*inner_state)
+
+ def combine(
+ self,
+ combine_input: CombineInput,
+ ) -> torch.Tensor:
+ self.combine_a(combine_input)
+ ret = self.combine_b()
+ return ret
+
+ def combine_a(
+ self,
+ combine_input: CombineInput,
+ ):
+ hidden_states, topk_ids, topk_weights = combine_input
+ self._update_stage(_Stage.AFTER_DISPATCH_B, _Stage.AFTER_COMBINE_A)
+ inner_state = self._get_impl().combine_a(
+ hidden_states=hidden_states,
+ topk_ids=topk_ids,
+ topk_weights=topk_weights,
+ )
+ self._combine_intermediate_state = inner_state
+
+ def combine_b(self):
+ self._update_stage(_Stage.AFTER_COMBINE_A, _Stage.INITIAL)
+ inner_state = self._combine_intermediate_state
+ del self._combine_intermediate_state
+ return self._get_impl().combine_b(*inner_state)
+
+ def _get_impl(self) -> _NixlEPDispatcherImplBase:
+ is_extend_in_batch = get_is_extend_in_batch()
+ resolved_deepep_mode = self.deepep_mode.resolve(is_extend_in_batch)
+ if resolved_deepep_mode == DeepEPMode.NORMAL:
+ raise NotImplementedError("Normal mode is not supported for Nixl EP yet.")
+ elif resolved_deepep_mode == DeepEPMode.LOW_LATENCY:
+ return self._low_latency_dispatcher
+ else:
+ raise ValueError(f"Invalid deepep_mode: {self.deepep_mode}")
+
+ def set_quant_config(self, quant_config: dict):
+ super().set_quant_config(quant_config)
+ if self.deepep_mode.enable_low_latency():
+ self._low_latency_dispatcher.set_quant_config(quant_config)
+
+ def set_overlap_args(self, combine_overlap_args, meta_overlap_args):
+ super().set_overlap_args(combine_overlap_args, meta_overlap_args)
+ if self.deepep_mode.enable_low_latency():
+ self._low_latency_dispatcher.set_overlap_args(
+ combine_overlap_args, meta_overlap_args
+ )
+
+ def _update_stage(self, old_stage, new_stage):
+ assert self._stage == old_stage
+ self._stage = new_stage
diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py
index 3a2c7f3e5..65da83c0e 100644
--- a/python/sglang/srt/layers/moe/utils.py
+++ b/python/sglang/srt/layers/moe/utils.py
@@ -22,6 +22,7 @@ class MoeA2ABackend(Enum):
NONE = "none"
DEEPEP = "deepep"
MOONCAKE = "mooncake"
+ NIXL = "nixl"
MORI = "mori"
ASCEND_FUSEEP = "ascend_fuseep"
FLASHINFER = "flashinfer"
@@ -44,6 +45,9 @@ class MoeA2ABackend(Enum):
def is_mooncake(self):
return self == MoeA2ABackend.MOONCAKE
+ def is_nixl(self):
+ return self == MoeA2ABackend.NIXL
+
def is_flashinfer(self):
return self == MoeA2ABackend.FLASHINFER
diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py
index a1c355174..57236c072 100644
--- a/python/sglang/srt/layers/quantization/fp8.py
+++ b/python/sglang/srt/layers/quantization/fp8.py
@@ -748,7 +748,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
return True
if moe_runner_backend.is_auto():
return deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and (
- get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake()
+ get_moe_a2a_backend().is_deepep()
+ or get_moe_a2a_backend().is_mooncake()
+ or get_moe_a2a_backend().is_nixl()
)
return False
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index e42dbc556..451fc56c6 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -808,6 +808,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
local_rank=self.gpu_id,
distributed_init_method=dist_init_method,
timeout=self.server_args.dist_timeout,
+ moe_a2a_backend=self.server_args.moe_a2a_backend,
)
initialize_model_parallel(
tensor_model_parallel_size=self.tp_size,
diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py
index e9a4f511b..8f0617142 100644
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -445,6 +445,7 @@ class DeepseekV2MoE(nn.Module):
dict(tp_rank=0, tp_size=1)
if get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
+ or get_moe_a2a_backend().is_nixl()
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep()
or get_moe_a2a_backend().is_flashinfer()
@@ -489,6 +490,7 @@ class DeepseekV2MoE(nn.Module):
if (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
+ or get_moe_a2a_backend().is_nixl()
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep()
):
@@ -510,6 +512,7 @@ class DeepseekV2MoE(nn.Module):
self._enable_a2a_moe = (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
+ or get_moe_a2a_backend().is_nixl()
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep()
or get_moe_a2a_backend().is_flashinfer()
diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py
index db8c1c7ce..85f13132c 100644
--- a/python/sglang/srt/models/glm4_moe.py
+++ b/python/sglang/srt/models/glm4_moe.py
@@ -420,7 +420,11 @@ class Glm4MoeSparseMoeBlock(nn.Module):
),
)
- if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
+ if (
+ get_moe_a2a_backend().is_deepep()
+ or get_moe_a2a_backend().is_mooncake()
+ or get_moe_a2a_backend().is_nixl()
+ ):
# TODO: we will support tp < ep in the future
self.ep_size = get_moe_expert_parallel_world_size()
self.num_experts = (
@@ -437,7 +441,9 @@ class Glm4MoeSparseMoeBlock(nn.Module):
)
self._enable_a2a_moe = (
- get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake()
+ get_moe_a2a_backend().is_deepep()
+ or get_moe_a2a_backend().is_mooncake()
+ or get_moe_a2a_backend().is_nixl()
)
def get_moe_weights(self):
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index f8be1b895..2be1c045a 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -191,6 +191,7 @@ MOE_A2A_BACKEND_CHOICES = [
"none",
"deepep",
"mooncake",
+ "nixl",
"mori",
"ascend_fuseep",
"flashinfer",
@@ -508,7 +509,7 @@ class ServerArgs:
# Expert parallelism
ep_size: int = 1
moe_a2a_backend: Literal[
- "none", "deepep", "mooncake", "mori", "ascend_fuseep", "flashinfer"
+ "none", "deepep", "mooncake", "nixl", "mori", "ascend_fuseep", "flashinfer"
] = "none"
moe_runner_backend: str = "auto"
flashinfer_mxfp4_moe_precision: Literal["default", "bf16"] = "default"
@@ -530,7 +531,7 @@ class ServerArgs:
enable_expert_distribution_metrics: bool = False
deepep_config: Optional[str] = None
moe_dense_tp_size: Optional[int] = None
- elastic_ep_backend: Literal[None, "mooncake"] = None
+ elastic_ep_backend: Literal[None, "mooncake", "nixl"] = None
enable_elastic_expert_backup: bool = False
mooncake_ib_device: Optional[str] = None
@@ -2558,6 +2559,12 @@ class ServerArgs:
f"Mooncake MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]."
)
+ if self.moe_a2a_backend == "nixl":
+ self.ep_size = self.tp_size
+ logger.warning(
+ f"Nixl MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]."
+ )
+
if self.moe_a2a_backend == "ascend_fuseep":
self.ep_size = self.tp_size
logger.warning(
@@ -2620,9 +2627,10 @@ class ServerArgs:
if self.enable_eplb:
if self.eplb_algorithm == "auto":
self.eplb_algorithm = "elasticity_aware"
- assert (
- self.eplb_algorithm == "elasticity_aware"
- ), "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware'."
+ assert self.eplb_algorithm in [
+ "elasticity_aware",
+ "elasticity_aware_hierarchical",
+ ], "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'."
if self.elastic_ep_backend == "mooncake":
self.mooncake_ib_device = self._validate_ib_devices(
@@ -4650,8 +4658,8 @@ class ServerArgs:
"--elastic-ep-backend",
type=str,
default=ServerArgs.elastic_ep_backend,
- choices=["none", "mooncake"],
- help="Specify the collective communication backend for elastic EP. Currently supports 'mooncake'.",
+ choices=["none", "mooncake", "nixl"],
+ help="Specify the collective communication backend for elastic EP. Supports 'mooncake' and 'nixl'.",
)
parser.add_argument(
"--enable-elastic-expert-backup",
diff --git a/test/manual/ep/test_nixl_ep.py b/test/manual/ep/test_nixl_ep.py
new file mode 100644
index 000000000..3a4be3b8c
--- /dev/null
+++ b/test/manual/ep/test_nixl_ep.py
@@ -0,0 +1,115 @@
+import os
+import time
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
+from sglang.test.server_fixtures.disaggregation_fixture import get_rdma_devices_args
+from sglang.test.test_utils import (
+ DEFAULT_MODEL_NAME_FOR_TEST_MLA,
+ DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+ DEFAULT_URL_FOR_TEST,
+ CustomTestCase,
+ popen_launch_server,
+)
+
+TEST_MODEL = os.environ.get("NIXL_EP_TEST_MODEL", DEFAULT_MODEL_NAME_FOR_TEST_MLA)
+os.environ.setdefault("SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK", "1024")
+
+ib_devices = get_rdma_devices_args()
+
+NIXL_COMMON = [
+ "--trust-remote-code",
+ "--moe-a2a-backend",
+ "nixl",
+ "--deepep-mode",
+ "low_latency",
+ "--tp",
+ "8",
+ "--mem-fraction-static",
+ "0.78",
+]
+DP_ATTN = ["--dp", "8", "--enable-dp-attention"]
+ELASTIC_NIXL = [
+ "--elastic-ep-backend",
+ "nixl",
+ "--enable-eplb",
+ "--ep-num-redundant-experts",
+ "24",
+]
+ELASTIC_MOONCAKE = [
+ "--elastic-ep-backend",
+ "mooncake",
+ "--mooncake-ib-device",
+ ib_devices,
+ "--enable-eplb",
+ "--ep-num-redundant-experts",
+ "24",
+]
+
+
+class _EPTestBase(CustomTestCase):
+ server_args: list[str] = []
+
+ @classmethod
+ def setUpClass(cls):
+ cls.model = TEST_MODEL
+ cls.base_url = DEFAULT_URL_FOR_TEST
+ cls.process = popen_launch_server(
+ cls.model,
+ cls.base_url,
+ timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+ other_args=cls.server_args,
+ )
+
+ @classmethod
+ def tearDownClass(cls):
+ kill_process_tree(cls.process.pid)
+ cls.process.wait(timeout=15)
+ time.sleep(2)
+
+ def _run_gsm8k(self):
+ args = SimpleNamespace(
+ num_shots=5,
+ data_path=None,
+ num_questions=200,
+ max_new_tokens=512,
+ parallel=128,
+ host="http://127.0.0.1",
+ port=int(self.base_url.split(":")[-1]),
+ )
+ metrics = run_eval_few_shot_gsm8k(args)
+ print(metrics)
+ return metrics
+
+ def test_gsm8k(self):
+ metrics = self._run_gsm8k()
+ self.assertGreater(metrics["accuracy"], 0.60)
+
+
+class TestNixlEPTP(_EPTestBase):
+ server_args = [*NIXL_COMMON]
+
+
+class TestNixlEPDPAttn(_EPTestBase):
+ server_args = [*NIXL_COMMON, *DP_ATTN]
+
+
+class TestNixlEPElasticEP(_EPTestBase):
+ server_args = [*NIXL_COMMON, *DP_ATTN, *ELASTIC_NIXL]
+
+
+class TestNixlMoeMooncakeElasticEP(_EPTestBase):
+ server_args = [*NIXL_COMMON, *DP_ATTN, *ELASTIC_MOONCAKE]
+
+ pkill_process_1 = "sglang::scheduler_DP1_TP8_EP8"
+
+ def test_gsm8k_fault_1(self):
+ os.system(f"pkill -f {self.pkill_process_1}")
+ metrics = self._run_gsm8k()
+ self.assertGreater(metrics["accuracy"], 0.60)
+
+
+if __name__ == "__main__":
+ unittest.main()