diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py new file mode 100644 index 000000000..a9213d83b --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py @@ -0,0 +1,353 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch + +from sglang.srt.layers.moe.moe_runner.base import ( + MoeQuantInfo, + MoeRunnerConfig, + register_fused_func, +) +from sglang.srt.utils.common import log_info_on_rank0, print_warning_once + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + StandardCombineInput, + StandardDispatchOutput, + ) + +logger = logging.getLogger(__name__) + +_FP4_SF_VEC_SIZE = 16 +_cutedsl_logged_scalarize: set = set() + + +# --------------------------------------------------------------------------- +# Weight / scale preparation utilities (called from modelopt_quant.py during +# process_weights_after_loading and lazy wrapper init) +# --------------------------------------------------------------------------- + + +def interleave_w13_halves( + tensor: torch.Tensor, group_size: int = 64, dim: int = 1 +) -> torch.Tensor: + """Interleave the two logical W13 halves for CuteDSL's SwiGLU GEMM1 layout. + + The caller is responsible for loading W13 in the expected two-half order. + This helper only rewrites the first and second halves into alternating + `group_size` chunks along `dim`. + """ + if tensor.shape[dim] % 2 != 0: + raise ValueError( + "Expected even size on interleave dimension for W13 half split." + ) + split = tensor.shape[dim] // 2 + if split % group_size != 0: + raise ValueError( + f"Expected split dim divisible by group_size={group_size}, got {split}." + ) + first_half = tensor.narrow(dim, 0, split) + second_half = tensor.narrow(dim, split, split) + first_half_groups = first_half.split(group_size, dim=dim) + second_half_groups = second_half.split(group_size, dim=dim) + interleaved = [ + item for pair in zip(first_half_groups, second_half_groups) for item in pair + ] + return torch.cat(interleaved, dim=dim) + + +def cutedsl_quant_scale_to_scalar( + quant_scale: torch.Tensor, + *, + name: str, +) -> torch.Tensor: + """Reduce per-expert quant-domain scale vector to a single scalar. + + The quant domain is the reciprocal of the raw checkpoint scale: + quant_scale = 1 / raw_scale + + Returns min(quant_scale) = 1/max(raw_scale), which is the TRTLLM CuteDSL + convention for global scalar activation scales (see TRTLLM quantization.py + lines 2137-2141: fc2_input_scale = tmp_fc2_input_scale.max().reciprocal()). + + If quant_scale is already scalar (numel==1), returns it unchanged. + """ + quant_scale = quant_scale.to(torch.float32) + if quant_scale.numel() == 0: + print_warning_once( + f"CuteDSL got empty {name}; using 1.0 fallback.", + ) + return torch.ones(1, device=quant_scale.device, dtype=torch.float32) + if quant_scale.numel() == 1: + return quant_scale.reshape(1) + if name not in _cutedsl_logged_scalarize: + log_info_on_rank0( + logger, + f"CuteDSL: reducing per-expert {name} to scalar via " + "min(quant_scale) = 1/max(raw_scale), matching TRTLLM convention.", + ) + _cutedsl_logged_scalarize.add(name) + return quant_scale.min().reshape(1) + + +def resolve_cutedsl_standard_scales( + layer: torch.nn.Module, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Resolve standard-path CuteDSL scales (baseline: scalar fc2/w13 input scales). + + Returns (w1_alpha, fc2_input_scale, w2_alpha, used_input_scale). + used_input_scale is the scalarized w13 input scale for FP4 quantize and GEMM1. + """ + + def _to_fp32_tensor(x: torch.Tensor | float, ref: torch.Tensor) -> torch.Tensor: + if not isinstance(x, torch.Tensor): + x = torch.tensor(x, device=ref.device) + return x.to(device=ref.device, dtype=torch.float32) + + def _align_scale_to_alpha( + scale: torch.Tensor, alpha: torch.Tensor, scale_name: str + ) -> torch.Tensor: + scale = scale.to(device=alpha.device, dtype=torch.float32) + alpha = alpha.to(torch.float32) + if scale.ndim == 0: + return scale + # Gated weight scales may be (num_experts, 2) with separate gate/up + # columns. Collapse to 1D by taking the first column (gate == up for + # well-formed checkpoints; mismatch is warned in process_weights_after_loading). + if scale.ndim == 2 and scale.shape[1] <= 2: + scale = scale[:, 0] + if scale.numel() == alpha.numel(): + return scale + if scale.numel() == 1: + return scale.reshape(()) + + # Some EP setups may carry global-per-expert scale vectors while alphas are + # local-per-expert vectors. Slice to this rank's local expert range. + num_local_experts = getattr(layer, "num_local_experts", None) + num_experts = getattr(layer, "num_experts", None) + moe_ep_rank = getattr(layer, "moe_ep_rank", 0) + if ( + num_local_experts is not None + and num_experts is not None + and scale.numel() == num_experts + and alpha.numel() == num_local_experts + ): + start = moe_ep_rank * num_local_experts + end = start + num_local_experts + return scale[start:end] + + raise ValueError( + f"Unable to align {scale_name} shape={tuple(scale.shape)} " + f"to alpha shape={tuple(alpha.shape)} for CuteDSL standard scale resolution." + ) + + def _resolve_w1_alpha_from_scalar_input_scale( + used_input_scale: torch.Tensor, + ) -> torch.Tensor: + """Resolve GEMM1 alpha consistent with scalarized activation quant scale. + + CuteDSL pre-quantizes x with a single scalar (used_input_scale), but + g1_alphas was derived with per-expert activation scales: + g1_alphas[e] = (1/w13_isq[e]) * w13_ws2[e] + Correct alpha for scalar quantization: + w1_alpha[e] = w13_ws2[e] / used_input_scale + = g1_alphas[e] * w13_isq[e] / used_input_scale + When w13_isq is already scalar, this is a no-op (ratio = 1). + """ + eps = 1e-12 + scalar = torch.clamp(used_input_scale.to(torch.float32).reshape(()), min=eps) + + if hasattr(layer, "w13_weight_scale_2"): + w13_weight_scale_2 = _align_scale_to_alpha( + layer.w13_weight_scale_2, layer.g1_alphas, "w13_weight_scale_2" + ) + return w13_weight_scale_2.to(torch.float32) / scalar + + w13_isq = _align_scale_to_alpha( + layer.w13_input_scale_quant, layer.g1_alphas, "w13_input_scale_quant" + ) + w13_isq = torch.clamp(_to_fp32_tensor(w13_isq, layer.g1_alphas), min=eps) + return (layer.g1_alphas.to(torch.float32) * w13_isq / scalar).to(torch.float32) + + def _resolve_w2_alpha_from_scalar_fc2_input_scale( + fc2_input_scale: torch.Tensor, + ) -> torch.Tensor: + """Resolve GEMM2 alpha consistent with scalarized FC2 input scale. + + CuteDSL standard path uses a scalar global scale for GEMM1 FP4 output + quantization (`fc2_input_scale`). GEMM2 alpha must use the same scalar + convention: alpha2 = w2_weight_scale_2 / fc2_input_scale. + """ + eps = 1e-12 + fc2_input_scale = fc2_input_scale.to(torch.float32) + fc2_scalar = torch.clamp(fc2_input_scale.reshape(-1)[:1], min=eps).reshape(()) + + if hasattr(layer, "w2_weight_scale_2"): + w2_weight_scale_2 = _align_scale_to_alpha( + layer.w2_weight_scale_2, layer.g2_alphas, "w2_weight_scale_2" + ) + w2_weight_scale_2 = w2_weight_scale_2.to(torch.float32) + return w2_weight_scale_2 / fc2_scalar + + w2_q_for_w2 = _align_scale_to_alpha( + layer.w2_input_scale_quant, layer.g2_alphas, "w2_input_scale_quant" + ) + w2_q_for_w2 = torch.clamp( + _to_fp32_tensor(w2_q_for_w2, layer.g2_alphas), min=eps + ) + w2_weight_scale_2 = layer.g2_alphas.to(torch.float32) * w2_q_for_w2 + return w2_weight_scale_2 / fc2_scalar + + fc2_input_scale = cutedsl_quant_scale_to_scalar( + layer.w2_input_scale_quant, + name="w2_input_scale_quant", + ) + w2_alpha = _resolve_w2_alpha_from_scalar_fc2_input_scale(fc2_input_scale) + used_input_scale = cutedsl_quant_scale_to_scalar( + layer.w13_input_scale_quant, + name="w13_input_scale_quant", + ) + w1_alpha = _resolve_w1_alpha_from_scalar_input_scale(used_input_scale) + return w1_alpha, fc2_input_scale, w2_alpha, used_input_scale + + +def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None: + """Lazily create CuteDslMoEWrapper and resolve scales on first forward. + + The wrapper is created lazily (not in __init__ / create_weights) because + it depends on final weight shapes and EP configuration. The wrapper's + CUDA-graph buffers are allocated inside CuteDslMoEWrapper.__init__, which + typically runs during the autotune dummy forward under inference_mode(). + We wrap the creation in inference_mode(False) so that those pre-allocated + buffers are normal tensors -- inference tensors cannot be inplace-updated + during later CUDA graph capture, which runs outside inference_mode. + """ + if getattr(layer, "_cutedsl_wrapper", None) is not None: + return + + try: + from flashinfer import CuteDslMoEWrapper + except ImportError as e: + raise ImportError( + "flashinfer_cutedsl backend requires FlashInfer with CuteDSL support. " + "Install with: pip install flashinfer" + ) from e + + from sglang.srt.server_args import get_global_server_args + + assert layer.intermediate_size_per_partition > 0, ( + f"CuteDSL MoE: intermediate_size_per_partition must be > 0, " + f"got {layer.intermediate_size_per_partition}. Check EP/TP configuration." + ) + + server_args = get_global_server_args() + use_cuda_graph = server_args is not None and not server_args.disable_cuda_graph + max_num_tokens = max( + getattr(server_args, "cuda_graph_max_bs", None) or 512, + getattr(server_args, "chunked_prefill_size", None) or 8192, + ) + top_k = layer.top_k if layer.top_k is not None else layer.moe_runner_config.top_k + # inference_mode(False) ensures the wrapper's pre-allocated CUDA-graph + # buffers are normal tensors. This call typically happens inside + # _dummy_run which runs under inference_mode(); inference tensors cannot + # be inplace-updated during later CUDA graph capture (which runs outside + # inference_mode), so we must opt out here. + with torch.inference_mode(False): + layer._cutedsl_wrapper = CuteDslMoEWrapper( + num_experts=layer.num_experts, + top_k=top_k, + hidden_size=layer.hidden_size, + intermediate_size=layer.intermediate_size_per_partition, + use_cuda_graph=use_cuda_graph, + max_num_tokens=max_num_tokens, + num_local_experts=layer.num_local_experts, + local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, + output_dtype=layer.moe_runner_config.params_dtype, + device=str(layer.w13_weight.device), + ) + + w1_alpha, fc2_input_scale, w2_alpha, used_input_scale = ( + resolve_cutedsl_standard_scales(layer) + ) + layer._cutedsl_scales = (w1_alpha, fc2_input_scale, w2_alpha) + layer._cutedsl_input_scale = used_input_scale + + +# --------------------------------------------------------------------------- +# Dataclass + fused function for moe_runner dispatch +# --------------------------------------------------------------------------- + + +@dataclass +class CuteDslFp4MoeQuantInfo(MoeQuantInfo): + """Quantization payload consumed by FlashInfer CuteDSL FP4 MoE kernels.""" + + # Lazily-created CuteDslMoEWrapper (stashed on layer) + wrapper: Any + + # Weights (uint8 FP4 packed) + w13_weight: torch.Tensor + w2_weight: torch.Tensor + + # Block-scale factors + w13_weight_sf: torch.Tensor + w2_weight_sf: torch.Tensor + + # Per-expert GEMM scales + w1_alpha: torch.Tensor + w2_alpha: torch.Tensor + + # Intermediate quantization scale (fc2 input) + fc2_input_scale: torch.Tensor + + # Activation quantization scale (scalarized) + input_scale: torch.Tensor + + +@register_fused_func("none", "flashinfer_cutedsl") +def fused_experts_none_to_flashinfer_cutedsl_fp4( + dispatch_output: StandardDispatchOutput, + quant_info: CuteDslFp4MoeQuantInfo, + runner_config: MoeRunnerConfig, +) -> StandardCombineInput: + from flashinfer import fp4_quantize + + from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput + from sglang.srt.layers.moe.topk import TopKOutputChecker + + assert runner_config.activation == "silu", "Only silu is supported for CuteDSL MoE." + + hidden_states = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + assert TopKOutputChecker.format_is_standard(topk_output) + + topk_ids = topk_output.topk_ids + topk_weights = topk_output.topk_weights + if topk_ids.dtype != torch.int32: + topk_ids = topk_ids.to(torch.int32) + + x_fp4, x_sf = fp4_quantize( + hidden_states, + quant_info.input_scale, + sf_vec_size=_FP4_SF_VEC_SIZE, + is_sf_swizzled_layout=False, + ) + + output = quant_info.wrapper.run( + x=x_fp4, + x_sf=x_sf, + token_selected_experts=topk_ids, + token_final_scales=topk_weights, + 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, + 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) diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index 758bf2e06..7e62a6424 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -54,6 +54,8 @@ class MoeRunner: or runner_backend.is_flashinfer_trtllm_routed() ): self.runner_core = None # FlashInfer TRT-LLM only supports fused path + elif runner_backend.is_flashinfer_cutedsl(): + self.runner_core = None # FlashInfer CuteDSL only supports fused path else: raise NotImplementedError(f"Unsupported runner backend: {runner_backend}") diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index ef78ca9f3..9e22d8092 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/standard.py @@ -88,8 +88,14 @@ class StandardDispatcher(BaseDispatcher): def __init__(self, moe_runner_config: MoeRunnerConfig): super().__init__() self.moe_ep_size = get_moe_expert_parallel_world_size() - self.enable_flashinfer_cutlass_moe = ( - get_moe_runner_backend().is_flashinfer_cutlass() + backend = get_moe_runner_backend() + self.enable_flashinfer_cutlass_moe = backend.is_flashinfer_cutlass() + # FlashInfer CUTLASS and CuteDSL handle EP internally with global expert IDs. + # Skip local expert mapping so topk_ids stay in global space. + self.skip_local_expert_mapping = ( + backend.is_flashinfer_cutlass() + or backend.is_flashinfer_cutedsl() + or backend.is_flashinfer_trtllm_routed() ) self.enable_flashinfer_trtllm_routed_moe = ( get_moe_runner_backend().is_flashinfer_trtllm_routed() @@ -149,8 +155,7 @@ class StandardDispatcher(BaseDispatcher): if ( self.moe_ep_size > 1 - and not self.enable_flashinfer_cutlass_moe - and not self.enable_flashinfer_trtllm_routed_moe + and not self.skip_local_expert_mapping and TopKOutputChecker.format_is_standard(topk_output) ): if self.local_expert_mapping is None: diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index d3b278c16..3115bf0a5 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -51,7 +51,6 @@ from sglang.srt.layers.quantization.utils import ( from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.utils import copy_or_rebind_param from sglang.srt.utils.common import ( - get_bool_env_var, is_cuda, is_sm120_supported, next_power_of_2, @@ -167,10 +166,6 @@ if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None): return -CUTEDSL_MOE_SCALAR_INPUT_SCALE = get_bool_env_var( - "SGLANG_CUTEDSL_MOE_SCALAR_INPUT_SCALE", "true" -) - # FP4 GEMM alignment constant - CUTLASS/FlashInfer kernels require dimensions divisible by 32 FP4_GEMM_ALIGNMENT = 32 @@ -993,9 +988,10 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): ) -> CombineInput: x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + from sglang.srt.layers.moe.topk import TopKOutputChecker # Fast path: TRT-LLM FP8 per-tensor MoE using BYPASSED TopK routing - from sglang.srt.layers.moe.topk import TopKOutputChecker if ( get_moe_runner_backend().is_flashinfer_trtllm() @@ -1089,8 +1085,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): activation_type=activation, )[0] - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - return StandardCombineInput(hidden_states=output) quant_info = TritonMoeQuantInfo( @@ -1547,9 +1541,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): @property def enable_flashinfer_cutedsl_moe(self) -> bool: + """Access the global enable_flashinfer_cutedsl_moe setting.""" from sglang.srt.layers.moe import get_moe_runner_backend - """Access the global enable_flashinfer_cutedsl_moe setting.""" return get_moe_runner_backend().is_flashinfer_cutedsl() def create_weights( @@ -1714,19 +1708,12 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): w13_input_scale = layer.w13_input_scale.max().to(torch.float32) w2_input_scale = layer.w2_input_scale.max().to(torch.float32) elif self.enable_flashinfer_cutedsl_moe: - # All-expert-one-input-scale is mathematically different from default per-expert-input-scale - # Thus we allow users to switch the flag to do thorough testing - if CUTEDSL_MOE_SCALAR_INPUT_SCALE: - w13_input_scale = ( - layer.w13_input_scale.max() - .to(torch.float32) - .repeat(layer.w13_input_scale.shape[0]) - ) - else: - w13_input_scale = layer.w13_input_scale.max(dim=1).values.to( - torch.float32 - ) - + # CuteDSL standard path uses a single scalar input scale (all experts). + w13_input_scale = ( + layer.w13_input_scale.max() + .to(torch.float32) + .repeat(layer.w13_input_scale.shape[0]) + ) w2_input_scale = layer.w2_input_scale def _slice_scale(w): @@ -1825,6 +1812,26 @@ 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. + from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( + interleave_w13_halves, + ) + + layer.w13_weight = Parameter( + interleave_w13_halves( + layer.w13_weight.view(torch.uint8), group_size=64, dim=1 + ).contiguous(), + requires_grad=False, + ) + layer.w13_weight_scale = Parameter( + interleave_w13_halves( + layer.w13_weight_scale, group_size=64, dim=1 + ).contiguous(), + requires_grad=False, + ) + # Process w13 weights w13_blockscale_swizzled = swizzle_blockscale(layer.w13_weight_scale) copy_or_rebind_param( @@ -1869,6 +1876,45 @@ 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. + from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout + + from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( + _FP4_SF_VEC_SIZE, + ) + + sf_vec_size = _FP4_SF_VEC_SIZE + num_local_experts = layer.w13_weight.shape[0] + w13_m = layer.w13_weight.shape[1] + w13_k = layer.w13_weight.shape[2] * 2 + w2_m = layer.w2_weight.shape[1] + w2_k = layer.w2_weight.shape[2] * 2 + layer.w13_blockscale_mma = Parameter( + convert_sf_to_mma_layout( + layer.w13_blockscale_swizzled.contiguous() + .view(torch.uint8) + .reshape(-1), + m=w13_m, + k=w13_k, + num_groups=num_local_experts, + sf_vec_size=sf_vec_size, + ), + requires_grad=False, + ) + layer.w2_blockscale_mma = Parameter( + convert_sf_to_mma_layout( + layer.w2_blockscale_swizzled.contiguous() + .view(torch.uint8) + .reshape(-1), + m=w2_m, + k=w2_k, + num_groups=num_local_experts, + sf_vec_size=sf_vec_size, + ), + requires_grad=False, + ) + # Both flashinfer cutlass and regular cutlass use same processing for w2 # Set up CUTLASS MoE parameters (reuse to keep CUDA graph stable) @@ -1894,27 +1940,34 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): @property def load_up_proj_weight_first(self) -> bool: - # FlashInfer CUTLASS kernel assumes [Up, Gate] Proj as W13 - return self.enable_flashinfer_cutlass_moe and self.moe_runner_config.is_gated + # Load W13 as [Up, Gate] for FlashInfer CUTLASS/CuteDSL kernels. + return self.moe_runner_config.is_gated and ( + self.enable_flashinfer_cutlass_moe or self.enable_flashinfer_cutedsl_moe + ) def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): self.moe_runner_config = moe_runner_config - if get_moe_runner_backend().is_flashinfer_trtllm(): - self.runner = MoeRunner( - MoeRunnerBackend.FLASHINFER_TRTLLM, moe_runner_config - ) - elif get_moe_runner_backend().is_flashinfer_trtllm_routed(): - self.runner = MoeRunner( - MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED, moe_runner_config - ) + moe_runner_backend = get_moe_runner_backend() + + if moe_runner_backend.is_auto(): + # TRTLLM is currently the most performant and tested FP4 MoE + # backend, so use it as the default. + moe_runner_backend = MoeRunnerBackend.FLASHINFER_TRTLLM + + if moe_runner_backend.is_flashinfer_cutedsl(): + import sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl # noqa: F401 – triggers @register_fused_func + + if not moe_runner_backend.is_flashinfer_cutlass(): + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) def apply( self, layer: FusedMoE, dispatch_output: StandardDispatchOutput, ) -> CombineInput: + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput x = dispatch_output.hidden_states x_sf = dispatch_output.hidden_states_scale @@ -1958,6 +2011,33 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): return self.runner.run(dispatch_output, quant_info) + if self.enable_flashinfer_cutedsl_moe: + from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( + CuteDslFp4MoeQuantInfo, + ensure_cutedsl_wrapper, + ) + + 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, + w1_alpha=w1_alpha, + w2_alpha=w2_alpha, + fc2_input_scale=fc2_input_scale, + input_scale=layer._cutedsl_input_scale, + ) + return self.runner.run(dispatch_output, quant_info) + if self.enable_flashinfer_cutlass_moe: from sglang.srt.layers.moe.token_dispatcher import DispatchOutputChecker @@ -1997,7 +2077,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): fc2_expert_weights=layer.w2_weight.view(torch.long), output_dtype=output_dtype, input_sf=x_sf, - # swizzled_input_sf=not get_moe_a2a_backend().is_flashinfer(), + # swizzled_input_sf intentionally omitted; not used for this path. quant_scales=[ layer.w13_input_scale_quant, layer.w13_blockscale_swizzled.view(torch.int32), @@ -2015,8 +2095,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): enable_alltoall=get_moe_a2a_backend().is_flashinfer(), )[0] - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - return StandardCombineInput(hidden_states=output) from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4 @@ -2038,8 +2116,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input, ).to(x.dtype) # Scale by routed_scaling_factor is fused into select_experts. - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - return StandardCombineInput(hidden_states=output) def apply_without_routing_weights( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index bde3c135e..76d9ded64 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2148,6 +2148,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): # TODO: Enable for flashinfer_trtllm_routed once https://github.com/flashinfer-ai/flashinfer/issues/2749 is fixed. # "flashinfer_trtllm_routed", "flashinfer_mxfp4", + "flashinfer_cutedsl", # TODO: flashinfer_cutlass will cause some flashinfer compilation errors. To be fixed. # "flashinfer_cutlass", ]: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 753d18c21..b6c5676cb 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2733,6 +2733,26 @@ class ServerArgs: self.tp_size, ], "The expert parallel size must be 1 or the same as the tensor parallel size" + if self.moe_runner_backend == "flashinfer_cutedsl": + assert self.quantization in [ + "modelopt_fp4" + ], f"Invalid quantization '{self.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4'." + assert self.ep_size in [ + 1, + self.tp_size, + ], "The expert parallel size must be 1 or the same as the tensor parallel size" + assert self.moe_a2a_backend in [ + "none", + "deepep", + ], ( + f"flashinfer_cutedsl supports moe_a2a_backend='none' (standard path) " + f"or 'deepep' (DeepEP low-latency path), got '{self.moe_a2a_backend}'." + ) + self.disable_shared_experts_fusion = True + logger.warning( + "FlashInfer CuteDSL MoE is enabled. --disable-shared-experts-fusion is automatically set." + ) + if self.moe_runner_backend == "flashinfer_trtllm": assert self.quantization in [ "modelopt_fp4", diff --git a/test/registered/backends/test_deepseek_v3_fp4_cutedsl_moe.py b/test/registered/backends/test_deepseek_v3_fp4_cutedsl_moe.py new file mode 100644 index 000000000..68e96593a --- /dev/null +++ b/test/registered/backends/test_deepseek_v3_fp4_cutedsl_moe.py @@ -0,0 +1,149 @@ +"""Backend tests for CuteDSL MoE (FusedMoE + moe_runner, moe_a2a=none). + +Exercises the CuteDSL moe_runner path with ModelOpt FP4 by launching a +server with --moe-runner-backend flashinfer_cutedsl. + +Two configurations are tested: + - EP=1, TP=4: each GPU holds all experts with TP-sharded intermediate dim + - EP=4, TP=4: each GPU holds 1/4 of experts at full intermediate width, + partial results combined via all-reduce (no A2A dispatch) + +Requires 4 GPUs. Run from repo root with: + python -m pytest test/registered/backends/test_deepseek_v3_fp4_cutedsl_moe.py -v -s +Or via the nightly suite: + python test/run_suite.py --hw cuda --suite nightly-4-gpu-b200 --nightly +""" + +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.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) + +register_cuda_ci(est_time=900, suite="nightly-4-gpu-b200", nightly=True) + +FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3-0324-FP4" +SERVER_LAUNCH_TIMEOUT = 1000 +GSM8K_ACCURACY_THRESHOLD = 0.935 + + +class TestDeepseekV3FP4CuteDSLMoE(CustomTestCase): + """CuteDSL standard moe_runner path: flashinfer_cutedsl + modelopt_fp4, EP=1.""" + + @classmethod + def setUpClass(cls): + cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = [ + "--tp", + "4", + "--ep", + "1", + "--mem-fraction-static", + "0.75", + "--attention-backend", + "trtllm_mla", + "--moe-runner-backend", + "flashinfer_cutedsl", + "--quantization", + "modelopt_fp4", + "--model-loader-extra-config", + '{"enable_multithread_load": true}', + ] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_a_gsm8k( + self, + ): # Append an "a" to make this test run first (alphabetically) to warm up the server + args = SimpleNamespace( + num_shots=8, + data_path=None, + num_questions=1319, + parallel=1319, + max_new_tokens=512, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + if is_in_ci(): + write_github_step_summary( + f"### test_gsm8k (deepseek-v3-fp4-cutedsl-moe)\n" + f'{metrics["accuracy"]=:.3f}\n' + ) + self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD) + + +class TestDeepseekV3FP4CuteDSLMoEEP4(CustomTestCase): + """CuteDSL standard moe_runner path: flashinfer_cutedsl + modelopt_fp4, EP=TP=4.""" + + @classmethod + def setUpClass(cls): + cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = [ + "--tp", + "4", + "--ep", + "4", + "--mem-fraction-static", + "0.75", + "--attention-backend", + "trtllm_mla", + "--moe-runner-backend", + "flashinfer_cutedsl", + "--moe-a2a-backend", + "none", + "--quantization", + "modelopt_fp4", + "--model-loader-extra-config", + '{"enable_multithread_load": true}', + ] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_a_gsm8k(self): + args = SimpleNamespace( + num_shots=8, + data_path=None, + num_questions=1319, + parallel=1319, + max_new_tokens=512, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + if is_in_ci(): + write_github_step_summary( + f"### test_gsm8k (deepseek-v3-fp4-cutedsl-moe-ep4)\n" + f'{metrics["accuracy"]=:.3f}\n' + ) + self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/moe/test_cutedsl_moe.py b/test/registered/moe/test_cutedsl_moe.py index 76b3b4a87..f8449c4c9 100644 --- a/test/registered/moe/test_cutedsl_moe.py +++ b/test/registered/moe/test_cutedsl_moe.py @@ -1,18 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 import unittest -from typing import Callable import torch from flashinfer import fp4_quantize, scaled_fp4_grouped_quantize from torch.nn import functional as F -from sglang.jit_kernel.nvfp4 import scaled_fp4_quant from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.moe.flashinfer_cutedsl_moe import flashinfer_cutedsl_moe_masked -from sglang.srt.layers.moe.topk import TopKConfig, select_experts from sglang.test.ci.ci_register import register_cuda_ci -register_cuda_ci(est_time=20, suite="stage-c-test-4-gpu-b200") +try: + from flashinfer import CuteDslMoEWrapper + from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout +except ImportError: + CuteDslMoEWrapper = None + convert_sf_to_mma_layout = None + +register_cuda_ci(est_time=300, suite="stage-c-test-4-gpu-b200") SKIP_TEST = torch.cuda.get_device_capability() < (10, 0) SKIP_REASON = "Nvfp4 Requires compute capability of 10 or above." @@ -78,6 +82,312 @@ def break_fp4_bytes(a, dtype): return values.reshape(m, n * 2).to(dtype=dtype) +def _interleave_w13_halves( + x: torch.Tensor, group_size: int = 64, dim: int = -1 +) -> torch.Tensor: + """Interleave the two logical W13 halves for the CuteDSL wrapper layout.""" + sizes = x.size() + dim = dim % x.dim() + assert sizes[dim] % (group_size * 2) == 0 + prev_sizes = sizes[:dim] + post_sizes = sizes[dim + 1 :] + x = x.view(*prev_sizes, 2, sizes[dim] // (group_size * 2), group_size, *post_sizes) + x = x.transpose(dim, dim + 1).contiguous().view(*sizes) + return x + + +def _create_cutedsl_wrapper_tensors( + num_tokens: int, + hidden_size: int, + intermediate_size: int, + num_experts: int, + top_k: int, + device: str = "cuda", + seed: int = 42, +): + """Create quantized tensors for CuteDslMoEWrapper.run() (MMA layout, same as production). + + Returns quantized inputs for the wrapper **and** the original bf16 weights + needed to compute a numerical reference. Scale values (w1_alpha, w2_alpha, + fc2_input_scale) are derived from weight magnitudes so that scale-contract + bugs are caught. + """ + assert CuteDslMoEWrapper is not None and convert_sf_to_mma_layout is not None + torch.manual_seed(seed) + sf_vec_size = 16 + + x_bf16 = ( + torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) / 10 + ) + a1_gs = torch.tensor([1.0], device=device, dtype=torch.float32) + x_quantized, x_sf = fp4_quantize( + x_bf16, + global_scale=a1_gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=False, + ) + x_sf = x_sf.unsqueeze(-1) + + router_logits = torch.randn(num_tokens, num_experts, device=device) + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) + routing_weights = routing_weights.float() + selected_experts = selected_experts.to(torch.int32) + + # --- GEMM1 weights --- + w1_bf16 = ( + torch.randn( + num_experts, + 2 * intermediate_size, + hidden_size, + dtype=torch.bfloat16, + device=device, + ) + / 10 + ) + w1_bf16_interleaved = _interleave_w13_halves(w1_bf16, group_size=64, dim=1) + w1_amax = w1_bf16.abs().amax(dim=(1, 2)).to(torch.float32) + w1_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax.mean() + w1_gs = w1_gs.unsqueeze(0) + w1_flat = w1_bf16_interleaved.view(num_experts * 2 * intermediate_size, hidden_size) + w1_q_flat, w1_sf_flat = fp4_quantize( + w1_flat, + global_scale=w1_gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=True, + ) + w1_q = w1_q_flat.view(num_experts, 2 * intermediate_size, hidden_size // 2) + w1_weight_sf = convert_sf_to_mma_layout( + w1_sf_flat, + m=2 * intermediate_size, + k=hidden_size, + num_groups=num_experts, + sf_vec_size=sf_vec_size, + ) + w1_alpha = 1.0 / (a1_gs * w1_gs).expand(num_experts) + + # --- GEMM2 weights --- + w2_bf16 = ( + torch.randn( + num_experts, + hidden_size, + intermediate_size, + dtype=torch.bfloat16, + device=device, + ) + / 10 + ) + w2_amax = w2_bf16.abs().amax(dim=(1, 2)).to(torch.float32) + w2_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax.mean() + w2_gs = w2_gs.unsqueeze(0) + w2_flat = w2_bf16.view(num_experts * hidden_size, intermediate_size) + w2_q_flat, w2_sf_flat = fp4_quantize( + w2_flat, + global_scale=w2_gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=True, + ) + w2_q = w2_q_flat.view(num_experts, hidden_size, intermediate_size // 2) + w2_weight_sf = convert_sf_to_mma_layout( + w2_sf_flat, + m=hidden_size, + k=intermediate_size, + num_groups=num_experts, + sf_vec_size=sf_vec_size, + ) + fc2_input_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax.mean() + fc2_input_scale = fc2_input_scale.unsqueeze(0) + w2_alpha = 1.0 / (fc2_input_scale * w2_gs).expand(num_experts) + + return { + "x": x_quantized, + "x_sf": x_sf, + "x_bf16": x_bf16, + "token_selected_experts": selected_experts, + "token_final_scales": routing_weights, + "w1_weight": w1_q, + "w1_weight_sf": w1_weight_sf, + "w1_weight_bf16": w1_bf16, + "w1_alpha": w1_alpha, + "fc2_input_scale": fc2_input_scale, + "w2_weight": w2_q, + "w2_weight_sf": w2_weight_sf, + "w2_weight_bf16": w2_bf16, + "w2_alpha": w2_alpha, + # Global scales needed by _quantize_local_expert_weights + "a1_gs": a1_gs, + "w1_gs": w1_gs, + "w2_gs": w2_gs, + } + + +def _quantize_local_expert_weights( + w1_bf16_local: torch.Tensor, + w2_bf16_local: torch.Tensor, + a1_gs: torch.Tensor, + w1_gs: torch.Tensor, + w2_gs: torch.Tensor, + fc2_input_scale: torch.Tensor, +): + """Independently quantize and MMA-convert a local expert weight shard. + + Mirrors the per-rank weight preprocessing that happens during model loading + in production (each rank holds [num_local_experts, ...] bf16 weights, + quantizes them, and calls convert_sf_to_mma_layout with + num_groups=num_local_experts). + """ + sf_vec_size = 16 + num_local_experts = w1_bf16_local.shape[0] + intermediate_size_2x = w1_bf16_local.shape[1] + hidden_size = w1_bf16_local.shape[2] + intermediate_size = w2_bf16_local.shape[2] + + # GEMM1: interleave -> quantize -> MMA layout + w1_interleaved = _interleave_w13_halves(w1_bf16_local, group_size=64, dim=1) + w1_flat = w1_interleaved.view(num_local_experts * intermediate_size_2x, hidden_size) + w1_q_flat, w1_sf_flat = fp4_quantize( + w1_flat, + global_scale=w1_gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=True, + ) + w1_q = w1_q_flat.view(num_local_experts, intermediate_size_2x, hidden_size // 2) + w1_sf = convert_sf_to_mma_layout( + w1_sf_flat, + m=intermediate_size_2x, + k=hidden_size, + num_groups=num_local_experts, + sf_vec_size=sf_vec_size, + ) + w1_alpha = 1.0 / (a1_gs * w1_gs).expand(num_local_experts) + + # GEMM2: quantize -> MMA layout + w2_flat = w2_bf16_local.view(num_local_experts * hidden_size, intermediate_size) + w2_q_flat, w2_sf_flat = fp4_quantize( + w2_flat, + global_scale=w2_gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=True, + ) + w2_q = w2_q_flat.view(num_local_experts, hidden_size, intermediate_size // 2) + w2_sf = convert_sf_to_mma_layout( + w2_sf_flat, + m=hidden_size, + k=intermediate_size, + num_groups=num_local_experts, + sf_vec_size=sf_vec_size, + ) + w2_alpha = 1.0 / (fc2_input_scale * w2_gs).expand(num_local_experts) + + return { + "w1_weight": w1_q, + "w1_weight_sf": w1_sf, + "w1_alpha": w1_alpha, + "w2_weight": w2_q, + "w2_weight_sf": w2_sf, + "w2_alpha": w2_alpha, + } + + +def _run_wrapper(wrapper, tensors, **overrides): + """Call wrapper.run() with the standard 11-arg dict from _create_cutedsl_wrapper_tensors.""" + kwargs = dict( + x=tensors["x"], + x_sf=tensors["x_sf"], + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + w1_weight=tensors["w1_weight"], + w1_weight_sf=tensors["w1_weight_sf"], + w1_alpha=tensors["w1_alpha"], + fc2_input_scale=tensors["fc2_input_scale"], + w2_weight=tensors["w2_weight"], + w2_weight_sf=tensors["w2_weight_sf"], + w2_alpha=tensors["w2_alpha"], + ) + kwargs.update(overrides) + return wrapper.run(**kwargs) + + +def _quant_dequant_fp4_reference( + tensor: torch.Tensor, + global_scale: torch.Tensor, + sf_vec_size: int = 16, +) -> torch.Tensor: + """Simulate FP4 quant-dequant roundtrip for reference computation.""" + from flashinfer.fp4_quantization import e2m1_and_ufp8sf_scale_to_float + + tensor_bf16 = tensor.to(torch.bfloat16) + fp4_packed, sf = fp4_quantize( + tensor_bf16, + global_scale=global_scale, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=False, + ) + sf_uint8 = sf.view(torch.uint8).reshape(-1) + dequantized = e2m1_and_ufp8sf_scale_to_float( + fp4_packed.cpu(), + sf_uint8.cpu(), + (1.0 / global_scale).cpu(), + sf_vec_size=sf_vec_size, + ufp8_type=1, + is_sf_swizzled_layout=False, + ).to(tensor.device) + return dequantized.float() + + +def _compute_reference_moe_fp4( + hidden_states: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm2_weights: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + fc2_input_scale: torch.Tensor, +) -> torch.Tensor: + """Pure-PyTorch MoE reference using bf16 weights (pre-interleave layout). + + gemm1_weights is [num_experts, 2*intermediate_size, hidden_size] with the + *original* (un-interleaved) layout: first half = linear, second half = gate. + """ + device = hidden_states.device + num_tokens = hidden_states.shape[0] + hidden_states = hidden_states.float() + gemm1_weights = gemm1_weights.float() + gemm2_weights = gemm2_weights.float() + + output = torch.zeros(num_tokens, hidden_size, dtype=torch.float32, device=device) + + for token_idx in range(num_tokens): + token_input = hidden_states[token_idx : token_idx + 1] + for k in range(top_k): + expert_idx = token_selected_experts[token_idx, k].item() + scale = token_final_scales[token_idx, k].item() + if expert_idx < 0 or expert_idx >= num_experts: + continue + + w1 = gemm1_weights[expert_idx] + gemm1_out = token_input @ w1.T + + linear = gemm1_out[:, :intermediate_size] + gate = gemm1_out[:, intermediate_size:] + swiglu_out = F.silu(gate) * linear + + if fc2_input_scale is not None: + swiglu_out = _quant_dequant_fp4_reference( + swiglu_out, fc2_input_scale, sf_vec_size=16 + ) + + w2 = gemm2_weights[expert_idx] + gemm2_out = swiglu_out @ w2.T + output[token_idx] += scale * gemm2_out.squeeze(0) + + return output + + def compute_routing(router_logits: torch.Tensor, top_k: int): routing_weights = torch.softmax(router_logits, dim=1, dtype=torch.float) routing_weights, selected_experts = torch.topk(routing_weights, top_k, dim=-1) @@ -109,20 +419,6 @@ def prepare_inputs( return hidden_states_3d, masked_m, topk_idx, routing_weights -MNK_FACTORS = [ - (2, 1024, 1024), - (2, 1024, 1536), - (2, 3072, 1024), - (2, 3072, 1536), - (64, 1024, 1024), - (64, 1024, 1536), - (64, 3072, 1024), - (64, 2048, 1024), - (224, 1024, 1024), - (224, 1024, 1536), -] - - # Reference implementation of torch_moe def torch_moe(a, w1, w2, score, topk, expert_map): B, D = a.shape @@ -158,7 +454,7 @@ def torch_moe_nvfp4(a, w1, w2, topk, topk_weight, topk_ids): if mask.sum(): m = w1[i].shape[0] assert m % 2 == 0 - # Note: w1 and w3 are swapped! + # The first and second W13 halves feed the two SwiGLU branches. w3_expert, w1_expert = w1[i][m // 2 :, :], w1[i][: m // 2, :] inter = F.silu(a[mask] @ w1_expert.t()) * (a[mask] @ w3_expert.t()) inter_gs = torch.tensor(1.0).cuda() @@ -177,128 +473,6 @@ def torch_moe_nvfp4(a, w1, w2, topk, topk_weight, topk_ids): ).sum(dim=1) -def check_moe( - m: int, - n: int, - k: int, - e: int, - topk: int, - dtype: torch.dtype, - moe_impl: Callable, - flip_w13: bool, -): - torch.manual_seed(7) - a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 - w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 - quant_blocksize = 16 - round_up = lambda x, y: (x + y - 1) // y * y - sf_w1_2n = round_up(2 * n, 128) - sf_w1_k = round_up(k // quant_blocksize, 4) - w1_blockscale = torch.empty( - (e, sf_w1_2n, sf_w1_k), device="cuda", dtype=torch.float8_e4m3fn - ) - - w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 - sf_w2_k = round_up(k, 128) - sf_w2_n = round_up(n // quant_blocksize, 4) - w2_blockscale = torch.empty( - (e, sf_w2_k, sf_w2_n), device="cuda", dtype=torch.float8_e4m3fn - ) - - w1_q = torch.empty((e, 2 * n, k // 2), device="cuda", dtype=torch.uint8) - w2_q = torch.empty((e, k, n // 2), device="cuda", dtype=torch.uint8) - w1_gs = torch.empty((e,), device="cuda", dtype=torch.float32) - w2_gs = torch.empty((e,), device="cuda", dtype=torch.float32) - - for expert in range(e): - w1_amax = torch.abs(w1).max().to(torch.float32) - w2_amax = torch.abs(w2).max().to(torch.float32) - w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax - w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax - - w1_q[expert], w1_blockscale[expert] = scaled_fp4_quant( - w1[expert], w1_gs[expert] - ) - - w2_q[expert], w2_blockscale[expert] = scaled_fp4_quant( - w2[expert], w2_gs[expert] - ) - - score = torch.randn((m, e), device="cuda", dtype=dtype) - - topk_output = select_experts( - hidden_states=a, - router_logits=score, - topk_config=TopKConfig(top_k=topk, renormalize=False), - ) - topk_weights, topk_ids, _ = topk_output - - a1_gs = torch.ones((e,), device="cuda", dtype=torch.float32) - a2_gs = torch.ones((e,), device="cuda", dtype=torch.float32) - test_output = moe_impl( - a=a, - topk_weights=topk_weights, - topk_ids=topk_ids, - w1_q=w1_q, - w2_q=w2_q, - a1_gs=a1_gs, - w1_blockscale=w1_blockscale, - w1_alphas=(1 / w1_gs), - a2_gs=a2_gs, - w2_blockscale=w2_blockscale, - w2_alphas=(1 / w2_gs), - ) - - # Reference check: - a_global_scale = ( - (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1) - ).to(torch.float32) - a_fp4, a_scale_interleaved = scaled_fp4_quant(a, a_global_scale) - _, m_k = a_fp4.shape - a_in_dtype = dequantize_nvfp4_to_dtype( - a_fp4, - a_scale_interleaved, - a_global_scale, - dtype=a.dtype, - device=a.device, - block_size=quant_blocksize, - ) - - w1_d = torch.empty((e, 2 * n, k), device="cuda", dtype=dtype) - w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype) - - for idx in range(0, e): - w1_d[idx] = dequantize_nvfp4_to_dtype( - w1_q[idx], - w1_blockscale[idx], - w1_gs[idx], - dtype=w1.dtype, - device=w1.device, - block_size=quant_blocksize, - ) - w2_d[idx] = dequantize_nvfp4_to_dtype( - w2_q[idx], - w2_blockscale[idx], - w2_gs[idx], - dtype=w2.dtype, - device=w2.device, - block_size=quant_blocksize, - ) - - if flip_w13: - dim = -2 - size = w1_d.size(dim) - assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}" - half = size // 2 - # Reorder weight - w1, w3 = w1_d.split(half, dim=dim) - w1_d = torch.cat([w3, w1], dim=dim).contiguous() - - torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk, None) - - torch.testing.assert_close(torch_output, test_output, atol=1e-1, rtol=1e-1) - - class TestFlashinferCutedslMoe(unittest.TestCase): @unittest.skipIf(SKIP_TEST, SKIP_REASON) def test_flashinfer_cutedsl_moe_masked(self): @@ -316,9 +490,6 @@ class TestFlashinferCutedslMoe(unittest.TestCase): with self.subTest( bs=bs, hidden_dim=hidden_dim, inter_dim=inter_dim, topk=topk ): - print( - f"Testing with bs={bs}, hidden_dim={hidden_dim}, inter_dim={inter_dim}, topk={topk}" - ) with torch.inference_mode(): torch.manual_seed(42) device = "cuda" @@ -476,8 +647,296 @@ class TestFlashinferCutedslMoe(unittest.TestCase): torch.testing.assert_close( out_weighted.cpu(), ref_output.cpu(), atol=5e-2, rtol=5e-2 ) - print( - f"Test passed with bs={bs}, hidden_dim={hidden_dim}, inter_dim={inter_dim}, topk={topk}" + + @unittest.skipIf(SKIP_TEST, SKIP_REASON) + @unittest.skipIf( + CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None, + "CuteDslMoEWrapper / convert_sf_to_mma_layout not available", + ) + def test_cutedsl_moe_wrapper_run(self): + """Call CuteDslMoEWrapper.run() with MMA-layout tensors and verify against reference.""" + test_cases = [ + # (num_tokens, hidden_size, intermediate_size, num_experts, top_k) + # Minimum dimensions match FlashInfer's test_wrapper_accuracy: + # num_experts >= 256, hidden_size >= 256, intermediate_size >= 512, + # num_tokens >= 128. The CuteDSL GEMM kernels have tile-size + # constraints that make smaller dimensions unreliable. + (128, 256, 512, 256, 2), + (128, 256, 512, 256, 8), + (256, 256, 512, 256, 4), + ] + + for ( + num_tokens, + hidden_size, + intermediate_size, + num_experts, + top_k, + ) in test_cases: + with self.subTest( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + top_k=top_k, + ): + tensors = _create_cutedsl_wrapper_tensors( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + top_k=top_k, + ) + + wrapper = CuteDslMoEWrapper( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + use_cuda_graph=False, + ) + + with torch.no_grad(): + out = _run_wrapper(wrapper, tensors) + + self.assertEqual(out.shape, (num_tokens, hidden_size)) + self.assertEqual(out.dtype, torch.bfloat16) + self.assertFalse( + torch.isnan(out).any().item() or torch.isinf(out).any().item(), + "Output contains NaN or Inf", + ) + + ref_output = _compute_reference_moe_fp4( + hidden_states=tensors["x_bf16"].float().cuda(), + gemm1_weights=tensors["w1_weight_bf16"].float().cuda(), + gemm2_weights=tensors["w2_weight_bf16"].float().cuda(), + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + fc2_input_scale=tensors["fc2_input_scale"], + ) + + out_f32 = out.float() + ref_f32 = ref_output.float() + output_scale = max(ref_f32.std().item(), 0.01) + atol = max(0.1, 3.0 * output_scale) + rtol = 0.85 + abs_diff = torch.abs(out_f32 - ref_f32) + rel_diff = abs_diff / (torch.abs(ref_f32) + 1e-8) + within_tol = (abs_diff < atol) | (rel_diff < rtol) + pct_within = within_tol.float().mean().item() + self.assertGreaterEqual( + pct_within, + 0.925, + f"Only {pct_within * 100:.2f}% of elements within tolerance " + f"(atol={atol:.4f})", + ) + + @unittest.skipIf(SKIP_TEST, SKIP_REASON) + @unittest.skipIf( + CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None, + "CuteDslMoEWrapper / convert_sf_to_mma_layout not available", + ) + def test_cutedsl_cuda_graph_parity(self): + """Verify non-graph and cuda_graph wrappers produce identical results. + + Also checks both match the pure-PyTorch reference, and that a second + cuda_graph pass reuses buffers deterministically (subsumes the former + cuda_graph_smoke test). + """ + test_cases = [ + # (num_tokens, hidden_size, intermediate_size, num_experts, top_k) + (128, 256, 512, 256, 2), + (256, 256, 512, 256, 4), + ] + + for ( + num_tokens, + hidden_size, + intermediate_size, + num_experts, + top_k, + ) in test_cases: + with self.subTest( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + top_k=top_k, + ): + tensors = _create_cutedsl_wrapper_tensors( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + top_k=top_k, + ) + + wrapper_args = dict( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + ) + wrapper_no_graph = CuteDslMoEWrapper( + **wrapper_args, use_cuda_graph=False + ) + wrapper_graph = CuteDslMoEWrapper( + **wrapper_args, + use_cuda_graph=True, + max_num_tokens=num_tokens, + ) + + with torch.no_grad(): + out_no_graph = _run_wrapper(wrapper_no_graph, tensors) + out_graph = _run_wrapper(wrapper_graph, tensors) + out_graph2 = _run_wrapper(wrapper_graph, tensors) + + torch.testing.assert_close( + out_no_graph, + out_graph, + atol=1e-2, + rtol=1e-2, + msg="non-graph vs cuda_graph wrapper outputs diverge", + ) + torch.testing.assert_close( + out_graph, + out_graph2, + atol=1e-5, + rtol=1e-5, + msg="second cuda_graph pass should reuse buffers identically", + ) + + ref_output = _compute_reference_moe_fp4( + hidden_states=tensors["x_bf16"].float().cuda(), + gemm1_weights=tensors["w1_weight_bf16"].float().cuda(), + gemm2_weights=tensors["w2_weight_bf16"].float().cuda(), + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + fc2_input_scale=tensors["fc2_input_scale"], + ) + + out_f32 = out_graph.float() + ref_f32 = ref_output.float() + output_scale = max(ref_f32.std().item(), 0.01) + atol = max(0.1, 3.0 * output_scale) + rtol = 0.85 + abs_diff = torch.abs(out_f32 - ref_f32) + rel_diff = abs_diff / (torch.abs(ref_f32) + 1e-8) + within_tol = (abs_diff < atol) | (rel_diff < rtol) + pct_within = within_tol.float().mean().item() + self.assertGreaterEqual( + pct_within, + 0.925, + f"graph vs reference: only {pct_within * 100:.2f}% within tol", + ) + + @unittest.skipIf(SKIP_TEST, SKIP_REASON) + @unittest.skipIf( + CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None, + "CuteDslMoEWrapper / convert_sf_to_mma_layout not available", + ) + def test_cutedsl_ep_sharded_allreduce(self): + """Verify EP-sharded execution: partial outputs from EP ranks sum to full result. + + Simulates the EP=TP all-reduce pattern used by the CuteDSL moe_runner when + ep_size > 1 and moe_a2a_backend=none. Each "rank" runs a wrapper with + num_local_experts < num_experts and a corresponding local_expert_offset, + receiving only the local slice of weights/scales/alphas — matching the + real runtime contract where each rank holds only its own expert partition. + The partial outputs are summed (simulating tensor_model_parallel_all_reduce) + and compared against a single wrapper processing all experts. + """ + test_cases = [ + # (num_tokens, hidden_size, intermediate_size, num_experts, top_k, ep_size) + # Dimensions match FlashInfer's minimum wrapper requirements. + (128, 256, 512, 256, 2, 2), + (128, 256, 512, 256, 2, 4), + (128, 256, 512, 256, 8, 8), + ] + + for ( + num_tokens, + hidden_size, + intermediate_size, + num_experts, + top_k, + ep_size, + ) in test_cases: + with self.subTest( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + top_k=top_k, + ep_size=ep_size, + ): + assert num_experts % ep_size == 0 + num_local_experts = num_experts // ep_size + + tensors = _create_cutedsl_wrapper_tensors( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + top_k=top_k, + ) + + # Full-expert baseline (EP=1): all experts on one "rank" + wrapper_full = CuteDslMoEWrapper( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + use_cuda_graph=False, + ) + with torch.no_grad(): + out_full = _run_wrapper(wrapper_full, tensors) + + # EP-sharded: each rank independently quantizes its local + # bf16 weight shard and calls convert_sf_to_mma_layout with + # num_groups=num_local_experts — matching the real per-rank + # weight preprocessing in the CuteDSL moe_runner path. + accumulated = torch.zeros_like(out_full) + for rank in range(ep_size): + lo = rank * num_local_experts + hi = lo + num_local_experts + + local_tensors = _quantize_local_expert_weights( + w1_bf16_local=tensors["w1_weight_bf16"][lo:hi], + w2_bf16_local=tensors["w2_weight_bf16"][lo:hi], + a1_gs=tensors["a1_gs"], + w1_gs=tensors["w1_gs"], + w2_gs=tensors["w2_gs"], + fc2_input_scale=tensors["fc2_input_scale"], + ) + + wrapper_shard = CuteDslMoEWrapper( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + use_cuda_graph=False, + num_local_experts=num_local_experts, + local_expert_offset=lo, + ) + with torch.no_grad(): + partial = _run_wrapper(wrapper_shard, tensors, **local_tensors) + accumulated += partial + + torch.testing.assert_close( + out_full, + accumulated, + atol=1e-2, + rtol=1e-2, + msg=( + f"EP-sharded all-reduce mismatch " + f"(ep_size={ep_size}, tokens={num_tokens})" + ), )