feat(moe): add fp8 megamoe fused path and nvfp4 layout guards

Wire MegaMoE FP8 through DeepGEMM's fused fp8_mega_moe path, preserve fallback runner layouts, and add explicit NVFP4 group-size guardrails for unsupported DeepGEMM scale transforms. Tested: PYTHONPYCACHEPREFIX=/private/tmp/sglang_pycache python3 -m py_compile python/sglang/srt/layers/moe/mega_moe.py python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py python/sglang/srt/layers/quantization/fp8.py python/sglang/srt/layers/quantization/modelopt_quant.py test/registered/unit/moe/test_glm_megamoe.py Tested: git diff --check Not-tested: python3 -m unittest test/registered/unit/moe/test_glm_megamoe.py (local Python environments do not have torch installed).
2026-07-08 11:42:51 +08:00
parent eb7f44a8ee
commit 89ba17ad05
5 changed files with 465 additions and 15 deletions
--- a/python/sglang/srt/layers/moe/mega_moe.py
+++ b/python/sglang/srt/layers/moe/mega_moe.py
@@ -17,12 +17,14 @@ from __future__ import annotations

 import os
 from contextlib import nullcontext
+from math import ceil
 from typing import TYPE_CHECKING, Optional

 import torch

 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
+from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.dp_attention import get_dp_global_num_tokens
 from sglang.srt.layers.moe.topk import select_experts
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
@@ -57,6 +59,7 @@ def _get_mega_moe_symm_buffer(
    num_topk: int,
    hidden: int,
    intermediate_hidden: int,
+    mma_type: str = "fp8xfp4",
 ) -> SymmBuffer:
    import deep_gemm

@@ -69,6 +72,7 @@ def _get_mega_moe_symm_buffer(
        num_topk,
        hidden,
        intermediate_hidden,
+        mma_type,
    )
    buf = _MEGA_MOE_SYMM_BUFFER.get(key)
    if buf is None:
@@ -80,6 +84,7 @@ def _get_mega_moe_symm_buffer(
            hidden,
            intermediate_hidden,
            use_fp8_dispatch=True,
+            mma_type=mma_type,
            activation="swiglu",
        )
        _MEGA_MOE_SYMM_BUFFER[key] = buf
@@ -91,8 +96,11 @@ def should_use_mega_moe(moe, hidden_states: torch.Tensor) -> bool:
        return False
    if not getattr(moe.experts, "_mega_moe_weights_built", False):
        return False
-    # This fork does not carry the DSV4 JIT FP8 pre-dispatch path yet. Keep the
-    # fast path gated to DeepGEMM's FP4 activation pre-dispatch API.
+    weight_format = getattr(moe.experts, "_mega_moe_weight_format", None)
+    if weight_format == "fp8":
+        return deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+    if weight_format != "fp4":
+        return False
    if not envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get():
        return False
    if get_is_capture_mode():
@@ -209,6 +217,18 @@ def _run_mega_routed(
    input_ids_global: Optional[torch.Tensor],
    num_tokens: int,
 ) -> torch.Tensor:
+    weight_format = getattr(moe.experts, "_mega_moe_weight_format", None)
+    if weight_format == "fp8":
+        return _run_mega_fp8_routed(
+            moe, hidden_states, forward_batch, input_ids_global, num_tokens
+        )
+    if weight_format != "fp4":
+        raise RuntimeError(
+            f"Unsupported MegaMoE weight format: {weight_format!r}. "
+            "Build weights with build_mega_moe_experts_weights() or "
+            "build_mega_moe_fp8_experts_weights()."
+        )
+
    import deep_gemm

    from sglang.srt.distributed.parallel_state import get_moe_ep_group
@@ -258,6 +278,7 @@ def _run_mega_routed(
        num_topk=top_k,
        hidden=hidden_size,
        intermediate_hidden=intermediate_size,
+        mma_type="fp8xfp4",
    )

    if num_tokens > 0:
@@ -304,6 +325,101 @@ def _run_mega_routed(
    return y


+def _run_mega_fp8_routed(
+    moe,
+    hidden_states: torch.Tensor,
+    forward_batch: Optional[ForwardBatch],
+    input_ids_global: Optional[torch.Tensor],
+    num_tokens: int,
+) -> torch.Tensor:
+    import deep_gemm
+
+    from sglang.srt.distributed.parallel_state import get_moe_ep_group
+
+    if num_tokens > 0:
+        router_logits = _call_gate(moe, hidden_states, forward_batch)
+        topk_output = _select_mega_topk(
+            moe, hidden_states, router_logits, forward_batch, input_ids_global
+        )
+        topk_ids = topk_output.topk_ids
+        topk_weights = topk_output.topk_weights
+    else:
+        topk_ids = None
+        topk_weights = None
+
+    experts = moe.experts
+    hidden_size = getattr(experts, "_mega_moe_hidden_size", hidden_states.shape[-1])
+    ep_group = get_moe_ep_group().device_group
+    num_experts = experts.num_experts
+    top_k = _get_mega_top_k(moe)
+    intermediate_size = getattr(
+        experts,
+        "_mega_moe_intermediate_size",
+        getattr(moe.config, "moe_intermediate_size"),
+    )
+    num_max_tokens_per_rank = (
+        envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get()
+    )
+    assert num_tokens <= num_max_tokens_per_rank, (
+        f"mega MoE: num_tokens={num_tokens} exceeds cap "
+        f"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK="
+        f"{num_max_tokens_per_rank}; raise the env var or shrink "
+        f"cuda_graph_max_bs / chunked_prefill_size accordingly"
+    )
+
+    buf = _get_mega_moe_symm_buffer(
+        ep_group,
+        num_experts=num_experts,
+        num_max_tokens_per_rank=num_max_tokens_per_rank,
+        num_topk=top_k,
+        hidden=hidden_size,
+        intermediate_hidden=intermediate_size,
+        mma_type="fp8xfp8",
+    )
+
+    if num_tokens > 0:
+        topk_ids_in = topk_ids.to(torch.int32)
+        topk_weights_in = topk_weights.to(torch.float32)
+    else:
+        topk_ids_in = hidden_states.new_empty((0, top_k), dtype=torch.int32)
+        topk_weights_in = hidden_states.new_empty((0, top_k), dtype=torch.float32)
+
+    deep_gemm.mega_moe_pre_dispatch(
+        hidden_states,
+        topk_ids_in,
+        topk_weights_in,
+        buf.x,
+        buf.x_sf,
+        buf.topk_idx,
+        buf.topk_weights,
+        num_tokens=num_tokens,
+        group_size=128,
+        use_fp4_acts=False,
+    )
+
+    y = torch.empty(
+        (max(num_tokens, 1), hidden_size),
+        dtype=torch.bfloat16,
+        device=hidden_states.device,
+    )
+    swiglu_limit = getattr(moe.config, "swiglu_limit", None)
+    deep_gemm.fp8_mega_moe(
+        y,
+        experts.mega_l1_weights,
+        experts.mega_l2_weights,
+        buf,
+        recipe=(128, 128, 128),
+        activation="swiglu",
+        activation_clamp=swiglu_limit,
+        fast_math=True,
+    )
+    y = y[:num_tokens]
+
+    if not getattr(experts, "should_fuse_routed_scaling_factor_in_topk", False):
+        y.mul_(getattr(moe, "routed_scaling_factor", 1.0))
+    return y
+
+
 def _as_deep_gemm_packed_fp4_weights(
    weights: tuple[torch.Tensor, torch.Tensor],
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -371,6 +487,37 @@ def _infer_mega_moe_scale_group_size(
    return k // num_scale_cols


+def _transform_mega_moe_sf(
+    transform_sf_into_required_layout,
+    sf: torch.Tensor,
+    *,
+    mn: int,
+    k: int,
+    group_size: int,
+    num_groups: int,
+    name: str,
+) -> torch.Tensor:
+    try:
+        return transform_sf_into_required_layout(
+            sf.to(torch.float32),
+            mn=mn,
+            k=k,
+            recipe=(1, group_size),
+            num_groups=num_groups,
+            disable_ue8m0_cast=False,
+        )
+    except RuntimeError as exc:
+        if "Unknown SF transformation" not in str(exc):
+            raise
+        raise RuntimeError(
+            f"DeepGEMM does not support the MegaMoE {name} scale layout "
+            f"recipe=(1,{group_size}) for mn={mn}, k={k}, "
+            f"num_groups={num_groups}. GLM NVFP4 checkpoints use group16 "
+            "scales, so this requires an sgl-deep-gemm/DeepGEMM build with "
+            "group16 MegaMoE SF transformation support."
+        ) from exc
+
+
 def build_mega_moe_experts_weights(
    experts,
    *,
@@ -397,21 +544,23 @@ def build_mega_moe_experts_weights(
            f"{group_k1=} and {group_k2=}."
        )

-    w13_sf = transform_sf_into_required_layout(
-        w13_sf_src.to(torch.float32),
+    w13_sf = _transform_mega_moe_sf(
+        transform_sf_into_required_layout,
+        w13_sf_src,
        mn=n1,
        k=k1,
-        recipe=(1, group_k1),
+        group_size=group_k1,
        num_groups=num_groups,
-        disable_ue8m0_cast=False,
+        name="w13",
    )
-    w2_sf = transform_sf_into_required_layout(
-        w2_sf_src.to(torch.float32),
+    w2_sf = _transform_mega_moe_sf(
+        transform_sf_into_required_layout,
+        w2_sf_src,
        mn=n2,
        k=k2,
-        recipe=(1, group_k2),
+        group_size=group_k2,
        num_groups=num_groups,
-        disable_ue8m0_cast=False,
+        name="w2",
    )

    if preserve_runner_layout:
@@ -455,4 +604,170 @@ def build_mega_moe_experts_weights(
    experts._mega_moe_hidden_size = k1
    experts._mega_moe_intermediate_size = k2
    experts._mega_moe_weight_group_size = group_k1
+    experts._mega_moe_weight_format = "fp4"
+    experts._mega_moe_weights_built = True
+
+
+def _swap_gate_up_halves(weight: torch.Tensor) -> torch.Tensor:
+    half = weight.shape[1] // 2
+    return torch.cat((weight[:, half:], weight[:, :half]), dim=1).contiguous()
+
+
+def _normalize_mega_moe_fp8_block_shape(block_shape) -> tuple[int, int]:
+    if block_shape is None:
+        return 128, 128
+    block_mn, block_k = int(block_shape[0]), int(block_shape[1])
+    if (block_mn, block_k) != (128, 128):
+        raise ValueError(
+            "MegaMoE FP8 fused path currently requires 128x128 weight "
+            f"block scales, got {block_shape}."
+        )
+    return block_mn, block_k
+
+
+def _expand_mega_moe_fp8_scale(
+    scale: torch.Tensor,
+    *,
+    mn: int,
+    k: int,
+    block_shape: tuple[int, int],
+    num_groups: int,
+    name: str,
+) -> torch.Tensor:
+    block_mn, block_k = block_shape
+    expected_shape = (num_groups, ceil(mn / block_mn), ceil(k / block_k))
+    if scale.dim() == 1:
+        if scale.shape[0] != num_groups:
+            raise ValueError(
+                f"MegaMoE FP8 {name} per-expert scale must have "
+                f"{num_groups} entries, got shape={tuple(scale.shape)}."
+            )
+        return scale.view(num_groups, 1, 1).expand(expected_shape).contiguous()
+    if tuple(scale.shape) != expected_shape:
+        raise ValueError(
+            f"MegaMoE FP8 {name} scale shape must be {expected_shape}, "
+            f"got {tuple(scale.shape)}."
+        )
+    return scale.contiguous()
+
+
+def _transform_mega_moe_fp8_sf(
+    transform_sf_into_required_layout,
+    sf: torch.Tensor,
+    *,
+    mn: int,
+    k: int,
+    block_shape: tuple[int, int],
+    num_groups: int,
+    name: str,
+) -> torch.Tensor:
+    block_mn, block_k = block_shape
+    sf = _expand_mega_moe_fp8_scale(
+        sf,
+        mn=mn,
+        k=k,
+        block_shape=block_shape,
+        num_groups=num_groups,
+        name=name,
+    )
+    try:
+        return transform_sf_into_required_layout(
+            sf.to(torch.float32),
+            mn=mn,
+            k=k,
+            recipe=(block_mn, block_k),
+            num_groups=num_groups,
+            disable_ue8m0_cast=False,
+        )
+    except RuntimeError as exc:
+        if "Unknown SF transformation" not in str(exc):
+            raise
+        raise RuntimeError(
+            f"DeepGEMM does not support the MegaMoE FP8 {name} scale layout "
+            f"recipe=({block_mn},{block_k}) for mn={mn}, k={k}, "
+            f"num_groups={num_groups}."
+        ) from exc
+
+
+def build_mega_moe_fp8_experts_weights(
+    experts,
+    *,
+    preserve_runner_layout: bool = True,
+    swap_w13_halves: bool = False,
+) -> None:
+    from deep_gemm import (
+        transform_sf_into_required_layout,
+        transform_weights_for_mega_moe,
+    )
+
+    if getattr(experts, "_mega_moe_weights_built", False):
+        return
+
+    w13_weight = experts.w13_weight.data
+    w2_weight = experts.w2_weight.data
+    if swap_w13_halves:
+        w13_weight = _swap_gate_up_halves(w13_weight)
+    elif preserve_runner_layout:
+        w13_weight = w13_weight.clone()
+    if preserve_runner_layout:
+        w2_weight = w2_weight.clone()
+
+    block_quant = hasattr(experts, "w13_weight_scale_inv") and hasattr(
+        experts, "w2_weight_scale_inv"
+    )
+    if block_quant:
+        w13_scale = experts.w13_weight_scale_inv.data
+        w2_scale = experts.w2_weight_scale_inv.data
+        if preserve_runner_layout:
+            w13_scale = w13_scale.clone()
+            w2_scale = w2_scale.clone()
+        quant_config = getattr(
+            getattr(experts, "quant_method", None), "quant_config", None
+        )
+        block_shape = getattr(quant_config, "weight_block_size", None)
+    else:
+        w13_scale = experts.w13_weight_scale.data
+        w2_scale = experts.w2_weight_scale.data
+        if preserve_runner_layout:
+            w13_scale = w13_scale.clone()
+            w2_scale = w2_scale.clone()
+        block_shape = [128, 128]
+
+    block_shape = _normalize_mega_moe_fp8_block_shape(block_shape)
+    num_groups, n1, k1 = w13_weight.shape
+    _, n2, k2 = w2_weight.shape
+    w13_sf = _transform_mega_moe_fp8_sf(
+        transform_sf_into_required_layout,
+        w13_scale,
+        mn=n1,
+        k=k1,
+        block_shape=block_shape,
+        num_groups=num_groups,
+        name="w13",
+    )
+    w2_sf = _transform_mega_moe_fp8_sf(
+        transform_sf_into_required_layout,
+        w2_scale,
+        mn=n2,
+        k=k2,
+        block_shape=block_shape,
+        num_groups=num_groups,
+        name="w2",
+    )
+    l1_pair, l2_pair = transform_weights_for_mega_moe(
+        (w13_weight, w13_sf), (w2_weight, w2_sf)
+    )
+
+    experts.mega_fp8_w13_weight = w13_weight
+    experts.mega_fp8_w2_weight = w2_weight
+    experts.mega_fp8_w13_scale = w13_scale
+    experts.mega_fp8_w2_scale = w2_scale
+    experts.mega_l1_weights = l1_pair
+    experts.mega_l2_weights = l2_pair
+    experts._mega_moe_fp8_block_quant = block_quant
+    experts._mega_moe_fp8_block_shape = block_shape
+    experts._mega_moe_preserve_runner_layout = preserve_runner_layout
+    experts._mega_moe_hidden_size = w13_weight.shape[2]
+    experts._mega_moe_intermediate_size = w2_weight.shape[2]
+    experts._mega_moe_weight_format = "fp8"
    experts._mega_moe_weights_built = True
--- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
@@ -1258,6 +1258,19 @@ def fused_experts_none_to_flashinfer_trtllm(
    )


+@register_fused_func("megamoe", "flashinfer_trtllm")
+def fused_experts_megamoe_to_flashinfer_trtllm(
+    dispatch_output: StandardDispatchOutput,
+    quant_info: MoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+) -> StandardCombineInput:
+    # MegaMoE uses the standard dispatcher in this fork. Keep FlashInfer TRTLLM
+    # available as the fallback path when the MegaMoE fast path is not selected.
+    return fused_experts_none_to_flashinfer_trtllm(
+        dispatch_output, quant_info, runner_config
+    )
+
+
@register_fused_func("none", "flashinfer_trtllm_routed")
 def fused_experts_none_to_flashinfer_trtllm_routed(
    dispatch_output: StandardDispatchOutput,
--- a/python/sglang/srt/layers/quantization/fp8.py
+++ b/python/sglang/srt/layers/quantization/fp8.py
@@ -27,7 +27,11 @@ from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
    FlashInferTrtllmFp8MoeQuantInfo,
 )
 from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
-from sglang.srt.layers.moe.utils import RoutingMethodType, get_moe_runner_backend
+from sglang.srt.layers.moe.utils import (
+    RoutingMethodType,
+    get_moe_a2a_backend,
+    get_moe_runner_backend,
+)
 from sglang.srt.layers.parameter import (
    BlockQuantScaleParameter,
    ModelWeightParameter,
@@ -757,6 +761,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
        moe_runner_backend = get_moe_runner_backend()
        if moe_runner_backend.is_deep_gemm():
            return True
+        if get_moe_a2a_backend().is_megamoe():
+            return deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
        if moe_runner_backend.is_auto():
            return deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and (
                get_moe_a2a_backend().is_deepep()
@@ -1038,9 +1044,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                and will_use_deepgemm
                and not layer.w13_weight_scale_inv.format_ue8m0
            ):
-                assert isinstance(
-                    layer, DeepEPMoE
-                ), "DeepGemm MoE is only supported with DeepEPMoE"
+                assert (
+                    isinstance(layer, DeepEPMoE)
+                    or get_moe_a2a_backend().is_megamoe()
+                ), (
+                    "DeepGemm MoE is only supported with DeepEPMoE or the "
+                    "MegaMoE fast path"
+                )
                weight_block_size = self.quant_config.weight_block_size
                requant_weight_ue8m0_inplace(
                    layer.w13_weight, layer.w13_weight_scale_inv, weight_block_size
@@ -1237,6 +1247,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):
        # Block quant doesn't need to process weights after loading
        if self.block_quant:
            self.process_weights_after_loading_block_quant(layer)
+            if get_moe_a2a_backend().is_megamoe():
+                from sglang.srt.layers.moe.mega_moe import (
+                    build_mega_moe_fp8_experts_weights,
+                )
+
+                build_mega_moe_fp8_experts_weights(
+                    layer, preserve_runner_layout=True
+                )
            return

        # If checkpoint is fp16 or bfloat16, quantize in place.
@@ -1267,6 +1285,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):

            if _is_hip:
                self.process_weights_hip_scale_padding(layer)
+            if get_moe_a2a_backend().is_megamoe():
+                from sglang.srt.layers.moe.mega_moe import (
+                    build_mega_moe_fp8_experts_weights,
+                )
+
+                build_mega_moe_fp8_experts_weights(
+                    layer, preserve_runner_layout=True
+                )
            return

        # If checkpoint is fp8, we need to handle that the
@@ -1351,6 +1377,15 @@ class Fp8MoEMethod(FusedMoEMethodBase):
            if _is_hip:
                self.process_weights_hip_scale_padding(layer)

+            if get_moe_a2a_backend().is_megamoe():
+                from sglang.srt.layers.moe.mega_moe import (
+                    build_mega_moe_fp8_experts_weights,
+                )
+
+                build_mega_moe_fp8_experts_weights(
+                    layer, preserve_runner_layout=True
+                )
+
            # Align FP8 weights to FlashInfer per-tensor kernel layout if enabled
            if get_moe_runner_backend().is_flashinfer_trtllm():
                from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
--- a/python/sglang/srt/layers/quantization/modelopt_quant.py
+++ b/python/sglang/srt/layers/quantization/modelopt_quant.py
@@ -958,6 +958,17 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
                layer.w2_input_scale.max(), requires_grad=False
            )

+        if get_moe_a2a_backend().is_megamoe():
+            from sglang.srt.layers.moe.mega_moe import (
+                build_mega_moe_fp8_experts_weights,
+            )
+
+            # ModelOpt FP8 stores W13 in [Up, Gate] order. Preserve a DeepGEMM
+            # compatible [Gate, Up] copy before FlashInfer mutates runner layout.
+            build_mega_moe_fp8_experts_weights(
+                layer, preserve_runner_layout=True, swap_w13_halves=True
+            )
+
        # Align FP8 weights to FlashInfer per-tensor kernel layout if enabled
        if get_moe_runner_backend().is_flashinfer_trtllm():
            from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (