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

[MoE Refactor] Migrate flashinfer_cutedsl + DeepEP to MoeRunner (#25525)

Browse Source 
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-17 14:48:17 -07:00
committed by GitHub
co-authored by Claude Opus 4.7
parent 89e501c5a8
commit 7158a255eb
5 changed files with 156 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/ep_moe/layer.py
+7 -29
View File
@@ -105,6 +105,12 @@ class DeepEPMoE(FusedMoE):
and envs.SGLANG_DEEPEP_BF16_DISPATCH.get()
):
self.deprecate_flag = True
elif (
get_moe_runner_backend().is_flashinfer_cutedsl()
and quant_config is not None
and quant_config.get_name() == "modelopt_fp4"
):
self.deprecate_flag = True
else:
self.deprecate_flag = False
@@ -134,15 +140,9 @@ class DeepEPMoE(FusedMoE):
self.deepep_mode.enable_low_latency()
and not _is_npu
and not _is_hip
and not (
get_moe_runner_backend().is_flashinfer_cutedsl()
and self.quant_config is not None
and self.quant_config.get_name() == "modelopt_fp4"
)
and quant_config is not None
):
# AMD HIP, NPU supports low_latency deepep without deepgemm
# NV FP4 quantization with flashinfer_cutedsl also supports low_latency deepep without deepgemm
# AMD HIP and NPU support low_latency DeepEP without DeepGEMM.
# Unquantized draft MoE uses BF16 DeepEP dispatch and a local fallback.
assert (
deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
@@ -220,12 +220,6 @@ class DeepEPMoE(FusedMoE):
elif DispatchOutputChecker.format_is_deepep_ll(dispatch_output):
if self.quant_config is None:
output = self.forward_unquantized_deepep_ll(dispatch_output)
elif (
get_moe_runner_backend().is_flashinfer_cutedsl()
and self.quant_config is not None
and self.quant_config.get_name() == "modelopt_fp4"
):
output = self.forward_flashinfer_cutedsl(dispatch_output)
elif self.use_w4afp8:
output = self.forward_cutlass_w4afp8_masked(dispatch_output)
else:
@@ -288,22 +282,6 @@ class DeepEPMoE(FusedMoE):
output = output + w2_bias.unsqueeze(1)
return output.masked_fill(~valid_mask, 0)
def forward_flashinfer_cutedsl(
self,
dispatch_output: DeepEPLLDispatchOutput,
):
hidden_states, hidden_states_scale, _, _, masked_m, _ = dispatch_output
assert self.quant_method is not None
assert self.moe_runner_config.activation == "silu"
output = self.quant_method.apply_without_routing_weights(
layer=self,
x=(hidden_states, hidden_states_scale),
masked_m=masked_m,
moe_runner_config=self.moe_runner_config,
)
return output
def forward_cutlass_w4afp8(
self,
dispatch_output: DeepEPNormalDispatchOutput,
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py
+103 -13
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, Optional
import torch
@@ -14,7 +14,10 @@ from sglang.srt.layers.moe.moe_runner.base import (
from sglang.srt.utils.common import log_info_on_rank0, print_warning_once
if TYPE_CHECKING:
from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
from sglang.srt.layers.moe.token_dispatcher import (
DeepEPLLCombineInput,
DeepEPLLDispatchOutput,
StandardCombineInput,
StandardDispatchOutput,
)
@@ -283,28 +286,49 @@ def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
@dataclass
class CuteDslFp4MoeQuantInfo(MoeQuantInfo):
"""Quantization payload consumed by FlashInfer CuteDSL FP4 MoE kernels."""
"""Quantization payload for FlashInfer CuteDSL FP4 MoE kernels.
# Lazily-created CuteDslMoEWrapper (stashed on layer)
wrapper: Any
Shared by the two CuteDSL runner entries:
# Weights (uint8 FP4 packed)
* "v2" standard path (a2a=``none``/``flashinfer``): consumed by the
``@register_fused_func("none", "flashinfer_cutedsl")`` entry, which
drives ``CuteDslMoEWrapper.run``. Weights are ``[Up, Gate]``
interleaved with MMA-layout blockscales. ``wrapper`` is set;
``w*_scale`` are scalarized.
* "v1" DeepEP low-latency path (a2a=``deepep``): consumed by the
``@register_fused_func("deepep", "flashinfer_cutedsl")`` entry,
which drives ``flashinfer_cutedsl_moe_masked``. Weights are
``[Gate, Up]`` non-interleaved with swizzled blockscales.
``wrapper`` is ``None``; ``w*_scale`` are per-expert.
"""
# FP4 packed weights (uint8)
w13_weight: torch.Tensor
w2_weight: torch.Tensor
# Block-scale factors
# Block-scale factors (MMA layout for v2, swizzled for v1)
w13_weight_sf: torch.Tensor
w2_weight_sf: torch.Tensor
# Per-expert GEMM scales
# Per-expert GEMM dequant alphas (scalarized for v2, per-expert for v1)
w1_alpha: torch.Tensor
w2_alpha: torch.Tensor
# Intermediate quantization scale (fc2 input)
fc2_input_scale: torch.Tensor
# Activation quant scales (1 / raw_input_scale).
# - a1_scale: quantizes hidden_states before GEMM1
# - a2_scale: quantizes GEMM1 output before GEMM2 (a.k.a. fc2 input)
a1_scale: torch.Tensor
a2_scale: torch.Tensor
# Activation quantization scale (scalarized)
input_scale: torch.Tensor
# v2 only: lazily-created CuteDslMoEWrapper (``None`` on the v1 path).
wrapper: Optional[Any] = None
# v1 only: ``True`` when DeepEP pre-quantizes activations to NVFP4.
use_nvfp4_dispatch: bool = False
# v1 only: SBO down-GEMM overlap args.
down_gemm_overlap_args: Optional["DownGemmOverlapArgs"] = None
@register_fused_func("none", "flashinfer_cutedsl")
@@ -318,6 +342,7 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4(
from sglang.srt.layers.quantization.fp4_utils import fp4_quantize
assert runner_config.activation == "silu", "Only silu is supported for CuteDSL MoE."
assert quant_info.wrapper is not None, "CuteDSL v2 path requires CuteDslMoEWrapper."
hidden_states = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
@@ -330,7 +355,7 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4(
x_fp4, x_sf = fp4_quantize(
hidden_states,
quant_info.input_scale,
quant_info.a1_scale,
sf_vec_size=_FP4_SF_VEC_SIZE,
is_sf_swizzled_layout=False,
)
@@ -343,10 +368,75 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4(
w1_weight=quant_info.w13_weight,
w1_weight_sf=quant_info.w13_weight_sf,
w1_alpha=quant_info.w1_alpha,
fc2_input_scale=quant_info.fc2_input_scale,
fc2_input_scale=quant_info.a2_scale,
w2_weight=quant_info.w2_weight,
w2_weight_sf=quant_info.w2_weight_sf,
w2_alpha=quant_info.w2_alpha,
)
return StandardCombineInput(hidden_states=output)
@register_fused_func("deepep", "flashinfer_cutedsl")
def fused_experts_deepep_to_flashinfer_cutedsl_fp4(
dispatch_output: DeepEPLLDispatchOutput,
quant_info: CuteDslFp4MoeQuantInfo,
runner_config: MoeRunnerConfig,
) -> DeepEPLLCombineInput:
from sglang.srt.layers.moe.flashinfer_cutedsl_moe import (
flashinfer_cutedsl_moe_masked,
)
from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPLLCombineInput
assert runner_config.activation == "silu", "Only silu is supported for CuteDSL MoE."
assert (
not runner_config.apply_router_weight_on_input
), "apply_router_weight_on_input is not supported for Flashinfer"
hidden_states, hidden_states_scale, _, _, masked_m, _ = dispatch_output
# flashinfer_cutedsl_moe_masked reinterprets scales as float8_e4m3fn.
# Same-dtype .view is a no-op; only wider dtypes (e.g. int32-packed
# UE8M0) need stride(-1)==1.
if (
quant_info.use_nvfp4_dispatch
and hidden_states_scale is not None
and hidden_states_scale.element_size() != 1
and hidden_states_scale.stride(-1) != 1
):
raise AssertionError(
f"NVFP4 dispatch scale has stride(-1)={hidden_states_scale.stride(-1)}, "
f"dtype={hidden_states_scale.dtype}; .view(float8_e4m3fn) requires stride(-1)==1. "
"Try SGLANG_MOE_NVFP4_DISPATCH=0 or check DeepEP version."
)
overlap = quant_info.down_gemm_overlap_args
output = flashinfer_cutedsl_moe_masked(
hidden_states=(hidden_states, hidden_states_scale),
input_global_scale=(
None if quant_info.use_nvfp4_dispatch else quant_info.a1_scale
),
w1=quant_info.w13_weight,
w1_blockscale=quant_info.w13_weight_sf,
w1_alpha=quant_info.w1_alpha,
w2=quant_info.w2_weight,
a2_global_scale=quant_info.a2_scale,
w2_blockscale=quant_info.w2_weight_sf,
w2_alpha=quant_info.w2_alpha,
masked_m=masked_m,
**(
dict(
down_sm_count=overlap.num_sms,
down_signals=overlap.signal,
down_start_event=overlap.start_event,
)
if overlap is not None
else {}
),
)
return DeepEPLLCombineInput(
hidden_states=output,
topk_ids=dispatch_output.topk_ids,
topk_weights=dispatch_output.topk_weights,
)
python/sglang/srt/layers/moe/moe_runner/runner.py
-2
View File
@@ -163,11 +163,9 @@ class MoeRunner:
def set_overlap_args(
self, down_gemm_overlap_args: DownGemmOverlapArgs, meta_overlap_args: dict
):
assert self.fused_func is None, "Fused func is not supported for overlap args"
self.down_gemm_overlap_args = down_gemm_overlap_args
self.meta_overlap_args = meta_overlap_args
def clear_overlap_args(self) -> None:
assert self.fused_func is None, "Fused func is not supported for overlap args"
self.down_gemm_overlap_args = None
self.meta_overlap_args = None
python/sglang/srt/layers/quantization/modelopt_quant.py
+44 -101
View File
@@ -67,7 +67,6 @@ from sglang.srt.utils.custom_op import register_custom_op
from sglang.srt.utils.patch_torch import register_fake_if_exists
if TYPE_CHECKING:
from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.token_dispatcher import (
CombineInput,
@@ -1551,24 +1550,24 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
# ----- CuteDSL v1 vs v2 path helpers -----
#
# "v1": cutedsl + deepep low-latency.
# - Bypasses MoeRunner entirely; calls apply_without_routing_weights ->
# flashinfer_cutedsl_moe_masked (grouped_gemm_nt_masked).
# - MoeRunner fused func calls flashinfer_cutedsl_moe_masked
# (grouped_gemm_nt_masked).
# - Expects W13 in default [Gate, Up] order, NOT interleaved.
# - Uses swizzled blockscales directly (w13_blockscale_swizzled).
#
# "v2" (standard): cutedsl + none/flashinfer a2a.
# - Uses MoeRunner with @register_fused_func CuteDslMoEWrapper kernels.
# - MoeRunner fused func calls CuteDslMoEWrapper kernels.
# - Expects W13 in [Up, Gate] order, interleaved in 64-row chunks.
# - Uses MMA-layout blockscales (w13_blockscale_mma).
@property
def _is_cutedsl_v1_deepep(self) -> bool:
"""CuteDSL v1 + DeepEP low-latency path (no MoeRunner)."""
"""CuteDSL v1 + DeepEP low-latency path (masked grouped GEMM)."""
return is_flashinfer_cutedsl_v1_path()
@property
def _is_cutedsl_v2_standard(self) -> bool:
"""New CuteDSL standard path (a2a=none or flashinfer, uses MoeRunner)."""
"""CuteDSL v2 standard path (a2a=none or flashinfer, uses CuteDslMoEWrapper)."""
return self.enable_flashinfer_cutedsl_moe and not self._is_cutedsl_v1_deepep
def create_weights(
@@ -1998,11 +1997,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
if moe_runner_backend.is_flashinfer_cutedsl():
import sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl # noqa: F401 – triggers @register_fused_func
# CuteDSL v1 (deepep) uses the apply_without_routing_weights
# path (flashinfer_cutedsl_moe_masked) and does not need a MoeRunner.
if self._is_cutedsl_v1_deepep:
return
if not moe_runner_backend.is_flashinfer_cutlass():
self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
@@ -2013,10 +2007,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
) -> CombineInput:
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
x = dispatch_output.hidden_states
x_sf = dispatch_output.hidden_states_scale
topk_output = dispatch_output.topk_output
# 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
# consume them.
activation = self.moe_runner_config.activation
assert (
@@ -2054,33 +2048,49 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
return self.runner.run(dispatch_output, quant_info)
# CuteDSL v2 standard path (a2a=none/flashinfer).
# The v1 (deepep) path never reaches apply(); it goes through
# apply_without_routing_weights instead.
if self.enable_flashinfer_cutedsl_moe:
from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
CuteDslFp4MoeQuantInfo,
ensure_cutedsl_wrapper,
)
if self._is_cutedsl_v1_deepep:
# v1 path: DeepEP low-latency + flashinfer_cutedsl_moe_masked.
# Weights are [Gate, Up] (non-interleaved) with swizzled blockscales.
quant_info = CuteDslFp4MoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
w13_weight_sf=layer.w13_blockscale_swizzled,
w2_weight_sf=layer.w2_blockscale_swizzled,
w1_alpha=layer.g1_alphas,
w2_alpha=layer.g2_alphas,
a1_scale=layer.w13_input_scale_quant,
a2_scale=layer.w2_input_scale_quant,
use_nvfp4_dispatch=MOE_NVFP4_DISPATCH,
down_gemm_overlap_args=getattr(
self.runner, "down_gemm_overlap_args", None
),
)
return self.runner.run(dispatch_output, quant_info)
# v2 standard path (a2a=none/flashinfer): uses CuteDslMoEWrapper
# with [Up, Gate] interleaved weights and MMA blockscales.
ensure_cutedsl_wrapper(layer)
w1_alpha, fc2_input_scale, w2_alpha = layer._cutedsl_scales
w1_weight_sf = getattr(
layer, "w13_blockscale_mma", layer.w13_blockscale_swizzled
)
w2_weight_sf = getattr(
layer, "w2_blockscale_mma", layer.w2_blockscale_swizzled
)
quant_info = CuteDslFp4MoeQuantInfo(
wrapper=layer._cutedsl_wrapper,
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
w13_weight_sf=w1_weight_sf,
w2_weight_sf=w2_weight_sf,
w13_weight_sf=getattr(
layer, "w13_blockscale_mma", layer.w13_blockscale_swizzled
),
w2_weight_sf=getattr(
layer, "w2_blockscale_mma", layer.w2_blockscale_swizzled
),
w1_alpha=w1_alpha,
w2_alpha=w2_alpha,
fc2_input_scale=fc2_input_scale,
input_scale=layer._cutedsl_input_scale,
a1_scale=layer._cutedsl_input_scale,
a2_scale=fc2_input_scale,
wrapper=layer._cutedsl_wrapper,
)
return self.runner.run(dispatch_output, quant_info)
@@ -2092,6 +2102,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
), "apply_router_weight_on_input is not supported for Flashinfer"
# TRTLLM Cutlass moe takes in activations in BF16/Half/nvfp4 precision
# and fp4 quantized weights loaded from the checkpoint
x = dispatch_output.hidden_states
x_sf = dispatch_output.hidden_states_scale
topk_output = dispatch_output.topk_output
topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
output_dtype = torch.bfloat16
@@ -2145,6 +2158,8 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
output = cutlass_moe_fp4(
a=x,
@@ -2163,75 +2178,3 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
).to(x.dtype)
# Scale by routed_scaling_factor is fused into select_experts.
return StandardCombineInput(hidden_states=output)
def apply_without_routing_weights(
self,
layer: FusedMoE,
x: tuple[torch.Tensor, Optional[torch.Tensor]],
masked_m: torch.Tensor,
moe_runner_config: MoeRunnerConfig,
) -> torch.Tensor:
"""CuteDSL v1 (deepep low-latency) path.
Called by the DeepEP dispatcher instead of apply(). Uses
flashinfer_cutedsl_moe_masked (grouped_gemm_nt_masked) directly,
bypassing MoeRunner. Weights must be in default [Gate, Up] order
and NOT interleaved -- see _is_cutedsl_v1_deepep guards in
process_weights_after_loading and load_up_proj_weight_first.
"""
assert (
moe_runner_config.activation == "silu"
), "Only SiLU activation is supported."
assert self.enable_flashinfer_cutedsl_moe, "only support flashinfer cutedsl moe"
assert (
not moe_runner_config.apply_router_weight_on_input
), "apply_router_weight_on_input is not supported for Flashinfer"
from sglang.srt.layers.moe.flashinfer_cutedsl_moe import (
flashinfer_cutedsl_moe_masked,
)
# flashinfer_cutedsl_moe_masked reinterprets scales as float8_e4m3fn.
# Same-dtype .view is a no-op; only wider dtypes (e.g. int32-packed
# UE8M0) need stride(-1)==1.
if (
MOE_NVFP4_DISPATCH
and x[1] is not None
and x[1].element_size() != 1
and x[1].stride(-1) != 1
):
raise AssertionError(
f"NVFP4 dispatch scale has stride(-1)={x[1].stride(-1)}, "
f"dtype={x[1].dtype}; .view(float8_e4m3fn) requires stride(-1)==1. "
"Try SGLANG_MOE_NVFP4_DISPATCH=0 or check DeepEP version."
)
down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = getattr(
layer, "down_gemm_overlap_args", None
)
out = flashinfer_cutedsl_moe_masked(
hidden_states=x,
input_global_scale=(
None if MOE_NVFP4_DISPATCH else layer.w13_input_scale_quant
),
w1=layer.w13_weight,
w1_blockscale=layer.w13_blockscale_swizzled,
w1_alpha=layer.g1_alphas,
w2=layer.w2_weight,
a2_global_scale=layer.w2_input_scale_quant,
w2_blockscale=layer.w2_blockscale_swizzled,
w2_alpha=layer.g2_alphas,
masked_m=masked_m,
**(
dict(
down_sm_count=down_gemm_overlap_args.num_sms,
down_signals=down_gemm_overlap_args.signal,
down_start_event=down_gemm_overlap_args.start_event,
)
if down_gemm_overlap_args is not None
else {}
),
)
return out