Support DSV4 shared expert fusion for DeepEP and MegaMOE (#27349)

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)

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
         )

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,

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 (

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)

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

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(

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,

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"]))

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",

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,

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()

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()