chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Feature] Integrate Elastic NIXL-EP into SGLang (#19248)

Browse Source 
Signed-off-by: Barak Biber <bbiber@nvidia.com>
Signed-off-by: Yoray Zack <yorayz@nvidia.com>
Signed-off-by: Itay Alroy <ialroy@nvidia.com>
Co-authored-by: Barak Biber <bbiber@nvidia.com>
This commit is contained in:
Yoray Zack
2026-03-11 17:37:43 +08:00
committed by GitHub
co-authored by Barak Biber
parent 680d9d98e4
commit 9991debde3
18 changed files with 752 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/batch_overlap/two_batch_overlap.py
+5
View File
@@ -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)
python/sglang/srt/distributed/parallel_state.py
+61
View File
@@ -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
python/sglang/srt/distributed/utils.py
+36
View File
@@ -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."""
python/sglang/srt/environ.py
+5
View File
@@ -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)
python/sglang/srt/eplb/eplb_algorithms/__init__.py
+8 -2
View File
@@ -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
python/sglang/srt/layers/moe/ep_moe/layer.py
+5 -1
View File
@@ -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
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+6 -1
View File
@@ -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
python/sglang/srt/layers/moe/token_dispatcher/__init__.py
+8
View File
@@ -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",
python/sglang/srt/layers/moe/token_dispatcher/nixl.py
+465
View File
@@ -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
python/sglang/srt/layers/moe/utils.py
+4
View File
@@ -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
python/sglang/srt/layers/quantization/fp8.py
+3 -1
View File
@@ -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
python/sglang/srt/model_executor/model_runner.py
+1
View File
@@ -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,
python/sglang/srt/models/deepseek_v2.py
+3
View File
@@ -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()
python/sglang/srt/models/glm4_moe.py
+8 -2
View File
@@ -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):
python/sglang/srt/server_args.py
+15 -7
View File
@@ -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",