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

DeepSeek V4 w4a4 MegaMoE (#25052)

Browse Source 
Co-authored-by: pranjalssh <adkz.photos@gmail.com>
This commit is contained in:
Baizhou Zhang
2026-05-13 18:35:32 -07:00
committed by GitHub
co-authored by pranjalssh
parent 34c0029f0a
commit b7f856df70
5 changed files with 212 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/pyproject.toml
+1 -1
View File
@@ -59,7 +59,7 @@ dependencies = [
"sentencepiece",
"setproctitle",
"flash-attn-4>=4.0.0b9",
"sgl-deep-gemm==0.0.1",
"sgl-deep-gemm==0.1.0",
"sglang-kernel==0.4.2.post1",
"soundfile==0.13.1",
"tiktoken",
python/sglang/srt/environ.py
+11
View File
@@ -595,6 +595,17 @@ class Envs:
# DeepGemm Mega MoE
SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE = EnvBool(False)
SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK = EnvInt(1024)
# 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)
SGLANG_OPT_FIX_MEGA_MOE_MEMORY = EnvBool(False)
# TopK
python/sglang/srt/layers/moe/mega_moe.py
+52 -10
View File
@@ -15,6 +15,7 @@
from __future__ import annotations
import os
from contextlib import nullcontext
from typing import TYPE_CHECKING, Optional
@@ -34,6 +35,26 @@ 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(
@@ -46,6 +67,8 @@ def _get_mega_moe_symm_buffer(
) -> SymmBuffer:
import deep_gemm
_apply_mega_moe_dg_env()
key = (
id(group),
num_max_tokens_per_rank,
@@ -188,16 +211,35 @@ def _run_mega_routed(
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)
mega_moe_pre_dispatch(
hidden_states,
topk_ids_in,
topk_weights_in,
buf.x,
buf.x_sf,
buf.topk_idx,
buf.topk_weights,
quant_group_size=32,
)
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
# handles the E2M1 packing variant. The jit implementation
# only emits FP8.
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=32,
use_fp4_acts=True,
)
else:
mega_moe_pre_dispatch(
hidden_states,
topk_ids_in,
topk_weights_in,
buf.x,
buf.x_sf,
buf.topk_idx,
buf.topk_weights,
quant_group_size=32,
)
# Allocate at least one row so y has a non-null CUDA data_ptr;
# the DeepGEMM tvm-ffi binding rejects nullptr in convert_to_torch_tensor().
test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py
-49
View File
@@ -31,14 +31,6 @@ _DEEPEP_ENV = {
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
_MEGAMOE_ENV = {
"SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE": "1",
"SGLANG_OPT_FIX_MEGA_MOE_MEMORY": "1",
"SGLANG_OPT_FIX_NEXTN_MEGA_MOE": "1",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "4096",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "0",
}
def _gsm8k_check(test_case):
args = SimpleNamespace(
@@ -138,46 +130,5 @@ class TestDSV4FlashFP4B200Balanced(ServerSanityMixin, CustomTestCase):
_gsm8k_check(self)
class TestDSV4FlashFP4B200MegaMoE(ServerSanityMixin, CustomTestCase):
"""Balanced recipe: TP=4, DP=4, MegaMoE."""
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(MODEL)
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
],
env=_MEGAMOE_ENV,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
_gsm8k_check(self)
if __name__ == "__main__":
unittest.main()
test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py
+148
View File
@@ -0,0 +1,148 @@
"""B200 per-commit CI: DeepSeek-V4-Flash FP4 (LowLatency recipe).
Launches TP=4 with flashinfer_mxfp4 MoE runner + EAGLE speculative decoding.
Runs 12 ServerSanity probes (correctness, streaming, concurrency, determinism)
plus a GSM8K accuracy gate.
Registry: stage-c-test-dsv4-4-gpu-b200 (per-commit, 4x B200)
"""
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.server_sanity_kit import ServerSanityMixin
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
try_cached_model,
)
register_cuda_ci(est_time=1800, suite="stage-c-test-dsv4-4-gpu-b200")
MODEL = "deepseek-ai/DeepSeek-V4-Flash"
SERVER_LAUNCH_TIMEOUT = 3600
_W4A8_MEGAMOE_ENV = {
"SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE": "1",
"SGLANG_OPT_FIX_MEGA_MOE_MEMORY": "1",
"SGLANG_OPT_FIX_NEXTN_MEGA_MOE": "1",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "4096",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "0",
}
_W4A4_MEGAMOE_ENV = {
"SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE": "1",
"SGLANG_OPT_FIX_MEGA_MOE_MEMORY": "1",
"SGLANG_OPT_FIX_NEXTN_MEGA_MOE": "1",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "4096",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "0",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS": "1",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND": "1",
}
def _gsm8k_check(test_case):
args = SimpleNamespace(
base_url=test_case.base_url,
model=test_case.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"[{type(test_case).__name__}] GSM8K {metrics=}")
test_case.assertGreater(metrics["score"], 0.93)
class TestDSV4FlashFP4B200W4A8MegaMoE(ServerSanityMixin, CustomTestCase):
"""Balanced recipe: TP=4, DP=4, MegaMoE."""
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(MODEL)
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
],
env=_W4A8_MEGAMOE_ENV,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
_gsm8k_check(self)
class TestDSV4FlashFP4B200W4A4MegaMoE(ServerSanityMixin, CustomTestCase):
"""Balanced recipe: TP=4, DP=4, MegaMoE."""
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(MODEL)
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
],
env=_W4A4_MEGAMOE_ENV,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
_gsm8k_check(self)
if __name__ == "__main__":
unittest.main()