minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[DeepSeek V4] Add W4A4 MegaMoE server flag (#35918)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-08-21 18:44:18 -07:00
committed by GitHub
parent 0be2a209ac
commit 3b5909de0e
15 changed files with 147 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/mega_moe_hook.py
+38
View File
@@ -0,0 +1,38 @@
from __future__ import annotations
import logging
import os
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
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:
if server_args.moe_runner_backend != "megamoe":
return
if server_args.moe_a2a_backend not in ("none", "megamoe"):
logger.warning(
"--moe-runner-backend megamoe is an alias for "
"--moe-a2a-backend megamoe; overriding "
"--moe-a2a-backend %s.",
server_args.moe_a2a_backend,
)
server_args.moe_runner_backend = "auto"
server_args.moe_a2a_backend = "megamoe"
def handle_w4a4_mxfp4_megamoe_env(server_args: ServerArgs) -> None:
if not server_args.enable_w4a4_mxfp4_megamoe:
return
os.environ["DG_USE_FP4_ACTS"] = "1"
os.environ["DG_USE_MXF4_KIND"] = "1"
python/sglang/srt/environ.py
+6 -10
View File
@@ -1052,16 +1052,6 @@ class Envs:
# DeepGEMM Mega MoE
# ===================================================================
SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK = EnvInt(8192)
# When set, the mega-MoE x slot is packed E2M1 (FP4) instead of FP8 E4M3.
# Halves symm-buffer footprint and unlocks the MXF4 mainloop downstream.
# Setting this also exports DG_USE_FP4_ACTS=1 so DeepGEMM's symm-buffer
# sizing + fp8_fp4_mega_moe pick up the FP4 layout.
SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS = EnvBool(False)
# Switches the L1+L2 mainloops from kind::mxf8f6f4 (K=32 with-padding) to
# kind::mxf4 (K=64 dense) inside fp8_fp4_mega_moe. No effect unless
# SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS is also set; DeepGEMM asserts
# this combination on the host side.
SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND = EnvBool(False)
# ===================================================================
# Top-k kernels
@@ -1631,6 +1621,12 @@ _DEPRECATED_ENVS: Dict[str, _DeprecatedEnv] = {
note="Please use '--moe-runner-backend=cutlass' and/or "
"'--speculative-moe-runner-backend=cutlass' instead."
),
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS": _DeprecatedEnv(
note="Please use '--enable-w4a4-mxfp4-megamoe' instead."
),
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND": _DeprecatedEnv(
note="Please use '--enable-w4a4-mxfp4-megamoe' instead."
),
"SGLANG_DFLASH_PREFILL_REFILL_TARGET": _DeprecatedEnv(
note="DFlash now auto-enables the min-free-slots delay; unset this env. "
"To override the threshold, use '--min-free-slots-delay'."
python/sglang/srt/layers/moe/mega_moe.py
+1 -23
View File
@@ -42,26 +42,6 @@ if TYPE_CHECKING:
_MEGA_MOE_SYMM_BUFFER: dict = {}
_MEGA_MOE_DG_ENV_APPLIED = False
def _apply_mega_moe_dg_env() -> None:
"""Forward sglang's FP4/MXF4 opt-in flags to DeepGEMM via env vars.
DeepGEMM reads `DG_USE_FP4_ACTS` (and `DG_USE_MXF4_KIND`) at host-function
call time — both `get_symm_buffer_for_mega_moe` and `fp8_fp4_mega_moe`.
Forwarding once at first use is sufficient (these are static config
flags, not per-request state) and matches the `setdefault` pattern so
explicit `DG_USE_*` overrides from outside still win.
"""
global _MEGA_MOE_DG_ENV_APPLIED
if _MEGA_MOE_DG_ENV_APPLIED:
return
if envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get():
os.environ.setdefault("DG_USE_FP4_ACTS", "1")
if envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND.get():
os.environ.setdefault("DG_USE_MXF4_KIND", "1")
_MEGA_MOE_DG_ENV_APPLIED = True
def _get_mega_moe_symm_buffer(
@@ -74,8 +54,6 @@ def _get_mega_moe_symm_buffer(
) -> SymmBuffer:
import deep_gemm
_apply_mega_moe_dg_env()
key = (
id(group),
num_max_tokens_per_rank,
@@ -232,7 +210,7 @@ def _run_mega_routed(
num_tokens,
)
use_fp4_acts = envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get()
use_fp4_acts = os.getenv("DG_USE_FP4_ACTS") == "1"
if use_fp4_acts:
# FP4 path goes through DeepGEMM's mega_moe_pre_dispatch which
# handles the E2M1 packing variant. The jit implementation
python/sglang/srt/server_args.py
+10 -15
View File
@@ -2390,6 +2390,13 @@ class ServerArgs:
),
NS("exec.moe"),
] = "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 "
"--moe-a2a-backend megamoe.",
NS("exec.moe"),
] = False
moe_runner_backend: A[
str,
Arg(
@@ -3651,7 +3658,9 @@ class ServerArgs:
# _handle_model_specific_adjustments never runs.
self._resolved_overrides = []
self._handle_moe_runner_backend_alias()
from sglang.srt.arg_groups.mega_moe_hook import handle_mega_moe
handle_mega_moe(self)
self._handle_return_hidden_states_mode()
self._handle_media_url_security()
self._handle_hicache_ratio_default()
@@ -3824,20 +3833,6 @@ class ServerArgs:
materialize_declarations(self)
def _handle_moe_runner_backend_alias(self):
if self.moe_runner_backend != "megamoe":
return
if self.moe_a2a_backend not in ("none", "megamoe"):
logger.warning(
"--moe-runner-backend megamoe is an alias for "
"--moe-a2a-backend megamoe; overriding "
"--moe-a2a-backend %s.",
self.moe_a2a_backend,
)
self.moe_runner_backend = "auto"
self.moe_a2a_backend = "megamoe"
def _handle_return_hidden_states_mode(self):
if self.return_hidden_states_mode not in (None, "last", "full"):
raise ValueError(