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 544c0494e..3d4cbfc44 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -45,7 +45,11 @@ from sglang.srt.layers.moe.topk import ( TopKOutput, TopKOutputChecker, ) -from sglang.srt.layers.moe.utils import RoutingMethodType, is_deepep_class_backend +from sglang.srt.layers.moe.utils import ( + RoutingMethodType, + has_per_rank_fused_shared_slots, + uses_per_rank_fused_shared_slots, +) from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, QuantizationConfig, @@ -54,6 +58,7 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsMxInt4MoE, ) from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod +from sglang.srt.layers.quantization.fp8_utils import quantize_block_fp8_weight_to_mxfp4 from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4FusedMoEMethod from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( @@ -197,10 +202,11 @@ class FusedMoE(torch.nn.Module): self.moe_tp_size = get_parallel().moe_tp_size self.moe_tp_rank = get_parallel().moe_tp_rank - # DeepEP: each rank has its own shared expert slot, so total shared - # weight slots = num_fused_shared_experts * ep_size. - # AMD/Standard: shared experts are global, slots = num_fused_shared_experts. - if num_fused_shared_experts > 0 and is_deepep_class_backend(): + # For fused shared experts, DeepEP-class and MegaMOE backends use + # per-rank physical shared slots, while other backends keep fused + # shared experts as global shared slots. When fusion is disabled, + # num_fused_shared_experts is 0 and no shared slots are added here. + if has_per_rank_fused_shared_slots(num_fused_shared_experts): num_shared_slots = num_fused_shared_experts * self.moe_ep_size else: num_shared_slots = num_fused_shared_experts @@ -210,6 +216,8 @@ class FusedMoE(torch.nn.Module): self._num_local_routed = self._num_global_routed // self.moe_ep_size self.num_local_experts = self._num_local_routed + num_fused_shared_experts self._has_fused_shared = num_fused_shared_experts > 0 + self._pending_fp8_shared_weights: dict[tuple[int, str], torch.Tensor] = {} + self._pending_fp8_shared_scales: dict[tuple[int, str], torch.Tensor] = {} assert intermediate_size % self.moe_tp_size == 0 self.intermediate_size_per_partition = intermediate_size // self.moe_tp_size @@ -570,6 +578,93 @@ class FusedMoE(torch.nn.Module): # w2, down_proj: Load into only logical weight of w2. expert_data.copy_(loaded_weight) + def _maybe_load_fp8_shared_expert_as_fp4( + self, + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + shard_dim: int, + tp_rank: int, + ) -> bool: + if ( + not self._has_fused_shared + or expert_id < self._num_local_routed + or self.quant_config is None + or not getattr(self.quant_config, "is_fp4_experts", False) + or shard_id not in ("w1", "w2", "w3") + ): + return False + + is_weight = ( + "weight" in weight_name + and "scale" not in weight_name + and loaded_weight.dtype == torch.float8_e4m3fn + ) + is_scale = "weight_scale_inv" in weight_name and loaded_weight.dtype in ( + torch.float8_e8m0fnu, + torch.float32, + ) + if not is_weight and not is_scale: + return False + + weight_param = self.w2_weight if shard_id == "w2" else self.w13_weight + scale_param = ( + self.w2_weight_scale_inv if shard_id == "w2" else self.w13_weight_scale_inv + ) + if param is not weight_param and param is not scale_param: + return False + + key = (expert_id, shard_id) + if is_weight: + fp8_weight = loaded_weight + fp8_scale = self._pending_fp8_shared_scales.pop(key, None) + if fp8_scale is None: + self._pending_fp8_shared_weights[key] = loaded_weight + return True + else: + fp8_weight = self._pending_fp8_shared_weights.pop(key, None) + fp8_scale = loaded_weight + if fp8_weight is None: + self._pending_fp8_shared_scales[key] = loaded_weight + return True + + logging.getLogger(__name__).warning_once( + "Loading FP8 shared expert weights into FP4 fused MoE weights. " + "The shared expert is quantized at load time and may differ " + "slightly from a checkpoint that stores shared experts directly " + "in FP4." + ) + + weight_block_size = getattr(self.quant_config, "weight_block_size", None) + if weight_block_size is None: + raise ValueError( + "Loading FP8 shared expert weights into FP4 fused MoE weights " + "requires block-FP8 weight_block_size." + ) + fp4_weight, fp4_scale = quantize_block_fp8_weight_to_mxfp4( + fp8_weight, fp8_scale, weight_block_size + ) + + weight_data = weight_param.data[expert_id] + scale_data = scale_param.data[expert_id] + self._load_model_weight_or_group_weight_scale( + shard_dim=shard_dim, + expert_data=weight_data, + shard_id=shard_id, + loaded_weight=fp4_weight, + tp_rank=tp_rank, + ) + self._load_model_weight_or_group_weight_scale( + shard_dim=shard_dim, + expert_data=scale_data, + shard_id=shard_id, + loaded_weight=fp4_scale, + tp_rank=tp_rank, + ) + return True + def _load_single_value( self, param: torch.nn.Parameter, loaded_weight: torch.Tensor, expert_id: int ): @@ -658,7 +753,7 @@ class FusedMoE(torch.nn.Module): if 0 <= shared_expert_id < self.num_fused_shared_experts: # Checkpoint shared experts start after logical routed experts, while # local fused MoE weights store them after physical routed experts. - if require_global_experts and is_deepep_class_backend(): + if require_global_experts and uses_per_rank_fused_shared_slots(): physical_expert_ids = [ rank * self.num_local_experts + self._num_local_routed @@ -849,6 +944,17 @@ class FusedMoE(torch.nn.Module): if is_transposed: shard_dim = int(not shard_dim) + if self._maybe_load_fp8_shared_expert_as_fp4( + param=param, + loaded_weight=loaded_weight, + weight_name=weight_name, + shard_id=shard_id, + expert_id=expert_id, + shard_dim=shard_dim, + tp_rank=tp_rank, + ): + return + # Case input scale: input_scale loading is only supported for fp8 if "input_scale" in weight_name: # INT4-FP8 (INT4 MoE Weight, FP8 Compute): Adjust input_scale for e4m3fnuz (AMD) diff --git a/python/sglang/srt/layers/moe/hash_topk.py b/python/sglang/srt/layers/moe/hash_topk.py index 55ec9ef6f..2be451e2d 100644 --- a/python/sglang/srt/layers/moe/hash_topk.py +++ b/python/sglang/srt/layers/moe/hash_topk.py @@ -16,9 +16,12 @@ from sglang.srt.eplb.expert_location_dispatch import ( ) from sglang.srt.layers.moe.topk import ( StandardTopKOutput, + TopKConfig, _mask_topk_ids_padded_region, _zero_topk_weights_padded_region, + remap_topk_for_per_rank_shared_slots, ) +from sglang.srt.layers.moe.utils import has_per_rank_fused_shared_slots from sglang.srt.utils import is_hip, is_npu logger = logging.getLogger(__name__) @@ -106,10 +109,18 @@ class HashTopK(nn.Module): topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device) topk_ids = torch.full((0, topk), -1, dtype=torch.int32, device=device) router_logits = torch.empty((0, topk), dtype=torch.float32, device=device) - return self._apply_deepep_waterfill( - StandardTopKOutput(topk_weights, topk_ids, router_logits), - num_tokens=0, - ) + topk_output = StandardTopKOutput(topk_weights, topk_ids, router_logits) + if has_per_rank_fused_shared_slots(self.num_fused_shared_experts): + n = self.num_fused_shared_experts + topk_output = topk_output._replace( + topk_ids=topk_output.topk_ids.new_empty( + (0, topk_output.topk_ids.shape[-1] + n) + ), + topk_weights=topk_output.topk_weights.new_empty( + (0, topk_output.topk_weights.shape[-1] + n) + ), + ) + return self._apply_deepep_waterfill(topk_output, num_tokens=0) def _apply_deepep_waterfill( self, topk_output: StandardTopKOutput, num_tokens: int @@ -198,6 +209,7 @@ class HashTopK(nn.Module): if self.apply_routed_scaling_factor_on_output: topk_weights = topk_weights * self.routed_scaling_factor + num_fused_shared_experts = self.num_fused_shared_experts log2phy_prob = None if ( expert_location_dispatch_info is not None @@ -212,14 +224,47 @@ class HashTopK(nn.Module): if lplb_solver is not None: log2phy_prob = lplb_solver.solve(topk_ids) - topk_ids = topk_ids_logical_to_physical( - topk_ids, expert_location_dispatch_info, log2phy_prob - ) + recorder_topk_ids = None + if has_per_rank_fused_shared_slots(num_fused_shared_experts): + shared_cols = topk_ids[:, -num_fused_shared_experts:] + routed_cols = topk_ids[:, :-num_fused_shared_experts] + routed_cols = topk_ids_logical_to_physical( + routed_cols, expert_location_dispatch_info, log2phy_prob + ) + topk_ids = torch.cat([routed_cols, shared_cols], dim=-1) + recorder_topk_ids = routed_cols + + num_physical_routed_experts = ( + expert_location_dispatch_info.num_physical_experts + if expert_location_dispatch_info is not None + else self.num_experts + ) + topk_ids, topk_weights = remap_topk_for_per_rank_shared_slots( + topk_ids, + topk_weights, + num_fused_shared_experts, + num_physical_routed_experts, + TopKConfig( + top_k=self.topk, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=self.routed_scaling_factor, + ), + ) + else: + topk_ids = topk_ids_logical_to_physical( + topk_ids, expert_location_dispatch_info, log2phy_prob + ) if is_hip(): _zero_topk_weights_padded_region(topk_weights, num_token_non_padded) else: _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) - get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids) + if recorder_topk_ids is not None: + _mask_topk_ids_padded_region(recorder_topk_ids, num_token_non_padded) + if recorder_topk_ids is None: + recorder_topk_ids = topk_ids + get_global_expert_distribution_recorder().on_select_experts( + topk_ids=recorder_topk_ids + ) topk_output = StandardTopKOutput( topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits ) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index ce423c47f..1b23554f7 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -100,7 +100,9 @@ from sglang.srt.eplb.expert_location_dispatch import ( ) from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe import get_moe_runner_backend -from sglang.srt.layers.moe.utils import is_deepep_class_backend +from sglang.srt.layers.moe.utils import ( + has_per_rank_fused_shared_slots, +) from sglang.srt.layers.utils import MultiPlatformOp from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer from sglang.srt.utils import ( @@ -596,7 +598,7 @@ class TopK(MultiPlatformOp): # FIXME: router_logits should be of size (0, num_experts) router_logits = torch.empty((0, topk), dtype=torch.float32, device=device) topk_output = StandardTopKOutput(topk_weights, topk_ids, router_logits) - if self.topk_config.num_fused_shared_experts > 0 and is_deepep_class_backend(): + if has_per_rank_fused_shared_slots(self.topk_config.num_fused_shared_experts): n = self.topk_config.num_fused_shared_experts topk_output = topk_output._replace( topk_ids=topk_output.topk_ids.new_empty( @@ -1470,7 +1472,8 @@ def biased_grouped_topk_gpu( if num_fused_shared_experts > 0: # Append shared expert columns: ID = num_experts (first shared slot), # weight = sum(routed) / scaling_factor (matching biased_grouped_topk_impl). - # DeepEP fusion will overwrite both in _remap_topk_ids_for_deepep_fusion. + # For DeepEP/MegaMOE per-rank shared-slot layout, post-process remaps + # this placeholder ID and overwrites the shared weight for the active scaling path. topk_ids = F.pad(topk_ids, (0, num_fused_shared_experts), value=num_experts) topk_weights = F.pad(topk_weights, (0, num_fused_shared_experts)) if routed_scaling_factor is not None: @@ -1686,18 +1689,18 @@ else: fused_topk_native = fused_topk_torch_native -def _remap_topk_for_deepep( +def remap_topk_for_per_rank_shared_slots( topk_ids: torch.Tensor, topk_weights: torch.Tensor, num_fused_shared_experts: int, num_physical_routed_experts: int, topk_config: TopKConfig, ) -> tuple[torch.Tensor, torch.Tensor]: - """Remap TopK output to DeepEP interleaved expert layout. + """Remap TopK IDs to a per-rank shared-slot layout. - DeepEP dispatch needs each rank's shared expert at a unique ID so tokens - route to the correct rank. The layout interleaves shared slots among - routed experts: [routed_0..L-1, shared, routed_L..2L-1, shared, ...]. + DeepEP and MegaMoE dispatch need each rank's shared expert at a unique ID + so tokens route to the correct rank. The layout is ordered by rank: + [rank0 routed..., rank0 shared, rank1 routed..., rank1 shared, ...]. Routed IDs: e -> e + e // num_local_routed Shared IDs: ep_rank * num_local_experts + num_local_routed @@ -1710,7 +1713,7 @@ def _remap_topk_for_deepep( ep_rank = get_parallel().moe_ep_rank # Static EPLB may add redundant physical experts. At this point routed # topk_ids have already been remapped from logical to physical ids, so the - # DeepEP interleaved layout must use the physical routed count. + # per-rank shared-slot layout must use the physical routed count. num_local_routed = num_physical_routed_experts // ep_size num_local_experts = num_local_routed + num_fused_shared_experts @@ -1786,6 +1789,9 @@ def _post_process_topk_ids( expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: num_fused_shared_experts = topk_config.num_fused_shared_experts + use_per_rank_shared_slots = has_per_rank_fused_shared_slots( + num_fused_shared_experts + ) fused_shared_experts_scaling_factor = ( topk_config.fused_shared_experts_scaling_factor ) @@ -1811,7 +1817,7 @@ def _post_process_topk_ids( topk_ids, expert_location_dispatch_info, log2phy_prob ) _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) - elif num_fused_shared_experts > 0 and is_deepep_class_backend(): + elif use_per_rank_shared_slots: # Shared experts appended as extra columns in topk_ids: their value # would be out-of-bounds for the logical-to-physical dispatch table, # so split, dispatch the routed cols, recombine. @@ -1822,8 +1828,8 @@ def _post_process_topk_ids( ) topk_ids = torch.cat([routed_cols, shared_cols], dim=-1) # ExpertDistributionRecorder tracks EPLB physical routed experts. - # DeepEP dispatch later inserts per-rank shared slots into topk_ids, - # so keep the routed physical ids separately for statistics. + # Per-rank shared-slot remap later adds shared slots to the topk ID + # space, so keep the routed physical ids separately for statistics. recorder_topk_ids = routed_cols else: topk_ids = _biased_grouped_topk_postprocess( @@ -1854,12 +1860,11 @@ def _post_process_topk_ids( recorder_topk_ids = topk_ids _aiter_append = num_fused_shared_experts > 0 and _use_aiter - _deepep_remap = num_fused_shared_experts > 0 and is_deepep_class_backend() - if _aiter_append and _deepep_remap: - # Fused path: append shared experts AND apply the DeepEP interleaved + if _aiter_append and use_per_rank_shared_slots: + # Fused path: append shared experts AND apply the per-rank shared-slot # remap in a single Triton kernel. This replaces the original - # fused_append_shared_experts() + eager _remap_topk_for_deepep() pair, + # fused_append_shared_experts() + eager per-rank shared-slot remap pair, # collapsing ~6 launch-bound elementwise kernels/layer (div_floor / add / # arange / fill / copy) into the one append kernel that already runs. # @@ -1867,7 +1872,7 @@ def _post_process_topk_ids( # aiter_biased_grouped_topk folds routed_scaling_factor into the routed # weights and forward_deepep skips the post-MoE multiply for _use_aiter, # so the always-on shared expert must contribute 1.0x. (The eager - # _remap_topk_for_deepep instead sets shared weight to + # per-rank shared-slot remap instead sets shared weight to # 1/routed_scaling_factor to compensate a post-MoE scale that the aiter # path does not apply; see PR #28237.) num_physical_routed_experts = ( @@ -1914,15 +1919,16 @@ def _post_process_topk_ids( scale_factor, N, # base id for shared experts ) - elif _deepep_remap: - # DeepEP: remap to interleaved expert layout where each rank's shared - # expert has a unique ID for dispatch routing. + + elif use_per_rank_shared_slots: + # DeepEP/MegaMOE: remap to per-rank shared-slot layout where each + # rank's shared expert has a unique ID for dispatch routing. num_physical_routed_experts = ( expert_location_dispatch_info.num_physical_experts if expert_location_dispatch_info is not None else router_logits.shape[1] ) - topk_ids, topk_weights = _remap_topk_for_deepep( + topk_ids, topk_weights = remap_topk_for_per_rank_shared_slots( topk_ids, topk_weights, num_fused_shared_experts, diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 3e1c62c4d..50ef3f7d4 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -369,6 +369,16 @@ def is_deepep_class_backend() -> bool: return b.is_deepep() or b.is_mooncake() or b.is_mori() +def uses_per_rank_fused_shared_slots() -> bool: + """Check whether fused shared experts use per-rank physical slots.""" + return is_deepep_class_backend() or get_moe_a2a_backend().is_megamoe() + + +def has_per_rank_fused_shared_slots(num_fused_shared_experts: int) -> bool: + """Check whether this layer has fused shared experts in per-rank slots.""" + return num_fused_shared_experts > 0 and uses_per_rank_fused_shared_slots() + + def is_flashinfer_cutedsl_v1_path() -> bool: """CuteDSL v1 + DeepEP low-latency path (no MoeRunner, no autotune).""" return ( diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 48fe2a83e..cd5129f7e 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -80,6 +80,18 @@ _AITER_GFX95_CK_W8A8_MAX_SAFE_M = { } +class _MXFP4QuantizedData(MXFP4QuantizeUtil): + def __init__( + self, + original_shape: torch.Size, + original_dtype: torch.dtype, + quantized_data: torch.Tensor, + ): + self.original_shape = original_shape + self.original_dtype = original_dtype + self.quantized_data = quantized_data + + # Force CK bpreshuffle (not Triton) for the dense w8a8-block GEMMs (MLA q/kv/o # projections), to match ATOM (CK preshuffle; Triton FP8 blockscale is slower). # Default OFF; DeepseekV4 enables it via set_force_ck_w8a8(True). The env var @@ -1283,6 +1295,30 @@ def block_quant_dequant( return (x_q_block.to(torch.float32) * x_scale_repeat).to(dtype) +def quantize_block_fp8_weight_to_mxfp4( + fp8_weight: torch.Tensor, + fp8_scale: torch.Tensor, + weight_block_size: List[int], + mxfp4_block_size: int = 32, +) -> Tuple[torch.Tensor, torch.Tensor]: + fp8_weight_dequant = block_quant_dequant( + fp8_weight, + fp8_scale.to(torch.float32), + weight_block_size, + torch.bfloat16, + ) + fp4_weight, fp4_scale = _MXFP4QuantizedData.quantize( + fp8_weight_dequant, block_size=mxfp4_block_size + ) + fp4_weight = fp4_weight.quantized_data + fp4_weight = fp4_weight.contiguous().view(torch.int8) + fp4_scale = fp4_scale.view( + *fp8_weight_dequant.shape[:-1], + fp8_weight_dequant.shape[-1] // mxfp4_block_size, + ) + return fp4_weight, fp4_scale.contiguous().view(torch.float8_e8m0fnu) + + def requant_weight_ue8m0_inplace(weight, weight_scale_inv, weight_block_size): assert isinstance(weight, torch.nn.Parameter) assert isinstance(weight_scale_inv, torch.nn.Parameter) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 096f026f7..4c49ad812 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -98,6 +98,7 @@ from sglang.srt.layers.moe.topk import BypassedTopKOutput, TopK, TopKOutputForma from sglang.srt.layers.moe.utils import ( RoutingMethodType, filter_moe_weight_param_global_expert, + has_per_rank_fused_shared_slots, is_deepep_class_backend, is_sbo_enabled, is_tbo_enabled, @@ -554,18 +555,18 @@ class DeepseekV2MoE(nn.Module): # mlp.shared_experts → mlp.experts.256 when > 0. self.num_fused_shared_experts = 0 if _fusion_disabled else n_shared_experts - # DeepEP shared expert fusion: shared expert is fused into the same MoE kernel - # as a local expert at the home EP rank. Expert layout is expanded from 256 - # routed to 256+EP_size (e.g. 272 for EP=16). TopK handles interleaving. - _is_deepep_fusion = ( - is_deepep_class_backend() and self.num_fused_shared_experts > 0 + # DeepEP and MegaMOE shared expert fusion: shared expert is fused into + # the same MoE kernel as a local expert at each EP rank. Expert layout + # is expanded from 256 routed to 256+EP_size (e.g. 272 for EP=16). + _uses_per_rank_shared_slots = has_per_rank_fused_shared_slots( + self.num_fused_shared_experts ) - if _is_deepep_fusion: + if _uses_per_rank_shared_slots: # 256 routed + EP_size shared slots = 272 experts total (for EP=16) num_experts_for_moe = config.n_routed_experts + self.moe_ep_size top_k_for_moe = config.num_experts_per_tok + 1 # 8 routed + 1 shared - # Interleaving for DeepEP dispatch is handled by TopK internally. + # Interleaving for DeepEP/MegaMOE dispatch is handled by TopK internally. else: num_experts_for_moe = ( config.n_routed_experts + self.num_fused_shared_experts @@ -604,13 +605,13 @@ class DeepseekV2MoE(nn.Module): ) # scaling factor for fused shared experts on AMD-platform. - # DeepEP doesn't need this: shared expert is only computed on home rank + # DeepEP/MegaMOE doesn't need this: shared expert is only computed on home rank # (not all-reduced), so no 1/ep_size correction is needed. fused_shared_experts_scaling_factor = None if ( self.moe_ep_size > 1 and self.num_fused_shared_experts > 0 - and not _is_deepep_fusion + and not _uses_per_rank_shared_slots ): # if enable_ep_moe tp_szie == ep_size, every gpu get shared experts gemm output # so we scale with 1 / self.moe_ep_size in ep mode which will make it equalation as in tp mode @@ -688,13 +689,13 @@ class DeepseekV2MoE(nn.Module): self.shared_experts_is_fp8 = False self.shared_experts_weight_block_size = None self._shared_expert_tp1 = False - # Shared experts: skip when fused into MoE kernel (self.num_fused_shared_experts > 0) - # or when DeepEP fusion is enabled (shared expert is local slot 16 in FusedMoE, no separate MLP). + # Shared experts: skip when fused into MoE kernel + # (self.num_fused_shared_experts > 0) or when DeepEP/MegaMOE fusion is enabled. if ( config.n_shared_experts is not None and config.n_shared_experts > 0 and self.num_fused_shared_experts == 0 - and not _is_deepep_fusion + and not _uses_per_rank_shared_slots ): intermediate_size = config.moe_intermediate_size * config.n_shared_experts # Disable TP for shared experts for A2A/FP4 allgather paths, or when @@ -1243,16 +1244,6 @@ class DeepseekV2MoE(nn.Module): topk_output = self.topk.empty_topk_output( hidden_states.device, layer_id=self.layer_id ) - if is_deepep_class_backend() and self.num_fused_shared_experts > 0: - n = self.num_fused_shared_experts - topk_output = topk_output._replace( - topk_ids=topk_output.topk_ids.new_empty( - (0, topk_output.topk_ids.shape[-1] + n) - ), - topk_weights=topk_output.topk_weights.new_empty( - (0, topk_output.topk_weights.shape[-1] + n) - ), - ) if sbo_overlap_dispatch_flag: shared_output = None diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index a5d50fe71..452cba6af 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -1865,28 +1865,25 @@ class DeepseekV4ForCausalLM(nn.Module): if get_global_server_args().disable_shared_experts_fusion: return - # Waterfill needs shared-experts fusion so it can dispatch shared - # expert tokens to least-loaded EP ranks. - if get_global_server_args().enable_deepep_waterfill: + disable_reason = None + if get_global_server_args().enforce_shared_experts_fusion: if self.config.n_shared_experts != 1: raise ValueError( - "DeepEP Waterfill for DeepSeek V4 expects exactly one shared " + "DeepSeek V4 shared-experts fusion expects exactly one shared " f"expert, but got n_shared_experts={self.config.n_shared_experts}." ) - self.num_fused_shared_experts = self.config.n_shared_experts + else: + disable_reason = "Config does not support fused shared expert(s)." + + if disable_reason is not None: + get_global_server_args().disable_shared_experts_fusion = True log_info_on_rank0( logger, - "DeepSeek V4: --enable-deepep-waterfill set; KEEP shared-experts " - "fusion enabled so waterfill can rebalance shared expert dispatch.", + f"{disable_reason} Shared experts fusion optimization is disabled.", ) return - get_global_server_args().disable_shared_experts_fusion = True - log_info_on_rank0( - logger, - "DeepSeek V4 requires different clamping for shared and routed experts. " - "Shared experts fusion optimization is disabled.", - ) + self.num_fused_shared_experts = self.config.n_shared_experts @torch.no_grad() def forward( diff --git a/test/registered/moe/test_fused_append_remap_deepep.py b/test/registered/moe/test_fused_append_remap_per_rank_shared_slots.py similarity index 92% rename from test/registered/moe/test_fused_append_remap_deepep.py rename to test/registered/moe/test_fused_append_remap_per_rank_shared_slots.py index 2825ed53f..e3a7c3472 100644 --- a/test/registered/moe/test_fused_append_remap_deepep.py +++ b/test/registered/moe/test_fused_append_remap_per_rank_shared_slots.py @@ -1,8 +1,9 @@ -"""Unit tests for the fused append + DeepEP-remap shared-experts Triton kernel. +"""Unit tests for fused append + per-rank shared-slot remap. Covers ``fused_append_remap_shared_experts_deepep``, which collapses -``fused_append_shared_experts()`` followed by ``_remap_topk_for_deepep()`` into a -single Triton launch on the aiter/DeepEP-class path. The kernel is GPU-only +``fused_append_shared_experts()`` followed by +``remap_topk_for_per_rank_shared_slots()`` into a +single Triton launch on the per-rank shared-slot path. The kernel is GPU-only (Triton), so these tests are skipped when no accelerator is present. """ @@ -14,7 +15,11 @@ from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels impo fused_append_remap_shared_experts_deepep, fused_append_shared_experts, ) -from sglang.srt.layers.moe.topk import TopKConfig, _remap_topk_for_deepep, _use_aiter +from sglang.srt.layers.moe.topk import ( + TopKConfig, + _use_aiter, + remap_topk_for_per_rank_shared_slots, +) from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci @@ -50,7 +55,7 @@ def _reference_append_remap( @unittest.skipUnless( torch.cuda.is_available(), "fused append+remap kernel requires a GPU" ) -class TestFusedAppendRemapDeepEP(CustomTestCase): +class TestFusedAppendRemapPerRankSharedSlots(CustomTestCase): # (m, k, num_physical_routed, ep_size, ep_rank, num_fused_shared_experts). # k and num_fused_shared_experts are kept powers of two (tl.arange constraint). CASES = [ @@ -107,7 +112,7 @@ class TestFusedAppendRemapDeepEP(CustomTestCase): self.assertTrue(torch.allclose(got_w, exp_w)) def test_equivalence_with_eager_append_then_remap(self): - """Fused kernel == fused_append_shared_experts() + _remap_topk_for_deepep(). + """Fused kernel == append shared experts + per-rank shared-slot remap. The eager remap overwrites the shared weight: 1.0 on the aiter/HIP path (routed_scaling_factor is pre-folded into the routed topk weights), else @@ -140,7 +145,7 @@ class TestFusedAppendRemapDeepEP(CustomTestCase): scale_factor, npr, # shared-expert base id (overwritten by the remap) ) - eager_ids, eager_w = _remap_topk_for_deepep( + eager_ids, eager_w = remap_topk_for_per_rank_shared_slots( eager_ids, eager_w, s, diff --git a/test/registered/moe/test_hash_topk.py b/test/registered/moe/test_hash_topk.py new file mode 100644 index 000000000..2ef89cfb1 --- /dev/null +++ b/test/registered/moe/test_hash_topk.py @@ -0,0 +1,145 @@ +import sys +from types import SimpleNamespace + +import pytest +import torch + +from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +from sglang.srt.layers.moe import hash_topk as hash_topk_module +from sglang.srt.layers.moe.hash_topk import HashTopK +from sglang.srt.layers.moe.topk import ( + StandardTopKOutput, +) +from sglang.srt.models.deepseek_v2 import DeepseekV2MoE +from sglang.srt.runtime_context import get_parallel +from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-b-test-cpu") + + +@pytest.fixture(autouse=True) +def _set_dummy_server_args(): + set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) + + +def test_hash_topk_remaps_per_rank_fused_shared_slots(monkeypatch): + monkeypatch.setattr( + hash_topk_module, "has_per_rank_fused_shared_slots", lambda *_args: True + ) + recorded = {} + + class FakeRecorder: + def on_select_experts(self, *, topk_ids): + recorded["topk_ids"] = topk_ids.clone() + + monkeypatch.setattr( + hash_topk_module, + "get_global_expert_distribution_recorder", + lambda: FakeRecorder(), + ) + + topk = HashTopK( + topk=3, + num_experts=256, + num_fused_shared_experts=1, + vocab_size=2, + scoring_func="sqrtsoftplus", + routed_scaling_factor=2.5, + ) + with torch.no_grad(): + topk.tid2eid.copy_(torch.tensor([[0, 65], [63, 127]], dtype=torch.int32)) + + info = ExpertLocationDispatchInfo( + ep_dispatch_algorithm="static", + partial_logical_to_rank_dispatch_physical_map=torch.arange( + 256, dtype=torch.int32 + ), + partial_logical_to_all_physical_map=torch.arange(256, dtype=torch.int32).view( + 256, 1 + ), + partial_logical_to_all_physical_map_num_valid=torch.ones( + 256, dtype=torch.int32 + ), + num_physical_experts=256, + ) + + with ( + get_parallel().override(moe_ep_size=4, moe_ep_rank=2), + hash_topk_module.envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.override(False), + ): + output = topk( + hidden_states=torch.empty(2, 4), + router_logits=torch.ones(2, 256), + input_ids=torch.tensor([0, 1], dtype=torch.int64), + expert_location_dispatch_info=info, + ) + + # Physical layout for EP=4 has 64 routed slots per rank plus one local + # shared slot: [0..63, shared, 64..127, shared, ...]. + assert output.topk_ids.tolist() == [[0, 66, 194], [63, 128, 194]] + assert torch.allclose(output.topk_weights[:, -1], torch.full((2,), 0.4)) + assert recorded["topk_ids"].tolist() == [[0, 65], [63, 127]] + + +def test_hash_topk_empty_output_keeps_per_rank_shared_slot(monkeypatch): + monkeypatch.setattr( + hash_topk_module, "has_per_rank_fused_shared_slots", lambda *_args: True + ) + + topk = HashTopK( + topk=7, + num_experts=256, + num_fused_shared_experts=1, + vocab_size=2, + scoring_func="softmax", + ) + + output = topk.empty_topk_output(torch.device("cpu")) + + assert output.topk_ids.shape == (0, 7) + assert output.topk_weights.shape == (0, 7) + assert output.router_logits.shape == (0, 6) + + +def test_deepep_empty_forward_does_not_append_shared_slot_twice(): + captured = {} + + class FakeTopK: + def empty_topk_output(self, device, *, layer_id=None): + return StandardTopKOutput( + topk_weights=torch.empty((0, 9), dtype=torch.float32, device=device), + topk_ids=torch.empty((0, 9), dtype=torch.int32, device=device), + router_logits=torch.empty((0, 8), dtype=torch.float32, device=device), + ) + + class FakeExperts: + should_fuse_routed_scaling_factor_in_topk = True + + def __call__(self, hidden_states, topk_output): + captured["topk_ids_shape"] = tuple(topk_output.topk_ids.shape) + captured["topk_weights_shape"] = tuple(topk_output.topk_weights.shape) + return hidden_states + + moe = SimpleNamespace( + _fuse_shared_experts_inside_sbo=False, + is_nextn=False, + num_fused_shared_experts=1, + layer_id=0, + topk=FakeTopK(), + experts=FakeExperts(), + alt_stream=None, + routed_scaling_factor=1.0, + ) + + hidden_states = torch.empty((0, 4), dtype=torch.float32) + forward_batch = SimpleNamespace(num_token_non_padded=None) + + DeepseekV2MoE.forward_deepep(moe, hidden_states, forward_batch) + + assert captured["topk_ids_shape"] == (0, 9) + assert captured["topk_weights_shape"] == (0, 9) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/eplb/test_deepep_waterfill_eplb.py b/test/registered/unit/eplb/test_deepep_waterfill_eplb.py index cb699eca4..7ddf00196 100644 --- a/test/registered/unit/eplb/test_deepep_waterfill_eplb.py +++ b/test/registered/unit/eplb/test_deepep_waterfill_eplb.py @@ -56,7 +56,7 @@ class TestDeepEPWaterfillEPLB(CustomTestCase): self.assertEqual(len(weights), 1) self.assertEqual(weights[0].shape, (experts.num_local_experts, 2)) - def test_topk_recorder_ids_exclude_deepep_fused_shared_slots(self): + def test_topk_recorder_ids_exclude_per_rank_fused_shared_slots(self): topk_ids = torch.tensor([[0, 33, 263, 256]], dtype=torch.int32) topk_weights = torch.ones_like(topk_ids, dtype=torch.float32) topk_config = TopKConfig( @@ -74,7 +74,9 @@ class TestDeepEPWaterfillEPLB(CustomTestCase): with ( patch.object(topk_module, "_is_cuda", True), patch.object(topk_module, "_use_aiter", False), - patch.object(topk_module, "is_deepep_class_backend", return_value=True), + patch.object( + topk_module, "has_per_rank_fused_shared_slots", return_value=True + ), get_parallel().override(moe_ep_size=8, moe_ep_rank=7), patch.object( topk_module, @@ -94,7 +96,7 @@ class TestDeepEPWaterfillEPLB(CustomTestCase): self.assertTrue(torch.equal(processed_ids, torch.tensor([[0, 34, 270, 271]]))) self.assertTrue(torch.equal(recorder_ids, torch.tensor([[0, 33, 263]]))) - def test_topk_recorder_ids_match_dispatch_ids_for_non_deepep_fusion(self): + def test_topk_recorder_ids_match_dispatch_ids_without_per_rank_shared_slots(self): topk_ids = torch.tensor([[0, 33, 263, 256]], dtype=torch.int32) topk_weights = torch.ones_like(topk_ids, dtype=torch.float32) topk_config = TopKConfig( @@ -112,7 +114,9 @@ class TestDeepEPWaterfillEPLB(CustomTestCase): with ( patch.object(topk_module, "_is_cuda", True), patch.object(topk_module, "_use_aiter", False), - patch.object(topk_module, "is_deepep_class_backend", return_value=False), + patch.object( + topk_module, "has_per_rank_fused_shared_slots", return_value=False + ), patch.object( topk_module, "_biased_grouped_topk_postprocess", diff --git a/test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py b/test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py index 321613ebb..4c570b9d3 100644 --- a/test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py +++ b/test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py @@ -1,10 +1,10 @@ -"""Unit tests for fused shared-expert weight scaling on the DeepEP layout. +"""Unit tests for fused shared-expert weight scaling on per-rank shared slots. -These tests pin the contract of ``_remap_topk_for_deepep`` for the fused shared -expert's topk weight on the two paths this fix covers: +These tests pin the contract of ``remap_topk_for_per_rank_shared_slots`` for +the fused shared expert's topk weight on the two paths this fix covers: * aiter (HIP) path: routed_scaling_factor is folded into the routed weights and - forward_deepep skips the post-MoE multiply, so the shared weight must be 1.0 + the post-MoE multiply is skipped, so the shared weight must be 1.0 for a net 1.0x contribution. * post-MoE scaling path (default): the whole MoE output is multiplied by routed_scaling_factor afterward, so the shared weight must be 1/rsf. @@ -56,7 +56,7 @@ class TestFusedSharedExpertScaling(CustomTestCase): ), ), ): - _out_ids, out_weights = topk_module._remap_topk_for_deepep( + _out_ids, out_weights = topk_module.remap_topk_for_per_rank_shared_slots( topk_ids.clone(), topk_weights.clone(), num_fused_shared_experts=1, @@ -100,7 +100,7 @@ class TestFusedSharedExpertScaling(CustomTestCase): ), ), ): - out_ids, _ = topk_module._remap_topk_for_deepep( + out_ids, _ = topk_module.remap_topk_for_per_rank_shared_slots( topk_ids.clone(), topk_weights.clone(), num_fused_shared_experts=1, diff --git a/test/registered/unit/layers/quantization/test_fp8_utils_mxfp4.py b/test/registered/unit/layers/quantization/test_fp8_utils_mxfp4.py new file mode 100644 index 000000000..dca230142 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_fp8_utils_mxfp4.py @@ -0,0 +1,50 @@ +import unittest + +import torch + +from sglang.srt.layers.quantization.fp8_utils import ( + quantize_block_fp8_weight_to_mxfp4, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=4, suite="base-a-test-cpu") + + +class TestFp8UtilsMxfp4(unittest.TestCase): + def test_quantize_block_fp8_weight_to_mxfp4_shapes_and_dtype(self): + fp8_weight = ( + torch.linspace(-2.0, 2.0, 32 * 32, dtype=torch.float32) + .reshape(32, 32) + .to(torch.float8_e4m3fn) + ) + fp8_scale = torch.ones(1, 1, dtype=torch.float8_e8m0fnu) + + fp4_weight, fp4_scale = quantize_block_fp8_weight_to_mxfp4( + fp8_weight, fp8_scale, [128, 128] + ) + + self.assertEqual(fp4_weight.dtype, torch.int8) + self.assertEqual(fp4_weight.shape, torch.Size([32, 16])) + self.assertEqual(fp4_scale.dtype, torch.float8_e8m0fnu) + self.assertEqual(fp4_scale.shape, torch.Size([32, 1])) + + def test_quantize_block_fp8_weight_to_mxfp4_grouped_weight(self): + fp8_weight = ( + torch.linspace(-2.0, 2.0, 2 * 32 * 32, dtype=torch.float32) + .reshape(2, 32, 32) + .to(torch.float8_e4m3fn) + ) + fp8_scale = torch.ones(2, 1, 1, dtype=torch.float8_e8m0fnu) + + fp4_weight, fp4_scale = quantize_block_fp8_weight_to_mxfp4( + fp8_weight, fp8_scale, [128, 128] + ) + + self.assertEqual(fp4_weight.dtype, torch.int8) + self.assertEqual(fp4_weight.shape, torch.Size([2, 32, 16])) + self.assertEqual(fp4_scale.dtype, torch.float8_e8m0fnu) + self.assertEqual(fp4_scale.shape, torch.Size([2, 32, 1])) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py new file mode 100644 index 000000000..48aa53ff5 --- /dev/null +++ b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py @@ -0,0 +1,50 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.srt.models import deepseek_v4 as deepseek_v4_module +from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=4, suite="base-a-test-cpu") + + +class TestDeepseekV4SharedExpertFusionPolicy(unittest.TestCase): + def _make_model(self, n_shared_experts=1): + return SimpleNamespace( + config=SimpleNamespace(n_shared_experts=n_shared_experts) + ) + + def test_disables_shared_fusion_without_enforce(self): + server_args = SimpleNamespace( + disable_shared_experts_fusion=False, + enforce_shared_experts_fusion=False, + ) + model = self._make_model() + + with patch.object( + deepseek_v4_module, "get_global_server_args", return_value=server_args + ): + DeepseekV4ForCausalLM.determine_num_fused_shared_experts(model) + + self.assertEqual(model.num_fused_shared_experts, 0) + self.assertTrue(server_args.disable_shared_experts_fusion) + + def test_enables_shared_fusion_when_enforced(self): + server_args = SimpleNamespace( + disable_shared_experts_fusion=False, + enforce_shared_experts_fusion=True, + ) + model = self._make_model() + + with patch.object( + deepseek_v4_module, "get_global_server_args", return_value=server_args + ): + DeepseekV4ForCausalLM.determine_num_fused_shared_experts(model) + + self.assertEqual(model.num_fused_shared_experts, 1) + self.assertFalse(server_args.disable_shared_experts_fusion) + + +if __name__ == "__main__": + unittest.main()