[Feature] Add DeepEPv2 (ElasticBuffer) MoE A2A backend (#35634)

Co-authored-by: menyu <menyu@nvidia.com>
Co-authored-by: Jinyan Chen <93358689+liz-badada@users.noreply.github.com>
Co-authored-by: Han Yu <helloyu0903@gmail.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
MengYu
2026-08-26 19:54:33 -07:00
committed by GitHub
co-authored by menyu Jinyan Chen Han Yu Cheng Wan
parent cbfe54fba8
commit a3ae667d67
20 changed files with 2414 additions and 35 deletions
+5 -1
View File
@@ -2806,7 +2806,7 @@ def _moe_runner_fusion_disable(view: Any) -> dict:
def _a2a_fusion_adjustments(view: Any) -> dict:
"""A2A-backend-driven shared-experts fusion adjustments, declared at the
legacy write slots in _handle_a2a_moe: Waterfill requires the
fusion enabled; FlashInfer A2A requires it disabled."""
fusion enabled; FlashInfer and DeepEP v2 A2A require it disabled."""
if view.moe_a2a_backend in ("deepep", "megamoe") and view.enable_waterfill:
if view.disable_shared_experts_fusion:
logger.warning(
@@ -2819,6 +2819,9 @@ def _a2a_fusion_adjustments(view: Any) -> dict:
"Flashinfer MoE A2A is enabled. --disable-shared-experts-fusion is automatically set."
)
return {"disable_shared_experts_fusion": True}
if view.moe_a2a_backend == "deepep_v2":
# Fused shared experts are not validated with DeepEP v2.
return {"disable_shared_experts_fusion": True}
return {}
@@ -2827,6 +2830,7 @@ _A2A_EP_SPANNING_BACKENDS = frozenset(
{
"megamoe",
"deepep",
"deepep_v2",
"mooncake",
"nixl",
"ascend_fuseep",
+4
View File
@@ -1043,6 +1043,10 @@ class Envs:
# read by several call sites; do not use in new code.
SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False)
SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
# Per-rank buffer capacity, not a model token limit.
SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
# 0 lets ElasticBuffer select its theoretical communication SM/QP counts.
SGLANG_DEEPEP_V2_NUM_SMS = EnvInt(0)
SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32)
SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False)
SGLANG_ENABLE_QWEN_DEEPEP_SHARED_OVERLAP = EnvBool(True)
+4 -1
View File
@@ -103,7 +103,9 @@ class DeepEPMoE(FusedMoE):
and quant_config is not None
and quant_config.get_name() == "humming"
)
if is_humming:
if get_moe_a2a_backend().is_deepep_v2():
self.deprecate_flag = True
elif is_humming:
self.deprecate_flag = True
elif _use_aiter:
self.deprecate_flag = True
@@ -354,6 +356,7 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]):
if (
get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_deepep_v2()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_nixl()
or get_moe_a2a_backend().is_pplx()
@@ -38,6 +38,7 @@ from sglang.srt.layers.moe.token_dispatcher.ascend_tp import (
AscendTPDispatcher,
)
from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher
from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import DeepEPv2Dispatcher
from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher
from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardDispatcher,
@@ -189,6 +190,15 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
async_finish=True,
return_recv_hook=True,
)
elif a2a_backend.is_deepep_v2():
return DeepEPv2Dispatcher(
group=get_tp_group().device_group,
router_topk=moe_runner_config.top_k,
num_experts=moe_runner_config.num_experts,
num_local_experts=moe_runner_config.num_local_experts,
hidden_size=moe_runner_config.hidden_size,
params_dtype=moe_runner_config.params_dtype,
)
elif a2a_backend.is_flashinfer():
return FlashinferDispatcher(
group=get_tp_group().device_group,
@@ -226,6 +236,34 @@ def _validate_hpc_ops_quant_method(quant_method) -> None:
)
def _validate_deepep_v2_quant_method(quant_method) -> None:
"""Validate the FP8 contract consumed by the DeepEP v2 adapter."""
if not get_moe_a2a_backend().is_deepep_v2():
return
config = (
quant_method.quant_config if isinstance(quant_method, Fp8MoEMethod) else None
)
reason = None
if not isinstance(quant_method, Fp8MoEMethod):
reason = f"selected {type(quant_method).__name__}"
elif quant_method.use_mxfp8:
reason = "selected MXFP8 weights"
elif quant_method.is_fp4_expert:
reason = "selected FP4 experts"
elif list(quant_method.weight_block_size or []) != [128, 128]:
reason = f"has weight_block_size={quant_method.weight_block_size}"
elif config.activation_scheme != "dynamic":
reason = f"has activation_scheme={config.activation_scheme!r}"
if reason is not None:
raise ValueError(
"--moe-a2a-backend deepep_v2 requires 128x128 blockwise FP8 "
f"experts with dynamic activation scaling, but this layer {reason}. "
"Use a compatible checkpoint or --moe-a2a-backend deepep."
)
class FusedMoE(torch.nn.Module):
"""FusedMoE layer for MoE models.
@@ -407,6 +445,7 @@ class FusedMoE(torch.nn.Module):
self.use_deep_gemm,
)
_validate_hpc_ops_quant_method(self.quant_method)
_validate_deepep_v2_quant_method(self.quant_method)
self.supports_deferred_finalize = (
envs.SGLANG_ENABLE_MOE_DEFERRED_FINALIZE.get()
and get_moe_runner_backend().is_flashinfer_trtllm()
@@ -49,6 +49,10 @@ if TYPE_CHECKING:
DeepEPNormalCombineInput,
DeepEPNormalDispatchOutput,
)
from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import (
DeepEPv2CombineInput,
DeepEPv2DispatchOutput,
)
from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardCombineInput,
StandardDispatchOutput,
@@ -206,6 +210,7 @@ class DeepGemmRunnerInput(RunnerInput):
masked_m: Optional[torch.Tensor] = None
expected_m: Optional[int] = None
m_indices: Optional[torch.Tensor] = None
hidden_states_scale_tma_aligned: bool = False
@property
def runner_backend(self) -> MoeRunnerBackend:
@@ -321,7 +326,10 @@ class DeepGemmRunnerCore(MoeRunnerCore):
device=hidden_states_device,
dtype=torch.bfloat16,
)
if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES:
if (
deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES
and not runner_input.hidden_states_scale_tma_aligned
):
hidden_states_scale = tma_align_input_scale(hidden_states_scale)
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(
@@ -1419,3 +1427,186 @@ def _apply_swiglu_limit(
out = torch.cat([gate, up], dim=-1)
assert out.shape == (num_tokens, hidden_size_x2)
return out
@register_pre_permute("deepep_v2", "deep_gemm")
def pre_permute_deepep_v2_to_deep_gemm(
dispatch_output: DeepEPv2DispatchOutput,
quant_info: DeepGemmMoeQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> DeepGemmRunnerInput:
from sglang.kernels.ops.moe.ep_moe_kernels import (
ep_expand_init_m_indices_from_psum,
ep_scatter_from_psum,
)
hidden_states = dispatch_output.hidden_states
hidden_states_scale = dispatch_output.hidden_states_scale
topk_ids = dispatch_output.topk_ids
topk_weights = dispatch_output.topk_weights
psum_num_recv_tokens_per_expert = dispatch_output.psum_num_recv_tokens_per_expert
is_expanded = dispatch_output.is_expanded
hidden_states_scale_tma_aligned = dispatch_output.hidden_states_scale_tma_aligned
deepep_v2_use_masked = dispatch_output.use_masked_gemm
deepep_v2_expected_m = dispatch_output.expected_m
deepep_v2_masked_max_m = dispatch_output.masked_max_m
deepep_v2_total_expanded = dispatch_output.total_expanded
deepep_v2_expert_alignment = dispatch_output.expert_alignment
if hidden_states_scale is None:
raise RuntimeError(
"DeepEP v2 -> DeepGEMM requires FP8 dispatch output with activation "
"scales, but the dispatch output carried none."
)
assert runner_config.activation == "silu"
if is_expanded:
if psum_num_recv_tokens_per_expert is None:
raise RuntimeError(
"DeepEP v2 requires the native expert prefix sums from the "
"ElasticBuffer dispatch handle."
)
all_tokens = hidden_states.shape[0]
running_state["all_tokens"] = all_tokens
running_state["hidden_states_shape"] = hidden_states.shape
running_state["hidden_states_device"] = hidden_states.device
running_state["hidden_states_dtype"] = hidden_states.dtype
running_state["topk_ids"] = None
running_state["topk_weights"] = topk_weights
running_state["deepep_v2_expanded"] = True
if deepep_v2_use_masked:
# masked_m bounds each expert independently of buffer capacity.
from sglang.kernels.ops.moe.ep_moe_kernels import expand_to_masked_slab
num_local_experts = psum_num_recv_tokens_per_expert.shape[0]
input_tensor, input_tensor_scale, masked_m = expand_to_masked_slab(
hidden_states,
hidden_states_scale,
psum_num_recv_tokens_per_expert,
num_local_experts,
deepep_v2_masked_max_m,
deepep_v2_expert_alignment,
)
running_state["deepep_v2_masked"] = True
running_state["deepep_v2_psum"] = psum_num_recv_tokens_per_expert
running_state["deepep_v2_total_expanded"] = deepep_v2_total_expanded
running_state["deepep_v2_expert_alignment"] = deepep_v2_expert_alignment
return DeepGemmRunnerInput(
hidden_states=input_tensor,
hidden_states_scale=input_tensor_scale,
use_masked_gemm=True,
masked_m=masked_m,
expected_m=deepep_v2_expected_m,
)
# Mark aligned expert rows and leave the unused receive tail at -1.
m_indices = torch.full(
(all_tokens,), -1, device=hidden_states.device, dtype=torch.int32
)
ep_expand_init_m_indices_from_psum(psum_num_recv_tokens_per_expert, m_indices)
return DeepGemmRunnerInput(
hidden_states=hidden_states,
hidden_states_scale=hidden_states_scale,
use_masked_gemm=False,
m_indices=m_indices,
hidden_states_scale_tma_aligned=hidden_states_scale_tma_aligned,
)
all_tokens = int(psum_num_recv_tokens_per_expert[-1].item())
K = hidden_states.shape[1]
running_state["all_tokens"] = all_tokens
running_state["hidden_states_shape"] = hidden_states.shape
running_state["hidden_states_device"] = hidden_states.device
running_state["hidden_states_dtype"] = hidden_states.dtype
running_state["topk_ids"] = topk_ids
running_state["topk_weights"] = topk_weights
input_tensor = torch.empty(
(all_tokens, K), device=hidden_states.device, dtype=hidden_states.dtype
)
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
# Packed UE8M0 scales require zero padding lanes.
input_tensor_scale = torch.zeros(
(ceil_div(K // 128, 4), all_tokens),
device=hidden_states.device,
dtype=torch.int,
).transpose(0, 1)
else:
input_tensor_scale = torch.empty(
(all_tokens, K // 128), device=hidden_states.device, dtype=torch.float32
)
m_indices = torch.empty(all_tokens, device=hidden_states.device, dtype=torch.int32)
output_index = torch.empty_like(topk_ids)
# Contiguous psum already includes the 128-row expert alignment.
expert_start_loc = torch.empty_like(psum_num_recv_tokens_per_expert)
ep_scatter_from_psum(
hidden_states,
hidden_states_scale,
topk_ids,
psum_num_recv_tokens_per_expert,
expert_start_loc,
input_tensor,
input_tensor_scale,
m_indices,
output_index,
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
)
dispose_tensor(hidden_states)
dispose_tensor(hidden_states_scale)
running_state["output_index"] = output_index
return DeepGemmRunnerInput(
hidden_states=input_tensor,
hidden_states_scale=input_tensor_scale,
use_masked_gemm=False,
m_indices=m_indices,
)
@register_post_permute("deep_gemm", "deepep_v2")
def post_permute_deep_gemm_to_deepep_v2(
runner_output: DeepGemmRunnerOutput,
quant_info: DeepGemmMoeQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> DeepEPv2CombineInput:
from sglang.kernels.ops.moe.ep_moe_kernels import ep_gather
from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import DeepEPv2CombineInput
if running_state.get("deepep_v2_expanded", False):
hidden_states = runner_output.hidden_states
topk_weights = running_state["topk_weights"]
if running_state.get("deepep_v2_masked", False):
# Expanded combine does not consume top-k weights.
from sglang.kernels.ops.moe.ep_moe_kernels import masked_slab_to_expand
hidden_states = masked_slab_to_expand(
hidden_states,
running_state["deepep_v2_psum"],
running_state["deepep_v2_total_expanded"],
running_state["deepep_v2_expert_alignment"],
topk_weights=topk_weights,
)
return DeepEPv2CombineInput(hidden_states, None)
if topk_weights is not None:
# Expanded combine does not consume top-k weights.
hidden_states = hidden_states * topk_weights.to(
hidden_states.dtype
).unsqueeze(-1)
return DeepEPv2CombineInput(hidden_states, None)
hidden_states = runner_output.hidden_states
topk_ids = running_state["topk_ids"]
topk_weights = running_state["topk_weights"]
output_index = running_state["output_index"]
gather_out = torch.empty(
running_state["hidden_states_shape"],
device=running_state["hidden_states_device"],
dtype=torch.bfloat16,
)
ep_gather(hidden_states, topk_ids, topk_weights, output_index, gather_out)
return DeepEPv2CombineInput(
hidden_states=gather_out,
topk_weights=topk_weights,
)
@@ -50,6 +50,15 @@ class MoeRunner:
"--moe-runner-backend hpc_ops for this model."
)
if get_moe_a2a_backend().is_deepep_v2() and not runner_backend.is_deep_gemm():
raise ValueError(
"--moe-a2a-backend deepep_v2 requires the deep_gemm MoE runner, "
f"but this MoE layer's quantization method selected the "
f"'{runner_backend.value}' runner. deepep_v2 dispatches FP8 "
"activations plus scales, which only deep_gemm consumes; use an "
"FP8 blockwise-quantized checkpoint, or --moe-a2a-backend deepep."
)
self.fused_func = None
if runner_backend.is_triton():
@@ -21,6 +21,11 @@ from sglang.srt.layers.moe.token_dispatcher.deepep import (
DeepEPNormalCombineInput,
DeepEPNormalDispatchOutput,
)
from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import (
DeepEPv2CombineInput,
DeepEPv2Dispatcher,
DeepEPv2DispatchOutput,
)
from sglang.srt.layers.moe.token_dispatcher.flashinfer import (
FlashinferDispatcher,
FlashinferDispatchOutput,
@@ -72,6 +77,9 @@ __all__ = [
"MoriEPLLDispatchOutput",
"MoriEPLLCombineInput",
"MoriEPDispatcher",
"DeepEPv2Dispatcher",
"DeepEPv2DispatchOutput",
"DeepEPv2CombineInput",
"NixlEPCombineInput",
"NixlEPDispatchOutput",
"NixlEPDispatcher",
@@ -27,6 +27,8 @@ if TYPE_CHECKING:
DeepEPLLDispatchOutput,
DeepEPNormalCombineInput,
DeepEPNormalDispatchOutput,
DeepEPv2CombineInput,
DeepEPv2DispatchOutput,
FlashinferCombineInput,
FlashinferDispatchOutput,
StandardCombineInput,
@@ -165,6 +167,12 @@ class DispatchOutputChecker:
) -> TypeGuard[FlashinferDispatchOutput]:
return dispatch_output.format.is_flashinfer()
@staticmethod
def format_is_deepep_v2(
dispatch_output: DispatchOutput,
) -> TypeGuard[DeepEPv2DispatchOutput]:
return dispatch_output.format.is_deepep_v2()
class DispatchOutputFormat(Enum):
@@ -172,6 +180,7 @@ class DispatchOutputFormat(Enum):
DEEPEP_NORMAL = "deepep_normal"
DEEPEP_LL = "deepep_ll"
FLASHINFER = "flashinfer"
DEEPEP_V2 = "deepep_v2"
ASCEND_TP = "ascend_tp"
def is_standard(self) -> bool:
@@ -195,6 +204,9 @@ class DispatchOutputFormat(Enum):
def is_flashinfer(self) -> bool:
return self == DispatchOutputFormat.FLASHINFER
def is_deepep_v2(self) -> bool:
return self == DispatchOutputFormat.DEEPEP_V2
@runtime_checkable
class DispatchOutput(Protocol):
@@ -249,12 +261,19 @@ class CombineInputChecker:
) -> TypeGuard[FlashinferCombineInput]:
return combine_input.format == CombineInputFormat.FLASHINFER
@staticmethod
def format_is_deepep_v2(
combine_input: CombineInput,
) -> TypeGuard[DeepEPv2CombineInput]:
return combine_input.format == CombineInputFormat.DEEPEP_V2
class CombineInputFormat(Enum):
STANDARD = "standard"
DEEPEP_NORMAL = "deepep_normal"
DEEPEP_LL = "deepep_ll"
FLASHINFER = "flashinfer"
DEEPEP_V2 = "deepep_v2"
ASCEND_TP = "ascend_tp"
@@ -0,0 +1,460 @@
from __future__ import annotations
import logging
import os
from typing import NamedTuple, Optional
import torch
import torch.distributed as dist
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import get_is_extend_in_batch
from sglang.srt.layers.moe.token_dispatcher.base import (
BaseDispatcher,
CombineInput,
CombineInputFormat,
DispatchOutput,
DispatchOutputFormat,
)
from sglang.srt.layers.moe.topk import TopKOutput
from sglang.srt.layers.moe.utils import (
DeepEPv2Fp8ScaleFormat,
get_deepep_v2_fp8_scale_format,
)
logger = logging.getLogger(__name__)
_SCALE_BLOCK_SIZE = 128
# Must match DeepGEMM's contiguous expert alignment.
_EXPERT_ALIGNMENT = 128
_deepep_v2_import_error: Optional[BaseException] = None
_fp8_quant_import_error: Optional[BaseException] = None
sglang_per_token_group_quant_fp8 = None
try:
from deep_ep import ElasticBuffer
use_deepep_v2 = True
except (ImportError, OSError) as exc:
use_deepep_v2 = False
_deepep_v2_import_error = exc
if use_deepep_v2:
try:
from sglang.kernels.ops.quantization.fp8_kernel import (
sglang_per_token_group_quant_fp8,
)
except (ImportError, OSError) as exc:
_fp8_quant_import_error = exc
class DeepEPv2DispatchOutput(NamedTuple):
hidden_states: torch.Tensor
hidden_states_scale: Optional[torch.Tensor]
topk_ids: Optional[torch.Tensor]
topk_weights: torch.Tensor
psum_num_recv_tokens_per_expert: Optional[torch.Tensor] = None
is_expanded: bool = False
hidden_states_scale_tma_aligned: bool = False
use_masked_gemm: bool = False
expected_m: int = 0
masked_max_m: int = 0
total_expanded: int = 0
expert_alignment: int = 128
@property
def format(self) -> DispatchOutputFormat:
return DispatchOutputFormat.DEEPEP_V2
class DeepEPv2CombineInput(NamedTuple):
hidden_states: torch.Tensor
topk_weights: Optional[torch.Tensor]
@property
def format(self) -> CombineInputFormat:
return CombineInputFormat.DEEPEP_V2
assert isinstance(DeepEPv2DispatchOutput, DispatchOutput)
assert isinstance(DeepEPv2CombineInput, CombineInput)
def _raise_deepep_v2_import_error() -> None:
detail = (
f" Original import error: {_deepep_v2_import_error}"
if _deepep_v2_import_error is not None
else ""
)
raise ImportError(
"DeepEP v2 (ElasticBuffer) is not available. Install DeepEP v2 from "
"https://github.com/deepseek-ai/DeepEP." + detail
)
def _ensure_deepep_v2_available() -> None:
if not use_deepep_v2:
_raise_deepep_v2_import_error()
def _ensure_fp8_quant_available() -> None:
_ensure_deepep_v2_available()
if sglang_per_token_group_quant_fp8 is None:
detail = (
f" Original import error: {_fp8_quant_import_error}"
if _fp8_quant_import_error is not None
else ""
)
raise ImportError(
"DeepEP v2 FP8 dispatch requires the SGLang FP8 quantization kernel."
+ detail
)
def _get_allow_hybrid_mode() -> bool:
from sglang.srt.runtime_context import get_exec
return get_exec().moe.deepep_v2_mode == "hybrid"
def _quantize_for_deepep_v2_dispatch(
hidden_states: torch.Tensor, scale_format: DeepEPv2Fp8ScaleFormat
):
_ensure_fp8_quant_available()
return sglang_per_token_group_quant_fp8(
hidden_states,
_SCALE_BLOCK_SIZE,
column_major_scales=scale_format.tma_aligned,
scale_tma_aligned=scale_format.tma_aligned,
scale_ue8m0=scale_format.ue8m0,
)
class DeepEPv2Buffer:
"""Facade for the process-wide ElasticBuffer stored in runtime resources."""
_STATE_KEY = "deepep_v2_ep_state"
@classmethod
def _state(cls):
from types import SimpleNamespace
from sglang.srt.runtime_context import get_resources
buffers = get_resources().buffers
state = buffers.get(cls._STATE_KEY)
if state is None:
state = SimpleNamespace(buffer=None, key=None)
buffers[cls._STATE_KEY] = state
return state
@classmethod
def get_buffer(
cls,
group: dist.ProcessGroup,
hidden_size: int,
router_topk: int,
num_max_dispatch_tokens_per_rank: int,
use_fp8_dispatch: bool,
allow_hybrid_mode: Optional[bool] = None,
) -> ElasticBuffer:
_ensure_deepep_v2_available()
if allow_hybrid_mode is None:
allow_hybrid_mode = _get_allow_hybrid_mode()
state = cls._state()
# A key change rebuilds ElasticBuffer collectively on every rank.
key = (
group,
hidden_size,
router_topk,
num_max_dispatch_tokens_per_rank,
use_fp8_dispatch,
allow_hybrid_mode,
dist.get_world_size(group),
)
if state.buffer is not None and state.key == key:
return state.buffer
# Native explicit teardown is unavailable unless explicitly_destroy=True.
cls.destroy()
# Communicator reuse requires a device-bound process group.
os.environ.setdefault("EP_REUSE_NCCL_COMM", "0")
buffer = ElasticBuffer(
group,
num_max_tokens_per_rank=num_max_dispatch_tokens_per_rank,
hidden=hidden_size,
num_topk=router_topk,
use_fp8_dispatch=use_fp8_dispatch,
allow_hybrid_mode=allow_hybrid_mode,
sl_idx=0,
prefer_overlap_with_compute=False,
)
# Publish only after collective construction succeeds.
state.buffer = buffer
state.key = key
logger.info(
"Initialized DeepEP v2 ElasticBuffer: world_size=%s hidden_size=%s "
"num_topk=%s max_dispatch_tokens_per_rank=%s use_fp8_dispatch=%s "
"allow_hybrid_mode=%s num_bytes=%s",
dist.get_world_size(group),
hidden_size,
router_topk,
num_max_dispatch_tokens_per_rank,
use_fp8_dispatch,
allow_hybrid_mode,
buffer.num_bytes,
)
return buffer
@classmethod
def destroy(cls) -> None:
state = cls._state()
state.buffer = None
state.key = None
class _DeepEPv2Impl:
def __init__(
self,
group: dist.ProcessGroup,
router_topk: int,
num_experts: int,
num_local_experts: int,
hidden_size: int,
scale_format: DeepEPv2Fp8ScaleFormat,
num_max_dispatch_tokens_per_rank: int,
):
self.group = group
self.router_topk = router_topk
self.num_experts = num_experts
self.num_local_experts = num_local_experts
self.hidden_size = hidden_size
self.scale_format = scale_format
self.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank
self.rank = dist.get_rank(group)
self._handle = None
self._pad_empty_combine = False
def _destroy_handle(self) -> None:
self._handle = None
def _get_buffer(self) -> ElasticBuffer:
return DeepEPv2Buffer.get_buffer(
self.group,
self.hidden_size,
self.router_topk,
self.num_max_dispatch_tokens_per_rank,
True,
)
def _validate_common(
self, hidden_states: torch.Tensor, topk_ids: torch.Tensor
) -> None:
if hidden_states.shape[0] > self.num_max_dispatch_tokens_per_rank:
raise ValueError(
f"DeepEP v2 dispatch input exceeds the per-rank buffer capacity "
f"{self.num_max_dispatch_tokens_per_rank}, got {hidden_states.shape[0]}. "
"Increase SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK or "
"lower the active prefill/decode batch limit."
)
if hidden_states.shape[1] != self.hidden_size:
raise ValueError(
f"DeepEP v2 hidden size mismatch: expected {self.hidden_size}, "
f"got {hidden_states.shape[1]}"
)
if self.hidden_size % _SCALE_BLOCK_SIZE != 0:
raise ValueError(
"DeepEP v2 FP8 dispatch requires hidden_size multiple of "
f"{_SCALE_BLOCK_SIZE}, got {self.hidden_size}"
)
if topk_ids.shape[1] != self.router_topk:
raise ValueError(
f"DeepEP v2 topk mismatch: expected {self.router_topk}, "
f"got {topk_ids.shape[1]}"
)
def dispatch(
self, hidden_states: torch.Tensor, topk_output: TopKOutput
) -> DeepEPv2DispatchOutput:
if self._handle is not None:
raise RuntimeError(
"DeepEP v2 dispatch called while the previous dispatch handle is "
"still unconsumed (missing combine)"
)
_ensure_deepep_v2_available()
topk_weights = topk_output.topk_weights
topk_ids = topk_output.topk_ids.to(torch.int64)
self._validate_common(hidden_states, topk_ids)
# Decode uses expanded/masked layout; extend uses contiguous in both modes.
use_expand_layout = not get_is_extend_in_batch()
use_masked = use_expand_layout
# CPU-synced dispatch needs a dummy token to notify from an idle rank.
self._pad_empty_combine = (not use_masked) and hidden_states.shape[0] == 0
if self._pad_empty_combine:
hidden_states = hidden_states.new_zeros((1, hidden_states.shape[-1]))
# Dummy routes need distinct expert ids; zero weights null the result.
topk_ids = torch.arange(
topk_ids.shape[-1], dtype=topk_ids.dtype, device=topk_ids.device
).unsqueeze(0)
topk_weights = topk_weights.new_zeros((1, topk_weights.shape[-1]))
_ensure_fp8_quant_available()
if use_masked:
_ue8m0 = self.scale_format.ue8m0
dispatch_x = sglang_per_token_group_quant_fp8(
hidden_states,
_SCALE_BLOCK_SIZE,
column_major_scales=_ue8m0,
scale_tma_aligned=_ue8m0,
scale_ue8m0=_ue8m0,
)
use_tma_aligned_col_major_sf = _ue8m0
else:
dispatch_x = _quantize_for_deepep_v2_dispatch(
hidden_states, self.scale_format
)
use_tma_aligned_col_major_sf = self.scale_format.tma_aligned
# This collective argument must not depend on a rank-local batch.
num_max_tokens = self.num_max_dispatch_tokens_per_rank
# Masked dispatch stays asynchronous for CUDA graph capture.
do_cpu_sync_val = True
if use_masked:
do_cpu_sync_val = False
buffer = self._get_buffer()
recv_x, recv_topk_idx, recv_topk_weights, handle, event = buffer.dispatch(
dispatch_x,
topk_idx=topk_ids,
topk_weights=topk_weights,
num_experts=self.num_experts,
num_max_tokens_per_rank=num_max_tokens,
expert_alignment=_EXPERT_ALIGNMENT,
num_sms=envs.SGLANG_DEEPEP_V2_NUM_SMS.get(),
use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf,
do_cpu_sync=do_cpu_sync_val,
do_expand=use_expand_layout,
)
self._handle = handle
local_tokens = hidden_states.shape[0]
if event.event is not None:
event.current_stream_wait()
if isinstance(recv_x, tuple):
recv_hidden_states, recv_hidden_states_scale = recv_x
else:
recv_hidden_states = recv_x
recv_hidden_states_scale = None
if use_expand_layout:
# Expanded combine uses handle metadata instead of recv_topk_idx.
local_topk_ids = None
else:
num_recv_tokens = int(
handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()
)
recv_topk_idx = recv_topk_idx[:num_recv_tokens]
recv_topk_weights = recv_topk_weights[:num_recv_tokens]
recv_hidden_states = recv_hidden_states[:num_recv_tokens]
if recv_hidden_states_scale is not None:
recv_hidden_states_scale = recv_hidden_states_scale[:num_recv_tokens]
local_topk_ids = recv_topk_idx
expected_m = 0
masked_max_m = 0
total_expanded = 0
if use_masked:
# expected_m is only a schedule hint; masked_m is the actual bound.
ep_group_size = max(1, self.num_experts // self.num_local_experts)
expected_m = max(
1,
(local_tokens * ep_group_size * self.router_topk + self.num_experts)
// self.num_experts,
)
# Account for the worst case where every rank targets one local expert.
masked_max_m = self.num_max_dispatch_tokens_per_rank * ep_group_size
total_expanded = recv_hidden_states.shape[0]
return DeepEPv2DispatchOutput(
recv_hidden_states,
recv_hidden_states_scale,
local_topk_ids,
recv_topk_weights,
handle.psum_num_recv_tokens_per_expert,
use_expand_layout,
use_tma_aligned_col_major_sf,
use_masked,
expected_m,
masked_max_m,
total_expanded,
_EXPERT_ALIGNMENT,
)
def combine(self, combine_input: DeepEPv2CombineInput) -> torch.Tensor:
if self._handle is None:
raise RuntimeError(
"DeepEP v2 combine called without a valid dispatch handle"
)
# Release the single-use handle even when combine fails.
try:
buffer = self._get_buffer()
combined_x, _, event = buffer.combine(
combine_input.hidden_states,
handle=self._handle,
topk_weights=combine_input.topk_weights,
)
if event.event is not None:
event.current_stream_wait()
if self._pad_empty_combine:
combined_x = combined_x[:0]
return combined_x
finally:
self._pad_empty_combine = False
self._destroy_handle()
class DeepEPv2Dispatcher(BaseDispatcher):
def __init__(
self,
group: dist.ProcessGroup,
router_topk: int,
num_experts: int,
num_local_experts: int,
hidden_size: int,
params_dtype: torch.dtype,
):
super().__init__()
if params_dtype != torch.bfloat16:
raise NotImplementedError(
"DeepEP v2 dispatch adapter currently expects BF16 model activations, "
f"got {params_dtype}"
)
scale_format = get_deepep_v2_fp8_scale_format()
self.num_max_dispatch_tokens_per_rank = (
envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get()
)
self._impl = _DeepEPv2Impl(
group=group,
router_topk=router_topk,
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden_size,
scale_format=scale_format,
num_max_dispatch_tokens_per_rank=self.num_max_dispatch_tokens_per_rank,
)
def dispatch(
self, hidden_states: torch.Tensor, topk_output: TopKOutput
) -> DispatchOutput:
return self._impl.dispatch(hidden_states, topk_output)
def combine(self, combine_input: CombineInput) -> torch.Tensor:
if combine_input.format != CombineInputFormat.DEEPEP_V2:
raise TypeError(
f"Expected DeepEP v2 combine input, got {combine_input.format}"
)
return self._impl.combine(combine_input)
+33 -2
View File
@@ -4,6 +4,7 @@ import logging
import os
from contextlib import contextmanager
from enum import Enum, IntEnum
from typing import NamedTuple
import torch
@@ -40,6 +41,7 @@ class MoeA2ABackend(Enum):
ASCEND_TP = "ascend_tp"
FLASHINFER = "flashinfer"
MEGAMOE = "megamoe"
DEEPEP_V2 = "deepep_v2"
PPLX = "pplx"
CUSTOMIZED = "customized"
@@ -79,6 +81,9 @@ class MoeA2ABackend(Enum):
def is_megamoe(self):
return self == MoeA2ABackend.MEGAMOE
def is_deepep_v2(self):
return self == MoeA2ABackend.DEEPEP_V2
def is_pplx(self):
return self == MoeA2ABackend.PPLX
@@ -178,6 +183,13 @@ class MoeRunnerBackend(Enum):
return self == MoeRunnerBackend.AITER
class DeepEPv2Fp8ScaleFormat(NamedTuple):
"""DeepGEMM FP8 activation-scale layout expected from DeepEP v2."""
tma_aligned: bool
ue8m0: bool
class DeepEPMode(Enum):
NORMAL = "normal"
@@ -311,6 +323,19 @@ def get_ascend_dispatcher_output_dtype(dispatcher):
return DispatcherOutputDtype.BF16
def get_deepep_v2_fp8_scale_format() -> DeepEPv2Fp8ScaleFormat:
"""Resolve the FP8 scale layout DeepEP v2 must pre-quantize into."""
from sglang.srt.layers import deep_gemm_wrapper
return DeepEPv2Fp8ScaleFormat(
tma_aligned=(
deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES
or deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
),
ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
)
def initialize_moe_config():
"""Seed the MoE runtime flags from the published configuration.
@@ -502,9 +527,15 @@ def is_sbo_enabled() -> bool:
def is_deepep_class_backend() -> bool:
"""Check if the MoE backend is DeepEP-family (DeepEP, Mooncake, Mori, or PPLX)."""
"""Return whether A2A combine occurs inside a DeepEP-family dispatcher."""
b = get_moe_a2a_backend()
return b.is_deepep() or b.is_mooncake() or b.is_mori() or b.is_pplx()
return (
b.is_deepep()
or b.is_deepep_v2()
or b.is_mooncake()
or b.is_mori()
or b.is_pplx()
)
def uses_per_rank_fused_shared_slots() -> bool:
+8 -4
View File
@@ -744,6 +744,7 @@ class DeepseekV2MoE(nn.Module):
or get_moe_a2a_backend().is_ascend_fuseep()
or get_moe_a2a_backend().is_flashinfer()
or get_moe_a2a_backend().is_megamoe()
or get_moe_a2a_backend().is_deepep_v2()
or should_use_flashinfer_cutlass_moe_fp4_allgather()
or envs.SGLANG_SHARED_EXPERT_TP1.get()
)
@@ -833,6 +834,7 @@ class DeepseekV2MoE(nn.Module):
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_deepep_v2()
):
# TODO: we will support tp < ep in the future
self.ep_size = get_parallel().moe_ep_size
@@ -855,6 +857,7 @@ class DeepseekV2MoE(nn.Module):
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep()
or get_moe_a2a_backend().is_flashinfer()
or get_moe_a2a_backend().is_deepep_v2()
)
self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo()
# SGLANG_OPT_MOE_QUANT_ONCE eligibility, resolved lazily on first
@@ -2757,10 +2760,11 @@ class DeepseekV2Model(nn.Module):
)
)
self.layers_to_capture = []
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
self.enable_a2a_moe = True
else:
self.enable_a2a_moe = False
self.enable_a2a_moe = (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_deepep_v2()
)
# llama_4_scaling: for supporting Mistral-Large-3 model
self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
+166
View File
@@ -303,10 +303,20 @@ MOE_A2A_BACKEND_CHOICES = [
"ascend_fuseep",
"flashinfer",
"megamoe",
"deepep_v2",
"pplx",
"ascend_tp",
]
# These architectures take the A2A MoE path and skip post-expert all-reduce.
_DEEPEP_V2_VALIDATED_ARCHITECTURES = frozenset(
{
"DeepseekV3ForCausalLM",
"DeepseekV4ForCausalLM",
"Qwen3MoeForCausalLM",
}
)
MXFP8_MOE_RUNNER_BACKEND_CHOICES = [
"cutlass",
"deep_gemm",
@@ -2443,6 +2453,8 @@ class ServerArgs:
"ascend_fuseep",
"flashinfer",
"megamoe",
"deepep_v2",
"ascend_tp",
"pplx",
],
Arg(
@@ -2459,6 +2471,15 @@ class ServerArgs:
"--moe-a2a-backend megamoe.",
NS("exec.moe"),
] = False
deepep_v2_mode: A[
Literal["direct", "hybrid"],
"DeepEP v2 ElasticBuffer communication topology, fixed at server init: "
"`direct` (single-node NVLink) or `hybrid` (multi-node scale-out). "
"Layout/grouped-GEMM and the decode CUDA graph are chosen per batch by "
"inference phase, independent of this knob; not equivalent to DeepEP v1 "
"normal/low_latency.",
NS("exec.moe"),
] = "direct"
moe_runner_backend: A[
str,
Arg(
@@ -4020,6 +4041,10 @@ class ServerArgs:
# time; last declarations of the resolution, mirroring that order.
self._handle_model_capability_adjustments()
# Validate after all batch-size declarations are visible.
self._validate_deepep_v2_speculative_draft()
self._validate_deepep_v2_dispatch_token_budget()
self._resolution_finished = True
def _handle_return_hidden_states_mode(self):
@@ -7415,6 +7440,93 @@ class ServerArgs:
f"(e.g. --max-prefill-tokens) to <= {max_cutedsl_tokens}."
)
def _validate_deepep_v2_dispatch_token_budget(self) -> None:
"""Check the configured prefill and decode-graph buffer bounds."""
view = resolved_view(self)
if view.moe_a2a_backend != "deepep_v2":
return
capacity = envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get()
if view.disaggregation_mode != "decode":
prefill_tokens = self.max_prefill_buffer_tokens() or (
view.max_prefill_tokens or 0
)
if prefill_tokens > capacity:
raise ValueError(
"DeepEP v2 per-rank prefill budget exceeds "
"SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK: "
f"required={prefill_tokens}, capacity={capacity}. Raise the "
"environment value or lower --chunked-prefill-size/"
"--max-prefill-tokens."
)
if view.disaggregation_mode == "prefill":
return
decode_config = getattr(view.cuda_graph_config, "decode", None)
if decode_config is None or decode_config.backend == Backend.DISABLED:
return
graph_bs = decode_config.max_bs or 0
if view.max_running_requests is not None:
attn_dp_size = view.dp_size if view.enable_dp_attention else 1
per_rank_pool_bs = max(1, view.max_running_requests // attn_dp_size)
graph_bs = min(graph_bs, per_rank_pool_bs)
tokens_per_req = (
self.max_speculative_num_draft_tokens or 1
if view.speculative_algorithm
else 1
)
graph_tokens = graph_bs * tokens_per_req
if graph_tokens > capacity:
raise ValueError(
"DeepEP v2 per-rank decode CUDA graph exceeds "
"SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK: "
f"required={graph_tokens}, capacity={capacity} "
f"(requests={graph_bs}, tokens/request={tokens_per_req}). Raise "
"the environment value or lower --cuda-graph-max-bs."
)
def _validate_deepep_v2_model_architecture(self) -> None:
"""Allow DeepEP v2 only where its model workflow is validated."""
if (
parse_connector_type(resolved_view(self).model_path)
== ConnectorType.INSTANCE
):
raise ValueError(
"DeepEP v2 MoE cannot validate a model loaded through an instance "
"connector. Load it from a model path or use "
"--moe-a2a-backend deepep."
)
architectures = (
getattr(self.get_model_config().hf_config, "architectures", None) or []
)
architecture = architectures[0] if architectures else None
if architecture not in _DEEPEP_V2_VALIDATED_ARCHITECTURES:
raise ValueError(
f"DeepEP v2 MoE is not validated for {architecture!r}; supported "
f"architectures are {sorted(_DEEPEP_V2_VALIDATED_ARCHITECTURES)}. "
"Other model workflows may require an all-reduce after A2A "
"combine. Use --moe-a2a-backend deepep."
)
def _validate_deepep_v2_speculative_draft(self) -> None:
"""Reject an explicit or inherited DeepEP v2 draft backend."""
view = resolved_view(self)
draft_backend = view.speculative_moe_a2a_backend
if draft_backend is None and view.speculative_algorithm:
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
algorithm = SpeculativeAlgorithm.from_string(view.speculative_algorithm)
if not algorithm.is_ngram():
draft_backend = view.moe_a2a_backend
if draft_backend == "deepep_v2":
raise ValueError(
"DeepEP v2 MoE is not validated as a speculative draft backend. "
"Select another --speculative-moe-a2a-backend."
)
def _handle_a2a_moe(self):
# The backend overrides and the ep_size=tp_size adjustments moved to
# the resolution pipeline (arg_groups/overrides.py:
@@ -7466,6 +7578,60 @@ class ServerArgs:
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
if a2a_backend == "deepep_v2":
self._validate_deepep_v2_model_architecture()
if resolved_view(self).enable_deterministic_inference:
raise ValueError(
"DeepEP v2 does not forward deterministic=True to "
"ElasticBuffer, so deterministic sorting remains disabled. "
"Disable --enable-deterministic-inference or use "
"--moe-a2a-backend deepep."
)
# ElasticBuffer requires CUMEM, but not NVLS or its preallocation.
os.environ.setdefault("NCCL_CUMEM_ENABLE", "1")
# Respect model-level runner declarations before resolving auto.
resolved_runner = resolved_view(self).moe_runner_backend
if resolved_runner == "auto":
self._declare("_handle_a2a_moe", moe_runner_backend="deep_gemm")
logger.warning(
"DeepEP v2 MoE: resolved --moe-runner-backend auto -> deep_gemm."
)
elif resolved_runner != "deep_gemm":
raise ValueError(
"DeepEP v2 MoE currently supports only "
f"--moe-runner-backend deep_gemm. Got {resolved_runner!r}. "
"Add a runner adapter before enabling DeepEP v2 with other "
"MoE runners."
)
if cfg.enable_two_batch_overlap or cfg.enable_single_batch_overlap:
raise ValueError(
"DeepEP v2 MoE has not implemented the TBO/SBO overlap hooks yet. "
"Disable --enable-two-batch-overlap and "
"--enable-single-batch-overlap when using --moe-a2a-backend deepep_v2."
)
if cfg.enforce_shared_experts_fusion:
raise ValueError(
"DeepEP v2 MoE has not validated fused shared experts yet. "
"Remove --enforce-shared-experts-fusion when using "
"--moe-a2a-backend deepep_v2."
)
# Prefill reads host counts and is not graph-capturable.
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
logger.warning(
f"DeepEP v2 MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{cfg.tp_size}]."
)
logger.warning(
"DeepEP v2 MoE is using deepep_v2_mode=%s. This controls "
"ElasticBuffer direct/hybrid mode and is independent from "
"--deepep-mode normal/low_latency. DeepEP v2 MoE enables the "
"decode CUDA graph on the masked decode path (any comm mode) "
"and disables shared expert fusion. "
"SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK is a "
"per-rank communication buffer capacity, not a model limit; "
"increase it for large prefill/chunked-prefill workloads.",
cfg.deepep_v2_mode,
)
# The resolving view, not the field: `_a2a_backend_overrides` may have
# moved this already (waterfill forces `deepep`).
a2a_now = resolved_view(self).moe_a2a_backend
@@ -20,6 +20,12 @@ from sglang.srt.runtime_context import (
from sglang.srt.state_capturer.base import BaseTopkCapturer
def _is_scattered_a2a_backend() -> bool:
"""Return whether routed tokens are scattered across attention-TP ranks."""
backend = get_moe_a2a_backend()
return backend.is_deepep() or backend.is_deepep_v2()
class RoutedExpertsCapturer(BaseTopkCapturer):
"""Capturer for routed experts with host buffer.
@@ -84,11 +90,8 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
device_topk_size=topk_size + num_fused_shared_experts,
)
# DeepEP a2a path: each attn-TP rank only sees its scattered slice of
# topk_ids. All-gather across attn-TP at capture time so device_cache
# holds the full batch and the existing _get_local_slice / D2H sync
# paths work unchanged. Pre-allocate the gather target.
if get_moe_a2a_backend().is_deepep():
# Rebuild the full token batch before routed-expert readback.
if _is_scattered_a2a_backend():
attn_tp_size = (
get_parallel().attn_tp_size if is_dp_attention_enabled() else 1
)
@@ -102,7 +105,7 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
)
def capture(self, layer_id: int, topk_indices: torch.Tensor):
if get_moe_a2a_backend().is_deepep():
if _is_scattered_a2a_backend():
local_topk = topk_indices
topk_indices = self.gather_buffer[
: local_topk.size(0) * get_parallel().attn_tp_size
@@ -116,10 +119,8 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
can_run_graph: bool,
cuda_graph_batch: Optional[int],
) -> torch.Tensor:
# Under DeepEP, capture() already attn_tp_all_gathered into the head of
# the per-rank buffer, so the local DP rank's data lives at [0:N_local]
# rather than at the global [start_pos:end_pos] offset.
if is_dp_attention_enabled() and not get_moe_a2a_backend().is_deepep():
# Gathered rows start at buffer offset zero on every DP rank.
if is_dp_attention_enabled() and not _is_scattered_a2a_backend():
# GPU->CPU sync would break overlap; operate on CPU directly.
local_start_pos, local_num_tokens = get_dp_local_slice_cpu(
forward_batch, can_run_graph, cuda_graph_batch