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

[MegaMoE] Wire Qwen MoE blocks to DeepGEMM MegaMoE (MXFP4 and NVFP4 experts) (#38080)

Browse Source
This commit is contained in:
YAMY
2026-09-19 16:00:38 -07:00
committed by GitHub
parent 3a64faa1f2
commit 9cc7da2ab0
15 changed files with 686 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/deepseek_v4_hook.py
-85
View File
@@ -20,91 +20,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def validate_deepseek_v4_mega_moe_token_budget(
server_args: ServerArgs,
) -> None:
"""Ensure the DSV4 prefill budget fits MegaMoE's per-rank buffer."""
cfg = resolving_view(server_args)
mega_moe_enabled = cfg.moe_a2a_backend == "megamoe"
if not mega_moe_enabled or cfg.disaggregation_mode == "decode":
# decode node will skip the check because decode bs is not relevant with --chunk-prefill-size
return
if cfg.pp_size > 1 and cfg.enable_dynamic_chunking:
return
if cfg.chunked_prefill_size is None or cfg.chunked_prefill_size <= 0:
raise ValueError(
"DeepSeekV4 with MegaMoE requires chunked prefill to be enabled. "
"Set --chunked-prefill-size to a positive value; "
"--chunked-prefill-size=-1 is unsafe because MegaMoE's per-rank "
"token requirement would not have a strict prefill-forward bound."
)
if cfg.enable_prefill_cp:
token_partition_size = cfg.attn_cp_size
token_partition_name = "attn_cp_size"
token_alignment = 1
local_chunked_prefill_size = (
cfg.chunked_prefill_size + token_partition_size - 1
) // token_partition_size
elif cfg.enable_dp_attention:
token_partition_size = cfg.dp_size
token_partition_name = "dp_size"
token_alignment = max(
cfg.tp_size // cfg.dp_size // cfg.attn_cp_size,
1,
)
local_chunked_prefill_size = cfg.chunked_prefill_size // token_partition_size
else:
# Pure TP and PP with static chunking are handled here.
token_partition_size = 1
token_partition_name = "none"
# global_num_tokens will ceil_align to attn_tp_size so the validation needs to do alignment as well
token_alignment = max(
cfg.tp_size // token_partition_size // cfg.attn_cp_size,
1,
)
local_chunked_prefill_size = cfg.chunked_prefill_size
if local_chunked_prefill_size <= 0:
raise ValueError(
"DeepSeekV4 with MegaMoE requires a positive effective per-rank "
"chunked prefill size. "
f"Current values: chunked_prefill_size="
f"{cfg.chunked_prefill_size}, "
f"token_partition={token_partition_name}, "
f"token_partition_size={token_partition_size}."
)
required_tokens_per_rank = (
(local_chunked_prefill_size + token_alignment - 1)
// token_alignment
* token_alignment
)
max_tokens_per_rank = (
envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get()
)
if max_tokens_per_rank < required_tokens_per_rank:
raise ValueError(
"DeepSeekV4 with MegaMoE requires "
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to "
"cover each rank's effective prefill token budget. "
f"Current values: chunked_prefill_size="
f"{cfg.chunked_prefill_size}, "
f"token_partition={token_partition_name}, "
f"token_partition_size={token_partition_size}, "
f"token_alignment={token_alignment}, "
f"required_per_rank={required_tokens_per_rank}, "
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK="
f"{max_tokens_per_rank}. Set "
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to at "
f"least {required_tokens_per_rank}, or lower "
"--chunked-prefill-size until the effective per-rank budget fits. "
"Otherwise MegaMoE falls back to the fused MoE path at runtime."
)
def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None:
"""Residual imperative arm of the DeepSeek V4 defaults.
python/sglang/srt/arg_groups/mega_moe_hook.py
+172
View File
@@ -8,14 +8,186 @@ if TYPE_CHECKING:
from sglang.srt.arg_groups.overrides import (
declare_resolution,
model_config_of,
resolving_view,
)
from sglang.srt.connector import ConnectorType
from sglang.srt.environ import envs
from sglang.srt.runtime_context import get_platform
from sglang.srt.utils.common import parse_connector_type
logger = logging.getLogger(__name__)
def handle_mega_moe(server_args: ServerArgs) -> None:
handle_moe_runner_backend_alias(server_args)
check_mega_moe_compat(server_args)
def check_mega_moe_compat(server_args: ServerArgs) -> None:
cfg = resolving_view(server_args)
if cfg.moe_a2a_backend != "megamoe":
return
if cfg.enable_two_batch_overlap or cfg.enable_single_batch_overlap:
# The fused kernel has no dispatch_a / combine_a split for TBO / SBO.
raise ValueError(
"--moe-a2a-backend megamoe has no two-batch / single-batch overlap "
"decomposition; disable --enable-two-batch-overlap and "
"--enable-single-batch-overlap."
)
platform = get_platform()
if not (platform.is_cuda and (platform.is_sm90 or platform.is_sm100)):
raise ValueError(
"--moe-a2a-backend megamoe needs a CUDA SM90 GPU (block-FP8 experts) "
"or an SM100-class GPU (MXFP4 / NVFP4 experts); it runs DeepGEMM "
"kernels over CUDA symmetric memory."
)
# MoE blocks with no fused-MoE fallback under megamoe.
MEGA_MOE_NO_FALLBACK_ARCHS = frozenset(
{
"DeepseekV4ForCausalLM",
"InternS2PreviewForConditionalGeneration",
"MellumForCausalLM",
"Qwen2MoeForCausalLM",
"Qwen3MoeForCausalLM",
"Qwen3NextForCausalLM",
"Qwen3VLMoeForConditionalGeneration",
"Qwen3_5MoeForCausalLM",
"Qwen3_5MoeForConditionalGeneration",
}
)
def mega_moe_needs_token_budget(
model_arch: str, hf_quant_config: dict, quantization: str | None
) -> bool:
# NVFP4 experts are repacked into the mega layout at load: no fallback.
nvfp4_experts = quantization == "modelopt_fp4" or "FP4" in str(
hf_quant_config.get("quant_algo", "")
)
return model_arch in MEGA_MOE_NO_FALLBACK_ARCHS or nvfp4_experts
def validate_mega_moe_token_budget_for_model(server_args: ServerArgs) -> None:
# Runs after speculative resolution so the decode / verify bound is final.
cfg = resolving_view(server_args)
if cfg.moe_a2a_backend != "megamoe":
return
if parse_connector_type(cfg.model_path) == ConnectorType.INSTANCE:
return
model_config = model_config_of(server_args)
model_arch = model_config.hf_config.architectures[0]
if mega_moe_needs_token_budget(
model_arch, model_config.hf_quant_config, cfg.quantization
):
validate_mega_moe_token_budget(server_args, model_arch)
def mega_moe_decode_tokens_per_rank(cfg) -> int:
cg_config = cfg.cuda_graph_config
decode_max_bs = (cg_config.decode.max_bs if cg_config is not None else 0) or 0
num_tokens_per_req = (
(cfg.speculative_num_draft_tokens or 1) if cfg.speculative_algorithm else 1
)
return decode_max_bs * num_tokens_per_req
def validate_mega_moe_token_budget(server_args: ServerArgs, model_label: str) -> None:
cfg = resolving_view(server_args)
if cfg.moe_a2a_backend != "megamoe":
return
max_tokens_per_rank = (
envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get()
)
if cfg.disaggregation_mode != "prefill":
decode_tokens = mega_moe_decode_tokens_per_rank(cfg)
if max_tokens_per_rank < decode_tokens:
raise ValueError(
f"{model_label} with MegaMoE requires "
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to cover the "
"largest decode / verify forward on one rank. Current values: "
f"decode cuda graph max bs x tokens per request = {decode_tokens}, "
f"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK="
f"{max_tokens_per_rank}. Raise the env var to at least "
f"{decode_tokens} or lower --cuda-graph-max-bs / "
"--speculative-num-draft-tokens."
)
if cfg.disaggregation_mode == "decode":
# decode node will skip the check because decode bs is not relevant with --chunk-prefill-size
return
if cfg.pp_size > 1 and cfg.enable_dynamic_chunking:
return
if cfg.chunked_prefill_size is None or cfg.chunked_prefill_size <= 0:
raise ValueError(
f"{model_label} with MegaMoE requires chunked prefill to be enabled. "
"Set --chunked-prefill-size to a positive value; "
"--chunked-prefill-size=-1 is unsafe because MegaMoE's per-rank "
"token requirement would not have a strict prefill-forward bound."
)
if cfg.enable_prefill_cp:
token_partition_size = cfg.attn_cp_size
token_partition_name = "attn_cp_size"
token_alignment = 1
local_chunked_prefill_size = (
cfg.chunked_prefill_size + token_partition_size - 1
) // token_partition_size
elif cfg.enable_dp_attention:
token_partition_size = cfg.dp_size
token_partition_name = "dp_size"
token_alignment = max(
cfg.tp_size // cfg.dp_size // cfg.attn_cp_size,
1,
)
local_chunked_prefill_size = cfg.chunked_prefill_size // token_partition_size
else:
# Pure TP and PP with static chunking are handled here.
token_partition_size = 1
token_partition_name = "none"
# global_num_tokens will ceil_align to attn_tp_size so the validation needs to do alignment as well
token_alignment = max(
cfg.tp_size // token_partition_size // cfg.attn_cp_size,
1,
)
local_chunked_prefill_size = cfg.chunked_prefill_size
if local_chunked_prefill_size <= 0:
raise ValueError(
f"{model_label} with MegaMoE requires a positive effective per-rank "
"chunked prefill size. "
f"Current values: chunked_prefill_size="
f"{cfg.chunked_prefill_size}, "
f"token_partition={token_partition_name}, "
f"token_partition_size={token_partition_size}."
)
required_tokens_per_rank = (
(local_chunked_prefill_size + token_alignment - 1)
// token_alignment
* token_alignment
)
if max_tokens_per_rank < required_tokens_per_rank:
raise ValueError(
f"{model_label} with MegaMoE requires "
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to "
"cover each rank's effective prefill token budget. "
f"Current values: chunked_prefill_size="
f"{cfg.chunked_prefill_size}, "
f"token_partition={token_partition_name}, "
f"token_partition_size={token_partition_size}, "
f"token_alignment={token_alignment}, "
f"required_per_rank={required_tokens_per_rank}, "
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK="
f"{max_tokens_per_rank}. Set "
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to at "
f"least {required_tokens_per_rank}, or lower "
"--chunked-prefill-size until the effective per-rank budget fits."
)
def handle_moe_runner_backend_alias(server_args: ServerArgs) -> None:
python/sglang/srt/arg_groups/model_hook.py
-2
View File
@@ -421,14 +421,12 @@ def handle_model_specific_adjustments(server_args: Any):
]:
from sglang.srt.arg_groups.deepseek_v4_hook import (
validate_deepseek_v4_cp,
validate_deepseek_v4_mega_moe_token_budget,
validate_deepseek_v41_features,
)
# Before the CP validation: V4.1 rejects CP outright, the actionable message.
validate_deepseek_v41_features(server_args)
validate_deepseek_v4_cp(server_args)
validate_deepseek_v4_mega_moe_token_budget(server_args)
if get_platform().is_sm120:
# FP8 wo_a stays opt-in on SM120: only recent DeepGEMM builds ship
python/sglang/srt/arg_groups/pipeline.py
+6
View File
@@ -337,6 +337,12 @@ def run_resolution_pipeline(server_args: Any) -> None:
# Validate the CuteDSL A2A token budget now that num_tokens_per_req is final.
run_hook(validate_cutedsl_a2a_token_budget, server_args)
from sglang.srt.arg_groups.mega_moe_hook import (
validate_mega_moe_token_budget_for_model,
)
run_hook(validate_mega_moe_token_budget_for_model, server_args)
# Handle model loading format.
run_hook(handle_load_format, server_args)
python/sglang/srt/arg_groups/resolution_hooks.py
+1
View File
@@ -107,6 +107,7 @@ _OVERRIDABLE_HOOKS: FrozenSet[str] = frozenset(
"handle_speculative_decoding",
"handle_layernorm_sp",
"validate_cutedsl_a2a_token_budget",
"validate_mega_moe_token_budget_for_model",
"handle_load_format",
"handle_encoder_disaggregation",
"handle_tokenizer_batching",
python/sglang/srt/configs/model_config.py
+1
View File
@@ -563,6 +563,7 @@ class ModelConfig:
_quant_config_to_dict(getattr(self.hf_config, "quantization_config", None))
or {}
)
self.hf_quant_config: dict = quantization_config
routed_experts_quant_method = quantization_config.get(
"routed_experts_quant_method"
)
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+3
View File
@@ -398,6 +398,9 @@ class FusedMoE(torch.nn.Module):
self._num_local_routed = self._num_global_routed // storage_ep_size
self.num_local_experts = self._num_local_routed + num_fused_shared_experts
self._has_fused_shared = num_fused_shared_experts > 0
# Set by the quant method when it repacks experts for MegaMoE.
self._mega_moe_weights_built = False
self._mega_moe_nvfp4 = False
self._pending_fp8_shared_weights: dict[tuple[int, str], torch.Tensor] = {}
self._pending_fp8_shared_scales: dict[tuple[int, str], torch.Tensor] = {}
python/sglang/srt/layers/moe/mega_moe.py
+105 -23
View File
@@ -45,7 +45,9 @@ if TYPE_CHECKING:
_MEGA_MOE_SYMM_BUFFER: dict = {}
def _mega_moe_mma_type() -> str:
def _mega_moe_mma_type(experts=None) -> str:
if experts is not None and experts._mega_moe_nvfp4:
return "nvfp4xnvfp4"
return "mxf4xmxf4" if get_exec().moe.enable_w4a4_mxfp4_megamoe else "fp8xfp4"
@@ -86,6 +88,20 @@ def _configure_mega_moe_deep_gemm_num_sms(deep_gemm):
deep_gemm.set_num_sms(current_num_sms)
def check_mega_moe_shapes(hidden: int, intermediate: int, mma_type: str) -> None:
# DeepGEMM keeps one scale row per token and needs 16-byte TMA alignment
# on it (layout/mega_moe.cuh), so both dims must be multiples of 16 * group.
scale_group = 16 if mma_type == "nvfp4xnvfp4" else 32
align = 16 * scale_group
if hidden % align != 0 or intermediate % align != 0:
raise ValueError(
f"DeepGEMM MegaMoE ({mma_type}) needs hidden_size and "
f"moe_intermediate_size to be multiples of {align}; got "
f"hidden_size={hidden}, moe_intermediate_size={intermediate}. "
"Use another --moe-a2a-backend for this model."
)
def _get_mega_moe_symm_buffer(
group,
num_experts: int,
@@ -93,10 +109,12 @@ def _get_mega_moe_symm_buffer(
num_topk: int,
hidden: int,
intermediate_hidden: int,
mma_type: Optional[str] = None,
) -> SymmBuffer:
import deep_gemm
mma_type = _mega_moe_mma_type()
if mma_type is None:
mma_type = _mega_moe_mma_type()
with _configure_mega_moe_deep_gemm_num_sms(deep_gemm):
key = (
id(group),
@@ -124,14 +142,20 @@ def _get_mega_moe_symm_buffer(
return buf
def is_mega_moe_experts_ready(experts) -> bool:
if not experts._mega_moe_weights_built:
return False
if _device_sm == 90:
return is_sm90_fp8_mega_moe_available(experts)
# The SM100 mega kernels exist for compute capability 10.x only.
return _device_sm // 10 == 10
def should_use_mega_moe(moe: DeepseekV2MoE, hidden_states: torch.Tensor) -> bool:
if not get_moe_a2a_backend().is_megamoe():
return False
if not getattr(moe.experts, "_mega_moe_weights_built", False):
if not is_mega_moe_experts_ready(moe.experts):
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
@@ -188,10 +212,6 @@ def _run_mega_routed(
input_ids_global: Optional[torch.Tensor],
num_tokens: int,
) -> torch.Tensor:
import deep_gemm
from sglang.srt.runtime_context import get_parallel
hidden_size = moe.config.hidden_size
if num_tokens > 0:
@@ -216,10 +236,44 @@ def _run_mega_routed(
topk_ids = None
topk_weights = None
return run_mega_routed_experts(
moe.experts,
hidden_states,
topk_ids,
topk_weights,
hidden_size=hidden_size,
intermediate_size=moe.config.moe_intermediate_size,
top_k=moe.config.num_experts_per_tok + moe.num_fused_shared_experts,
num_tokens=num_tokens,
activation_clamp=moe.experts.moe_runner_config.swiglu_limit,
routed_scaling_factor=(
1.0
if moe.experts.should_fuse_routed_scaling_factor_in_topk
else float(moe.routed_scaling_factor)
),
)
def run_mega_routed_experts(
experts,
hidden_states: torch.Tensor,
topk_ids: Optional[torch.Tensor],
topk_weights: Optional[torch.Tensor],
*,
hidden_size: int,
intermediate_size: int,
top_k: int,
num_tokens: int,
activation_clamp: Optional[float] = None,
routed_scaling_factor: float = 1.0,
) -> torch.Tensor:
# Rows are this rank's tokens; the returned rows are fully combined.
import deep_gemm
from sglang.srt.runtime_context import get_parallel
ep_group = get_parallel().moe_ep_group.device_group
num_experts = moe.experts.num_experts
top_k = moe.config.num_experts_per_tok + moe.num_fused_shared_experts
intermediate_size = moe.config.moe_intermediate_size
num_experts = experts.num_experts
num_max_tokens_per_rank = (
envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get()
)
@@ -230,6 +284,7 @@ def _run_mega_routed(
f"--cuda-graph-max-bs-decode / --chunked-prefill-size accordingly"
)
mma_type = _mega_moe_mma_type(experts)
buf = _get_mega_moe_symm_buffer(
ep_group,
num_experts=num_experts,
@@ -237,6 +292,7 @@ def _run_mega_routed(
num_topk=top_k,
hidden=hidden_size,
intermediate_hidden=intermediate_size,
mma_type=mma_type,
)
if num_tokens > 0:
@@ -248,16 +304,42 @@ def _run_mega_routed(
if _device_sm == 90:
return run_sm90_mega_routed(
moe,
experts,
hidden_states,
topk_ids_in,
topk_weights_in,
buf,
num_tokens,
hidden_size=hidden_size,
activation_clamp=activation_clamp,
routed_scaling_factor=routed_scaling_factor,
)
mma_type = _mega_moe_mma_type()
if mma_type == "mxf4xmxf4":
mega_kwargs = {"recipe": (1, 1, 32)}
if mma_type == "nvfp4xnvfp4":
# Per-token outer scales go to buf.x_scales; the kernel folds them with
# the per-expert alphas in the L1 / L2 epilogues.
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=16,
mma_type=mma_type,
buf_x_scales=buf.x_scales,
)
mega_kwargs = {
"recipe": (1, 1, 16),
"use_x_scales": True,
"l1_alphas": experts.mega_l1_alphas,
"l2_alphas": experts.mega_l2_alphas,
"l2_act_scales": experts.mega_l2_act_scales,
}
elif 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.
@@ -292,22 +374,21 @@ def _run_mega_routed(
dtype=torch.bfloat16,
device=hidden_states.device,
)
swiglu_limit = getattr(moe.config, "swiglu_limit", None)
with _configure_mega_moe_deep_gemm_num_sms(deep_gemm):
deep_gemm.fp8_fp4_mega_moe(
y,
moe.experts.mega_l1_weights,
moe.experts.mega_l2_weights,
experts.mega_l1_weights,
experts.mega_l2_weights,
buf,
recipe=(1, 1, 32),
activation="swiglu",
activation_clamp=swiglu_limit,
activation_clamp=activation_clamp,
fast_math=True,
**mega_kwargs,
)
y = y[:num_tokens]
if not moe.experts.should_fuse_routed_scaling_factor_in_topk:
y.mul_(moe.routed_scaling_factor)
if routed_scaling_factor != 1.0:
y.mul_(routed_scaling_factor)
return y
@@ -366,6 +447,7 @@ def build_mega_moe_experts_weights(experts) -> None:
num_groups, n1, half_k1 = w13.shape
k1 = half_k1 * 2
check_mega_moe_shapes(hidden=k1, intermediate=n1 // 2, mma_type=mma_type)
_, n2, half_k2 = w2.shape
k2 = half_k2 * 2
python/sglang/srt/layers/moe/mega_moe_sm90.py
+10 -13
View File
@@ -15,7 +15,7 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Optional
import torch
@@ -24,8 +24,6 @@ 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:
@@ -42,20 +40,19 @@ def is_sm90_fp8_mega_moe_available(experts) -> bool:
def run_sm90_mega_routed(
moe: DeepseekV2MoE,
experts,
hidden_states: torch.Tensor,
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
buf: SymmBuffer,
num_tokens: int,
*,
hidden_size: int,
activation_clamp: Optional[float] = None,
routed_scaling_factor: float = 1.0,
) -> 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,
@@ -70,18 +67,18 @@ def run_sm90_mega_routed(
)
y = torch.empty(
(max(num_tokens, 1), moe.config.hidden_size),
(max(num_tokens, 1), 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,
experts.mega_l1_weights,
experts.mega_l2_weights,
buf,
recipe=(128, 128, 128),
activation="swiglu",
activation_clamp=getattr(moe.config, "swiglu_limit", None),
activation_clamp=activation_clamp,
fast_math=True,
)
y = y[:num_tokens]
python/sglang/srt/layers/quantization/modelopt_quant.py
+68
View File
@@ -2539,6 +2539,59 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
w2_input_scale._sglang_require_global_experts = True
layer.register_parameter("w2_input_scale", w2_input_scale)
def _build_mega_moe_weights(self, layer: torch.nn.Module) -> None:
# Activations are quantized per token at dispatch, so w13_input_scale
# is not used.
import deep_gemm
from deep_gemm.utils.math import transform_ue4m3_sf_into_required_layout
from sglang.srt.layers.moe.mega_moe import check_mega_moe_shapes
assert layer.moe_runner_config.is_gated, "MegaMoE NVFP4 needs a gated MLP"
n1 = layer.w13_weight.shape[1]
n2 = layer.w2_weight.shape[1]
check_mega_moe_shapes(
hidden=layer.w13_weight.shape[2] * 2,
intermediate=n1 // 2,
mma_type="nvfp4xnvfp4",
)
w13_sf = transform_ue4m3_sf_into_required_layout(
layer.w13_weight_scale.data.view(torch.float8_e4m3fn), n1
)
w2_sf = transform_ue4m3_sf_into_required_layout(
layer.w2_weight_scale.data.view(torch.float8_e4m3fn), n2
)
l1_pair, l2_pair = deep_gemm.transform_weights_for_mega_moe(
(layer.w13_weight.data.view(torch.int8), w13_sf),
(layer.w2_weight.data.view(torch.int8), w2_sf),
mma_type="nvfp4xnvfp4",
)
num_local = layer.num_local_experts
ones = torch.ones(num_local, dtype=torch.float32, device=l1_pair[0].device)
g1_gate, g1_up = _compute_gemm1_alphas(layer.w13_weight_scale_2, ones, True)
l1_alphas = torch.stack([g1_gate, g1_up], dim=1).contiguous()
w2_input_scale = _input_scale_to_local_experts(
layer.w2_input_scale, num_local, layer.num_experts, layer.moe_ep_rank
)
l2_act_scales = (1.0 / w2_input_scale).contiguous()
l2_alphas = (
w2_input_scale * layer.w2_weight_scale_2.to(torch.float32)
).contiguous()
layer.mega_l1_weights = l1_pair
layer.mega_l2_weights = l2_pair
layer.mega_l1_alphas = l1_alphas
layer.mega_l2_alphas = l2_alphas
layer.mega_l2_act_scales = l2_act_scales
# Free the checkpoint layout.
layer.w13_weight.data = l1_pair[0]
layer.w13_weight_scale.data = l1_pair[1]
layer.w2_weight.data = l2_pair[0]
layer.w2_weight_scale.data = l2_pair[1]
layer._mega_moe_nvfp4 = True
layer._mega_moe_weights_built = True
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""Transform packed FP4 MoE weights and scales for the selected backend."""
if getattr(layer, "inference_moe_w13_interleaved", False) and not getattr(
@@ -2553,6 +2606,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
)
layer._w13_deinterleaved = True
if get_moe_a2a_backend().is_megamoe():
self._build_mega_moe_weights(layer)
return
# GEMM1 scale processing is deferred until the input scale is known;
# see _compute_gemm1_alphas, which splits w13's gate/up weight scales.
moe_runner_backend = getattr(
@@ -2922,6 +2979,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
self._moe_runner_backend = moe_runner_backend
if get_moe_a2a_backend().is_megamoe():
# FusedMoE.forward is never reached under megamoe.
self.runner = None
return
if moe_runner_backend.is_flashinfer_cutedsl():
import sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl # noqa: F401 – triggers @register_fused_func
@@ -2972,6 +3034,12 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
layer: FusedMoE,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
if layer._mega_moe_nvfp4:
raise RuntimeError(
"NVFP4 MegaMoE experts cannot fall back to the fused MoE path; "
"the model block must route every forward through "
"run_mega_routed_experts (check the MegaMoE token budget)."
)
# Note: dispatch_output may be a DeepEPLLDispatchOutput (no topk_output
# attribute -- topk_ids/topk_weights live directly on the dispatch
# tuple). Defer per-attribute access to the branches that actually
python/sglang/srt/layers/quantization/mxfp4.py
+9 -1
View File
@@ -677,9 +677,17 @@ 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
from sglang.srt.layers.moe.mega_moe import (
_mega_moe_mma_type,
check_mega_moe_shapes,
)
mma_type = _mega_moe_mma_type()
check_mega_moe_shapes(
hidden=layer.w13_weight.shape[2] * 2,
intermediate=layer.w13_weight.shape[1] // 2,
mma_type=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),
python/sglang/srt/models/qwen2_moe.py
+71
View File
@@ -377,6 +377,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
or get_moe_a2a_backend().is_deepep_v2()
or get_moe_a2a_backend().is_flashinfer()
or get_moe_a2a_backend().is_flashinfer_megamoe()
or get_moe_a2a_backend().is_megamoe()
)
else {}
),
@@ -406,6 +407,10 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
)
self.top_k = config.num_experts_per_tok
self.is_nextn = is_nextn
self._use_mega_moe = get_moe_a2a_backend().is_megamoe()
self._mega_top_k = config.num_experts_per_tok + self.num_fused_shared_experts
self._mega_intermediate_size = config.moe_intermediate_size
self._mega_hidden_size = config.hidden_size
def get_moe_weights(self):
return [
@@ -625,6 +630,69 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
return final_hidden_states
def _forward_mega_moe(
self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch]
) -> torch.Tensor:
# Same contract as _forward_deepep: combined rows, no TP all-reduce.
from sglang.srt.layers.moe.mega_moe import (
is_mega_moe_experts_ready,
run_mega_routed_experts,
)
if not is_mega_moe_experts_ready(self.experts):
raise RuntimeError(
"moe_a2a_backend=megamoe needs MegaMoE expert weights on this "
"model: on SM100 load a checkpoint with MXFP4 or NVFP4 routed "
"experts; on SM90 load a block-FP8 checkpoint with a DeepGEMM "
"that ships fp8_mega_moe."
)
num_tokens = hidden_states.shape[0]
shared_output = None
topk_ids = None
topk_weights = None
if num_tokens > 0:
router_logits, _ = self.gate(hidden_states)
shared_output = self._forward_shared_experts(hidden_states)
topk_output = self.topk(
hidden_states,
router_logits,
num_token_non_padded=(
forward_batch.num_token_non_padded
if forward_batch is not None
else None
),
expert_location_dispatch_info=(
ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id)
if not self.is_nextn
else None
),
)
if self.enable_shared_expert_fusion:
topk_output = self._append_shared_to_topk_output(
topk_output, hidden_states
)
assert TopKOutputChecker.format_is_standard(topk_output), (
"MegaMoE pre-dispatch consumes raw topk ids/weights; "
"pick a MoE runner backend that emits standard TopK output"
)
topk_ids = topk_output.topk_ids
topk_weights = topk_output.topk_weights
final_hidden_states = run_mega_routed_experts(
self.experts,
hidden_states,
topk_ids,
topk_weights,
hidden_size=self._mega_hidden_size,
intermediate_size=self._mega_intermediate_size,
top_k=self._mega_top_k,
num_tokens=num_tokens,
)
if shared_output is not None:
final_hidden_states.add_(shared_output)
return final_hidden_states
@property
def supports_deferred_finalize(self) -> bool:
return bool(
@@ -743,6 +811,9 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
if defer_finalize and num_tokens == 0:
raise RuntimeError("Qwen deferred finalize does not support M=0")
if self._use_mega_moe:
return self._forward_mega_moe(hidden_states, forward_batch)
if (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_deepep_v2()
python/sglang/srt/models/qwen3_moe.py
+60 -1
View File
@@ -43,7 +43,7 @@ from sglang.srt.layers.moe import (
)
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.topk import TopK
from sglang.srt.layers.moe.topk import TopK, TopKOutputChecker
from sglang.srt.layers.moe.utils import (
RoutingMethodType,
filter_moe_weight_param_global_expert,
@@ -291,12 +291,19 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
)
self.top_k = config.num_experts_per_tok
self._use_mega_moe = get_moe_a2a_backend().is_megamoe()
self._mega_top_k = config.num_experts_per_tok
self._mega_intermediate_size = config.moe_intermediate_size
self._mega_hidden_size = config.hidden_size
def forward(
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
) -> torch.Tensor:
if self._use_mega_moe:
return self._forward_mega_moe(hidden_states, forward_batch)
if (
not is_deepep_class_backend()
and not get_moe_a2a_backend().is_ascend_fuseep()
@@ -356,6 +363,58 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
)
return final_hidden_states
def _forward_mega_moe(
self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch]
) -> torch.Tensor:
# Same contract as forward_deepep: combined rows, no TP all-reduce.
from sglang.srt.layers.moe.mega_moe import (
is_mega_moe_experts_ready,
run_mega_routed_experts,
)
if not is_mega_moe_experts_ready(self.experts):
raise RuntimeError(
"moe_a2a_backend=megamoe needs MegaMoE expert weights on this "
"model: on SM100 load a checkpoint with MXFP4 or NVFP4 routed "
"experts; on SM90 load a block-FP8 checkpoint with a DeepGEMM "
"that ships fp8_mega_moe."
)
num_tokens = hidden_states.shape[0]
topk_ids = None
topk_weights = None
if num_tokens > 0:
router_logits, _ = self.gate(hidden_states)
topk_output = self.topk(
hidden_states,
router_logits,
num_token_non_padded=(
forward_batch.num_token_non_padded
if forward_batch is not None
else None
),
expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
layer_id=self.layer_id,
),
)
assert TopKOutputChecker.format_is_standard(topk_output), (
"MegaMoE pre-dispatch consumes raw topk ids/weights; "
"pick a MoE runner backend that emits standard TopK output"
)
topk_ids = topk_output.topk_ids
topk_weights = topk_output.topk_weights
return run_mega_routed_experts(
self.experts,
hidden_states,
topk_ids,
topk_weights,
hidden_size=self._mega_hidden_size,
intermediate_size=self._mega_intermediate_size,
top_k=self._mega_top_k,
num_tokens=num_tokens,
)
def op_gate(self, state):
if is_non_idle_and_non_empty(
state.forward_batch.forward_mode, state.hidden_states_mlp_input