From 9991debde37b0ea1ecca8e13f0aef24c6b498d17 Mon Sep 17 00:00:00 2001 From: Yoray Zack Date: Wed, 11 Mar 2026 11:37:43 +0200 Subject: [PATCH] [Feature] Integrate Elastic NIXL-EP into SGLang (#19248) Signed-off-by: Barak Biber Signed-off-by: Yoray Zack Signed-off-by: Itay Alroy Co-authored-by: Barak Biber --- docs/advanced_features/expert_parallelism.md | 5 +- docs/advanced_features/server_arguments.md | 2 +- .../srt/batch_overlap/two_batch_overlap.py | 5 + .../sglang/srt/distributed/parallel_state.py | 61 +++ python/sglang/srt/distributed/utils.py | 36 ++ python/sglang/srt/environ.py | 5 + .../srt/eplb/eplb_algorithms/__init__.py | 10 +- python/sglang/srt/layers/moe/ep_moe/layer.py | 6 +- .../srt/layers/moe/fused_moe_triton/layer.py | 7 +- .../layers/moe/token_dispatcher/__init__.py | 8 + .../srt/layers/moe/token_dispatcher/nixl.py | 465 ++++++++++++++++++ python/sglang/srt/layers/moe/utils.py | 4 + python/sglang/srt/layers/quantization/fp8.py | 4 +- .../sglang/srt/model_executor/model_runner.py | 1 + python/sglang/srt/models/deepseek_v2.py | 3 + python/sglang/srt/models/glm4_moe.py | 10 +- python/sglang/srt/server_args.py | 22 +- test/manual/ep/test_nixl_ep.py | 115 +++++ 18 files changed, 752 insertions(+), 17 deletions(-) create mode 100644 python/sglang/srt/layers/moe/token_dispatcher/nixl.py create mode 100644 test/manual/ep/test_nixl_ep.py 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()