diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 3696cff43..789427edb 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -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, diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py index be2ae6b42..5e1f69978 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py @@ -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, + ) diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index 98ce75c7f..392534517 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index faeec7bf7..bab2e6f59 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -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 diff --git a/test/registered/moe/test_cutedsl_moe.py b/test/registered/moe/test_cutedsl_moe.py index 08ddd7db7..90205fd1b 100644 --- a/test/registered/moe/test_cutedsl_moe.py +++ b/test/registered/moe/test_cutedsl_moe.py @@ -780,7 +780,8 @@ class TestCuteDslV2(unittest.TestCase): class TestCuteDslV1(unittest.TestCase): """Correctness tests for the CuteDSL v1 (deepep) path. - The v1 path (apply_without_routing_weights -> flashinfer_cutedsl_moe_masked) + The v1 path (flashinfer_cutedsl_moe_masked, dispatched via the + @register_fused_func("deepep", "flashinfer_cutedsl") MoeRunner entry) is used when --moe-runner-backend flashinfer_cutedsl and --moe-a2a-backend deepep are combined. It expects: - W13 in default [Gate, Up] order (load_up_proj_weight_first = False)