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

fix legacy deepep path for flashinfer_cutedsl (#22925)

Browse Source
This commit is contained in:
Lee Nau
2026-04-20 11:49:33 -07:00
committed by GitHub
parent 4698f4cd10
commit b4bb036b73
5 changed files with 664 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+18 -5
View File
@@ -620,9 +620,25 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
input_global_scale = self.quant_config.get("input_global_scale", None)
if input_global_scale is not None:
use_nvfp4 = True
else:
elif not get_moe_runner_backend().is_flashinfer_cutedsl():
# flashinfer_cutedsl expects BF16 dispatch when NVFP4 dispatch is
# off; its kernel quantizes to NVFP4 internally.
use_fp8 = True
# round_scale / use_ue8m0 are FP8-DeepGEMM specific; they cause DeepEP
# to return int32-packed UE8M0 scales that don't feed the flashinfer
# cutedsl kernel.
fp8_deepgemm_scale_opts = (
dict(
round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
and deep_gemm_wrapper.DEEPGEMM_BLACKWELL,
use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
and deep_gemm_wrapper.DEEPGEMM_BLACKWELL,
)
if use_fp8
else dict()
)
buffer = self._get_buffer()
_deepep_precompile_tp_barrier()
packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = (
@@ -640,10 +656,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
),
async_finish=not self.return_recv_hook,
return_recv_hook=self.return_recv_hook,
round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
and deep_gemm_wrapper.DEEPGEMM_BLACKWELL,
use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
and deep_gemm_wrapper.DEEPGEMM_BLACKWELL,
**fp8_deepgemm_scale_opts,
)
)
return packed_recv_hidden, self.packed_recv_count, event, hook
python/sglang/srt/layers/moe/utils.py
+8
View File
@@ -263,6 +263,14 @@ def is_deepep_class_backend() -> bool:
return b.is_deepep() or b.is_mooncake() or b.is_mori()
def is_flashinfer_cutedsl_v1_path() -> bool:
"""CuteDSL v1 + DeepEP low-latency path (no MoeRunner, no autotune)."""
return (
get_moe_runner_backend().is_flashinfer_cutedsl()
and get_moe_a2a_backend().is_deepep()
)
def get_tbo_token_distribution_threshold() -> float:
global TBO_TOKEN_DISTRIBUTION_THRESHOLD
if TBO_TOKEN_DISTRIBUTION_THRESHOLD is None:
python/sglang/srt/layers/quantization/modelopt_quant.py
+70 -8
View File
@@ -24,7 +24,10 @@ from sglang.srt.layers.moe import (
)
from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
from sglang.srt.layers.moe.utils import should_use_flashinfer_cutlass_moe_fp4_allgather
from sglang.srt.layers.moe.utils import (
is_flashinfer_cutedsl_v1_path,
should_use_flashinfer_cutlass_moe_fp4_allgather,
)
from sglang.srt.layers.parameter import ModelWeightParameter, PerTensorScaleParameter
from sglang.srt.layers.quantization.base_config import (
FusedMoEMethodBase,
@@ -1546,6 +1549,29 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
return get_moe_runner_backend().is_flashinfer_cutedsl()
# ----- 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).
# - 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.
# - 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)."""
return is_flashinfer_cutedsl_v1_path()
@property
def _is_cutedsl_v2_standard(self) -> bool:
"""New CuteDSL standard path (a2a=none or flashinfer, uses MoeRunner)."""
return self.enable_flashinfer_cutedsl_moe and not self._is_cutedsl_v1_deepep
def create_weights(
self,
layer: torch.nn.Module,
@@ -1812,9 +1838,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
else:
# CUTLASS processing - handle w13 and w2 separately
if self.enable_flashinfer_cutedsl_moe and layer.moe_runner_config.is_gated:
# For the CuteDSL FP4 path, interleave the two logical W13 halves
# in 64-row chunks before swizzling the block-scales.
if self._is_cutedsl_v2_standard and layer.moe_runner_config.is_gated:
# CuteDSL v2 only: interleave the two logical W13 halves in
# 64-row chunks for the fused SwiGLU GEMM1 layout expected by
# CuteDslMoEWrapper. The v1 (deepep) path uses
# grouped_gemm_nt_masked which expects plain contiguous halves.
from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
interleave_w13_halves,
)
@@ -1876,8 +1904,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
layer, "w2_blockscale_swizzled", w2_blockscale_swizzled
)
if self.enable_flashinfer_cutedsl_moe:
# CuteDSL expects MMA layout for weight scales. Convert from swizzled bytes.
if self._is_cutedsl_v2_standard:
# CuteDSL v2 only: convert blockscales to MMA layout for
# CuteDslMoEWrapper. The v1 (deepep) path uses the
# swizzled blockscales directly via flashinfer_cutedsl_moe_masked.
from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
@@ -1940,9 +1970,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
@property
def load_up_proj_weight_first(self) -> bool:
# Load W13 as [Up, Gate] for FlashInfer CUTLASS/CuteDSL kernels.
# Load W13 as [Up, Gate] for FlashInfer CUTLASS and CuteDSL v2 kernels.
# The CuteDSL v1 (deepep) path uses [Gate, Up] -- do NOT flip.
return self.moe_runner_config.is_gated and (
self.enable_flashinfer_cutlass_moe or self.enable_flashinfer_cutedsl_moe
self.enable_flashinfer_cutlass_moe or self._is_cutedsl_v2_standard
)
def create_moe_runner(
@@ -1959,6 +1990,11 @@ 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)
@@ -2010,6 +2046,9 @@ 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,
@@ -2124,6 +2163,14 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
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."
@@ -2137,6 +2184,21 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
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
)
python/sglang/srt/model_executor/model_runner.py
+11
View File
@@ -2193,6 +2193,17 @@ class ModelRunner(ModelRunnerKVCacheMixin):
if self.server_args.disable_flashinfer_autotune:
return False
# CuteDSL v1 (cutedsl runner + deepep a2a) bypasses MoeRunner and must not
# be autotuned -- its _dummy_run would dispatch more tokens per rank than
# SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK, tripping a DeepEP assert.
# Read server_args directly to avoid depending on initialize_moe_config()
# having already populated the MoE backend globals.
if (
self.server_args.moe_runner_backend == "flashinfer_cutedsl"
and self.server_args.moe_a2a_backend == "deepep"
):
return False
backend_str = self.server_args.moe_runner_backend
# TODO smor- support other cases for flashinfer autotune, such as, mamba backend