From a3ae667d67a4ca37e9a4d1348dd3eaa09062a757 Mon Sep 17 00:00:00 2001 From: MengYu Date: Thu, 27 Aug 2026 10:54:33 +0800 Subject: [PATCH] [Feature] Add DeepEPv2 (ElasticBuffer) MoE A2A backend (#35634) Co-authored-by: menyu Co-authored-by: Jinyan Chen <93358689+liz-badada@users.noreply.github.com> Co-authored-by: Han Yu Co-authored-by: Cheng Wan --- .../sglang/kernels/ops/moe/ep_moe_kernels.py | 363 ++++++++++++++ python/sglang/srt/arg_groups/overrides.py | 6 +- python/sglang/srt/environ.py | 4 + python/sglang/srt/layers/moe/ep_moe/layer.py | 5 +- .../srt/layers/moe/fused_moe_triton/layer.py | 39 ++ .../srt/layers/moe/moe_runner/deep_gemm.py | 193 +++++++- .../srt/layers/moe/moe_runner/runner.py | 9 + .../layers/moe/token_dispatcher/__init__.py | 8 + .../srt/layers/moe/token_dispatcher/base.py | 19 + .../layers/moe/token_dispatcher/deepep_v2.py | 460 ++++++++++++++++++ python/sglang/srt/layers/moe/utils.py | 35 +- python/sglang/srt/models/deepseek_v2.py | 12 +- python/sglang/srt/server_args.py | 166 +++++++ .../srt/state_capturer/routed_experts.py | 21 +- .../ep/test_routed_experts_dp_readback.py | 194 ++++++++ .../moe/test_deepep_v2_buffer_lifecycle.py | 154 ++++++ .../layers/moe/test_deepep_v2_masked_slab.py | 244 ++++++++++ .../layers/moe/test_hpc_ops_runner_guard.py | 109 ++++- .../unit/server_args/test_server_args.py | 322 ++++++++++++ .../test_routed_experts_scattered_a2a.py | 86 ++++ 20 files changed, 2414 insertions(+), 35 deletions(-) create mode 100644 python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py create mode 100644 test/registered/ep/test_routed_experts_dp_readback.py create mode 100644 test/registered/unit/layers/moe/test_deepep_v2_buffer_lifecycle.py create mode 100644 test/registered/unit/layers/moe/test_deepep_v2_masked_slab.py create mode 100644 test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 9268572e3..23ba56439 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/python/sglang/kernels/ops/moe/ep_moe_kernels.py @@ -1221,6 +1221,142 @@ 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): + # Mask the tail because this expert is packed against the next one. + 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, @@ -2021,6 +2157,233 @@ def fp8_per_token_to_per_tensor_quant_triton( ) +# Expanded psum starts each expert at align(psum[e-1]); contiguous psum already +# includes the alignment padding. + +_DEEPEP_V2_REPACK_WORKERS_PER_EXPERT = 64 + + +# Scale strides support row-major FP32 and packed column-major UE8M0. +@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, +): + # A fixed worker grid makes conservative max_m values 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: + # Graph capture omits this host-visible overflow flag. + 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, + ) + # Write physical [E, SCALE_HIDDEN, MAX_M] for an mn-major view. + 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() + # Dummy pointer under capture; CHECK_OVERFLOW compiles out every store to it. + 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] + # Store [E, sh, max_m] so the returned transpose is 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, + ) + # Capture relies on max_m = cap * ep_group_size; eager also checks counts. + 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: + 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, +): + 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, +): + """Convert masked-GEMM output to expanded order. + + Only real rows are written; combine ignores uninitialized padding through the + handle. Optional top-k weights are fused into the copy over real rows only. + """ + 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 + 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_rows( inputs: torch.Tensor, topk_ids: torch.Tensor, diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 880de9a7a..6163dfa58 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -2806,7 +2806,7 @@ def _moe_runner_fusion_disable(view: Any) -> dict: def _a2a_fusion_adjustments(view: Any) -> dict: """A2A-backend-driven shared-experts fusion adjustments, declared at the legacy write slots in _handle_a2a_moe: Waterfill requires the - fusion enabled; FlashInfer A2A requires it disabled.""" + fusion enabled; FlashInfer and DeepEP v2 A2A require it disabled.""" if view.moe_a2a_backend in ("deepep", "megamoe") and view.enable_waterfill: if view.disable_shared_experts_fusion: logger.warning( @@ -2819,6 +2819,9 @@ def _a2a_fusion_adjustments(view: Any) -> dict: "Flashinfer MoE A2A is enabled. --disable-shared-experts-fusion is automatically set." ) return {"disable_shared_experts_fusion": True} + if view.moe_a2a_backend == "deepep_v2": + # Fused shared experts are not validated with DeepEP v2. + return {"disable_shared_experts_fusion": True} return {} @@ -2827,6 +2830,7 @@ _A2A_EP_SPANNING_BACKENDS = frozenset( { "megamoe", "deepep", + "deepep_v2", "mooncake", "nixl", "ascend_fuseep", diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 4e14098ac..1ab737a1f 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1043,6 +1043,10 @@ class Envs: # read by several call sites; do not use in new code. SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False) SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + # Per-rank buffer capacity, not a model token limit. + SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + # 0 lets ElasticBuffer select its theoretical communication SM/QP counts. + SGLANG_DEEPEP_V2_NUM_SMS = EnvInt(0) SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32) SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False) SGLANG_ENABLE_QWEN_DEEPEP_SHARED_OVERLAP = EnvBool(True) diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 513ba0aaa..ba81c9f1d 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -103,7 +103,9 @@ class DeepEPMoE(FusedMoE): and quant_config is not None and quant_config.get_name() == "humming" ) - if is_humming: + if get_moe_a2a_backend().is_deepep_v2(): + self.deprecate_flag = True + elif is_humming: self.deprecate_flag = True elif _use_aiter: self.deprecate_flag = True @@ -354,6 +356,7 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]): if ( get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_deepep_v2() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_pplx() diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 97240babb..fb2a63bfb 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -38,6 +38,7 @@ from sglang.srt.layers.moe.token_dispatcher.ascend_tp import ( AscendTPDispatcher, ) from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher +from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import DeepEPv2Dispatcher from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher from sglang.srt.layers.moe.token_dispatcher.standard import ( StandardDispatcher, @@ -189,6 +190,15 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: async_finish=True, return_recv_hook=True, ) + elif a2a_backend.is_deepep_v2(): + return DeepEPv2Dispatcher( + group=get_tp_group().device_group, + router_topk=moe_runner_config.top_k, + num_experts=moe_runner_config.num_experts, + num_local_experts=moe_runner_config.num_local_experts, + hidden_size=moe_runner_config.hidden_size, + params_dtype=moe_runner_config.params_dtype, + ) elif a2a_backend.is_flashinfer(): return FlashinferDispatcher( group=get_tp_group().device_group, @@ -226,6 +236,34 @@ def _validate_hpc_ops_quant_method(quant_method) -> None: ) +def _validate_deepep_v2_quant_method(quant_method) -> None: + """Validate the FP8 contract consumed by the DeepEP v2 adapter.""" + if not get_moe_a2a_backend().is_deepep_v2(): + return + + config = ( + quant_method.quant_config if isinstance(quant_method, Fp8MoEMethod) else None + ) + reason = None + if not isinstance(quant_method, Fp8MoEMethod): + reason = f"selected {type(quant_method).__name__}" + elif quant_method.use_mxfp8: + reason = "selected MXFP8 weights" + elif quant_method.is_fp4_expert: + reason = "selected FP4 experts" + elif list(quant_method.weight_block_size or []) != [128, 128]: + reason = f"has weight_block_size={quant_method.weight_block_size}" + elif config.activation_scheme != "dynamic": + reason = f"has activation_scheme={config.activation_scheme!r}" + + if reason is not None: + raise ValueError( + "--moe-a2a-backend deepep_v2 requires 128x128 blockwise FP8 " + f"experts with dynamic activation scaling, but this layer {reason}. " + "Use a compatible checkpoint or --moe-a2a-backend deepep." + ) + + class FusedMoE(torch.nn.Module): """FusedMoE layer for MoE models. @@ -407,6 +445,7 @@ class FusedMoE(torch.nn.Module): self.use_deep_gemm, ) _validate_hpc_ops_quant_method(self.quant_method) + _validate_deepep_v2_quant_method(self.quant_method) self.supports_deferred_finalize = ( envs.SGLANG_ENABLE_MOE_DEFERRED_FINALIZE.get() and get_moe_runner_backend().is_flashinfer_trtllm() diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index fd9c0211a..c183762c9 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -49,6 +49,10 @@ if TYPE_CHECKING: DeepEPNormalCombineInput, DeepEPNormalDispatchOutput, ) + from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import ( + DeepEPv2CombineInput, + DeepEPv2DispatchOutput, + ) from sglang.srt.layers.moe.token_dispatcher.standard import ( StandardCombineInput, StandardDispatchOutput, @@ -206,6 +210,7 @@ class DeepGemmRunnerInput(RunnerInput): masked_m: Optional[torch.Tensor] = None expected_m: Optional[int] = None m_indices: Optional[torch.Tensor] = None + hidden_states_scale_tma_aligned: bool = False @property def runner_backend(self) -> MoeRunnerBackend: @@ -321,7 +326,10 @@ class DeepGemmRunnerCore(MoeRunnerCore): device=hidden_states_device, dtype=torch.bfloat16, ) - if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES: + if ( + deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES + and not runner_input.hidden_states_scale_tma_aligned + ): hidden_states_scale = tma_align_input_scale(hidden_states_scale) deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig( @@ -1419,3 +1427,186 @@ def _apply_swiglu_limit( out = torch.cat([gate, up], dim=-1) assert out.shape == (num_tokens, hidden_size_x2) return out + + +@register_pre_permute("deepep_v2", "deep_gemm") +def pre_permute_deepep_v2_to_deep_gemm( + dispatch_output: DeepEPv2DispatchOutput, + quant_info: DeepGemmMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> DeepGemmRunnerInput: + from sglang.kernels.ops.moe.ep_moe_kernels import ( + ep_expand_init_m_indices_from_psum, + ep_scatter_from_psum, + ) + + hidden_states = dispatch_output.hidden_states + hidden_states_scale = dispatch_output.hidden_states_scale + topk_ids = dispatch_output.topk_ids + topk_weights = dispatch_output.topk_weights + psum_num_recv_tokens_per_expert = dispatch_output.psum_num_recv_tokens_per_expert + is_expanded = dispatch_output.is_expanded + hidden_states_scale_tma_aligned = dispatch_output.hidden_states_scale_tma_aligned + deepep_v2_use_masked = dispatch_output.use_masked_gemm + deepep_v2_expected_m = dispatch_output.expected_m + deepep_v2_masked_max_m = dispatch_output.masked_max_m + deepep_v2_total_expanded = dispatch_output.total_expanded + deepep_v2_expert_alignment = dispatch_output.expert_alignment + if hidden_states_scale is None: + raise RuntimeError( + "DeepEP v2 -> DeepGEMM requires FP8 dispatch output with activation " + "scales, but the dispatch output carried none." + ) + assert runner_config.activation == "silu" + + if is_expanded: + if psum_num_recv_tokens_per_expert is None: + raise RuntimeError( + "DeepEP v2 requires the native expert prefix sums from the " + "ElasticBuffer dispatch handle." + ) + all_tokens = hidden_states.shape[0] + running_state["all_tokens"] = all_tokens + running_state["hidden_states_shape"] = hidden_states.shape + running_state["hidden_states_device"] = hidden_states.device + running_state["hidden_states_dtype"] = hidden_states.dtype + running_state["topk_ids"] = None + running_state["topk_weights"] = topk_weights + running_state["deepep_v2_expanded"] = True + + if deepep_v2_use_masked: + # masked_m bounds each expert independently of buffer capacity. + from sglang.kernels.ops.moe.ep_moe_kernels import expand_to_masked_slab + + num_local_experts = psum_num_recv_tokens_per_expert.shape[0] + input_tensor, input_tensor_scale, masked_m = expand_to_masked_slab( + hidden_states, + hidden_states_scale, + psum_num_recv_tokens_per_expert, + num_local_experts, + deepep_v2_masked_max_m, + deepep_v2_expert_alignment, + ) + running_state["deepep_v2_masked"] = True + running_state["deepep_v2_psum"] = psum_num_recv_tokens_per_expert + running_state["deepep_v2_total_expanded"] = deepep_v2_total_expanded + running_state["deepep_v2_expert_alignment"] = deepep_v2_expert_alignment + return DeepGemmRunnerInput( + hidden_states=input_tensor, + hidden_states_scale=input_tensor_scale, + use_masked_gemm=True, + masked_m=masked_m, + expected_m=deepep_v2_expected_m, + ) + + # Mark aligned expert rows and leave the unused receive tail at -1. + m_indices = torch.full( + (all_tokens,), -1, device=hidden_states.device, dtype=torch.int32 + ) + ep_expand_init_m_indices_from_psum(psum_num_recv_tokens_per_expert, m_indices) + return DeepGemmRunnerInput( + hidden_states=hidden_states, + hidden_states_scale=hidden_states_scale, + use_masked_gemm=False, + m_indices=m_indices, + hidden_states_scale_tma_aligned=hidden_states_scale_tma_aligned, + ) + + all_tokens = int(psum_num_recv_tokens_per_expert[-1].item()) + K = hidden_states.shape[1] + running_state["all_tokens"] = all_tokens + running_state["hidden_states_shape"] = hidden_states.shape + running_state["hidden_states_device"] = hidden_states.device + running_state["hidden_states_dtype"] = hidden_states.dtype + running_state["topk_ids"] = topk_ids + running_state["topk_weights"] = topk_weights + + input_tensor = torch.empty( + (all_tokens, K), device=hidden_states.device, dtype=hidden_states.dtype + ) + if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + # Packed UE8M0 scales require zero padding lanes. + input_tensor_scale = torch.zeros( + (ceil_div(K // 128, 4), all_tokens), + device=hidden_states.device, + dtype=torch.int, + ).transpose(0, 1) + else: + input_tensor_scale = torch.empty( + (all_tokens, K // 128), device=hidden_states.device, dtype=torch.float32 + ) + m_indices = torch.empty(all_tokens, device=hidden_states.device, dtype=torch.int32) + output_index = torch.empty_like(topk_ids) + # Contiguous psum already includes the 128-row expert alignment. + expert_start_loc = torch.empty_like(psum_num_recv_tokens_per_expert) + ep_scatter_from_psum( + hidden_states, + hidden_states_scale, + topk_ids, + psum_num_recv_tokens_per_expert, + expert_start_loc, + input_tensor, + input_tensor_scale, + m_indices, + output_index, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + dispose_tensor(hidden_states) + dispose_tensor(hidden_states_scale) + running_state["output_index"] = output_index + + return DeepGemmRunnerInput( + hidden_states=input_tensor, + hidden_states_scale=input_tensor_scale, + use_masked_gemm=False, + m_indices=m_indices, + ) + + +@register_post_permute("deep_gemm", "deepep_v2") +def post_permute_deep_gemm_to_deepep_v2( + runner_output: DeepGemmRunnerOutput, + quant_info: DeepGemmMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> DeepEPv2CombineInput: + from sglang.kernels.ops.moe.ep_moe_kernels import ep_gather + from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import DeepEPv2CombineInput + + if running_state.get("deepep_v2_expanded", False): + hidden_states = runner_output.hidden_states + topk_weights = running_state["topk_weights"] + if running_state.get("deepep_v2_masked", False): + # Expanded combine does not consume top-k weights. + from sglang.kernels.ops.moe.ep_moe_kernels import masked_slab_to_expand + + hidden_states = masked_slab_to_expand( + hidden_states, + running_state["deepep_v2_psum"], + running_state["deepep_v2_total_expanded"], + running_state["deepep_v2_expert_alignment"], + topk_weights=topk_weights, + ) + return DeepEPv2CombineInput(hidden_states, None) + if topk_weights is not None: + # Expanded combine does not consume top-k weights. + hidden_states = hidden_states * topk_weights.to( + hidden_states.dtype + ).unsqueeze(-1) + return DeepEPv2CombineInput(hidden_states, None) + + hidden_states = runner_output.hidden_states + topk_ids = running_state["topk_ids"] + topk_weights = running_state["topk_weights"] + output_index = running_state["output_index"] + gather_out = torch.empty( + running_state["hidden_states_shape"], + device=running_state["hidden_states_device"], + dtype=torch.bfloat16, + ) + ep_gather(hidden_states, topk_ids, topk_weights, output_index, gather_out) + return DeepEPv2CombineInput( + hidden_states=gather_out, + topk_weights=topk_weights, + ) diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index 6aff7365e..fe8458b21 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -50,6 +50,15 @@ class MoeRunner: "--moe-runner-backend hpc_ops for this model." ) + if get_moe_a2a_backend().is_deepep_v2() and not runner_backend.is_deep_gemm(): + raise ValueError( + "--moe-a2a-backend deepep_v2 requires the deep_gemm MoE runner, " + f"but this MoE layer's quantization method selected the " + f"'{runner_backend.value}' runner. deepep_v2 dispatches FP8 " + "activations plus scales, which only deep_gemm consumes; use an " + "FP8 blockwise-quantized checkpoint, or --moe-a2a-backend deepep." + ) + self.fused_func = None if runner_backend.is_triton(): diff --git a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py index 7f2c0942f..806dabf7d 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py @@ -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", diff --git a/python/sglang/srt/layers/moe/token_dispatcher/base.py b/python/sglang/srt/layers/moe/token_dispatcher/base.py index 1ff2beb5b..e718c0085 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/base.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/base.py @@ -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" diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py new file mode 100644 index 000000000..6f56cd305 --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py @@ -0,0 +1,460 @@ +from __future__ import annotations + +import logging +import os +from typing import NamedTuple, Optional + +import torch +import torch.distributed as dist + +from sglang.srt.environ import envs +from sglang.srt.layers.dp_attention import get_is_extend_in_batch +from sglang.srt.layers.moe.token_dispatcher.base import ( + BaseDispatcher, + CombineInput, + CombineInputFormat, + DispatchOutput, + DispatchOutputFormat, +) +from sglang.srt.layers.moe.topk import TopKOutput +from sglang.srt.layers.moe.utils import ( + DeepEPv2Fp8ScaleFormat, + get_deepep_v2_fp8_scale_format, +) + +logger = logging.getLogger(__name__) + +_SCALE_BLOCK_SIZE = 128 +# Must match DeepGEMM's contiguous expert alignment. +_EXPERT_ALIGNMENT = 128 +_deepep_v2_import_error: Optional[BaseException] = None +_fp8_quant_import_error: Optional[BaseException] = None +sglang_per_token_group_quant_fp8 = None + +try: + from deep_ep import ElasticBuffer + + use_deepep_v2 = True +except (ImportError, OSError) as exc: + use_deepep_v2 = False + _deepep_v2_import_error = exc + +if use_deepep_v2: + try: + from sglang.kernels.ops.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8, + ) + except (ImportError, OSError) as exc: + _fp8_quant_import_error = exc + + +class DeepEPv2DispatchOutput(NamedTuple): + hidden_states: torch.Tensor + hidden_states_scale: Optional[torch.Tensor] + topk_ids: Optional[torch.Tensor] + topk_weights: torch.Tensor + psum_num_recv_tokens_per_expert: Optional[torch.Tensor] = None + is_expanded: bool = False + hidden_states_scale_tma_aligned: bool = False + use_masked_gemm: bool = False + expected_m: int = 0 + masked_max_m: int = 0 + total_expanded: int = 0 + expert_alignment: int = 128 + + @property + def format(self) -> DispatchOutputFormat: + return DispatchOutputFormat.DEEPEP_V2 + + +class DeepEPv2CombineInput(NamedTuple): + hidden_states: torch.Tensor + topk_weights: Optional[torch.Tensor] + + @property + def format(self) -> CombineInputFormat: + return CombineInputFormat.DEEPEP_V2 + + +assert isinstance(DeepEPv2DispatchOutput, DispatchOutput) +assert isinstance(DeepEPv2CombineInput, CombineInput) + + +def _raise_deepep_v2_import_error() -> None: + detail = ( + f" Original import error: {_deepep_v2_import_error}" + if _deepep_v2_import_error is not None + else "" + ) + raise ImportError( + "DeepEP v2 (ElasticBuffer) is not available. Install DeepEP v2 from " + "https://github.com/deepseek-ai/DeepEP." + detail + ) + + +def _ensure_deepep_v2_available() -> None: + if not use_deepep_v2: + _raise_deepep_v2_import_error() + + +def _ensure_fp8_quant_available() -> None: + _ensure_deepep_v2_available() + if sglang_per_token_group_quant_fp8 is None: + detail = ( + f" Original import error: {_fp8_quant_import_error}" + if _fp8_quant_import_error is not None + else "" + ) + raise ImportError( + "DeepEP v2 FP8 dispatch requires the SGLang FP8 quantization kernel." + + detail + ) + + +def _get_allow_hybrid_mode() -> bool: + from sglang.srt.runtime_context import get_exec + + return get_exec().moe.deepep_v2_mode == "hybrid" + + +def _quantize_for_deepep_v2_dispatch( + hidden_states: torch.Tensor, scale_format: DeepEPv2Fp8ScaleFormat +): + _ensure_fp8_quant_available() + return sglang_per_token_group_quant_fp8( + hidden_states, + _SCALE_BLOCK_SIZE, + column_major_scales=scale_format.tma_aligned, + scale_tma_aligned=scale_format.tma_aligned, + scale_ue8m0=scale_format.ue8m0, + ) + + +class DeepEPv2Buffer: + """Facade for the process-wide ElasticBuffer stored in runtime resources.""" + + _STATE_KEY = "deepep_v2_ep_state" + + @classmethod + def _state(cls): + from types import SimpleNamespace + + from sglang.srt.runtime_context import get_resources + + buffers = get_resources().buffers + state = buffers.get(cls._STATE_KEY) + if state is None: + state = SimpleNamespace(buffer=None, key=None) + buffers[cls._STATE_KEY] = state + return state + + @classmethod + def get_buffer( + cls, + group: dist.ProcessGroup, + hidden_size: int, + router_topk: int, + num_max_dispatch_tokens_per_rank: int, + use_fp8_dispatch: bool, + allow_hybrid_mode: Optional[bool] = None, + ) -> ElasticBuffer: + _ensure_deepep_v2_available() + + if allow_hybrid_mode is None: + allow_hybrid_mode = _get_allow_hybrid_mode() + state = cls._state() + # A key change rebuilds ElasticBuffer collectively on every rank. + key = ( + group, + hidden_size, + router_topk, + num_max_dispatch_tokens_per_rank, + use_fp8_dispatch, + allow_hybrid_mode, + dist.get_world_size(group), + ) + if state.buffer is not None and state.key == key: + return state.buffer + + # Native explicit teardown is unavailable unless explicitly_destroy=True. + cls.destroy() + + # Communicator reuse requires a device-bound process group. + os.environ.setdefault("EP_REUSE_NCCL_COMM", "0") + buffer = ElasticBuffer( + group, + num_max_tokens_per_rank=num_max_dispatch_tokens_per_rank, + hidden=hidden_size, + num_topk=router_topk, + use_fp8_dispatch=use_fp8_dispatch, + allow_hybrid_mode=allow_hybrid_mode, + sl_idx=0, + prefer_overlap_with_compute=False, + ) + # Publish only after collective construction succeeds. + state.buffer = buffer + state.key = key + logger.info( + "Initialized DeepEP v2 ElasticBuffer: world_size=%s hidden_size=%s " + "num_topk=%s max_dispatch_tokens_per_rank=%s use_fp8_dispatch=%s " + "allow_hybrid_mode=%s num_bytes=%s", + dist.get_world_size(group), + hidden_size, + router_topk, + num_max_dispatch_tokens_per_rank, + use_fp8_dispatch, + allow_hybrid_mode, + buffer.num_bytes, + ) + return buffer + + @classmethod + def destroy(cls) -> None: + state = cls._state() + state.buffer = None + state.key = None + + +class _DeepEPv2Impl: + def __init__( + self, + group: dist.ProcessGroup, + router_topk: int, + num_experts: int, + num_local_experts: int, + hidden_size: int, + scale_format: DeepEPv2Fp8ScaleFormat, + num_max_dispatch_tokens_per_rank: int, + ): + self.group = group + self.router_topk = router_topk + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.hidden_size = hidden_size + self.scale_format = scale_format + self.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank + self.rank = dist.get_rank(group) + self._handle = None + self._pad_empty_combine = False + + def _destroy_handle(self) -> None: + self._handle = None + + def _get_buffer(self) -> ElasticBuffer: + return DeepEPv2Buffer.get_buffer( + self.group, + self.hidden_size, + self.router_topk, + self.num_max_dispatch_tokens_per_rank, + True, + ) + + def _validate_common( + self, hidden_states: torch.Tensor, topk_ids: torch.Tensor + ) -> None: + if hidden_states.shape[0] > self.num_max_dispatch_tokens_per_rank: + raise ValueError( + f"DeepEP v2 dispatch input exceeds the per-rank buffer capacity " + f"{self.num_max_dispatch_tokens_per_rank}, got {hidden_states.shape[0]}. " + "Increase SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK or " + "lower the active prefill/decode batch limit." + ) + if hidden_states.shape[1] != self.hidden_size: + raise ValueError( + f"DeepEP v2 hidden size mismatch: expected {self.hidden_size}, " + f"got {hidden_states.shape[1]}" + ) + if self.hidden_size % _SCALE_BLOCK_SIZE != 0: + raise ValueError( + "DeepEP v2 FP8 dispatch requires hidden_size multiple of " + f"{_SCALE_BLOCK_SIZE}, got {self.hidden_size}" + ) + if topk_ids.shape[1] != self.router_topk: + raise ValueError( + f"DeepEP v2 topk mismatch: expected {self.router_topk}, " + f"got {topk_ids.shape[1]}" + ) + + def dispatch( + self, hidden_states: torch.Tensor, topk_output: TopKOutput + ) -> DeepEPv2DispatchOutput: + if self._handle is not None: + raise RuntimeError( + "DeepEP v2 dispatch called while the previous dispatch handle is " + "still unconsumed (missing combine)" + ) + _ensure_deepep_v2_available() + topk_weights = topk_output.topk_weights + topk_ids = topk_output.topk_ids.to(torch.int64) + self._validate_common(hidden_states, topk_ids) + # Decode uses expanded/masked layout; extend uses contiguous in both modes. + use_expand_layout = not get_is_extend_in_batch() + use_masked = use_expand_layout + + # CPU-synced dispatch needs a dummy token to notify from an idle rank. + self._pad_empty_combine = (not use_masked) and hidden_states.shape[0] == 0 + if self._pad_empty_combine: + hidden_states = hidden_states.new_zeros((1, hidden_states.shape[-1])) + # Dummy routes need distinct expert ids; zero weights null the result. + topk_ids = torch.arange( + topk_ids.shape[-1], dtype=topk_ids.dtype, device=topk_ids.device + ).unsqueeze(0) + topk_weights = topk_weights.new_zeros((1, topk_weights.shape[-1])) + + _ensure_fp8_quant_available() + if use_masked: + _ue8m0 = self.scale_format.ue8m0 + dispatch_x = sglang_per_token_group_quant_fp8( + hidden_states, + _SCALE_BLOCK_SIZE, + column_major_scales=_ue8m0, + scale_tma_aligned=_ue8m0, + scale_ue8m0=_ue8m0, + ) + use_tma_aligned_col_major_sf = _ue8m0 + else: + dispatch_x = _quantize_for_deepep_v2_dispatch( + hidden_states, self.scale_format + ) + use_tma_aligned_col_major_sf = self.scale_format.tma_aligned + + # This collective argument must not depend on a rank-local batch. + num_max_tokens = self.num_max_dispatch_tokens_per_rank + # Masked dispatch stays asynchronous for CUDA graph capture. + do_cpu_sync_val = True + if use_masked: + do_cpu_sync_val = False + + buffer = self._get_buffer() + recv_x, recv_topk_idx, recv_topk_weights, handle, event = buffer.dispatch( + dispatch_x, + topk_idx=topk_ids, + topk_weights=topk_weights, + num_experts=self.num_experts, + num_max_tokens_per_rank=num_max_tokens, + expert_alignment=_EXPERT_ALIGNMENT, + num_sms=envs.SGLANG_DEEPEP_V2_NUM_SMS.get(), + use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf, + do_cpu_sync=do_cpu_sync_val, + do_expand=use_expand_layout, + ) + self._handle = handle + local_tokens = hidden_states.shape[0] + if event.event is not None: + event.current_stream_wait() + + if isinstance(recv_x, tuple): + recv_hidden_states, recv_hidden_states_scale = recv_x + else: + recv_hidden_states = recv_x + recv_hidden_states_scale = None + + if use_expand_layout: + # Expanded combine uses handle metadata instead of recv_topk_idx. + local_topk_ids = None + else: + num_recv_tokens = int( + handle.psum_num_recv_tokens_per_scaleup_rank[-1].item() + ) + recv_topk_idx = recv_topk_idx[:num_recv_tokens] + recv_topk_weights = recv_topk_weights[:num_recv_tokens] + recv_hidden_states = recv_hidden_states[:num_recv_tokens] + if recv_hidden_states_scale is not None: + recv_hidden_states_scale = recv_hidden_states_scale[:num_recv_tokens] + + local_topk_ids = recv_topk_idx + + expected_m = 0 + masked_max_m = 0 + total_expanded = 0 + if use_masked: + # expected_m is only a schedule hint; masked_m is the actual bound. + ep_group_size = max(1, self.num_experts // self.num_local_experts) + expected_m = max( + 1, + (local_tokens * ep_group_size * self.router_topk + self.num_experts) + // self.num_experts, + ) + # Account for the worst case where every rank targets one local expert. + masked_max_m = self.num_max_dispatch_tokens_per_rank * ep_group_size + total_expanded = recv_hidden_states.shape[0] + + return DeepEPv2DispatchOutput( + recv_hidden_states, + recv_hidden_states_scale, + local_topk_ids, + recv_topk_weights, + handle.psum_num_recv_tokens_per_expert, + use_expand_layout, + use_tma_aligned_col_major_sf, + use_masked, + expected_m, + masked_max_m, + total_expanded, + _EXPERT_ALIGNMENT, + ) + + def combine(self, combine_input: DeepEPv2CombineInput) -> torch.Tensor: + if self._handle is None: + raise RuntimeError( + "DeepEP v2 combine called without a valid dispatch handle" + ) + # Release the single-use handle even when combine fails. + try: + buffer = self._get_buffer() + combined_x, _, event = buffer.combine( + combine_input.hidden_states, + handle=self._handle, + topk_weights=combine_input.topk_weights, + ) + if event.event is not None: + event.current_stream_wait() + if self._pad_empty_combine: + combined_x = combined_x[:0] + return combined_x + finally: + self._pad_empty_combine = False + self._destroy_handle() + + +class DeepEPv2Dispatcher(BaseDispatcher): + def __init__( + self, + group: dist.ProcessGroup, + router_topk: int, + num_experts: int, + num_local_experts: int, + hidden_size: int, + params_dtype: torch.dtype, + ): + super().__init__() + if params_dtype != torch.bfloat16: + raise NotImplementedError( + "DeepEP v2 dispatch adapter currently expects BF16 model activations, " + f"got {params_dtype}" + ) + scale_format = get_deepep_v2_fp8_scale_format() + self.num_max_dispatch_tokens_per_rank = ( + envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() + ) + self._impl = _DeepEPv2Impl( + group=group, + router_topk=router_topk, + num_experts=num_experts, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + scale_format=scale_format, + num_max_dispatch_tokens_per_rank=self.num_max_dispatch_tokens_per_rank, + ) + + def dispatch( + self, hidden_states: torch.Tensor, topk_output: TopKOutput + ) -> DispatchOutput: + return self._impl.dispatch(hidden_states, topk_output) + + def combine(self, combine_input: CombineInput) -> torch.Tensor: + if combine_input.format != CombineInputFormat.DEEPEP_V2: + raise TypeError( + f"Expected DeepEP v2 combine input, got {combine_input.format}" + ) + return self._impl.combine(combine_input) diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 2a844cd67..42a653cf7 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -4,6 +4,7 @@ import logging import os from contextlib import contextmanager from enum import Enum, IntEnum +from typing import NamedTuple import torch @@ -40,6 +41,7 @@ class MoeA2ABackend(Enum): ASCEND_TP = "ascend_tp" FLASHINFER = "flashinfer" MEGAMOE = "megamoe" + DEEPEP_V2 = "deepep_v2" PPLX = "pplx" CUSTOMIZED = "customized" @@ -79,6 +81,9 @@ class MoeA2ABackend(Enum): def is_megamoe(self): return self == MoeA2ABackend.MEGAMOE + def is_deepep_v2(self): + return self == MoeA2ABackend.DEEPEP_V2 + def is_pplx(self): return self == MoeA2ABackend.PPLX @@ -178,6 +183,13 @@ class MoeRunnerBackend(Enum): return self == MoeRunnerBackend.AITER +class DeepEPv2Fp8ScaleFormat(NamedTuple): + """DeepGEMM FP8 activation-scale layout expected from DeepEP v2.""" + + tma_aligned: bool + ue8m0: bool + + class DeepEPMode(Enum): NORMAL = "normal" @@ -311,6 +323,19 @@ def get_ascend_dispatcher_output_dtype(dispatcher): return DispatcherOutputDtype.BF16 +def get_deepep_v2_fp8_scale_format() -> DeepEPv2Fp8ScaleFormat: + """Resolve the FP8 scale layout DeepEP v2 must pre-quantize into.""" + from sglang.srt.layers import deep_gemm_wrapper + + return DeepEPv2Fp8ScaleFormat( + tma_aligned=( + deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES + or deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + ), + ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + + def initialize_moe_config(): """Seed the MoE runtime flags from the published configuration. @@ -502,9 +527,15 @@ def is_sbo_enabled() -> bool: def is_deepep_class_backend() -> bool: - """Check if the MoE backend is DeepEP-family (DeepEP, Mooncake, Mori, or PPLX).""" + """Return whether A2A combine occurs inside a DeepEP-family dispatcher.""" b = get_moe_a2a_backend() - return b.is_deepep() or b.is_mooncake() or b.is_mori() or b.is_pplx() + return ( + b.is_deepep() + or b.is_deepep_v2() + or b.is_mooncake() + or b.is_mori() + or b.is_pplx() + ) def uses_per_rank_fused_shared_slots() -> bool: diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 8bafc13e0..bde04d145 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -744,6 +744,7 @@ class DeepseekV2MoE(nn.Module): or get_moe_a2a_backend().is_ascend_fuseep() or get_moe_a2a_backend().is_flashinfer() or get_moe_a2a_backend().is_megamoe() + or get_moe_a2a_backend().is_deepep_v2() or should_use_flashinfer_cutlass_moe_fp4_allgather() or envs.SGLANG_SHARED_EXPERT_TP1.get() ) @@ -833,6 +834,7 @@ class DeepseekV2MoE(nn.Module): or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_ascend_fuseep() + or get_moe_a2a_backend().is_deepep_v2() ): # TODO: we will support tp < ep in the future self.ep_size = get_parallel().moe_ep_size @@ -855,6 +857,7 @@ class DeepseekV2MoE(nn.Module): or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_ascend_fuseep() or get_moe_a2a_backend().is_flashinfer() + or get_moe_a2a_backend().is_deepep_v2() ) self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo() # SGLANG_OPT_MOE_QUANT_ONCE eligibility, resolved lazily on first @@ -2757,10 +2760,11 @@ class DeepseekV2Model(nn.Module): ) ) self.layers_to_capture = [] - if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake(): - self.enable_a2a_moe = True - else: - self.enable_a2a_moe = False + self.enable_a2a_moe = ( + get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_mooncake() + or get_moe_a2a_backend().is_deepep_v2() + ) # llama_4_scaling: for supporting Mistral-Large-3 model self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4a7820f28..aa4a13580 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -303,10 +303,20 @@ MOE_A2A_BACKEND_CHOICES = [ "ascend_fuseep", "flashinfer", "megamoe", + "deepep_v2", "pplx", "ascend_tp", ] +# These architectures take the A2A MoE path and skip post-expert all-reduce. +_DEEPEP_V2_VALIDATED_ARCHITECTURES = frozenset( + { + "DeepseekV3ForCausalLM", + "DeepseekV4ForCausalLM", + "Qwen3MoeForCausalLM", + } +) + MXFP8_MOE_RUNNER_BACKEND_CHOICES = [ "cutlass", "deep_gemm", @@ -2443,6 +2453,8 @@ class ServerArgs: "ascend_fuseep", "flashinfer", "megamoe", + "deepep_v2", + "ascend_tp", "pplx", ], Arg( @@ -2459,6 +2471,15 @@ class ServerArgs: "--moe-a2a-backend megamoe.", NS("exec.moe"), ] = False + deepep_v2_mode: A[ + Literal["direct", "hybrid"], + "DeepEP v2 ElasticBuffer communication topology, fixed at server init: " + "`direct` (single-node NVLink) or `hybrid` (multi-node scale-out). " + "Layout/grouped-GEMM and the decode CUDA graph are chosen per batch by " + "inference phase, independent of this knob; not equivalent to DeepEP v1 " + "normal/low_latency.", + NS("exec.moe"), + ] = "direct" moe_runner_backend: A[ str, Arg( @@ -4020,6 +4041,10 @@ class ServerArgs: # time; last declarations of the resolution, mirroring that order. self._handle_model_capability_adjustments() + # Validate after all batch-size declarations are visible. + self._validate_deepep_v2_speculative_draft() + self._validate_deepep_v2_dispatch_token_budget() + self._resolution_finished = True def _handle_return_hidden_states_mode(self): @@ -7415,6 +7440,93 @@ class ServerArgs: f"(e.g. --max-prefill-tokens) to <= {max_cutedsl_tokens}." ) + def _validate_deepep_v2_dispatch_token_budget(self) -> None: + """Check the configured prefill and decode-graph buffer bounds.""" + view = resolved_view(self) + if view.moe_a2a_backend != "deepep_v2": + return + + capacity = envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() + if view.disaggregation_mode != "decode": + prefill_tokens = self.max_prefill_buffer_tokens() or ( + view.max_prefill_tokens or 0 + ) + if prefill_tokens > capacity: + raise ValueError( + "DeepEP v2 per-rank prefill budget exceeds " + "SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK: " + f"required={prefill_tokens}, capacity={capacity}. Raise the " + "environment value or lower --chunked-prefill-size/" + "--max-prefill-tokens." + ) + + if view.disaggregation_mode == "prefill": + return + decode_config = getattr(view.cuda_graph_config, "decode", None) + if decode_config is None or decode_config.backend == Backend.DISABLED: + return + + graph_bs = decode_config.max_bs or 0 + if view.max_running_requests is not None: + attn_dp_size = view.dp_size if view.enable_dp_attention else 1 + per_rank_pool_bs = max(1, view.max_running_requests // attn_dp_size) + graph_bs = min(graph_bs, per_rank_pool_bs) + tokens_per_req = ( + self.max_speculative_num_draft_tokens or 1 + if view.speculative_algorithm + else 1 + ) + graph_tokens = graph_bs * tokens_per_req + if graph_tokens > capacity: + raise ValueError( + "DeepEP v2 per-rank decode CUDA graph exceeds " + "SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK: " + f"required={graph_tokens}, capacity={capacity} " + f"(requests={graph_bs}, tokens/request={tokens_per_req}). Raise " + "the environment value or lower --cuda-graph-max-bs." + ) + + def _validate_deepep_v2_model_architecture(self) -> None: + """Allow DeepEP v2 only where its model workflow is validated.""" + if ( + parse_connector_type(resolved_view(self).model_path) + == ConnectorType.INSTANCE + ): + raise ValueError( + "DeepEP v2 MoE cannot validate a model loaded through an instance " + "connector. Load it from a model path or use " + "--moe-a2a-backend deepep." + ) + + architectures = ( + getattr(self.get_model_config().hf_config, "architectures", None) or [] + ) + + architecture = architectures[0] if architectures else None + if architecture not in _DEEPEP_V2_VALIDATED_ARCHITECTURES: + raise ValueError( + f"DeepEP v2 MoE is not validated for {architecture!r}; supported " + f"architectures are {sorted(_DEEPEP_V2_VALIDATED_ARCHITECTURES)}. " + "Other model workflows may require an all-reduce after A2A " + "combine. Use --moe-a2a-backend deepep." + ) + + def _validate_deepep_v2_speculative_draft(self) -> None: + """Reject an explicit or inherited DeepEP v2 draft backend.""" + view = resolved_view(self) + draft_backend = view.speculative_moe_a2a_backend + if draft_backend is None and view.speculative_algorithm: + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + algorithm = SpeculativeAlgorithm.from_string(view.speculative_algorithm) + if not algorithm.is_ngram(): + draft_backend = view.moe_a2a_backend + if draft_backend == "deepep_v2": + raise ValueError( + "DeepEP v2 MoE is not validated as a speculative draft backend. " + "Select another --speculative-moe-a2a-backend." + ) + def _handle_a2a_moe(self): # The backend overrides and the ep_size=tp_size adjustments moved to # the resolution pipeline (arg_groups/overrides.py: @@ -7466,6 +7578,60 @@ class ServerArgs: cfg.cuda_graph_config.decode.backend = Backend.DISABLED cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + if a2a_backend == "deepep_v2": + self._validate_deepep_v2_model_architecture() + if resolved_view(self).enable_deterministic_inference: + raise ValueError( + "DeepEP v2 does not forward deterministic=True to " + "ElasticBuffer, so deterministic sorting remains disabled. " + "Disable --enable-deterministic-inference or use " + "--moe-a2a-backend deepep." + ) + # ElasticBuffer requires CUMEM, but not NVLS or its preallocation. + os.environ.setdefault("NCCL_CUMEM_ENABLE", "1") + # Respect model-level runner declarations before resolving auto. + resolved_runner = resolved_view(self).moe_runner_backend + if resolved_runner == "auto": + self._declare("_handle_a2a_moe", moe_runner_backend="deep_gemm") + logger.warning( + "DeepEP v2 MoE: resolved --moe-runner-backend auto -> deep_gemm." + ) + elif resolved_runner != "deep_gemm": + raise ValueError( + "DeepEP v2 MoE currently supports only " + f"--moe-runner-backend deep_gemm. Got {resolved_runner!r}. " + "Add a runner adapter before enabling DeepEP v2 with other " + "MoE runners." + ) + if cfg.enable_two_batch_overlap or cfg.enable_single_batch_overlap: + raise ValueError( + "DeepEP v2 MoE has not implemented the TBO/SBO overlap hooks yet. " + "Disable --enable-two-batch-overlap and " + "--enable-single-batch-overlap when using --moe-a2a-backend deepep_v2." + ) + if cfg.enforce_shared_experts_fusion: + raise ValueError( + "DeepEP v2 MoE has not validated fused shared experts yet. " + "Remove --enforce-shared-experts-fusion when using " + "--moe-a2a-backend deepep_v2." + ) + # Prefill reads host counts and is not graph-capturable. + cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + logger.warning( + f"DeepEP v2 MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{cfg.tp_size}]." + ) + logger.warning( + "DeepEP v2 MoE is using deepep_v2_mode=%s. This controls " + "ElasticBuffer direct/hybrid mode and is independent from " + "--deepep-mode normal/low_latency. DeepEP v2 MoE enables the " + "decode CUDA graph on the masked decode path (any comm mode) " + "and disables shared expert fusion. " + "SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK is a " + "per-rank communication buffer capacity, not a model limit; " + "increase it for large prefill/chunked-prefill workloads.", + cfg.deepep_v2_mode, + ) + # The resolving view, not the field: `_a2a_backend_overrides` may have # moved this already (waterfill forces `deepep`). a2a_now = resolved_view(self).moe_a2a_backend diff --git a/python/sglang/srt/state_capturer/routed_experts.py b/python/sglang/srt/state_capturer/routed_experts.py index 15a200cd0..fce0fbcaa 100644 --- a/python/sglang/srt/state_capturer/routed_experts.py +++ b/python/sglang/srt/state_capturer/routed_experts.py @@ -20,6 +20,12 @@ from sglang.srt.runtime_context import ( from sglang.srt.state_capturer.base import BaseTopkCapturer +def _is_scattered_a2a_backend() -> bool: + """Return whether routed tokens are scattered across attention-TP ranks.""" + backend = get_moe_a2a_backend() + return backend.is_deepep() or backend.is_deepep_v2() + + class RoutedExpertsCapturer(BaseTopkCapturer): """Capturer for routed experts with host buffer. @@ -84,11 +90,8 @@ class RoutedExpertsCapturer(BaseTopkCapturer): device_topk_size=topk_size + num_fused_shared_experts, ) - # DeepEP a2a path: each attn-TP rank only sees its scattered slice of - # topk_ids. All-gather across attn-TP at capture time so device_cache - # holds the full batch and the existing _get_local_slice / D2H sync - # paths work unchanged. Pre-allocate the gather target. - if get_moe_a2a_backend().is_deepep(): + # Rebuild the full token batch before routed-expert readback. + if _is_scattered_a2a_backend(): attn_tp_size = ( get_parallel().attn_tp_size if is_dp_attention_enabled() else 1 ) @@ -102,7 +105,7 @@ class RoutedExpertsCapturer(BaseTopkCapturer): ) def capture(self, layer_id: int, topk_indices: torch.Tensor): - if get_moe_a2a_backend().is_deepep(): + if _is_scattered_a2a_backend(): local_topk = topk_indices topk_indices = self.gather_buffer[ : local_topk.size(0) * get_parallel().attn_tp_size @@ -116,10 +119,8 @@ class RoutedExpertsCapturer(BaseTopkCapturer): can_run_graph: bool, cuda_graph_batch: Optional[int], ) -> torch.Tensor: - # Under DeepEP, capture() already attn_tp_all_gathered into the head of - # the per-rank buffer, so the local DP rank's data lives at [0:N_local] - # rather than at the global [start_pos:end_pos] offset. - if is_dp_attention_enabled() and not get_moe_a2a_backend().is_deepep(): + # Gathered rows start at buffer offset zero on every DP rank. + if is_dp_attention_enabled() and not _is_scattered_a2a_backend(): # GPU->CPU sync would break overlap; operate on CPU directly. local_start_pos, local_num_tokens = get_dp_local_slice_cpu( forward_batch, can_run_graph, cuda_graph_batch diff --git a/test/registered/ep/test_routed_experts_dp_readback.py b/test/registered/ep/test_routed_experts_dp_readback.py new file mode 100644 index 000000000..cb21cc7bb --- /dev/null +++ b/test/registered/ep/test_routed_experts_dp_readback.py @@ -0,0 +1,194 @@ +"""DP>1 routed-expert readback parity for DeepEP-family A2A backends.""" + +import concurrent.futures +import json +import os +import unittest + +import numpy as np +import pybase64 +import requests +import torch + +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="4-gpu-h100") + +_MODEL = os.environ.get("SGLANG_ROUTED_EXPERTS_TEST_MODEL", "deepseek-ai/DeepSeek-V3") +_NUM_EXPERTS = 24 +_NUM_LAYERS = 1 +_TOPK = 8 + +_DUMMY_WEIGHT_ENV = { + "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) + + +def _deep_ep_nccl_compatible() -> bool: + try: + version = torch.cuda.nccl.version() + except (AttributeError, RuntimeError): + return False + return version is not None and version >= (2, 30, 7) + + +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", + # Keep the startup budget within the test's 256-token buffer. + "--chunked-prefill-size", + "256", + "--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): + if getattr(cls, "process", None): + 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): + solo = [self._one_request(i) for i in range(self._N_REQ)] + + 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" +) +@unittest.skipUnless( + _deep_ep_nccl_compatible(), "DeepEP v2 requires NCCL runtime >= 2.30.7" +) +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() diff --git a/test/registered/unit/layers/moe/test_deepep_v2_buffer_lifecycle.py b/test/registered/unit/layers/moe/test_deepep_v2_buffer_lifecycle.py new file mode 100644 index 000000000..91258cd88 --- /dev/null +++ b/test/registered/unit/layers/moe/test_deepep_v2_buffer_lifecycle.py @@ -0,0 +1,154 @@ +"""CPU-only tests for the DeepEP v2 ElasticBuffer ownership facade.""" + +import unittest +from unittest.mock import patch + +import torch + +from sglang.srt.layers.moe.token_dispatcher import deepep_v2 +from sglang.srt.runtime_context import get_resources, reset_context +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class _FakeGroup: + pass + + +class _FakeBuffer: + instances = [] + + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + self.num_bytes = 1 << 20 + type(self).instances.append(self) + + +class TestDeepEPv2BufferLifecycle(CustomTestCase): + def setUp(self): + reset_context() + _FakeBuffer.instances = [] + self._patches = [ + patch.object(deepep_v2, "use_deepep_v2", True), + patch.object(deepep_v2, "ElasticBuffer", _FakeBuffer, create=True), + patch.object(deepep_v2.dist, "get_world_size", return_value=8), + ] + for item in self._patches: + item.start() + + def tearDown(self): + reset_context() + for item in reversed(self._patches): + item.stop() + + def _get(self, group=None, **overrides): + kwargs = { + "group": group or _FakeGroup(), + "hidden_size": 4096, + "router_topk": 8, + "num_max_dispatch_tokens_per_rank": 128, + "use_fp8_dispatch": True, + "allow_hybrid_mode": False, + } + kwargs.update(overrides) + return deepep_v2.DeepEPv2Buffer.get_buffer(**kwargs) + + def test_same_key_reuses_buffer(self): + group = _FakeGroup() + first = self._get(group) + second = self._get(group) + self.assertIs(first, second) + self.assertEqual(len(_FakeBuffer.instances), 1) + + def test_constructor_inputs_participate_in_key(self): + group = _FakeGroup() + first = self._get(group) + second = self._get(group, num_max_dispatch_tokens_per_rank=256) + third = self._get( + group, + num_max_dispatch_tokens_per_rank=256, + allow_hybrid_mode=True, + ) + self.assertIsNot(first, second) + self.assertIsNot(second, third) + self.assertEqual(len(_FakeBuffer.instances), 3) + + def test_key_keeps_process_group_object(self): + group = _FakeGroup() + self._get(group) + state = get_resources().buffers[deepep_v2.DeepEPv2Buffer._STATE_KEY] + self.assertIs(state.key[0], group) + + def test_distinct_process_group_rebuilds(self): + first = self._get(_FakeGroup()) + second = self._get(_FakeGroup()) + self.assertIsNot(first, second) + self.assertEqual(len(_FakeBuffer.instances), 2) + + def test_state_lives_in_runtime_resources(self): + self._get() + self.assertIn( + deepep_v2.DeepEPv2Buffer._STATE_KEY, + get_resources().buffers, + ) + + def test_reset_context_drops_state_and_rebuilds(self): + group = _FakeGroup() + self._get(group) + reset_context() + self.assertNotIn( + deepep_v2.DeepEPv2Buffer._STATE_KEY, + get_resources().buffers, + ) + self._get(group) + self.assertEqual(len(_FakeBuffer.instances), 2) + + def test_failed_constructor_is_not_published(self): + class _FailingBuffer: + def __init__(self, *args, **kwargs): + raise RuntimeError("construct failed") + + with patch.object(deepep_v2, "ElasticBuffer", _FailingBuffer): + with self.assertRaisesRegex(RuntimeError, "construct failed"): + self._get() + + state = get_resources().buffers[deepep_v2.DeepEPv2Buffer._STATE_KEY] + self.assertIsNone(state.buffer) + self.assertIsNone(state.key) + self._get() + self.assertEqual(len(_FakeBuffer.instances), 1) + + def test_destroy_clears_facade_state(self): + group = _FakeGroup() + first = self._get(group) + deepep_v2.DeepEPv2Buffer.destroy() + state = get_resources().buffers[deepep_v2.DeepEPv2Buffer._STATE_KEY] + self.assertIsNone(state.buffer) + self.assertIsNone(state.key) + second = self._get(group) + self.assertIsNot(first, second) + + def test_unavailable_deepep_fails_before_state_creation(self): + with patch.object(deepep_v2, "use_deepep_v2", False): + with self.assertRaisesRegex(ImportError, "github.com/deepseek-ai/DeepEP"): + self._get() + self.assertNotIn( + deepep_v2.DeepEPv2Buffer._STATE_KEY, + get_resources().buffers, + ) + + def test_dispatch_capacity_guard_uses_actual_input_rows(self): + impl = object.__new__(deepep_v2._DeepEPv2Impl) + impl.num_max_dispatch_tokens_per_rank = 4 + impl.hidden_size = 128 + impl.router_topk = 2 + impl._validate_common(torch.empty(4, 128), torch.zeros(4, 2)) + with self.assertRaisesRegex(ValueError, "per-rank buffer capacity"): + impl._validate_common(torch.empty(5, 128), torch.zeros(5, 2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/moe/test_deepep_v2_masked_slab.py b/test/registered/unit/layers/moe/test_deepep_v2_masked_slab.py new file mode 100644 index 000000000..0cc11cb64 --- /dev/null +++ b/test/registered/unit/layers/moe/test_deepep_v2_masked_slab.py @@ -0,0 +1,244 @@ +"""Tests for the DeepEP v2 expanded/masked repack kernels.""" + +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 synthetic expanded-layout buffers for per-expert counts.""" + 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) + + # Vary rows and columns to expose broadcast or stride errors. + 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) + # Vary scale columns to expose pack-dimension stride errors. + 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 + ) + + self.assertEqual(masked_m.tolist(), list(counts)) + self.assertEqual(tuple(masked_x.shape), (E, self.MAX_M, self.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() + ) + if with_scale: + torch.testing.assert_close( + masked_x_scale[e, j].float(), scale[s + j].float() + ) + + 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): + self._check_expand_roundtrip( + [self.MAX_M, 1, self.MAX_M], torch.bfloat16, with_scale=False + ) + + def test_overflow_fails_fast(self): + 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): + """Build expanded rows with the production packed UE8M0 quantizer.""" + from sglang.kernels.ops.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8, + ) + + # hidden=1024 ensures the packed scale has multiple columns. + 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): + 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): + # Exercise replay with the production packed 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 dispatch/combine handle guards.""" + + @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() diff --git a/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py b/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py index e5fcaf0f5..6c9e441e2 100644 --- a/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py +++ b/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py @@ -1,16 +1,13 @@ -"""The hpc_ops MoE runner backend makes the standard dispatcher keep global -expert ids, so a quant method that silently falls back to another runner -(e.g. an unquantized MoE) would misroute tokens under EP>1. MoeRunner must -reject that combination loudly at startup. -""" +"""Startup guards for MoE runner and dispatcher quantization contracts.""" import sys +from types import SimpleNamespace import pytest from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.runner import MoeRunner -from sglang.srt.layers.moe.utils import MoeRunnerBackend +from sglang.srt.layers.moe.utils import MoeA2ABackend, MoeRunnerBackend from sglang.srt.runtime_context import get_flags from sglang.test.ci.ci_register import register_cpu_ci @@ -18,27 +15,27 @@ register_cpu_ci(est_time=6, suite="base-c-test-cpu") @pytest.fixture -def _runner_backend_flag(): +def _moe_flags(): moe = get_flags().moe - saved = moe.runner_backend + saved = (moe.runner_backend, moe.a2a_backend) yield moe - moe.runner_backend = saved + moe.runner_backend, moe.a2a_backend = saved -def test_non_hpc_runner_rejected_when_hpc_ops_requested(_runner_backend_flag): - _runner_backend_flag.runner_backend = MoeRunnerBackend.HPC_OPS +def test_non_hpc_runner_rejected_when_hpc_ops_requested(_moe_flags): + _moe_flags.runner_backend = MoeRunnerBackend.HPC_OPS with pytest.raises(ValueError, match="hpc_ops"): MoeRunner(MoeRunnerBackend.TRITON, MoeRunnerConfig()) -def test_triton_runner_allowed_without_hpc_ops(_runner_backend_flag): - _runner_backend_flag.runner_backend = MoeRunnerBackend.TRITON +def test_triton_runner_allowed_without_hpc_ops(_moe_flags): + _moe_flags.runner_backend = MoeRunnerBackend.TRITON runner = MoeRunner(MoeRunnerBackend.TRITON, MoeRunnerConfig()) assert runner.runner_core is not None def test_direct_kernel_quant_method_rejected_when_hpc_ops_requested( - _runner_backend_flag, + _moe_flags, ): # W4AFp8MoEMethod never constructs a MoeRunner (apply() calls its kernel # directly), so it bypasses the MoeRunner-level guard; the layer-level @@ -49,15 +46,95 @@ def test_direct_kernel_quant_method_rejected_when_hpc_ops_requested( from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod from sglang.srt.layers.quantization.w4afp8 import W4AFp8MoEMethod - _runner_backend_flag.runner_backend = MoeRunnerBackend.HPC_OPS + _moe_flags.runner_backend = MoeRunnerBackend.HPC_OPS with pytest.raises(ValueError, match="hpc_ops"): _validate_hpc_ops_quant_method(object.__new__(W4AFp8MoEMethod)) # The FP8 method (the one the hpc_ops runner supports) passes. _validate_hpc_ops_quant_method(object.__new__(Fp8MoEMethod)) # Without hpc_ops requested, any quant method passes. - _runner_backend_flag.runner_backend = MoeRunnerBackend.TRITON + _moe_flags.runner_backend = MoeRunnerBackend.TRITON _validate_hpc_ops_quant_method(object.__new__(W4AFp8MoEMethod)) +def _fp8_method(**overrides): + from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod + + values = { + "activation_scheme": "dynamic", + "weight_block_size": (128, 128), + "use_mxfp8": False, + "is_fp4_expert": False, + } + values.update(overrides) + method = object.__new__(Fp8MoEMethod) + method.quant_config = SimpleNamespace( + activation_scheme=values["activation_scheme"], + ) + method.weight_block_size = values["weight_block_size"] + method.use_mxfp8 = values["use_mxfp8"] + method.is_fp4_expert = values["is_fp4_expert"] + return method + + +def test_deepep_v2_quant_contract_accepts_blockwise_fp8(_moe_flags): + from sglang.srt.layers.moe.fused_moe_triton.layer import ( + _validate_deepep_v2_quant_method, + ) + + _moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2 + _validate_deepep_v2_quant_method(_fp8_method(weight_block_size=[128, 128])) + + +@pytest.mark.parametrize( + ("overrides", "expected"), + [ + ({"activation_scheme": "static"}, "activation_scheme"), + ({"weight_block_size": None}, "weight_block_size"), + ({"weight_block_size": (1, 32), "use_mxfp8": True}, "MXFP8"), + ({"is_fp4_expert": True}, "FP4 experts"), + ], +) +def test_deepep_v2_quant_contract_rejects_incompatible_fp8( + _moe_flags, overrides, expected +): + from sglang.srt.layers.moe.fused_moe_triton.layer import ( + _validate_deepep_v2_quant_method, + ) + + _moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2 + with pytest.raises(ValueError, match=expected): + _validate_deepep_v2_quant_method(_fp8_method(**overrides)) + + +def test_deepep_v2_quant_contract_rejects_incompatible_methods(_moe_flags): + from sglang.srt.layers.moe.fused_moe_triton.layer import ( + _validate_deepep_v2_quant_method, + ) + from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod + from sglang.srt.layers.quantization.w4afp8 import W4AFp8MoEMethod + + _moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2 + for method_type in (UnquantizedFusedMoEMethod, W4AFp8MoEMethod): + with pytest.raises(ValueError, match=method_type.__name__): + _validate_deepep_v2_quant_method(object.__new__(method_type)) + + +def test_deepep_v2_quant_contract_does_not_affect_other_backends(_moe_flags): + from sglang.srt.layers.moe.fused_moe_triton.layer import ( + _validate_deepep_v2_quant_method, + ) + from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod + + _moe_flags.a2a_backend = MoeA2ABackend.DEEPEP + _validate_deepep_v2_quant_method(object.__new__(UnquantizedFusedMoEMethod)) + + +def test_deepep_v2_runner_backstop(_moe_flags): + _moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2 + with pytest.raises(ValueError, match="deep_gemm"): + MoeRunner(MoeRunnerBackend.TRITON, MoeRunnerConfig()) + assert MoeRunner(MoeRunnerBackend.DEEP_GEMM, MoeRunnerConfig()).runner_core + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 9407f3b57..e8890e2af 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -2106,6 +2106,328 @@ class TestSamplingBackendTokenOracleEnvGate(CustomTestCase): self.assertEqual(parsed.sampling_backend, "token_oracle") +class TestDeepEPv2Args(CustomTestCase): + """DeepEP v2 server-argument resolution and validation.""" + + def _args(self, **overrides): + server_args = ServerArgs(model_path="dummy", moe_a2a_backend="deepep_v2") + server_args.model_config = SimpleNamespace( + hf_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]) + ) + # The dummy path does not initialize phase configs. + server_args.cuda_graph_config = CudaGraphConfig( + decode=PhaseConfig(backend=Backend.FULL, max_bs=512), + prefill=PhaseConfig(backend=Backend.FULL, max_bs=512), + ) + server_args._resolved_overrides = [] + valid = {f.name for f in dataclasses.fields(ServerArgs)} + for key, value in overrides.items(): + # Reject stale field names before setattr silently accepts them. + assert key in valid, f"{key} is not a ServerArgs field" + setattr(server_args, key, value) + return server_args + + def test_validated_architectures_allowed(self): + for architecture in ( + "DeepseekV3ForCausalLM", + "DeepseekV4ForCausalLM", + "Qwen3MoeForCausalLM", + ): + args = self._args(moe_runner_backend="deep_gemm") + args.model_config.hf_config.architectures = [architecture] + args._handle_a2a_moe() + + def test_unvalidated_and_missing_architectures_rejected(self): + for architectures in ( + ["Qwen2MoeForCausalLM"], + ["Qwen3_5MoeForCausalLM"], + [], + None, + ): + args = self._args(moe_runner_backend="deep_gemm") + args.model_config.hf_config.architectures = architectures + with self.assertRaisesRegex(ValueError, "not validated"): + args._handle_a2a_moe() + + def test_instance_connector_rejected(self): + args = self._args( + model_path="instance://worker/model", + moe_runner_backend="deep_gemm", + ) + with self.assertRaisesRegex(ValueError, "instance connector"): + args._handle_a2a_moe() + + def test_deterministic_inference_rejected(self): + args = self._args( + moe_runner_backend="deep_gemm", + enable_deterministic_inference=True, + ) + with self.assertRaisesRegex(ValueError, "deterministic sorting"): + args._handle_a2a_moe() + + def test_rl_on_policy_deterministic_inference_rejected(self): + args = self._args( + moe_runner_backend="deep_gemm", + rl_on_policy_target="fsdp", + ) + args.model_config.hf_config.architectures = ["Qwen3MoeForCausalLM"] + with ( + envs.SGLANG_VLM_CACHE_SIZE_MB.override(envs.SGLANG_VLM_CACHE_SIZE_MB.get()), + envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.override( + envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get() + ), + ): + args._handle_deterministic_inference() + with self.assertRaisesRegex(ValueError, "deterministic sorting"): + args._handle_a2a_moe() + + def test_deterministic_inference_does_not_affect_legacy_deepep(self): + args = self._args( + moe_a2a_backend="deepep", + moe_runner_backend="deep_gemm", + enable_deterministic_inference=True, + ) + args._handle_a2a_moe() + + def test_runner_restored_by_declaration_fails_fast(self): + # Validate the declaration-resolved runner rather than the raw field. + 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_resolve_ep_size_and_fusion(self): + from sglang.srt.arg_groups.overrides import resolved_view + + args = self._args(moe_runner_backend="auto", tp_size=2) + args._handle_a2a_moe() + self.assertEqual(resolved_view(args).ep_size, args.tp_size) + self.assertTrue(resolved_view(args).disable_shared_experts_fusion) + + def test_auto_runner_defaults_to_deep_gemm(self): + from sglang.srt.arg_groups.overrides import resolved_view + + args = self._args(moe_runner_backend="auto") + args._handle_a2a_moe() + self.assertEqual(resolved_view(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): + 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): + 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() + + def test_speculative_draft_backend_rejected(self): + for main_backend in ("none", "deepep", "deepep_v2"): + args = self._args( + moe_a2a_backend=main_backend, + moe_runner_backend="deep_gemm", + speculative_moe_a2a_backend="deepep_v2", + ) + with self.assertRaisesRegex(ValueError, "speculative draft backend"): + args._validate_deepep_v2_speculative_draft() + + def test_inherited_speculative_draft_backend_rejected(self): + args = self._args( + moe_runner_backend="deep_gemm", + speculative_algorithm="EAGLE", + ) + with self.assertRaisesRegex(ValueError, "speculative draft backend"): + args._validate_deepep_v2_speculative_draft() + + def test_ngram_does_not_inherit_a_draft_backend(self): + args = self._args( + moe_runner_backend="deep_gemm", + speculative_algorithm="NGRAM", + ) + args._validate_deepep_v2_speculative_draft() + + def test_explicit_legacy_speculative_backend_allowed(self): + args = self._args( + moe_runner_backend="deep_gemm", + speculative_algorithm="EAGLE", + speculative_moe_a2a_backend="deepep", + ) + args._validate_deepep_v2_speculative_draft() + + def test_resolved_legacy_speculative_backend_allowed(self): + args = self._args( + moe_runner_backend="deep_gemm", + speculative_algorithm="EAGLE", + ) + args._resolved_overrides = [ + ( + "test_speculative_backend", + {"speculative_moe_a2a_backend": "deepep"}, + ) + ] + args._validate_deepep_v2_speculative_draft() + + def test_prefill_chunk_exceeding_cap_rejected(self): + args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=2048) + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1024): + with self.assertRaisesRegex(ValueError, "NUM_MAX_DISPATCH_TOKENS_PER_RANK"): + args._validate_deepep_v2_dispatch_token_budget() + + def test_prefill_chunk_at_cap_boundary_accepted(self): + args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=1024) + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1024): + args._validate_deepep_v2_dispatch_token_budget() + + def test_dynamic_chunking_probe_is_included(self): + args = self._args( + chunked_prefill_size=1024, + max_prefill_tokens=1024, + enable_dynamic_chunking=True, + pp_size=2, + disaggregation_mode="prefill", + ) + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1024): + with self.assertRaisesRegex(ValueError, "required=1280"): + args._validate_deepep_v2_dispatch_token_budget() + + def test_disabled_chunking_uses_max_prefill_tokens(self): + for disabled in (None, 0, -1): + args = self._args( + chunked_prefill_size=disabled, + max_prefill_tokens=1024, + disaggregation_mode="prefill", + ) + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128): + with self.assertRaisesRegex(ValueError, "required=1024"): + args._validate_deepep_v2_dispatch_token_budget() + + def test_decode_role_skips_prefill_capacity(self): + args = self._args( + chunked_prefill_size=4096, + disaggregation_mode="decode", + max_running_requests=32, + dp_size=1, + ) + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128): + args._validate_deepep_v2_dispatch_token_budget() + + def test_decode_graph_capacity_boundaries(self): + for max_bs, raises in ((128, False), (129, True)): + args = self._args( + disaggregation_mode="decode", + max_running_requests=None, + ) + args.cuda_graph_config.decode.max_bs = max_bs + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128): + if raises: + with self.assertRaisesRegex(ValueError, "decode CUDA graph"): + args._validate_deepep_v2_dispatch_token_budget() + else: + args._validate_deepep_v2_dispatch_token_budget() + + def test_dp_attention_divides_max_running_requests_per_rank(self): + args = self._args( + disaggregation_mode="decode", + max_running_requests=256, + tp_size=8, + dp_size=8, + enable_dp_attention=True, + ) + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128): + args._validate_deepep_v2_dispatch_token_budget() + + def test_tp_only_max_running_requests_is_not_divided(self): + args = self._args( + disaggregation_mode="decode", + max_running_requests=256, + tp_size=8, + dp_size=1, + enable_dp_attention=False, + ) + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128): + with self.assertRaisesRegex(ValueError, "decode CUDA graph"): + args._validate_deepep_v2_dispatch_token_budget() + + def test_memory_derived_eager_pool_remains_runtime_validated(self): + args = self._args( + disaggregation_mode="decode", + max_running_requests=None, + ) + args.cuda_graph_config.decode.backend = Backend.DISABLED + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1): + args._validate_deepep_v2_dispatch_token_budget() + + def test_speculative_decode_width_is_included(self): + args = self._args( + disaggregation_mode="decode", + speculative_algorithm="EAGLE", + speculative_num_draft_tokens=8, + max_running_requests=256, + dp_size=8, + enable_dp_attention=True, + ) + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128): + with self.assertRaisesRegex(ValueError, "tokens/request=8"): + args._validate_deepep_v2_dispatch_token_budget() + + def test_adaptive_speculative_uses_widest_candidate(self): + args = self._args( + disaggregation_mode="decode", + speculative_algorithm="EAGLE", + speculative_num_draft_tokens=4, + speculative_adaptive=True, + max_running_requests=128, + dp_size=8, + enable_dp_attention=True, + ) + with patch.object( + ServerArgs, + "max_speculative_num_draft_tokens", + new=property(lambda _self: 16), + ): + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128): + with self.assertRaisesRegex(ValueError, "tokens/request=16"): + args._validate_deepep_v2_dispatch_token_budget() + + def test_prefill_role_skips_decode_capacity(self): + args = self._args( + disaggregation_mode="prefill", + chunked_prefill_size=64, + max_running_requests=8192, + ) + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128): + args._validate_deepep_v2_dispatch_token_budget() + + def test_other_backend_skips_capacity_validation(self): + args = self._args( + moe_a2a_backend="deepep", + chunked_prefill_size=4096, + max_running_requests=4096, + ) + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1): + args._validate_deepep_v2_dispatch_token_budget() + + def test_capacity_validation_uses_resolved_backend(self): + args = self._args(chunked_prefill_size=4096) + args._resolved_overrides = [("test", {"moe_a2a_backend": "deepep"})] + with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1): + args._validate_deepep_v2_dispatch_token_budget() + + class TestHandleCrashDumpEnv(CustomTestCase): _COREDUMP_ENV_KEYS = ( "CUDA_ENABLE_COREDUMP_ON_EXCEPTION", diff --git a/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py b/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py new file mode 100644 index 000000000..37b558789 --- /dev/null +++ b/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py @@ -0,0 +1,86 @@ +"""DeepEP-family backend recognition in RoutedExpertsCapturer.""" + +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_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestScatteredA2ABackendHelper(CustomTestCase): + def test_classification(self): + 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 + + 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()