[Feature] Add DeepEPv2 (ElasticBuffer) MoE A2A backend (#29525)
Co-authored-by: menyu <menyu@nvidia.com>
This commit is contained in:
MengYu
@@ -1221,6 +1221,144 @@ def ep_scatter(
|
||||
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
|
||||
def _fwd_kernel_ep_gather(
|
||||
total_token_num,
|
||||
@@ -2003,6 +2141,270 @@ 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(
|
||||
inputs: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
|
||||
@@ -2546,7 +2546,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(
|
||||
@@ -2559,6 +2559,10 @@ 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":
|
||||
# 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 {}
|
||||
|
||||
|
||||
@@ -2567,6 +2571,7 @@ _A2A_EP_SPANNING_BACKENDS = frozenset(
|
||||
{
|
||||
"megamoe",
|
||||
"deepep",
|
||||
"deepep_v2",
|
||||
"mooncake",
|
||||
"nixl",
|
||||
"ascend_fuseep",
|
||||
|
||||
@@ -1014,6 +1014,12 @@ 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)
|
||||
# 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_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False)
|
||||
# Force dynamic Waterfill with runtime EP all-reduce instead of the default
|
||||
|
||||
@@ -103,7 +103,12 @@ 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():
|
||||
# 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
|
||||
elif _use_aiter:
|
||||
self.deprecate_flag = True
|
||||
@@ -354,6 +359,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,
|
||||
@@ -166,6 +167,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,
|
||||
|
||||
@@ -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,212 @@ 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-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,
|
||||
)
|
||||
|
||||
@@ -50,6 +50,22 @@ class MoeRunner:
|
||||
"--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
|
||||
|
||||
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,548 @@
|
||||
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)
|
||||
@@ -4,7 +4,7 @@ import logging
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from enum import Enum, IntEnum
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, NamedTuple
|
||||
|
||||
import torch
|
||||
|
||||
@@ -12,7 +12,13 @@ from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.dp_attention import (
|
||||
is_dp_attention_enabled,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_exec, get_flags, get_forward, get_parallel
|
||||
from sglang.srt.runtime_context import (
|
||||
get_exec,
|
||||
get_flags,
|
||||
get_forward,
|
||||
get_parallel,
|
||||
get_server_args,
|
||||
)
|
||||
from sglang.srt.utils import is_cuda, is_npu
|
||||
|
||||
_is_npu = is_npu()
|
||||
@@ -20,7 +26,6 @@ _is_npu = is_npu()
|
||||
if TYPE_CHECKING:
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -37,6 +42,7 @@ class MoeA2ABackend(Enum):
|
||||
ASCEND_TP = "ascend_tp"
|
||||
FLASHINFER = "flashinfer"
|
||||
MEGAMOE = "megamoe"
|
||||
DEEPEP_V2 = "deepep_v2"
|
||||
PPLX = "pplx"
|
||||
CUSTOMIZED = "customized"
|
||||
|
||||
@@ -76,6 +82,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
|
||||
|
||||
@@ -175,6 +184,20 @@ class MoeRunnerBackend(Enum):
|
||||
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):
|
||||
|
||||
NORMAL = "normal"
|
||||
@@ -308,6 +331,23 @@ 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.
|
||||
|
||||
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):
|
||||
moe = get_flags().moe
|
||||
moe.a2a_backend = MoeA2ABackend(server_args.moe_a2a_backend)
|
||||
|
||||
@@ -732,6 +732,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()
|
||||
)
|
||||
@@ -811,6 +812,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
|
||||
@@ -833,6 +835,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
|
||||
@@ -2704,7 +2707,10 @@ class DeepseekV2Model(nn.Module):
|
||||
for i in range(len(self.layers)):
|
||||
if isinstance(self.layers[i].mlp, DeepseekV2MoE):
|
||||
# tp_size = get_parallel().tp_size
|
||||
is_a2a_moe = is_deepep_class_backend()
|
||||
is_a2a_moe = (
|
||||
is_deepep_class_backend()
|
||||
or get_moe_a2a_backend().is_deepep_v2()
|
||||
)
|
||||
tp_size = 1 if is_a2a_moe else get_parallel().tp_size
|
||||
intermediate_size = (
|
||||
config.moe_intermediate_size * config.n_shared_experts
|
||||
@@ -2724,10 +2730,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)
|
||||
|
||||
@@ -280,6 +280,7 @@ MOE_A2A_BACKEND_CHOICES = [
|
||||
"ascend_fuseep",
|
||||
"flashinfer",
|
||||
"megamoe",
|
||||
"deepep_v2",
|
||||
"pplx",
|
||||
"ascend_tp",
|
||||
]
|
||||
@@ -2354,6 +2355,8 @@ class ServerArgs:
|
||||
"ascend_fuseep",
|
||||
"flashinfer",
|
||||
"megamoe",
|
||||
"deepep_v2",
|
||||
"ascend_tp",
|
||||
"pplx",
|
||||
],
|
||||
Arg(
|
||||
@@ -2363,6 +2366,15 @@ class ServerArgs:
|
||||
),
|
||||
NS("exec.moe"),
|
||||
] = "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[
|
||||
str,
|
||||
Arg(
|
||||
@@ -6964,6 +6976,95 @@ class ServerArgs:
|
||||
self.cuda_graph_config.decode.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 (
|
||||
self.moe_a2a_backend == "none" and is_npu()
|
||||
) or self.moe_a2a_backend == "ascend_tp":
|
||||
|
||||
@@ -20,6 +20,15 @@ from sglang.srt.runtime_context import (
|
||||
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):
|
||||
"""Capturer for routed experts with host buffer.
|
||||
|
||||
@@ -84,11 +93,11 @@ 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():
|
||||
# DeepEP-class 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 _is_scattered_a2a_backend():
|
||||
attn_tp_size = (
|
||||
get_parallel().attn_tp_size if is_dp_attention_enabled() else 1
|
||||
)
|
||||
@@ -102,7 +111,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 +125,11 @@ 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():
|
||||
# Under DeepEP-class backends, 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 _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
|
||||
|
||||
@@ -0,0 +1,207 @@
|
||||
"""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()
|
||||
@@ -0,0 +1,277 @@
|
||||
"""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()
|
||||
@@ -1898,6 +1898,120 @@ class TestSamplingBackendTokenOracleEnvGate(CustomTestCase):
|
||||
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):
|
||||
_COREDUMP_ENV_KEYS = (
|
||||
"CUDA_ENABLE_COREDUMP_ON_EXCEPTION",
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user