Bump sgl-deep-gemm to 0.1.7 (#37279)

2026-08-31 17:16:34 -07:00
parent 455232de6e
commit 07c8f7294d
8 changed files with 301 additions and 29 deletions
@@ -1,7 +1,6 @@
 from __future__ import annotations

 import logging
-import os
 from typing import TYPE_CHECKING

 if TYPE_CHECKING:
@@ -17,7 +16,6 @@ logger = logging.getLogger(__name__)

 def handle_mega_moe(server_args: ServerArgs) -> None:
    handle_moe_runner_backend_alias(server_args)
-    handle_w4a4_mxfp4_megamoe_env(server_args)


 def handle_moe_runner_backend_alias(server_args: ServerArgs) -> None:
@@ -38,12 +36,3 @@ def handle_moe_runner_backend_alias(server_args: ServerArgs) -> None:
        moe_runner_backend="auto",
        moe_a2a_backend="megamoe",
    )
-
-
-def handle_w4a4_mxfp4_megamoe_env(server_args: ServerArgs) -> None:
-    cfg = resolving_view(server_args)
-    if not cfg.enable_w4a4_mxfp4_megamoe:
-        return
-
-    os.environ["DG_USE_FP4_ACTS"] = "1"
-    os.environ["DG_USE_MXF4_KIND"] = "1"
@@ -16,7 +16,6 @@
 from __future__ import annotations

 import functools
-import os
 from contextlib import contextmanager, nullcontext
 from typing import TYPE_CHECKING, Optional

@@ -34,6 +33,7 @@ from sglang.srt.layers.moe.mega_moe_sm90 import (
 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
+from sglang.srt.runtime_context import get_exec

 if TYPE_CHECKING:
    from deep_gemm import SymmBuffer
@@ -45,6 +45,10 @@ if TYPE_CHECKING:
 _MEGA_MOE_SYMM_BUFFER: dict = {}


+def _mega_moe_mma_type() -> str:
+    return "mxf4xmxf4" if get_exec().moe.enable_w4a4_mxfp4_megamoe else "fp8xfp4"
+
+
@functools.lru_cache(maxsize=1)
 def _mega_moe_max_num_sms() -> Optional[int]:
    if _device_sm < 100:
@@ -92,6 +96,7 @@ def _get_mega_moe_symm_buffer(
 ) -> SymmBuffer:
    import deep_gemm

+    mma_type = _mega_moe_mma_type()
    key = (
        id(group),
        num_max_tokens_per_rank,
@@ -99,6 +104,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:
@@ -109,7 +115,7 @@ def _get_mega_moe_symm_buffer(
            num_topk,
            hidden,
            intermediate_hidden,
-            use_fp8_dispatch=True,
+            mma_type=mma_type,
            activation="swiglu",
        )
        _MEGA_MOE_SYMM_BUFFER[key] = buf
@@ -248,8 +254,8 @@ def _run_mega_routed(
            num_tokens,
        )

-    use_fp4_acts = os.getenv("DG_USE_FP4_ACTS") == "1"
-    if use_fp4_acts:
+    mma_type = _mega_moe_mma_type()
+    if mma_type == "mxf4xmxf4":
        # FP4 path goes through DeepGEMM's mega_moe_pre_dispatch which
        # handles the E2M1 packing variant. The jit implementation
        # only emits FP8.
@@ -263,7 +269,7 @@ def _run_mega_routed(
            buf.topk_weights,
            num_tokens=num_tokens,
            group_size=32,
-            use_fp4_acts=True,
+            mma_type=mma_type,
        )
    else:
        mega_moe_pre_dispatch(
@@ -304,22 +310,30 @@ def _run_mega_routed(


 def _interleave_mega_moe_gate_up(t: torch.Tensor, gran: int = 8) -> torch.Tensor:
-    # Match DeepGEMM's L1 gate/up layout:
-    # [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...].
+    # Match DeepGEMM's L1 gate/up layouts. FP8 activations use contiguous
+    # gran-8 chunks; packed MXFP4 activations use even/odd gran-16 chunks.
    num_groups, n, *rest = t.shape
    half = n // 2
    gate = t[:, :half].reshape(num_groups, half // gran, gran, *rest)
    up = t[:, half:].reshape(num_groups, half // gran, gran, *rest)
-    result = torch.stack([gate, up], dim=2).reshape(num_groups, n, *rest)
+    if gran == 16:
+        result = torch.cat(
+            [gate[:, :, 0::2], up[:, :, 0::2], gate[:, :, 1::2], up[:, :, 1::2]],
+            dim=2,
+        ).reshape(num_groups, n, *rest)
+    else:
+        result = torch.stack([gate, up], dim=2).reshape(num_groups, n, *rest)
    return torch.empty_like(t).copy_(result)


 def _interleave_mega_moe_l1_weights(
    l1_weights: tuple[torch.Tensor, torch.Tensor],
+    mma_type: str,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    gran = 16 if mma_type == "mxf4xmxf4" else 8
    return (
-        _interleave_mega_moe_gate_up(l1_weights[0]),
-        _interleave_mega_moe_gate_up(l1_weights[1]),
+        _interleave_mega_moe_gate_up(l1_weights[0], gran=gran),
+        _interleave_mega_moe_gate_up(l1_weights[1], gran=gran),
    )


@@ -342,6 +356,7 @@ def build_mega_moe_experts_weights(experts) -> None:
    if getattr(experts, "_mega_moe_weights_built", False):
        return

+    mma_type = _mega_moe_mma_type()
    w13 = experts.w13_weight.data
    w13_sf_fp32 = experts.w13_weight_scale_inv.data
    w2 = experts.w2_weight.data
@@ -375,7 +390,9 @@ def build_mega_moe_experts_weights(experts) -> None:
    # the deep-ep path consumes the non-transposed interleaved scale and a
    # swizzle-aware activation kernel. L2 weight is untouched by the mega
    # transform, so the existing `w2_weight.data` is shared directly.
-    w13_interleaved, w13_sf_interleaved = _interleave_mega_moe_l1_weights((w13, w13_sf))
+    w13_interleaved, w13_sf_interleaved = _interleave_mega_moe_l1_weights(
+        (w13, w13_sf), mma_type
+    )
    w13_sf_utccp = _transpose_mega_moe_sf_for_utccp(w13_sf_interleaved)
    w2_sf_utccp = _transpose_mega_moe_sf_for_utccp(w2_sf)

@@ -666,9 +666,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                # (keeping both layouts OOMs: 92 layers double the experts).
                from deep_gemm import transform_weights_for_mega_moe

+                from sglang.srt.layers.moe.mega_moe import _mega_moe_mma_type
+
+                mma_type = _mega_moe_mma_type()
                l1_pair, l2_pair = transform_weights_for_mega_moe(
                    (layer.w13_weight.data, layer.w13_weight_scale.data),
                    (layer.w2_weight.data, layer.w2_weight_scale.data),
+                    mma_type=mma_type,
                )
                layer.mega_l1_weights = l1_pair
                layer.mega_l2_weights = l2_pair
@@ -2389,8 +2389,8 @@ class ServerArgs:
    ] = "none"
    enable_w4a4_mxfp4_megamoe: A[
        bool,
-        "Enable the W4A4 MXFP4 MegaMoE path by setting DeepGEMM's "
-        "DG_USE_FP4_ACTS=1 and DG_USE_MXF4_KIND=1. Use with "
+        "Enable the W4A4 MXFP4 MegaMoE path with DeepGEMM's "
+        "mxf4xmxf4 MMA type. Use with "
        "--moe-a2a-backend megamoe.",
        NS("exec.moe"),
    ] = False