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