minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Revert "[Feature] Add DeepEPv2 (ElasticBuffer) MoE A2A backend" (#35568)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-19 14:14:38 -07:00
committed by GitHub
parent a6bc0532c9
commit 1270204d2c
18 changed files with 21 additions and 2112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/moe/ep_moe_kernels.py
-402
View File
@@ -1221,144 +1221,6 @@ def ep_scatter(
return return
@triton.jit
def _fwd_kernel_ep_scatter_psum_init(
psum_num_recv_tokens_per_expert,
expert_start_loc,
m_indices,
BLOCK_E: tl.constexpr,
):
cur_expert = tl.program_id(0)
cur_end = tl.load(psum_num_recv_tokens_per_expert + cur_expert)
cur_start = tl.load(
psum_num_recv_tokens_per_expert + cur_expert - 1,
mask=cur_expert > 0,
other=0,
)
cur_token_num = cur_end - cur_start
tl.store(expert_start_loc + cur_expert, cur_start)
off_expert = tl.arange(0, BLOCK_E)
for start_m in tl.range(0, cur_token_num, BLOCK_E, num_stages=4):
# cur_token_num need not be a multiple of BLOCK_E; mask the tail block so
# the final partial iteration does not write past this expert's region
# (which is packed right up against the next expert) and corrupt it.
idx = cur_start + start_m + off_expert
tl.store(m_indices + idx, cur_expert, mask=idx < cur_end)
@torch.no_grad()
def ep_scatter_from_psum(
recv_x: torch.Tensor,
recv_x_scale: torch.Tensor,
recv_topk: torch.Tensor,
psum_num_recv_tokens_per_expert: torch.Tensor,
expert_start_loc: torch.Tensor,
output_tensor: torch.Tensor,
output_tensor_scale: torch.Tensor,
m_indices: torch.Tensor,
output_index: torch.Tensor,
scale_ue8m0: bool = False,
):
BLOCK_E = 128
BLOCK_D = 128
num_warps = 8
num_experts = psum_num_recv_tokens_per_expert.shape[0]
hidden_size = recv_x.shape[1]
scale_hidden_size = hidden_size // BLOCK_D
if scale_ue8m0:
scale_hidden_size = ceil_div(scale_hidden_size, 4)
assert m_indices.shape[0] % BLOCK_E == 0
is_fp8 = recv_x_scale is not None and recv_x.dtype != torch.bfloat16
if is_fp8:
assert recv_x_scale.dtype == output_tensor_scale.dtype
assert (
recv_x_scale.shape[1] == output_tensor_scale.shape[1] == scale_hidden_size
)
_fwd_kernel_ep_scatter_psum_init[(num_experts,)](
psum_num_recv_tokens_per_expert,
expert_start_loc,
m_indices,
num_warps=num_warps,
BLOCK_E=BLOCK_E,
)
grid = min(recv_topk.shape[0], 1024 * 8)
_fwd_kernel_ep_scatter_2[(grid,)](
recv_topk.shape[0],
expert_start_loc,
recv_x,
recv_x.stride(0),
recv_x.stride(1),
recv_x_scale,
recv_x_scale.stride(0) if is_fp8 else 0,
recv_x_scale.stride(1) if is_fp8 else 0,
recv_topk,
recv_topk.stride(0),
recv_topk.stride(1),
output_tensor,
output_tensor.stride(0),
output_tensor.stride(1),
output_tensor_scale,
output_tensor_scale.stride(0) if is_fp8 else 0,
output_tensor_scale.stride(1) if is_fp8 else 0,
output_index,
output_index.stride(0),
output_index.stride(1),
topk_num=recv_topk.shape[1],
num_warps=num_warps,
HIDDEN_SIZE=hidden_size,
HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
SCALE_HIDDEN_SIZE=scale_hidden_size,
SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size),
ATOMIC_ADD_SEM=None if not _is_musa else "relaxed",
IS_FP8=is_fp8,
)
return
@triton.jit
def _fwd_kernel_ep_expand_m_indices_init(
psum_num_recv_tokens_per_expert,
m_indices,
BLOCK_E: tl.constexpr,
):
cur_expert = tl.program_id(0)
cur_end = tl.load(psum_num_recv_tokens_per_expert + cur_expert)
prev_end = tl.load(
psum_num_recv_tokens_per_expert + cur_expert - 1,
mask=cur_expert > 0,
other=0,
)
cur_start = ((prev_end + BLOCK_E - 1) // BLOCK_E) * BLOCK_E
aligned_end = ((cur_end + BLOCK_E - 1) // BLOCK_E) * BLOCK_E
off_expert = tl.arange(0, BLOCK_E)
for start_m in tl.range(0, aligned_end - cur_start, BLOCK_E, num_stages=4):
idx = cur_start + start_m + off_expert
tl.store(m_indices + idx, cur_expert, mask=idx < aligned_end)
@torch.no_grad()
def ep_expand_init_m_indices_from_psum(
psum_num_recv_tokens_per_expert: torch.Tensor,
m_indices: torch.Tensor,
):
BLOCK_E = 128
num_warps = 8
num_experts = psum_num_recv_tokens_per_expert.shape[0]
assert m_indices.shape[0] % BLOCK_E == 0
_fwd_kernel_ep_expand_m_indices_init[(num_experts,)](
psum_num_recv_tokens_per_expert,
m_indices,
num_warps=num_warps,
BLOCK_E=BLOCK_E,
)
return
@triton.jit @triton.jit
def _fwd_kernel_ep_gather( def _fwd_kernel_ep_gather(
total_token_num, total_token_num,
@@ -2141,270 +2003,6 @@ def fp8_per_token_to_per_tensor_quant_triton(
) )
# ---------------------------------------------------------------------------
# DeepEP v2 decode masked-GEMM bridge: repack the expanded expert-packed
# dispatch buffer into a regular [E_local, max_m, hidden] slab so DeepGEMM's
# *masked* grouped GEMM can bound compute by per-expert real counts (masked_m)
# instead of the dispatch capacity. All-GPU, static shapes -> cuda-graph safe.
# Expanded psum semantics (DeepEP v2): psum[e] = align(psum[e-1], ALIGN) + count_e,
# so expert e occupies recv rows [align(psum[e-1]) : psum[e]); count_e real tokens.
# Non-expand (contiguous) psum semantics differ: psum[e] is the inclusive prefix
# sum of alignment-PADDED counts, so every psum[e] is a multiple of ALIGN and
# psum[e-1] is expert e's aligned group start (consumed by ep_scatter_from_psum).
# ---------------------------------------------------------------------------
_DEEPEP_V2_REPACK_WORKERS_PER_EXPERT = 64
# recv_x_scale_stride1 carries the scale pack-dim stride so the repack reads
# both row-major (Hopper fp32) and column-major packed UE8M0 (Blackwell int32)
# dispatch-scale layouts correctly.
@triton.jit
def _fwd_kernel_expand_to_masked_slab(
psum_ptr,
recv_x_ptr,
recv_x_stride0,
recv_x_scale_ptr,
recv_x_scale_stride0,
recv_x_scale_stride1,
output_tensor_ptr,
output_tensor_stride0,
output_tensor_scale_ptr,
masked_m_ptr,
overflow_ptr,
MAX_M: tl.constexpr,
ALIGN: tl.constexpr,
HIDDEN: tl.constexpr,
HIDDEN_PAD: tl.constexpr,
SCALE_HIDDEN: tl.constexpr,
SCALE_HIDDEN_PAD: tl.constexpr,
IS_FP8: tl.constexpr,
CHECK_OVERFLOW: tl.constexpr,
NUM_WORKERS: tl.constexpr,
):
# Keep a fixed worker pool per expert and let each worker walk only real rows.
# This avoids launching cdiv(MAX_M, BLOCK_M) programs for a conservative
# max_m when decode traffic contains only a few rows per expert. The grid is
# still static and therefore cuda-graph safe.
e = tl.program_id(0)
worker = tl.program_id(1)
prev_end = tl.load(psum_ptr + e - 1, mask=e > 0, other=0)
start = ((prev_end + ALIGN - 1) // ALIGN) * ALIGN
end = tl.load(psum_ptr + e)
raw_count = end - start
count = tl.minimum(raw_count, MAX_M)
if worker == 0:
tl.store(masked_m_ptr + e, count)
if CHECK_OVERFLOW:
# Eager execution reports an invalid bound instead of truncating.
# Graph replay uses the proven cap * ep_size upper bound and omits
# this host-observable flag and its per-layer reset kernel.
ovf = tl.arange(0, 1)
tl.store(overflow_ptr + ovf, 1, mask=raw_count > MAX_M)
off = tl.arange(0, HIDDEN_PAD)
mask = off < HIDDEN
off_s = tl.arange(0, SCALE_HIDDEN_PAD)
mask_s = off_s < SCALE_HIDDEN
for j in tl.range(worker, count, NUM_WORKERS):
src = (start + j).to(tl.int64)
dst = (e * MAX_M + j).to(tl.int64)
v = tl.load(recv_x_ptr + src * recv_x_stride0 + off, mask=mask)
tl.store(output_tensor_ptr + dst * output_tensor_stride0 + off, v, mask=mask)
if IS_FP8:
vs = tl.load(
recv_x_scale_ptr
+ src * recv_x_scale_stride0
+ off_s * recv_x_scale_stride1,
mask=mask_s,
)
# mn-major write: physical layout [E, SCALE_HIDDEN, MAX_M], element
# (e, s, j). Viewed as [E, MAX_M, SCALE_HIDDEN] this is the mn-major
# TMA-aligned layout deep_gemm wants, so the GEMM-side transpose
# (get_mn_major_tma_aligned_tensor) becomes a no-op.
tl.store(
output_tensor_scale_ptr + e * SCALE_HIDDEN * MAX_M + off_s * MAX_M + j,
vs,
mask=mask_s,
)
@torch.no_grad()
def expand_to_masked_slab(
recv_x: torch.Tensor,
recv_x_scale,
psum_num_recv_tokens_per_expert: torch.Tensor,
num_local_experts: int,
max_m: int,
expert_alignment: int,
):
"""expanded [total, hidden] -> ([E_local, max_m, hidden], [E_local, max_m, sh] or None, masked_m[E_local])."""
hidden = recv_x.shape[1]
is_fp8 = recv_x_scale is not None and recv_x.dtype != torch.bfloat16
output_tensor = torch.empty(
(num_local_experts * max_m, hidden), device=recv_x.device, dtype=recv_x.dtype
)
masked_m = torch.empty(
(num_local_experts,), device=recv_x.device, dtype=torch.int32
)
check_overflow = not torch.cuda.is_current_stream_capturing()
overflow = (
torch.zeros((1,), device=recv_x.device, dtype=torch.int32)
if check_overflow
else masked_m
)
if is_fp8:
sh = recv_x_scale.shape[1]
# mn-major scale: store physically as [E, sh, max_m] (contiguous),
# return a [E, max_m, sh] view with mn-major stride. This matches
# deep_gemm's mn-major TMA-aligned scale layout, so the per-layer
# get_mn_major_tma_aligned_tensor call on the GEMM side is a no-op.
# On Hopper that call still runs and would transpose if the layout ever
# failed to match. On Blackwell (DEEPGEMM_SCALE_UE8M0) it does not:
# _run_masked_gemm takes the packed-ue8m0 branch and consumes this scale
# as-is, so correctness there does depend on this write being mn-major.
output_tensor_scale = torch.empty(
(num_local_experts * sh, max_m),
device=recv_x.device,
dtype=recv_x_scale.dtype,
)
scale_arg = recv_x_scale
scale_s0 = recv_x_scale.stride(0)
scale_s1 = recv_x_scale.stride(1)
else:
sh = 1
output_tensor_scale = None
scale_arg = recv_x
scale_s0 = 0
scale_s1 = 0
num_workers = min(max_m, _DEEPEP_V2_REPACK_WORKERS_PER_EXPERT)
_fwd_kernel_expand_to_masked_slab[(num_local_experts, num_workers)](
psum_num_recv_tokens_per_expert,
recv_x,
recv_x.stride(0),
scale_arg,
scale_s0,
scale_s1,
output_tensor,
output_tensor.stride(0),
output_tensor_scale if is_fp8 else scale_arg,
masked_m,
overflow,
MAX_M=max_m,
ALIGN=expert_alignment,
HIDDEN=hidden,
HIDDEN_PAD=triton.next_power_of_2(hidden),
SCALE_HIDDEN=sh,
SCALE_HIDDEN_PAD=triton.next_power_of_2(sh),
IS_FP8=is_fp8,
CHECK_OVERFLOW=check_overflow,
NUM_WORKERS=num_workers,
num_warps=4,
)
# Outside cuda graph capture, fail fast on slab overflow rather than return a
# silently truncated result. During capture we skip the host read to keep the
# path graph-safe; the eager warmup forward validates representative shapes.
# Safety under graph replay therefore relies on the static upper bound
# max_m = cap * ep_group_size holding: each rank sends at most `cap` tokens
# (enforced by the dispatch-entry assert) and a token contributes at most once
# per local expert, so no expert can exceed max_m. If those invariants change,
# graph replay would NOT fail-fast on overflow — re-validate before relying on it.
if check_overflow and int(overflow.item()) != 0:
raise RuntimeError(
f"DeepEP v2 masked slab overflow: an expert received more than max_m="
f"{max_m} tokens; increase "
f"SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK."
)
output_tensor = output_tensor.view(num_local_experts, max_m, hidden)
if is_fp8:
# physical [E, sh, max_m] -> [E, max_m, sh] view with mn-major stride (no copy)
output_tensor_scale = output_tensor_scale.view(
num_local_experts, sh, max_m
).transpose(1, 2)
return output_tensor, output_tensor_scale, masked_m
@triton.jit
def _fwd_kernel_masked_slab_to_expand(
psum_ptr,
input_tensor_ptr,
input_tensor_stride0,
output_tensor_ptr,
output_tensor_stride0,
weight_ptr,
MAX_M: tl.constexpr,
ALIGN: tl.constexpr,
HIDDEN: tl.constexpr,
HIDDEN_PAD: tl.constexpr,
HAS_W: tl.constexpr,
NUM_WORKERS: tl.constexpr,
):
# Fixed worker pool; see _fwd_kernel_expand_to_masked_slab. cuda-graph safe.
e = tl.program_id(0)
worker = tl.program_id(1)
prev_end = tl.load(psum_ptr + e - 1, mask=e > 0, other=0)
start = ((prev_end + ALIGN - 1) // ALIGN) * ALIGN
end = tl.load(psum_ptr + e)
count = end - start
count = tl.minimum(count, MAX_M)
off = tl.arange(0, HIDDEN_PAD)
mask = off < HIDDEN
for j in tl.range(worker, count, NUM_WORKERS):
src = (e * MAX_M + j).to(tl.int64)
dst = (start + j).to(tl.int64)
v = tl.load(input_tensor_ptr + src * input_tensor_stride0 + off, mask=mask)
if HAS_W:
w = tl.load(weight_ptr + dst)
v = (v.to(tl.float32) * w).to(v.dtype)
tl.store(output_tensor_ptr + dst * output_tensor_stride0 + off, v, mask=mask)
@torch.no_grad()
def masked_slab_to_expand(
input_tensor: torch.Tensor,
psum_num_recv_tokens_per_expert: torch.Tensor,
total_expanded_tokens: int,
expert_alignment: int,
topk_weights=None,
):
"""[E_local, max_m, hidden] masked-GEMM output -> [total, hidden] expanded order.
Only real rows are written; padding rows are uninitialized (the output is
torch.empty) and are never read -- combine consumes only real rows via handle
metadata. When topk_weights is given ([total_expanded], per expanded row), the
top-k weight is fused into the copy so the weighted-combine multiply happens
only on real rows (not the worst-case buffer).
"""
num_local_experts, max_m, hidden = input_tensor.shape
output_tensor = torch.empty(
(total_expanded_tokens, hidden),
device=input_tensor.device,
dtype=input_tensor.dtype,
)
input_tensor2d = input_tensor.view(num_local_experts * max_m, hidden)
has_w = topk_weights is not None
if has_w:
weight_arg = topk_weights.reshape(-1).to(torch.float32).contiguous()
else:
weight_arg = input_tensor2d # dummy, unused
num_workers = min(max_m, _DEEPEP_V2_REPACK_WORKERS_PER_EXPERT)
_fwd_kernel_masked_slab_to_expand[(num_local_experts, num_workers)](
psum_num_recv_tokens_per_expert,
input_tensor2d,
input_tensor2d.stride(0),
output_tensor,
output_tensor.stride(0),
weight_arg,
MAX_M=max_m,
ALIGN=expert_alignment,
HIDDEN=hidden,
HIDDEN_PAD=triton.next_power_of_2(hidden),
HAS_W=has_w,
NUM_WORKERS=num_workers,
num_warps=4,
)
return output_tensor
def moe_permute( def moe_permute(
inputs: torch.Tensor, inputs: torch.Tensor,
topk_ids: torch.Tensor, topk_ids: torch.Tensor,
python/sglang/srt/arg_groups/overrides.py
+1 -6
View File
@@ -2546,7 +2546,7 @@ def _moe_runner_fusion_disable(view: Any) -> dict:
def _a2a_fusion_adjustments(view: Any) -> dict: def _a2a_fusion_adjustments(view: Any) -> dict:
"""A2A-backend-driven shared-experts fusion adjustments, declared at the """A2A-backend-driven shared-experts fusion adjustments, declared at the
legacy write slots in _handle_a2a_moe: Waterfill requires the legacy write slots in _handle_a2a_moe: Waterfill requires the
fusion enabled; FlashInfer and DeepEP v2 A2A require it disabled.""" fusion enabled; FlashInfer A2A requires it disabled."""
if view.moe_a2a_backend in ("deepep", "megamoe") and view.enable_waterfill: if view.moe_a2a_backend in ("deepep", "megamoe") and view.enable_waterfill:
if view.disable_shared_experts_fusion: if view.disable_shared_experts_fusion:
logger.warning( logger.warning(
@@ -2559,10 +2559,6 @@ def _a2a_fusion_adjustments(view: Any) -> dict:
"Flashinfer MoE A2A is enabled. --disable-shared-experts-fusion is automatically set." "Flashinfer MoE A2A is enabled. --disable-shared-experts-fusion is automatically set."
) )
return {"disable_shared_experts_fusion": True} return {"disable_shared_experts_fusion": True}
if view.moe_a2a_backend == "deepep_v2":
# DeepEP v2 has not validated fused shared experts yet; the handler
# rejects an explicit --enforce-shared-experts-fusion.
return {"disable_shared_experts_fusion": True}
return {} return {}
@@ -2571,7 +2567,6 @@ _A2A_EP_SPANNING_BACKENDS = frozenset(
{ {
"megamoe", "megamoe",
"deepep", "deepep",
"deepep_v2",
"mooncake", "mooncake",
"nixl", "nixl",
"ascend_fuseep", "ascend_fuseep",
python/sglang/srt/environ.py
-6
View File
@@ -1014,12 +1014,6 @@ class Envs:
# read by several call sites; do not use in new code. # read by several call sites; do not use in new code.
SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False) SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False)
SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
# DeepEP v2 per-rank communication buffer capacity. This is not a model
# semantic token limit; large prefill/chunked-prefill workloads may need a
# larger value.
SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
# 0 lets DeepEP v2 ElasticBuffer choose the communication SM count.
SGLANG_DEEPEP_V2_NUM_SMS = EnvInt(0)
SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32) SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32)
SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False) SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False)
# Force dynamic Waterfill with runtime EP all-reduce instead of the default # Force dynamic Waterfill with runtime EP all-reduce instead of the default
python/sglang/srt/layers/moe/ep_moe/layer.py
+1 -7
View File
@@ -103,12 +103,7 @@ class DeepEPMoE(FusedMoE):
and quant_config is not None and quant_config is not None
and quant_config.get_name() == "humming" and quant_config.get_name() == "humming"
) )
if get_moe_a2a_backend().is_deepep_v2(): if is_humming:
# deepep_v2 runs on the base FusedMoE forward via its own
# DeepEPv2Dispatcher, so always delegate (never use DeepEPMoE's
# v1-specific dispatch/run_moe_core path).
self.deprecate_flag = True
elif is_humming:
self.deprecate_flag = True self.deprecate_flag = True
elif _use_aiter: elif _use_aiter:
self.deprecate_flag = True self.deprecate_flag = True
@@ -359,7 +354,6 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]):
if ( if (
get_moe_a2a_backend().is_mori() get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_deepep() 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_mooncake()
or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_nixl()
or get_moe_a2a_backend().is_pplx() or get_moe_a2a_backend().is_pplx()
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
-10
View File
@@ -38,7 +38,6 @@ from sglang.srt.layers.moe.token_dispatcher.ascend_tp import (
AscendTPDispatcher, AscendTPDispatcher,
) )
from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher 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.flashinfer import FlashinferDispatcher
from sglang.srt.layers.moe.token_dispatcher.standard import ( from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardDispatcher, StandardDispatcher,
@@ -167,15 +166,6 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
async_finish=True, async_finish=True,
return_recv_hook=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(): elif a2a_backend.is_flashinfer():
return FlashinferDispatcher( return FlashinferDispatcher(
group=get_tp_group().device_group, group=get_tp_group().device_group,
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
+1 -218
View File
@@ -49,10 +49,6 @@ if TYPE_CHECKING:
DeepEPNormalCombineInput, DeepEPNormalCombineInput,
DeepEPNormalDispatchOutput, DeepEPNormalDispatchOutput,
) )
from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import (
DeepEPv2CombineInput,
DeepEPv2DispatchOutput,
)
from sglang.srt.layers.moe.token_dispatcher.standard import ( from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardCombineInput, StandardCombineInput,
StandardDispatchOutput, StandardDispatchOutput,
@@ -210,7 +206,6 @@ class DeepGemmRunnerInput(RunnerInput):
masked_m: Optional[torch.Tensor] = None masked_m: Optional[torch.Tensor] = None
expected_m: Optional[int] = None expected_m: Optional[int] = None
m_indices: Optional[torch.Tensor] = None m_indices: Optional[torch.Tensor] = None
hidden_states_scale_tma_aligned: bool = False
@property @property
def runner_backend(self) -> MoeRunnerBackend: def runner_backend(self) -> MoeRunnerBackend:
@@ -326,10 +321,7 @@ class DeepGemmRunnerCore(MoeRunnerCore):
device=hidden_states_device, device=hidden_states_device,
dtype=torch.bfloat16, dtype=torch.bfloat16,
) )
if ( if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES:
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) hidden_states_scale = tma_align_input_scale(hidden_states_scale)
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig( deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(
@@ -1427,212 +1419,3 @@ def _apply_swiglu_limit(
out = torch.cat([gate, up], dim=-1) out = torch.cat([gate, up], dim=-1)
assert out.shape == (num_tokens, hidden_size_x2) assert out.shape == (num_tokens, hidden_size_x2)
return out 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-GEMM bridge: see expand_to_masked_slab -- bounds compute by
# per-expert masked_m instead of the dispatch capacity, cuda-graph safe.
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,
)
# do_cpu_sync=False -> the recv buffer is worst-case sized. ep_expand_init
# labels each expert's rows up to its 128-aligned end (the contiguous layout
# needs a whole 128-row tile to share one expert id) but never touches the
# tail past the last expert, so pre-fill with -1 to make the GEMM skip it.
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,
)
# ElasticBuffer always populates the handle's per-expert prefix sum, so the
# contiguous path never needs a host-side count list.
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
# Match the legacy deepep_normal adapter (same ep_scatter + grouped GEMM):
# ep_scatter writes only real-token rows and the post-permute ep_gather reads
# them back via output_index, so the alignment padding rows are never consumed
# and need no zero-init -- except under deterministic inference, where pad
# garbage would leak batch-dependent values into the grouped GEMM. The ue8m0
# packed-scale layout always keeps zeros (its in-int32 padding lanes must be 0).
deterministic = get_exec().deterministic.enable_deterministic_inference
buffer_init = torch.zeros if deterministic else torch.empty
input_tensor = buffer_init(
(all_tokens, K), device=hidden_states.device, dtype=hidden_states.dtype
)
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
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 = buffer_init(
(all_tokens, K // 128), device=hidden_states.device, dtype=torch.float32
)
m_indices = buffer_init(all_tokens, device=hidden_states.device, dtype=torch.int32)
output_index = torch.empty_like(topk_ids)
# Contiguous-path alignment contract: this psum comes from ElasticBuffer
# dispatch(do_expand=False, expert_alignment=_EXPERT_ALIGNMENT), and DeepEP
# documents the non-expand psum as the inclusive prefix sum of
# alignment-PADDED per-expert counts (deep_ep/buffers/elastic.py). The
# dispatcher pins that alignment to 128 ==
# get_m_alignment_for_contiguous_layout(), so psum[e-1] is a valid
# 128-aligned group start for the contiguous grouped GEMM. Do NOT re-align
# here: an align_up would silently mask an upstream contract break.
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):
# Masked path: GEMM output is the [E_local, max_m, hidden] slab.
# Repack it back to expanded row order (only real rows written;
# padding left uninitialized and never read) and fold in the
# top-k weights before combine (expanded combine does not
# consume them).
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, so apply them to
# each expert slot before combine. Keep this out-of-place until the
# runner/communication buffer reuse contract is explicitly audited.
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,
)
python/sglang/srt/layers/moe/moe_runner/runner.py
-16
View File
@@ -50,22 +50,6 @@ class MoeRunner:
"--moe-runner-backend hpc_ops for this model." "--moe-runner-backend hpc_ops for this model."
) )
# deepep_v2 only registers permute adapters for the deep_gemm runner.
# --moe-runner-backend is validated at server start, but the runner is
# picked per layer by the quant method and several of them hard-select
# Triton regardless (blockwise_int8, moe_wna16, w8a8_*, modelopt,
# unquant, ...). Without this, such a model reaches the permute
# registry and dies on a bare assert inside the MoE forward, after the
# weights are already loaded.
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 self.fused_func = None
if runner_backend.is_triton(): if runner_backend.is_triton():
python/sglang/srt/layers/moe/token_dispatcher/__init__.py
-8
View File
@@ -21,11 +21,6 @@ from sglang.srt.layers.moe.token_dispatcher.deepep import (
DeepEPNormalCombineInput, DeepEPNormalCombineInput,
DeepEPNormalDispatchOutput, DeepEPNormalDispatchOutput,
) )
from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import (
DeepEPv2CombineInput,
DeepEPv2Dispatcher,
DeepEPv2DispatchOutput,
)
from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( from sglang.srt.layers.moe.token_dispatcher.flashinfer import (
FlashinferDispatcher, FlashinferDispatcher,
FlashinferDispatchOutput, FlashinferDispatchOutput,
@@ -77,9 +72,6 @@ __all__ = [
"MoriEPLLDispatchOutput", "MoriEPLLDispatchOutput",
"MoriEPLLCombineInput", "MoriEPLLCombineInput",
"MoriEPDispatcher", "MoriEPDispatcher",
"DeepEPv2Dispatcher",
"DeepEPv2DispatchOutput",
"DeepEPv2CombineInput",
"NixlEPCombineInput", "NixlEPCombineInput",
"NixlEPDispatchOutput", "NixlEPDispatchOutput",
"NixlEPDispatcher", "NixlEPDispatcher",
python/sglang/srt/layers/moe/token_dispatcher/base.py
-19
View File
@@ -27,8 +27,6 @@ if TYPE_CHECKING:
DeepEPLLDispatchOutput, DeepEPLLDispatchOutput,
DeepEPNormalCombineInput, DeepEPNormalCombineInput,
DeepEPNormalDispatchOutput, DeepEPNormalDispatchOutput,
DeepEPv2CombineInput,
DeepEPv2DispatchOutput,
FlashinferCombineInput, FlashinferCombineInput,
FlashinferDispatchOutput, FlashinferDispatchOutput,
StandardCombineInput, StandardCombineInput,
@@ -167,12 +165,6 @@ class DispatchOutputChecker:
) -> TypeGuard[FlashinferDispatchOutput]: ) -> TypeGuard[FlashinferDispatchOutput]:
return dispatch_output.format.is_flashinfer() 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): class DispatchOutputFormat(Enum):
@@ -180,7 +172,6 @@ class DispatchOutputFormat(Enum):
DEEPEP_NORMAL = "deepep_normal" DEEPEP_NORMAL = "deepep_normal"
DEEPEP_LL = "deepep_ll" DEEPEP_LL = "deepep_ll"
FLASHINFER = "flashinfer" FLASHINFER = "flashinfer"
DEEPEP_V2 = "deepep_v2"
ASCEND_TP = "ascend_tp" ASCEND_TP = "ascend_tp"
def is_standard(self) -> bool: def is_standard(self) -> bool:
@@ -204,9 +195,6 @@ class DispatchOutputFormat(Enum):
def is_flashinfer(self) -> bool: def is_flashinfer(self) -> bool:
return self == DispatchOutputFormat.FLASHINFER return self == DispatchOutputFormat.FLASHINFER
def is_deepep_v2(self) -> bool:
return self == DispatchOutputFormat.DEEPEP_V2
@runtime_checkable @runtime_checkable
class DispatchOutput(Protocol): class DispatchOutput(Protocol):
@@ -261,19 +249,12 @@ class CombineInputChecker:
) -> TypeGuard[FlashinferCombineInput]: ) -> TypeGuard[FlashinferCombineInput]:
return combine_input.format == CombineInputFormat.FLASHINFER 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): class CombineInputFormat(Enum):
STANDARD = "standard" STANDARD = "standard"
DEEPEP_NORMAL = "deepep_normal" DEEPEP_NORMAL = "deepep_normal"
DEEPEP_LL = "deepep_ll" DEEPEP_LL = "deepep_ll"
FLASHINFER = "flashinfer" FLASHINFER = "flashinfer"
DEEPEP_V2 = "deepep_v2"
ASCEND_TP = "ascend_tp" ASCEND_TP = "ascend_tp"
python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py
-548
View File
@@ -1,548 +0,0 @@
from __future__ import annotations
import logging
import os
from typing import NamedTuple, Optional, Tuple
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
# Per-expert row alignment requested from ElasticBuffer. Must equal DeepGEMM's
# get_m_alignment_for_contiguous_layout() so the non-expand psum doubles as a
# valid group-start table for the contiguous grouped GEMM.
_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:
# direct/hybrid is a communication-topology knob resolved from ServerArgs.
# Callers without a running server (synthetic/unit tests) pass
# allow_hybrid_mode to DeepEPv2Buffer.get_buffer instead (get_server_args()
# raises when the process-wide ServerArgs is not set).
from sglang.srt.runtime_context import get_server_args
return get_server_args().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:
_buffer: Optional[ElasticBuffer] = None
_buffer_key: Optional[Tuple] = None
@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()
key = (
id(group),
hidden_size,
router_topk,
num_max_dispatch_tokens_per_rank,
use_fp8_dispatch,
allow_hybrid_mode,
dist.get_world_size(group),
)
if cls._buffer is not None and cls._buffer_key == key:
return cls._buffer
if cls._buffer is not None:
cls.destroy()
# DeepEP reuses the torch process group's internal NCCL communicator
# when EP_REUSE_NCCL_COMM=1 (its default). That path requires the group
# to be device-bound at init_process_group time (eager comm init),
# which SGLang's shared init does not do -- reusing then reads an
# uninitialized communicator and ElasticBuffer sizing segfaults in
# ncclTeamWorld. Default to letting DeepEP create its own communicator
# (it binds to the already-set current device); setdefault keeps any
# explicit user override.
os.environ.setdefault("EP_REUSE_NCCL_COMM", "0")
cls._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,
)
cls._buffer_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,
cls._buffer.num_bytes,
)
return cls._buffer
@classmethod
def destroy(cls) -> None:
cls._buffer = None
cls._buffer_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, # deepep_v2 always dispatches FP8 activations + scales
)
def _resolve_num_sms_qps(self, buffer: ElasticBuffer) -> Tuple[int, int]:
# num_sms/num_qps are NOT auto-resolved by ElasticBuffer when left at 0
# (0 means "0 SMs / 0 QPs"). Resolve both from the theoretical helpers
# (matches the DeepEP elastic test harness): num_sms from
# SGLANG_DEEPEP_V2_NUM_SMS or get_theoretical_num_sms, and num_qps always
# from get_theoretical_num_qps(num_sms). Multi-node RDMA dispatch needs
# the real QPs; single-node NVLink is unaffected by the extra QPs.
# Both helpers are host-only: get_theoretical_num_sms is cached in DeepEP
# for the fixed inputs here (and first runs during eager warmup), and
# get_theoretical_num_qps is plain arithmetic. On the CUDA-graph decode
# path this costs no device work, so it is capture-safe.
num_sms = envs.SGLANG_DEEPEP_V2_NUM_SMS.get()
if num_sms == 0:
num_sms = buffer.get_theoretical_num_sms(self.num_experts, self.router_topk)
num_qps = buffer.get_theoretical_num_qps(num_sms)
return num_sms, num_qps
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."
)
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:
# Guard-first (before the import check) so misuse is reportable without
# DeepEP installed.
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)
# DeepEP v2's native expanded layout is profitable for decode-like DeepGEMM
# FP8 workloads but regresses prefill-like ones, so layout is chosen by
# inference PHASE, independently of the comm mode (direct/hybrid is a topology
# knob fixed at server init): decode (non-extend) -> native expanded layout;
# prefill/extend -> non-expanded contiguous layout. This decouples the
# masked-GEMM + CUDA-graph decode fast path from the comm mode, so it is
# available under multi-node `hybrid` too.
use_expand_layout = not get_is_extend_in_batch()
# masked GEMM is built from the expanded layout (expand_to_masked_slab), so
# masked <=> expanded. async dispatch (cpu_sync=False) gives a static
# capturable recv shape; the masked GEMM bounds compute by masked_m, so the
# full (safe) cap costs no extra GEMM.
use_masked = use_expand_layout
# ElasticBuffer requires >=1 token per rank on the non-masked (contiguous /
# extend) path: DeepEP's own ElasticBuffer test pads every rank to
# `max(1, num_tokens)` (tests/elastic/test_ep.py). An idle DP rank with 0
# tokens never fires the dispatch notify / scale-up-reduction warps, so no
# rank's recv count becomes "ready" and the do_cpu_sync CPU readback times
# out ("Dispatch CPU wait", buffer.hpp:1032). Pad an empty local batch to a
# single dummy token; combine() slices that row back off so this rank's
# output is empty again. The masked decode path tolerates empty
# (do_cpu_sync=False), so it is left untouched.
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]))
# A token's top-k experts must be DISTINCT valid ids: duplicates (e.g.
# all-zero -> expert 0 repeated) fault the dispatch kernel. Route the
# dummy to experts [0, 1, ..., topk-1] with zero weights so it
# contributes nothing even before combine() slices it off.
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:
# Follow the hardware scale format (DEEPGEMM_SCALE_UE8M0 via
# scale_format.ue8m0). Hopper (False): plain row-major fp32 scale,
# and _run_masked_gemm does its own e8m0/tma-major alignment.
# Blackwell (True): pre-quantize the activation against a col-major
# UE8M0 scale so it already matches the layout the masked GEMM
# consumes.
_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
# num_max_tokens_per_rank is a COLLECTIVE dispatch arg (ElasticBuffer
# requires the same value on all ranks). Keep it at the fixed buffer cap
# (class-level, cross-rank-consistent), matching DeepEP LL which uses a
# fixed _num_max_dispatch_tokens_per_rank rather than a per-forward token
# count. Do NOT derive it from the local hidden_states.shape[0]: under
# ragged DP load (or TP attention) the ranks would disagree on this
# collective arg.
num_max_tokens = self.num_max_dispatch_tokens_per_rank
# Non-masked (extend/prefill) path reads exact per-expert
# recv
# counts on the CPU, so it must wait for the GPU to finish writing them
# (matches the DeepEP elastic test which passes do_cpu_sync=1). Leaving
# it None lets the CPU read zeros on multi-node (scaleup) dispatch. Only
# the masked decode path keeps do_cpu_sync=False for graph capturability.
do_cpu_sync_val = True
if use_masked:
do_cpu_sync_val = False
buffer = self._get_buffer()
_num_sms, _num_qps = self._resolve_num_sms_qps(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=_num_sms,
num_qps=_num_qps,
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]
# event.current_stream_wait() is a GPU stream dependency (not a CPU
# sync); the do_cpu_sync=False masked decode path stays CUDA-graph
# capturable.
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 layout already has one row per local expert slot. There is
# no recv_topk_idx tensor in this native layout; combine uses handle
# metadata and expects top-k weights to be applied before combine.
# Avoid exact-count CPU reads that are only needed by non-expanded
# slicing/scatter paths.
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]
# Elastic dispatch epilogue already converts global expert ids to local
# expert ids and marks non-local choices as -1. Keep it on-GPU and avoid
# an unnecessary max().item() synchronization in the decode path.
local_topk_ids = recv_topk_idx
expected_m = 0
masked_max_m = 0
total_expanded = 0
if use_masked:
# expected_m: average tokens-per-expert across the EP group, a
# per-rank-local schedule hint for the masked GEMM (NOT a hard bound;
# the real per-expert bound is masked_m on the GPU). Derive it from
# the actual local batch * EP group size, matching DeepEP LL
# (deepep.py dispatch_a uses hidden_states.shape[0]). Per-rank-local,
# so the actual batch is safe here even under ragged DP. group size
# == ep world size == num_experts // num_local_experts.
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,
)
# Size the masked slab to the FIXED worst case cap * ep_group_size,
# matching DeepEP LL's fixed buffer. A local expert receives the sum
# over all ranks of the tokens routed to it; each rank sends at most
# `cap` tokens (enforced by the dispatch-entry assert), so the count
# is bounded by cap * ep_group_size regardless of DP padding mode
# (MAX_LEN / SUM_LEN / skewed). Using the local batch for the slab
# would be unsafe: under skewed SUM_LEN decode another rank's larger
# batch could overflow this rank's slab.
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:
# Guard-first (before any DeepEP work) so misuse is reportable without
# DeepEP installed.
if self._handle is None:
raise RuntimeError(
"DeepEP v2 combine called without a valid dispatch handle"
)
# The handle is single-use: release it whether combine succeeds or
# raises, so a failed step cannot poison the next dispatch.
try:
buffer = self._get_buffer()
_num_sms, _num_qps = self._resolve_num_sms_qps(buffer)
combined_x, _, event = buffer.combine(
combine_input.hidden_states,
handle=self._handle,
topk_weights=combine_input.topk_weights,
num_sms=_num_sms,
num_qps=_num_qps,
)
# Stream dependency, not a CPU sync (graph-safe).
if event.event is not None:
event.current_stream_wait()
if self._pad_empty_combine:
# Drop the dummy token padded onto an empty local batch in
# dispatch so this idle rank's combined output is empty again.
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,
)
# This backend intentionally exposes only single-shot dispatch()/combine():
# TBO/SBO are rejected at server start, and our overlap PoC showed the naive
# two-phase split cannot overlap anyway (ElasticBuffer.dispatch is
# host-blocking); a split API will land together with real TBO support.
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)
python/sglang/srt/layers/moe/utils.py
+3 -43
View File
@@ -4,7 +4,7 @@ import logging
import os import os
from contextlib import contextmanager from contextlib import contextmanager
from enum import Enum, IntEnum from enum import Enum, IntEnum
from typing import TYPE_CHECKING, NamedTuple from typing import TYPE_CHECKING
import torch import torch
@@ -12,13 +12,7 @@ from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
is_dp_attention_enabled, is_dp_attention_enabled,
) )
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import get_exec, get_flags, get_forward, get_parallel
get_exec,
get_flags,
get_forward,
get_parallel,
get_server_args,
)
from sglang.srt.utils import is_cuda, is_npu from sglang.srt.utils import is_cuda, is_npu
_is_npu = is_npu() _is_npu = is_npu()
@@ -26,6 +20,7 @@ _is_npu = is_npu()
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
from sglang.srt.runtime_context import get_server_args
from sglang.srt.utils.common import log_info_on_rank0 from sglang.srt.utils.common import log_info_on_rank0
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -42,7 +37,6 @@ class MoeA2ABackend(Enum):
ASCEND_TP = "ascend_tp" ASCEND_TP = "ascend_tp"
FLASHINFER = "flashinfer" FLASHINFER = "flashinfer"
MEGAMOE = "megamoe" MEGAMOE = "megamoe"
DEEPEP_V2 = "deepep_v2"
PPLX = "pplx" PPLX = "pplx"
CUSTOMIZED = "customized" CUSTOMIZED = "customized"
@@ -82,9 +76,6 @@ class MoeA2ABackend(Enum):
def is_megamoe(self): def is_megamoe(self):
return self == MoeA2ABackend.MEGAMOE return self == MoeA2ABackend.MEGAMOE
def is_deepep_v2(self):
return self == MoeA2ABackend.DEEPEP_V2
def is_pplx(self): def is_pplx(self):
return self == MoeA2ABackend.PPLX return self == MoeA2ABackend.PPLX
@@ -184,20 +175,6 @@ class MoeRunnerBackend(Enum):
return self == MoeRunnerBackend.AITER return self == MoeRunnerBackend.AITER
class DeepEPv2Fp8ScaleFormat(NamedTuple):
"""
Layout of the FP8 activation scales DeepEP v2 dispatches to DeepGEMM.
Both fields come from the DeepGEMM JIT configuration and therefore vary by
HARDWARE, not by runner: Hopper wants row-major fp32, Blackwell wants
column-major packed UE8M0. Resolving them here keeps the dispatcher from
importing deep_gemm_wrapper and reading JIT flags itself.
"""
tma_aligned: bool
ue8m0: bool
class DeepEPMode(Enum): class DeepEPMode(Enum):
NORMAL = "normal" NORMAL = "normal"
@@ -331,23 +308,6 @@ def get_ascend_dispatcher_output_dtype(dispatcher):
return DispatcherOutputDtype.BF16 return DispatcherOutputDtype.BF16
def get_deepep_v2_fp8_scale_format() -> DeepEPv2Fp8ScaleFormat:
"""Resolve the FP8 scale layout DeepEP v2 must pre-quantize into.
deepep_v2 dispatches FP8 activations plus scales, which only the deep_gemm
runner consumes; MoeRunner rejects any other runner for this backend.
"""
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(server_args: ServerArgs): def initialize_moe_config(server_args: ServerArgs):
moe = get_flags().moe moe = get_flags().moe
moe.a2a_backend = MoeA2ABackend(server_args.moe_a2a_backend) moe.a2a_backend = MoeA2ABackend(server_args.moe_a2a_backend)
python/sglang/srt/models/deepseek_v2.py
+5 -12
View File
@@ -732,7 +732,6 @@ class DeepseekV2MoE(nn.Module):
or get_moe_a2a_backend().is_ascend_fuseep() or get_moe_a2a_backend().is_ascend_fuseep()
or get_moe_a2a_backend().is_flashinfer() or get_moe_a2a_backend().is_flashinfer()
or get_moe_a2a_backend().is_megamoe() or get_moe_a2a_backend().is_megamoe()
or get_moe_a2a_backend().is_deepep_v2()
or should_use_flashinfer_cutlass_moe_fp4_allgather() or should_use_flashinfer_cutlass_moe_fp4_allgather()
or envs.SGLANG_SHARED_EXPERT_TP1.get() or envs.SGLANG_SHARED_EXPERT_TP1.get()
) )
@@ -812,7 +811,6 @@ class DeepseekV2MoE(nn.Module):
or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_nixl()
or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep() 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 # TODO: we will support tp < ep in the future
self.ep_size = get_parallel().moe_ep_size self.ep_size = get_parallel().moe_ep_size
@@ -835,7 +833,6 @@ class DeepseekV2MoE(nn.Module):
or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep() or get_moe_a2a_backend().is_ascend_fuseep()
or get_moe_a2a_backend().is_flashinfer() 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() self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo()
# SGLANG_OPT_MOE_QUANT_ONCE eligibility, resolved lazily on first # SGLANG_OPT_MOE_QUANT_ONCE eligibility, resolved lazily on first
@@ -2707,10 +2704,7 @@ class DeepseekV2Model(nn.Module):
for i in range(len(self.layers)): for i in range(len(self.layers)):
if isinstance(self.layers[i].mlp, DeepseekV2MoE): if isinstance(self.layers[i].mlp, DeepseekV2MoE):
# tp_size = get_parallel().tp_size # tp_size = get_parallel().tp_size
is_a2a_moe = ( is_a2a_moe = is_deepep_class_backend()
is_deepep_class_backend()
or get_moe_a2a_backend().is_deepep_v2()
)
tp_size = 1 if is_a2a_moe else get_parallel().tp_size tp_size = 1 if is_a2a_moe else get_parallel().tp_size
intermediate_size = ( intermediate_size = (
config.moe_intermediate_size * config.n_shared_experts config.moe_intermediate_size * config.n_shared_experts
@@ -2730,11 +2724,10 @@ class DeepseekV2Model(nn.Module):
) )
) )
self.layers_to_capture = [] self.layers_to_capture = []
self.enable_a2a_moe = ( if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
get_moe_a2a_backend().is_deepep() self.enable_a2a_moe = True
or get_moe_a2a_backend().is_mooncake() else:
or get_moe_a2a_backend().is_deepep_v2() self.enable_a2a_moe = False
)
# llama_4_scaling: for supporting Mistral-Large-3 model # llama_4_scaling: for supporting Mistral-Large-3 model
self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None) self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
python/sglang/srt/server_args.py
-101
View File
@@ -280,7 +280,6 @@ MOE_A2A_BACKEND_CHOICES = [
"ascend_fuseep", "ascend_fuseep",
"flashinfer", "flashinfer",
"megamoe", "megamoe",
"deepep_v2",
"pplx", "pplx",
"ascend_tp", "ascend_tp",
] ]
@@ -2360,8 +2359,6 @@ class ServerArgs:
"ascend_fuseep", "ascend_fuseep",
"flashinfer", "flashinfer",
"megamoe", "megamoe",
"deepep_v2",
"ascend_tp",
"pplx", "pplx",
], ],
Arg( Arg(
@@ -2371,15 +2368,6 @@ class ServerArgs:
), ),
NS("exec.moe"), NS("exec.moe"),
] = "none" ] = "none"
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[ moe_runner_backend: A[
str, str,
Arg( Arg(
@@ -6982,95 +6970,6 @@ class ServerArgs:
self.cuda_graph_config.decode.backend = Backend.DISABLED self.cuda_graph_config.decode.backend = Backend.DISABLED
self.cuda_graph_config.prefill.backend = Backend.DISABLED self.cuda_graph_config.prefill.backend = Backend.DISABLED
if a2a_backend == "deepep_v2":
if self.moe_runner_backend == "auto":
# The generic auto -> runner resolution above only fires for
# moe_a2a_backend "none", so deepep_v2 would otherwise reach the
# check below still holding "auto" and fail. deepep_v2 dispatches
# FP8 activations plus scales, which only deep_gemm consumes.
self.moe_runner_backend = "deep_gemm"
logger.warning(
"DeepEP v2 MoE: resolved --moe-runner-backend auto -> deep_gemm."
)
# Validate the FINAL resolved runner, not the raw field. A model
# declaration (e.g. mxfp8 + auto -> flashinfer_trtllm) is
# materialized after this handler, so self.moe_runner_backend set
# above is not necessarily what the runtime will use. resolved_view
# reflects those pending declarations: validate and drive the graph
# decision off it, so an unsupported resolved runner fails fast here
# instead of being silently restored at materialize time.
resolved_runner = resolved_view(self).moe_runner_backend
if 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 self.enable_two_batch_overlap or self.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 self.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 capacity pre-check: the ElasticBuffer per-rank capacity
# must cover the largest extend forward, which is bounded by the
# chunked prefill budget. self.chunked_prefill_size is already the
# per-rank value here (_handle_data_parallelism divides the CLI
# value by dp_size under DP attention and runs before this
# handler). Without this check the server boots and only fails at
# the first full prefill chunk (the dispatcher's runtime capacity
# guard), which small smoke traffic may never trigger. Decode does
# not need a boot check: with CUDA graphs the padded capture batch
# goes through the same runtime guard during startup, and without
# graphs the guard still fails fast at runtime. Mirrors the MoRI and
# pplx chunk checks later in this handler, and the CuteDSL
# token-budget check in its own __post_init__ slot.
if (
self.chunked_prefill_size
and self.chunked_prefill_size > 0
and (self.disaggregation_mode != "decode")
):
deepep_v2_cap = (
envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get()
)
if self.chunked_prefill_size > deepep_v2_cap:
raise ValueError(
"DeepEP v2 MoE: the per-rank chunked prefill budget "
f"({self.chunked_prefill_size} tokens; the CLI "
"--chunked-prefill-size is divided by dp_size under DP "
"attention) exceeds the per-rank dispatch buffer "
"capacity SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_"
f"RANK={deepep_v2_cap}. Raise the env (it sizes the "
"communication buffer) or lower --chunked-prefill-size."
)
# The decode graph stays enabled under ANY comm mode (direct or
# hybrid): the masked layout is chosen per batch by inference phase
# (decode), not by the comm mode, giving static shapes with no host
# readback. The prefill/extend contiguous path reads exact per-expert
# counts back on the host, so it is never capturable.
self.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[{self.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.",
self.deepep_v2_mode,
)
if ( if (
self.moe_a2a_backend == "none" and is_npu() self.moe_a2a_backend == "none" and is_npu()
) or self.moe_a2a_backend == "ascend_tp": ) or self.moe_a2a_backend == "ascend_tp":
python/sglang/srt/state_capturer/routed_experts.py
+10 -20
View File
@@ -20,15 +20,6 @@ from sglang.srt.runtime_context import (
from sglang.srt.state_capturer.base import BaseTopkCapturer from sglang.srt.state_capturer.base import BaseTopkCapturer
def _is_scattered_a2a_backend() -> bool:
"""True for a2a backends whose MoE layer sees only this attn-TP rank's
slice of topk_ids (see the gather in capture()). DeepEP v2 shares legacy
DeepEP's token topology; classifying it as a TP-MoE backend would make
dp_rank > 0 read unwritten buffer rows."""
backend = get_moe_a2a_backend()
return backend.is_deepep() or backend.is_deepep_v2()
class RoutedExpertsCapturer(BaseTopkCapturer): class RoutedExpertsCapturer(BaseTopkCapturer):
"""Capturer for routed experts with host buffer. """Capturer for routed experts with host buffer.
@@ -93,11 +84,11 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
device_topk_size=topk_size + num_fused_shared_experts, device_topk_size=topk_size + num_fused_shared_experts,
) )
# DeepEP-class a2a path: each attn-TP rank only sees its scattered # DeepEP a2a path: each attn-TP rank only sees its scattered slice of
# slice of topk_ids. All-gather across attn-TP at capture time so # topk_ids. All-gather across attn-TP at capture time so device_cache
# device_cache holds the full batch and the existing _get_local_slice / # holds the full batch and the existing _get_local_slice / D2H sync
# D2H sync paths work unchanged. Pre-allocate the gather target. # paths work unchanged. Pre-allocate the gather target.
if _is_scattered_a2a_backend(): if get_moe_a2a_backend().is_deepep():
attn_tp_size = ( attn_tp_size = (
get_parallel().attn_tp_size if is_dp_attention_enabled() else 1 get_parallel().attn_tp_size if is_dp_attention_enabled() else 1
) )
@@ -111,7 +102,7 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
) )
def capture(self, layer_id: int, topk_indices: torch.Tensor): def capture(self, layer_id: int, topk_indices: torch.Tensor):
if _is_scattered_a2a_backend(): if get_moe_a2a_backend().is_deepep():
local_topk = topk_indices local_topk = topk_indices
topk_indices = self.gather_buffer[ topk_indices = self.gather_buffer[
: local_topk.size(0) * get_parallel().attn_tp_size : local_topk.size(0) * get_parallel().attn_tp_size
@@ -125,11 +116,10 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
can_run_graph: bool, can_run_graph: bool,
cuda_graph_batch: Optional[int], cuda_graph_batch: Optional[int],
) -> torch.Tensor: ) -> torch.Tensor:
# Under DeepEP-class backends, capture() already attn_tp_all_gathered # Under DeepEP, capture() already attn_tp_all_gathered into the head of
# into the head of the per-rank buffer, so the local DP rank's data # the per-rank buffer, so the local DP rank's data lives at [0:N_local]
# lives at [0:N_local] rather than at the global [start_pos:end_pos] # rather than at the global [start_pos:end_pos] offset.
# offset. if is_dp_attention_enabled() and not get_moe_a2a_backend().is_deepep():
if is_dp_attention_enabled() and not _is_scattered_a2a_backend():
# GPU->CPU sync would break overlap; operate on CPU directly. # GPU->CPU sync would break overlap; operate on CPU directly.
local_start_pos, local_num_tokens = get_dp_local_slice_cpu( local_start_pos, local_num_tokens = get_dp_local_slice_cpu(
forward_batch, can_run_graph, cuda_graph_batch forward_batch, can_run_graph, cuda_graph_batch
test/registered/ep/test_routed_experts_dp_readback.py
-207
View File
@@ -1,207 +0,0 @@
"""DP>1 readback of routed experts over DeepEP-class a2a backends.
With DP attention + a DeepEP-class a2a backend, the MoE layer sees only the
attention rank's DP-local tokens, so RoutedExpertsCapturer must gather at
capture time and read back from the buffer head. If the backend is not
recognized, requests owned by dp_rank > 0 read unwritten buffer rows and
silently return garbage expert ids (dp_rank 0 sits at offset 0 and looks
correct, which is why a DP>1 test is required).
Oracle: solo-vs-concurrent consistency. A request served alone is correct
even on a misclassifying tree (with the other rank empty, the global offset
degenerates to 0), so its per-token expert sets form a valid baseline. The
same prompts served concurrently must reproduce those sets; a misclassified
backend instead reads whatever the offset region holds (often well-formed
rows belonging to other tokens or graph warmup, which per-row validity
checks cannot catch). Radix cache is disabled so the concurrent phase cannot
serve cached prefix rows written by the solo phase.
Uses a dummy-weight single-layer 24-expert DeepSeek-V3 so each server boots
in seconds (same pattern as test_deepseek_v3_cutedsl_4gpu.py); generation
quality is irrelevant — only the capture/readback plumbing is under test.
"""
import concurrent.futures
import json
import os
import unittest
import numpy as np
import pybase64
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=900, stage="base-c", runner_config="deepep-8-gpu-h200")
_MODEL = os.environ.get("SGLANG_ROUTED_EXPERTS_TEST_MODEL", "deepseek-ai/DeepSeek-V3")
_NUM_EXPERTS = 24
_NUM_LAYERS = 1
_TOPK = 8 # DeepSeek-V3 num_experts_per_tok
_DUMMY_WEIGHT_ENV = {
# Dummy random weights legitimately produce NaN logits; sanitize instead
# of crashing (same rationale as test_deepseek_v3_cutedsl_4gpu.py).
"SGLANG_ENABLE_ASYNC_ASSERT": "0",
"SGLANG_SANITIZE_NAN_LOGITS": "1",
"SGLANG_CUDA_COREDUMP": "0",
"CUDA_ENABLE_COREDUMP_ON_EXCEPTION": "0",
"SGLANG_CUDA_COREDUMP_BEFORE_CRASH": "0",
}
def _deep_ep_has(attr: str) -> bool:
try:
import deep_ep # noqa: F401
except ImportError:
return False
return hasattr(deep_ep, attr)
class _ReadbackMixin:
backend_args: list
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--load-format",
"dummy",
"--json-model-override-args",
json.dumps(
{
"num_hidden_layers": _NUM_LAYERS,
"first_k_dense_replace": 0,
"n_routed_experts": _NUM_EXPERTS,
}
),
"--tp",
"2",
"--dp",
"2",
"--ep",
"2",
"--enable-dp-attention",
"--enable-return-routed-experts",
"--disable-cuda-graph",
"--disable-radix-cache",
"--mem-fraction-static",
"0.5",
*cls.backend_args,
]
cls.process = popen_launch_server(
_MODEL,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
env={
**os.environ,
**_DUMMY_WEIGHT_ENV,
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
"SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
},
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def _one_request(self, i: int):
resp = requests.post(
self.base_url + "/generate",
json={
"text": f"{self._WORDS[i]} is item number {i}. Describe it in detail.",
"sampling_params": {"max_new_tokens": 24, "temperature": 0},
"return_routed_experts": True,
},
timeout=300,
)
self.assertEqual(resp.status_code, 200)
meta = resp.json()["meta_info"]
self.assertIn("routed_experts", meta)
arr = np.frombuffer(pybase64.b64decode(meta["routed_experts"]), dtype=np.int32)
self.assertEqual(
arr.size % (_NUM_LAYERS * _TOPK),
0,
f"req{i}: payload size {arr.size} not a multiple of layers*topk",
)
rows = arr.reshape(-1, _NUM_LAYERS, _TOPK)
self.assertGreater(rows.shape[0], 0)
self.assertTrue(
bool(((rows >= 0) & (rows < _NUM_EXPERTS)).all()),
f"req{i}: expert id out of range [{rows.min()}, {rows.max()}]",
)
return rows
_WORDS = ["Alpha", "Bravo", "Charlie", "Delta", "Echo", "Foxtrot"]
_N_REQ = 6
def test_dp2_readback(self):
# Phase 1 — solo baselines: sequential requests leave the other DP
# rank empty, the global offset degenerates to 0, and the readback is
# correct even when the backend is misclassified.
solo = [self._one_request(i) for i in range(self._N_REQ)]
# Phase 2 — the same prompts concurrently: joint forward batches give
# dp_rank > 0 requests a non-zero global offset, which is exactly the
# path a misclassified backend gets wrong.
with concurrent.futures.ThreadPoolExecutor(max_workers=self._N_REQ) as ex:
conc = list(ex.map(self._one_request, range(self._N_REQ)))
for i in range(self._N_REQ):
a, b = solo[i], conc[i]
n = min(a.shape[0], b.shape[0])
total = match = 0
for t in range(n):
for layer in range(_NUM_LAYERS):
total += 1
if set(a[t, layer].tolist()) == set(b[t, layer].tolist()):
match += 1
frac = match / max(1, total)
self.assertGreaterEqual(
frac,
0.9,
f"req{i}: only {frac:.1%} of per-token expert sets match the "
"solo baseline — the capturer is reading rows that belong to "
"other tokens (DeepEP-class backend misclassification)",
)
@unittest.skipUnless(_deep_ep_has("Buffer"), "DeepEP (v1 Buffer) not installed")
class TestRoutedExpertsReadbackDeepEP(_ReadbackMixin, CustomTestCase):
backend_args = [
"--moe-a2a-backend",
"deepep",
"--deepep-mode",
"low_latency",
"--deepep-dispatcher-output-dtype",
"fp8",
"--moe-runner-backend",
"deep_gemm",
]
@unittest.skipUnless(
_deep_ep_has("ElasticBuffer"), "DeepEP v2 (ElasticBuffer) not installed"
)
class TestRoutedExpertsReadbackDeepEPv2(_ReadbackMixin, CustomTestCase):
backend_args = [
"--moe-a2a-backend",
"deepep_v2",
"--deepep-v2-mode",
"direct",
"--moe-runner-backend",
"deep_gemm",
]
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/moe/test_deepep_v2_masked_slab.py
-277
View File
@@ -1,277 +0,0 @@
"""Unit tests for the DeepEP v2 masked-masked_x repack Triton kernels.
Covers the corner cases flagged in review: empty expert, single hot expert,
per-expert count near / over max_m (overflow -> fail-fast, not silent truncation),
top-k weight fusion on real rows only, expanded<->masked_x round-trip layout, and the
fp8 activation+scale path.
"""
import unittest
import torch
from sglang.kernels.ops.moe.ep_moe_kernels import (
expand_to_masked_slab,
masked_slab_to_expand,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")
DEVICE = "cuda"
def _build_layout(counts, align, hidden, dtype, with_scale=False, scale_hidden=4):
"""Build (recv_x, recv_x_scale, psum, starts, total) for given per-expert counts.
Mirrors the DeepEP v2 expanded layout: expert e occupies rows
[align(psum[e-1]), psum[e]) with psum[-1] == 0.
"""
starts, psum = [], []
prev_end = 0
for c in counts:
start = ((prev_end + align - 1) // align) * align
end = start + c
starts.append(start)
psum.append(end)
prev_end = end
total = max(((prev_end + align - 1) // align) * align, 1)
# Unique value per real row (kept small so the x2 column stays inside the
# e4m3 range), alternating x1/x2 along hidden so a kernel that broadcast one
# column across the row, or mis-strided the hidden offset, is caught.
base = torch.zeros((total, hidden), dtype=torch.float32, device=DEVICE)
col_gain = 1.0 + (torch.arange(hidden, device=DEVICE) % 2).float()
for s, c in zip(starts, counts):
for j in range(c):
base[s + j] = float((s + j) % 200 + 1) * col_gain
recv_x = base.to(dtype)
scale = None
if with_scale:
scale = torch.zeros((total, scale_hidden), dtype=torch.float32, device=DEVICE)
# Distinct value per scale column: a row constant along scale_hidden
# cannot catch a kernel that reads column 0 for every column.
col = torch.arange(scale_hidden, dtype=torch.float32, device=DEVICE)
for s, c in zip(starts, counts):
for j in range(c):
scale[s + j] = float((s + j) % 50 + 1) * 0.5 + col
psum_t = torch.tensor(psum, dtype=torch.int32, device=DEVICE)
return recv_x, scale, psum_t, starts, total
def _real_rows(starts, counts):
rows = []
for s, c in zip(starts, counts):
rows.extend(range(s, s + c))
return rows
class TestDeepEPv2MaskedSlab(CustomTestCase):
ALIGN = 16
HIDDEN = 8
MAX_M = 32
def _check_expand_roundtrip(self, counts, dtype, with_scale, topk=False):
recv_x, scale, psum, starts, total = _build_layout(
counts, self.ALIGN, self.HIDDEN, dtype, with_scale=with_scale
)
E = len(counts)
masked_x, masked_x_scale, masked_m = expand_to_masked_slab(
recv_x, scale, psum, E, self.MAX_M, self.ALIGN
)
# masked_m == real per-expert count
self.assertEqual(masked_m.tolist(), list(counts))
self.assertEqual(tuple(masked_x.shape), (E, self.MAX_M, self.HIDDEN))
# masked_x real rows == source expanded rows
for e, (s, c) in enumerate(zip(starts, counts)):
for j in range(c):
torch.testing.assert_close(
masked_x[e, j].float(), recv_x[s + j].float()
)
if with_scale:
torch.testing.assert_close(
masked_x_scale[e, j].float(), scale[s + j].float()
)
# round-trip back to expanded order
weights = None
if topk:
weights = torch.zeros(total, dtype=torch.float32, device=DEVICE)
for r in _real_rows(starts, counts):
weights[r] = 0.25 + (r % 7) * 0.1
out = masked_slab_to_expand(
masked_x, psum, total, self.ALIGN, topk_weights=weights
)
self.assertEqual(tuple(out.shape), (total, self.HIDDEN))
for e, (s, c) in enumerate(zip(starts, counts)):
for j in range(c):
expected = masked_x[e, j].float()
if topk:
expected = (expected * weights[s + j]).to(masked_x.dtype).float()
torch.testing.assert_close(out[s + j].float(), expected)
def test_roundtrip_bf16(self):
self._check_expand_roundtrip([3, 0, 5, 1], torch.bfloat16, with_scale=False)
def test_roundtrip_bf16_with_topk_weight(self):
self._check_expand_roundtrip(
[2, 4, 0, 7], torch.bfloat16, with_scale=False, topk=True
)
def test_roundtrip_fp8_with_scale(self):
self._check_expand_roundtrip([3, 1, 6, 2], torch.float8_e4m3fn, with_scale=True)
def test_empty_experts(self):
self._check_expand_roundtrip([0, 0, 0, 0], torch.bfloat16, with_scale=False)
def test_single_hot_expert(self):
self._check_expand_roundtrip(
[0, self.MAX_M, 0, 0], torch.bfloat16, with_scale=False, topk=True
)
def test_count_at_max_m_boundary(self):
# exactly max_m must be kept (no overflow, no truncation)
self._check_expand_roundtrip(
[self.MAX_M, 1, self.MAX_M], torch.bfloat16, with_scale=False
)
def test_overflow_fails_fast(self):
# one expert exceeds max_m -> must raise, not silently truncate
counts = [self.MAX_M + 1, 2]
recv_x, scale, psum, starts, total = _build_layout(
counts, self.ALIGN, self.HIDDEN, torch.bfloat16
)
with self.assertRaises(RuntimeError):
expand_to_masked_slab(
recv_x, None, psum, len(counts), self.MAX_M, self.ALIGN
)
def _production_packed_ue8m0_layout(self, counts):
"""Expanded FP8 rows + scales built by the PRODUCTION quantizer with the
Blackwell flags (packed ue8m0, column-major): scale is int32 with
pack-dim stride != 1, unlike Hopper's row-major fp32."""
from sglang.kernels.ops.quantization.fp8_kernel import (
sglang_per_token_group_quant_fp8,
)
# 1024 = 8 quant groups of 128 -> the packed scale has ceil(8/4) = 2 int32
# columns. At hidden <= 512 it collapses to a single column: the pack-dim
# offset is always 0 and the pack-dim stride is never exercised.
hidden = 1024
raw, _, psum, starts, total = _build_layout(
counts, self.ALIGN, hidden, torch.bfloat16
)
recv_x, recv_x_scale = sglang_per_token_group_quant_fp8(
raw,
128,
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
)
self.assertEqual(recv_x_scale.dtype, torch.int32)
self.assertGreater(recv_x_scale.shape[1], 1, "pack dim must be indexed")
self.assertNotEqual(recv_x_scale.stride(1), 1)
return recv_x, recv_x_scale, psum, starts, total, hidden
def test_fp8_packed_ue8m0_scale_from_production_quantizer(self):
# The Blackwell dispatch scale is packed ue8m0 in column-major layout,
# so the repack must honor the scale pack-dim stride; the row-major
# fp32 cases above (stride(1) == 1) cannot regress it.
counts = [3, 1, 6, 2]
recv_x, recv_x_scale, psum, starts, _, hidden = (
self._production_packed_ue8m0_layout(counts)
)
E = len(counts)
masked_x, masked_x_scale, masked_m = expand_to_masked_slab(
recv_x, recv_x_scale, psum, E, self.MAX_M, self.ALIGN
)
self.assertEqual(masked_m.tolist(), list(counts))
self.assertEqual(tuple(masked_x.shape), (E, self.MAX_M, hidden))
for e, (s, c) in enumerate(zip(starts, counts)):
for j in range(c):
torch.testing.assert_close(
masked_x[e, j].float(), recv_x[s + j].float()
)
torch.testing.assert_close(masked_x_scale[e, j], recv_x_scale[s + j])
def test_expand_under_cuda_graph_capture(self):
# The masked repack runs inside the captured decode CUDA graph, so it
# must be capture-safe (no host sync) and correct after replay, with the
# production packed ue8m0 scale layout.
counts = [3, 1, 6, 2]
recv_x, recv_x_scale, psum, starts, _, _ = self._production_packed_ue8m0_layout(
counts
)
E = len(counts)
warm = torch.cuda.Stream()
warm.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warm):
expand_to_masked_slab(recv_x, recv_x_scale, psum, E, self.MAX_M, self.ALIGN)
torch.cuda.current_stream().wait_stream(warm)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
masked_x, masked_x_scale, masked_m = expand_to_masked_slab(
recv_x, recv_x_scale, psum, E, self.MAX_M, self.ALIGN
)
graph.replay()
torch.cuda.synchronize()
self.assertEqual(masked_m.tolist(), list(counts))
for e, (s, c) in enumerate(zip(starts, counts)):
for j in range(c):
torch.testing.assert_close(
masked_x[e, j].float(), recv_x[s + j].float()
)
torch.testing.assert_close(masked_x_scale[e, j], recv_x_scale[s + j])
class TestDeepEPv2HandleLifecycle(CustomTestCase):
"""CPU-only guards of the dispatch/combine handle lifecycle.
The guards are ordered before any DeepEP work, so misuse is testable
without deep_ep installed and without a GPU. The positive dispatch ->
combine path needs real ElasticBuffer communication and is covered by the
GPU accuracy runs instead.
"""
@staticmethod
def _bare_impl():
from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import _DeepEPv2Impl
impl = object.__new__(_DeepEPv2Impl)
impl._handle = None
impl._pad_empty_combine = False
return impl
def test_combine_without_dispatch_raises(self):
impl = self._bare_impl()
with self.assertRaisesRegex(RuntimeError, "without a valid dispatch handle"):
impl.combine(None)
def test_dispatch_with_unconsumed_handle_raises(self):
impl = self._bare_impl()
impl._handle = object()
with self.assertRaisesRegex(RuntimeError, "unconsumed"):
impl.dispatch(None, None)
def test_handle_cleared_when_combine_fails(self):
impl = self._bare_impl()
impl._handle = object()
impl._pad_empty_combine = True
def _boom():
raise RuntimeError("boom")
impl._get_buffer = _boom
with self.assertRaisesRegex(RuntimeError, "boom"):
impl.combine(None)
self.assertIsNone(impl._handle)
self.assertFalse(impl._pad_empty_combine)
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_server_args.py
-114
View File
@@ -1917,120 +1917,6 @@ class TestSamplingBackendTokenOracleEnvGate(CustomTestCase):
self.assertEqual(parsed.sampling_backend, "token_oracle") self.assertEqual(parsed.sampling_backend, "token_oracle")
class TestDeepEPv2Args(CustomTestCase):
"""DeepEP v2 server-args resolution + validation. The dummy-model path
short-circuits __post_init__, so _handle_a2a_moe() is invoked directly."""
def _args(self, **overrides):
server_args = ServerArgs(model_path="dummy", moe_a2a_backend="deepep_v2")
# The deepep_v2 branch mutates cuda_graph_config.{decode,prefill}.backend,
# so it must exist (the dummy path leaves it unset otherwise).
server_args.cuda_graph_config = CudaGraphConfig(
decode=PhaseConfig(backend=Backend.FULL, max_bs=512),
prefill=PhaseConfig(backend=Backend.FULL, max_bs=512),
)
valid = {f.name for f in dataclasses.fields(ServerArgs)}
for key, value in overrides.items():
# ServerArgs has no __slots__, so setattr of a stale field name would
# silently succeed and leave the test asserting nothing.
assert key in valid, f"{key} is not a ServerArgs field"
setattr(server_args, key, value)
return server_args
def test_runner_restored_by_declaration_fails_fast(self):
# mxfp8 + auto: a model declaration restores an unsupported runner at
# materialize time, which runs after this handler. The handler must
# validate the declaration-resolved runner, not the raw value it just set.
args = self._args(moe_runner_backend="auto")
args._resolved_overrides = [
("test_mxfp8", {"moe_runner_backend": "flashinfer_trtllm"})
]
with self.assertRaises(ValueError):
args._handle_a2a_moe()
def test_declarations_materialize_ep_size_and_fusion(self):
from sglang.srt.arg_groups.overrides import materialize_declarations
args = self._args(moe_runner_backend="auto", tp_size=2)
args._handle_a2a_moe()
# ep_size / shared-experts fusion are declared by the a2a passes and land
# on the fields only at materialization, like every other a2a backend.
materialize_declarations(args)
self.assertEqual(args.ep_size, args.tp_size)
self.assertTrue(args.disable_shared_experts_fusion)
def test_auto_runner_defaults_to_deep_gemm(self):
args = self._args(moe_runner_backend="auto")
args._handle_a2a_moe()
self.assertEqual(args.moe_runner_backend, "deep_gemm")
def test_unsupported_runner_rejected(self):
args = self._args(moe_runner_backend="flashinfer_trtllm")
with self.assertRaises(ValueError):
args._handle_a2a_moe()
def test_triton_runner_rejected(self):
# deepep_v2 registers permute adapters for deep_gemm only. Rejecting
# triton here is what keeps a user from reaching the permute registry
# and dying on a bare assert inside the MoE forward.
args = self._args(moe_runner_backend="triton")
with self.assertRaises(ValueError):
args._handle_a2a_moe()
def test_decode_graph_stays_enabled_in_both_comm_modes(self):
# Capturability follows the inference phase (masked decode), not the
# comm mode, so neither direct nor hybrid may disable the decode graph.
for mode in ("direct", "hybrid"):
args = self._args(moe_runner_backend="deep_gemm", deepep_v2_mode=mode)
args._handle_a2a_moe()
self.assertEqual(args.cuda_graph_config.decode.backend, Backend.FULL)
self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.DISABLED)
def test_two_batch_overlap_rejected(self):
args = self._args(moe_runner_backend="deep_gemm", enable_two_batch_overlap=True)
with self.assertRaises(ValueError):
args._handle_a2a_moe()
# --- prefill capacity pre-check (per-rank chunk vs dispatch buffer cap) ---
_CAP_ENV = "SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK"
def test_prefill_chunk_exceeding_cap_rejected(self):
args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=2048)
with patch.dict(os.environ, {self._CAP_ENV: "1024"}):
with self.assertRaisesRegex(ValueError, "NUM_MAX_DISPATCH_TOKENS_PER_RANK"):
args._handle_a2a_moe()
def test_prefill_chunk_at_cap_boundary_accepted(self):
# chunk == cap is the documented (and currently benchmarked) edge; the
# guard must be strict-greater-than.
args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=1024)
with patch.dict(os.environ, {self._CAP_ENV: "1024"}):
args._handle_a2a_moe()
self.assertEqual(args.moe_runner_backend, "deep_gemm")
def test_prefill_chunk_rejected_under_default_cap(self):
# Default cap is 128: a typical 1024-token per-rank chunk must be
# rejected at boot instead of at the first full prefill chunk.
args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=1024)
with self.assertRaisesRegex(ValueError, "chunked prefill budget"):
args._handle_a2a_moe()
def test_prefill_chunk_check_skipped_for_decode_disaggregation(self):
args = self._args(
moe_runner_backend="deep_gemm",
chunked_prefill_size=4096,
disaggregation_mode="decode",
)
args._handle_a2a_moe()
def test_prefill_chunk_check_skipped_when_chunking_disabled(self):
for disabled in (None, 0, -1):
args = self._args(
moe_runner_backend="deep_gemm", chunked_prefill_size=disabled
)
args._handle_a2a_moe()
class TestHandleCrashDumpEnv(CustomTestCase): class TestHandleCrashDumpEnv(CustomTestCase):
_COREDUMP_ENV_KEYS = ( _COREDUMP_ENV_KEYS = (
"CUDA_ENABLE_COREDUMP_ON_EXCEPTION", "CUDA_ENABLE_COREDUMP_ON_EXCEPTION",
test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py
-98
View File
@@ -1,98 +0,0 @@
"""DeepEP-class backend recognition in RoutedExpertsCapturer.
The capturer keys its buffer layout on the a2a backend: DeepEP-class
dispatchers hand the MoE layer only the attention rank's DP-local tokens, so
``capture()`` must attn-TP-gather and ``_get_local_slice()`` must read the
buffer head instead of the global DP offset. These tests pin that DeepEP v2
is classified like DeepEP (it shares that token topology); a miss makes
dp_rank > 0 read unwritten rows (silent wrong data), see the DP>1 readback
test in test/registered/ep/test_routed_experts_dp_readback.py.
"""
import unittest
from types import SimpleNamespace
from unittest import mock
import torch
from sglang.srt.layers.moe.utils import MoeA2ABackend
from sglang.srt.state_capturer import routed_experts as re_mod
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")
class TestScatteredA2ABackendHelper(CustomTestCase):
def test_classification(self):
# deepep_v2 shares DeepEP's scattered token topology. Other backends
# keep their existing classification (mooncake/mori are deliberately
# not reclassified here).
expected = {
"deepep": True,
"deepep_v2": True,
"none": False,
"mooncake": False,
}
for value, exp in expected.items():
with mock.patch.object(
re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend(value)
):
self.assertEqual(
re_mod._is_scattered_a2a_backend(), exp, f"backend={value}"
)
class TestGetLocalSliceBackendBranch(CustomTestCase):
T, L, K = 16, 3, 4 # buffer tokens, layers, top-k
def _capturer(self):
cap = object.__new__(re_mod.RoutedExpertsCapturer)
buf = torch.arange(self.T * self.L * self.K, dtype=torch.int32).reshape(
self.T, self.L, self.K
)
cap.device_cache = SimpleNamespace(buffer=buf)
cap.topk_size = self.K
return cap, buf
def _slice(self, cap, n_local):
fb = SimpleNamespace(out_cache_loc=torch.empty(n_local))
return cap._get_local_slice(fb, can_run_graph=False, cuda_graph_batch=None)
def test_deepep_v2_reads_buffer_head(self):
cap, buf = self._capturer()
with mock.patch.object(
re_mod, "is_dp_attention_enabled", return_value=True
), mock.patch.object(
re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend("deepep_v2")
):
out = self._slice(cap, n_local=5)
self.assertTrue(torch.equal(out, buf[0:5, :, : self.K]))
def test_deepep_v2_matches_deepep(self):
cap, _ = self._capturer()
outs = []
for backend in ("deepep", "deepep_v2"):
with mock.patch.object(
re_mod, "is_dp_attention_enabled", return_value=True
), mock.patch.object(
re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend(backend)
):
outs.append(self._slice(cap, n_local=7))
self.assertTrue(torch.equal(outs[0], outs[1]))
def test_tp_moe_reads_global_offset(self):
cap, buf = self._capturer()
with mock.patch.object(
re_mod, "is_dp_attention_enabled", return_value=True
), mock.patch.object(
re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend("none")
), mock.patch.object(
re_mod, "get_dp_local_slice_cpu", return_value=(6, 4)
):
out = self._slice(cap, n_local=999)
self.assertTrue(torch.equal(out, buf[6:10, :, : self.K]))
if __name__ == "__main__":
unittest.main()