Add SM90 FP8 MegaMoE support for DeepSeek-V4 (#29016)

Co-authored-by: yinding <yinding@bytedance.com>

python/sglang/srt/layers/moe/mega_moe.py

@@ -26,8 +26,13 @@ from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.dp_attention import get_dp_global_num_tokens
+from sglang.srt.layers.moe.mega_moe_sm90 import (
+    is_sm90_fp8_mega_moe_available,
+    run_sm90_mega_routed,
+)
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
 from sglang.srt.model_executor.runner import get_is_capture_mode
+from sglang.srt.models.deepseek_common.utils import _device_sm
 
 if TYPE_CHECKING:
     from deep_gemm import SymmBuffer
@@ -100,6 +105,9 @@ def should_use_mega_moe(moe: DeepseekV2MoE, hidden_states: torch.Tensor) -> bool
         return False
     if not getattr(moe.experts, "_mega_moe_weights_built", False):
         return False
+    if _device_sm == 90:
+        if not is_sm90_fp8_mega_moe_available(moe.experts):
+            return False
     if get_is_capture_mode():
         return True
 
@@ -214,6 +222,16 @@ def _run_mega_routed(
         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)
 
+    if _device_sm == 90:
+        return run_sm90_mega_routed(
+            moe,
+            hidden_states,
+            topk_ids_in,
+            topk_weights_in,
+            buf,
+            num_tokens,
+        )
+
     use_fp4_acts = envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get()
     if use_fp4_acts:
         # FP4 path goes through DeepGEMM's mega_moe_pre_dispatch which

python/sglang/srt/layers/moe/mega_moe_sm90.py

@@ -0,0 +1,179 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SM90 FP8 Mega-MoE forward path and expert-weight prep."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.models.deepseek_common.utils import _device_sm
+
+if TYPE_CHECKING:
+    from deep_gemm import SymmBuffer
+
+    from sglang.srt.models.deepseek_v2 import DeepseekV2MoE
+
+
+def is_sm90_fp8_mega_moe_available(experts) -> bool:
+    if _device_sm != 90:
+        return False
+    try:
+        import deep_gemm
+    except ImportError:
+        return False
+    return (
+        hasattr(deep_gemm, "fp8_mega_moe")
+        and hasattr(deep_gemm, "mega_moe_pre_dispatch_sm90")
+        and getattr(experts, "_mega_moe_sm90_fp8_weights", False)
+    )
+
+
+def run_sm90_mega_routed(
+    moe: DeepseekV2MoE,
+    hidden_states: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    buf: SymmBuffer,
+    num_tokens: int,
+) -> torch.Tensor:
+    import deep_gemm
+
+    if moe.experts.should_fuse_routed_scaling_factor_in_topk:
+        routed_scaling_factor = 1.0
+    else:
+        routed_scaling_factor = float(moe.routed_scaling_factor)
+
+    deep_gemm.mega_moe_pre_dispatch_sm90(
+        hidden_states,
+        topk_ids,
+        topk_weights,
+        buf.x,
+        buf.x_sf,
+        buf.topk_idx,
+        buf.topk_weights,
+        num_tokens=num_tokens,
+        group_size=128,
+        routed_scaling_factor=routed_scaling_factor,
+    )
+
+    y = torch.empty(
+        (max(num_tokens, 1), moe.config.hidden_size),
+        dtype=torch.bfloat16,
+        device=hidden_states.device,
+    )
+    deep_gemm.fp8_mega_moe(
+        y,
+        moe.experts.mega_l1_weights,
+        moe.experts.mega_l2_weights,
+        buf,
+        recipe=(128, 128, 128),
+        activation="swiglu",
+        activation_clamp=getattr(moe.config, "swiglu_limit", None),
+        fast_math=True,
+    )
+    y = y[:num_tokens]
+
+    return y
+
+
+def _interleave_l1_weight_only(weight: torch.Tensor, gran: int = 8) -> torch.Tensor:
+    num_groups, n, *rest = weight.shape
+    half = n // 2
+    gate = weight[:, :half].reshape(num_groups, half // gran, gran, *rest)
+    up = weight[:, half:].reshape(num_groups, half // gran, gran, *rest)
+    return torch.stack([gate, up], dim=2).reshape(num_groups, n, *rest)
+
+
+def build_sm90_mega_moe_experts_weights(experts) -> None:
+    if getattr(experts, "_mega_moe_weights_built", False):
+        return
+
+    w13 = experts.w13_weight.data
+    w13_sf_fp32 = experts.w13_weight_scale_inv.data
+    w2 = experts.w2_weight.data
+    w2_sf_fp32 = experts.w2_weight_scale_inv.data
+
+    assert w13.dtype == torch.float8_e4m3fn
+    assert w2.dtype == torch.float8_e4m3fn
+
+    num_groups, n1, k1 = w13.shape
+    _, n2, k2 = w2.shape
+    scale_group_mn, scale_group_k = 128, 128
+
+    assert k1 % scale_group_k == 0 and k2 % scale_group_k == 0, (
+        f"invalid SM90 mega-moe K/group_size: k1={k1}, k2={k2}, "
+        f"group_k={scale_group_k}"
+    )
+    expected_n_groups_1 = (n1 + scale_group_mn - 1) // scale_group_mn
+    expected_n_groups_2 = (n2 + scale_group_mn - 1) // scale_group_mn
+    expected_k_groups_1 = k1 // scale_group_k
+    expected_k_groups_2 = k2 // scale_group_k
+    assert w13_sf_fp32.shape[1] == expected_n_groups_1, (
+        f"w13 scale N groups mismatch: got {w13_sf_fp32.shape[1]}, "
+        f"expected {expected_n_groups_1} (n1={n1}, group_mn={scale_group_mn})"
+    )
+    assert w2_sf_fp32.shape[1] == expected_n_groups_2, (
+        f"w2 scale N groups mismatch: got {w2_sf_fp32.shape[1]}, "
+        f"expected {expected_n_groups_2} (n2={n2}, group_mn={scale_group_mn})"
+    )
+    assert w13_sf_fp32.shape[2] == expected_k_groups_1, (
+        f"w13 scale K groups mismatch: got {w13_sf_fp32.shape[2]}, "
+        f"expected {expected_k_groups_1} (k1={k1}, group_k={scale_group_k})"
+    )
+    assert w2_sf_fp32.shape[2] == expected_k_groups_2, (
+        f"w2 scale K groups mismatch: got {w2_sf_fp32.shape[2]}, "
+        f"expected {expected_k_groups_2} (k2={k2}, group_k={scale_group_k})"
+    )
+
+    if envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.get():
+        w13_interleaved = _interleave_l1_weight_only(w13)
+        experts.w13_weight.data = w13_interleaved
+        experts.mega_l1_weights = (
+            experts.w13_weight.data,
+            experts.w13_weight_scale_inv.data,
+        )
+        experts.mega_l2_weights = (
+            experts.w2_weight.data,
+            experts.w2_weight_scale_inv.data,
+        )
+    else:
+        import deep_gemm
+
+        w13_sf = deep_gemm.transform_sf_into_required_layout(
+            w13_sf_fp32,
+            mn=n1,
+            k=k1,
+            recipe=(128, 128),
+            num_groups=num_groups,
+            disable_ue8m0_cast=True,
+        )
+        w2_sf = deep_gemm.transform_sf_into_required_layout(
+            w2_sf_fp32,
+            mn=n2,
+            k=k2,
+            recipe=(128, 128),
+            num_groups=num_groups,
+            disable_ue8m0_cast=True,
+        )
+        l1_pair, l2_pair = deep_gemm.transform_weights_for_mega_moe_sm90(
+            (w13, w13_sf), (w2, w2_sf)
+        )
+        experts.mega_l1_weights = l1_pair
+        experts.mega_l2_weights = l2_pair
+
+    experts._mega_moe_sm90_fp8_weights = True
+    experts._mega_moe_weights_built = True

python/sglang/srt/layers/quantization/fp8.py

@@ -1663,6 +1663,15 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                     layer.w13_weight_scale_inv.format_ue8m0 = True
                     layer.w2_weight_scale_inv.format_ue8m0 = True
 
+            if get_moe_a2a_backend().is_megamoe() and is_sm90_supported():
+                from sglang.srt.layers.moe.mega_moe_sm90 import (
+                    build_sm90_mega_moe_experts_weights,
+                )
+
+                assert not self.is_fp4_expert
+                build_sm90_mega_moe_experts_weights(layer)
+                return
+
             if not self.is_fp4_expert:
                 weight_block_size = self.quant_config.weight_block_size
                 if requant_block_scale_ue8m0_for_deepgemm(