diff --git a/tests/kernels/moe/test_flashinfer_cutedsl_layout.py b/tests/kernels/moe/test_flashinfer_cutedsl_layout.py new file mode 100644 index 0000000000..8d71e3eaf5 --- /dev/null +++ b/tests/kernels/moe/test_flashinfer_cutedsl_layout.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Weight-layout normalization for the FlashInfer CuTeDSL NVFP4 MoE backend.""" + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( + reorder_w13_to_w31_for_flashinfer_cutedsl, +) + +_GATE = torch.tensor([[[1], [2], [3], [4]]]) +_UP = torch.tensor([[[10], [20], [30], [40]]]) +_EXPECTED = torch.cat([_UP, _GATE], dim=1) + + +def test_reorder_w13_swigluoai_interleaved(): + """gpt-oss w13 is [gate0, up0, gate1, ...] rather than packed [gate; up].""" + w13 = torch.empty(1, 8, 1, dtype=_GATE.dtype) + w13[:, 0::2] = _GATE + w13[:, 1::2] = _UP + + out, out_scale = reorder_w13_to_w31_for_flashinfer_cutedsl( + MoEActivation.SWIGLUOAI, w13, w13 + 100 + ) + + torch.testing.assert_close(out, _EXPECTED) + torch.testing.assert_close(out_scale, _EXPECTED + 100) + + +@pytest.mark.parametrize( + "activation", [MoEActivation.SILU, MoEActivation.SWIGLUOAI_UNINTERLEAVE] +) +def test_reorder_w13_packed_layouts(activation: MoEActivation): + w13 = torch.cat([_GATE, _UP], dim=1) + + out, out_scale = reorder_w13_to_w31_for_flashinfer_cutedsl( + activation, w13, w13 + 100 + ) + + torch.testing.assert_close(out, _EXPECTED) + torch.testing.assert_close(out_scale, _EXPECTED + 100) diff --git a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py index a7a7c5251b..1688c4c648 100644 --- a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py +++ b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py @@ -13,9 +13,13 @@ from tests.kernels.quantization.nvfp4_utils import ( FLOAT8_E4M3_MAX, break_fp4_bytes, ) -from tests.kernels.utils import torch_moe from vllm import _custom_ops as ops from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.activation import ( + SiluAndMul, + SiluAndMulWithClamp, + SwigluOAIAndMul, +) from vllm.model_executor.layers.fused_moe import fused_topk from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -46,6 +50,77 @@ if not has_flashinfer_cutedsl_moe_nvfp4() or not ( allow_module_level=True, ) +_SWIGLU_ALPHA = 1.702 +_SWIGLU_BETA = 1.0 +# The GEMM1 output is O(0.5) at these tensor scales, so the OAI default of 7.0 +# would never clamp anything. Pick a limit that actually bites. +_SWIGLU_LIMIT = 0.3 + +_ACT_CASES = [ + pytest.param(MoEActivation.SILU, None, None, None, id="silu"), + pytest.param(MoEActivation.SILU, None, None, _SWIGLU_LIMIT, id="silu-clamp"), + pytest.param(MoEActivation.RELU2_NO_MUL, None, None, None, id="relu2_no_mul"), + pytest.param( + MoEActivation.SWIGLUOAI, + _SWIGLU_ALPHA, + _SWIGLU_BETA, + _SWIGLU_LIMIT, + id="swigluoai", + ), + pytest.param( + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + _SWIGLU_ALPHA, + _SWIGLU_BETA, + _SWIGLU_LIMIT, + id="swigluoai_uninterleave", + ), +] + + +def _reference_activation( + activation: MoEActivation, + alpha: float | None, + beta: float | None, + limit: float | None, +): + """vLLM's own op for this activation, so the reference is not re-derived.""" + if activation == MoEActivation.RELU2_NO_MUL: + return lambda x: torch.square(torch.relu(x)) + if activation == MoEActivation.SWIGLUOAI: + # SwigluOAIAndMul hardcodes beta=1 and reads gate/up interleaved. + assert beta == 1.0 + return SwigluOAIAndMul(alpha=alpha, limit=limit) + if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: + return SiluAndMulWithClamp(limit, alpha, beta, compile_native=False) + if limit is not None: + return SiluAndMulWithClamp(limit, compile_native=False) + return SiluAndMul() + + +def _torch_moe_reference( + a: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + score: torch.Tensor, + topk: int, + act_fn, +) -> torch.Tensor: + m = a.shape[0] + weights, ids = torch.topk(torch.softmax(score, dim=-1, dtype=torch.float32), topk) + x = a.view(m, 1, -1).repeat(1, topk, 1).reshape(m * topk, -1) + flat_ids = ids.reshape(-1) + out = torch.zeros(m * topk, w2.shape[1], dtype=a.dtype, device=a.device) + for expert in range(w1.shape[0]): + mask = flat_ids == expert + if mask.any(): + acc = act_fn(x[mask] @ w1[expert].transpose(0, 1)).to(a.dtype) + out[mask] = acc @ w2[expert].transpose(0, 1) + return ( + (out.view(m, topk, -1).float() * weights.view(m, topk, 1)) + .sum(dim=1) + .to(a.dtype) + ) + def _quantize_nvfp4_linear( weight: torch.Tensor, @@ -85,14 +160,19 @@ def _dequantize_nvfp4_linear( @pytest.mark.parametrize("m,n,k,e,topk", [(16, 128, 512, 4, 2)]) +@pytest.mark.parametrize("activation,alpha,beta,limit", _ACT_CASES) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @torch.inference_mode() -def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul( +def test_flashinfer_cutedsl_fp4_moe( m: int, n: int, k: int, e: int, topk: int, + activation: MoEActivation, + alpha: float | None, + beta: float | None, + limit: float | None, dtype: torch.dtype, workspace_init, ): @@ -102,7 +182,8 @@ def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul( ): hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10 - w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 15 + w1_rows = 2 * n if activation.is_gated else n + w1 = torch.randn((e, w1_rows, k), device="cuda", dtype=dtype) / 15 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15 w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1) w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2) @@ -112,7 +193,6 @@ def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul( hidden_states, score, topk, renormalize=False ) - activation = MoEActivation.RELU2_NO_MUL fake_layer = SimpleNamespace(activation=activation) a1_scale = torch.ones(1, device="cuda", dtype=torch.float32) a2_scale = torch.ones(1, device="cuda", dtype=torch.float32) @@ -144,6 +224,11 @@ def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul( w1_scale=w1_scale_cutedsl, w2_scale=w2_scale_cutedsl, is_scale_swizzled=False, + # Unset params must be omitted rather than forwarded as None into + # the kernel's float-typed SwiGLU arguments. + gemm1_alpha=alpha, + gemm1_beta=beta, + gemm1_clamp_limit=limit, ) moe_config = FusedMoEConfig( num_experts=e, @@ -198,7 +283,7 @@ def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul( dtype=dtype, ) - w1_d = torch.empty((e, n, k), device="cuda", dtype=dtype) + w1_d = torch.empty((e, w1_rows, k), device="cuda", dtype=dtype) w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype) for idx in range(e): w1_d[idx] = _dequantize_nvfp4_linear( @@ -214,17 +299,25 @@ def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul( dtype=dtype, ) - torch_output = torch_moe( + torch_output = _torch_moe_reference( a_in_dtype, w1_d, w2_d, score, topk, - activation=activation, + _reference_activation(activation, alpha, beta, limit), ) torch.testing.assert_close( torch_output, cutedsl_output, - atol=2e-1, + atol=3e-2, rtol=2e-1, ) + # Outputs here are O(1e-2) while NVFP4 noise is O(1e-3), so an absolute + # tolerance loose enough for the quantization error also accepts a zero + # tensor. Compare direction too, which dropped SwiGLU params or a wrong + # w13 layout would break. + cosine = torch.nn.functional.cosine_similarity( + cutedsl_output.flatten().float(), torch_output.flatten().float(), dim=0 + ) + assert cosine > 0.99, f"cosine similarity {cosine:.4f} below 0.99" diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 0c633090ad..054240271f 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -814,6 +814,8 @@ def nvfp4_moe_quant_config( w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, is_scale_swizzled: bool = True, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: """ @@ -833,6 +835,8 @@ def nvfp4_moe_quant_config( per_out_ch_quant=False, block_shape=None, is_scale_swizzled=is_scale_swizzled, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, gemm1_clamp_limit=gemm1_clamp_limit, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index 823ae43b5e..e54a62124a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -59,6 +59,9 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): self.global_num_experts = moe_config.num_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank self.local_expert_offset = self.ep_rank * self.local_num_experts + self.gemm1_alpha = quant_config.gemm1_alpha + self.gemm1_beta = quant_config.gemm1_beta + self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) @@ -93,7 +96,12 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation in (MoEActivation.SILU, MoEActivation.RELU2_NO_MUL) + return activation in ( + MoEActivation.SILU, + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + MoEActivation.RELU2_NO_MUL, + ) @staticmethod def _supports_parallel_config( @@ -149,6 +157,22 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): # The functional API expects x_sf with trailing dim: (M, K//16, 1). x_sf = a1q_scale.unsqueeze(-1) + # The kernel defaults swiglu_{alpha,beta,limit} to the plain-SwiGLU + # values, so only forward the ones the model actually sets. + swiglu_params: dict[str, float | None] = {} + if activation == MoEActivation.SILU: + swiglu_params = {"swiglu_limit": self.gemm1_clamp_limit} + elif activation in ( + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ): + swiglu_params = { + "swiglu_alpha": self.gemm1_alpha, + "swiglu_beta": self.gemm1_beta, + "swiglu_limit": self.gemm1_clamp_limit, + } + swiglu_kwargs = {k: v for k, v in swiglu_params.items() if v is not None} + flashinfer_cute_dsl_fused_moe_nvfp4( x=hidden_states, x_sf=x_sf, @@ -167,4 +191,5 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): local_expert_offset=self.local_expert_offset, moe_output=output, activation_type=activation_to_flashinfer_int(activation), + **swiglu_kwargs, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 8439490d30..30c8458cbd 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -190,6 +190,7 @@ def select_nvfp4_moe_backend( NVFP4_BACKENDS_WITH_CLAMP = { NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTLASS, + NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.MARLIN, NvFp4MoeBackend.HUMMING, @@ -261,8 +262,8 @@ def select_nvfp4_moe_backend( f"Model sets swiglu_limit={config.swiglu_limit}, but the " f"explicitly requested moe_backend={runner_backend!r} does " f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm', " - f"'flashinfer_cutlass', 'cutlass', 'marlin', or 'humming' " - f"instead." + f"'flashinfer_cutlass', 'flashinfer_cutedsl', 'cutlass', " + f"'marlin', or 'humming' instead." ) return _return_or_raise( requested_backend, config, weight_key, activation_key, activation_format @@ -471,6 +472,8 @@ def make_nvfp4_moe_quant_config( a13_scale: torch.Tensor, a2_scale: torch.Tensor, swiglu_limit: float | None = None, + swiglu_alpha: float | None = None, + swiglu_beta: float | None = None, layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig: if backend == NvFp4MoeBackend.HUMMING: @@ -497,6 +500,8 @@ def make_nvfp4_moe_quant_config( a2_gscale=a2_scale, w1_scale=w13_scale, w2_scale=w2_scale, + gemm1_alpha=swiglu_alpha, + gemm1_beta=swiglu_beta, gemm1_clamp_limit=swiglu_limit, ) @@ -521,6 +526,8 @@ def make_nvfp4_moe_quant_config( NvFp4MoeBackend.FLASHINFER_CUTEDSL, ) ), + gemm1_alpha=swiglu_alpha, + gemm1_beta=swiglu_beta, gemm1_clamp_limit=swiglu_limit, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index c184d20689..9ae807fcb6 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -251,6 +251,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, swiglu_limit=getattr(layer, "swiglu_limit", None), + swiglu_alpha=getattr(layer, "swiglu_alpha", None), + swiglu_beta=getattr(layer, "swiglu_beta", None), layer=layer, ) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 05f486493d..8e22a6b73e 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1591,6 +1591,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, swiglu_limit=getattr(layer, "swiglu_limit", None), + swiglu_alpha=getattr(layer, "swiglu_alpha", None), + swiglu_beta=getattr(layer, "swiglu_beta", None), layer=layer, ) diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index f9468cf5da..508c715aec 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1659,6 +1659,9 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod): w2_scale_2=layer.w2_weight_scale_2, a13_scale=layer.w13_input_scale_2, a2_scale=layer.w2_input_scale_2, + swiglu_limit=getattr(layer, "swiglu_limit", None), + swiglu_alpha=getattr(layer, "swiglu_alpha", None), + swiglu_beta=getattr(layer, "swiglu_beta", None), layer=layer, ) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 1b513e9c0c..8fe1a22371 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING import torch from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( align_fp4_moe_weights_for_fi, align_trtllm_fp4_moe_hidden_dim_for_fi, @@ -80,6 +81,28 @@ def interleave_linear_and_gate( return x +def reorder_w13_to_w31_for_flashinfer_cutedsl( + activation: MoEActivation, + w13: torch.Tensor, + w13_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Normalize gated w13 rows to the [up; gate] order used by FlashInfer.""" + if activation == MoEActivation.SWIGLUOAI: + # gpt-oss checkpoints store w13 interleaved as [gate0, up0, gate1, ...]. + gate, up = w13[:, 0::2], w13[:, 1::2] + gate_scale, up_scale = w13_scale[:, 0::2], w13_scale[:, 1::2] + return ( + torch.cat([up, gate], dim=1).contiguous(), + torch.cat([up_scale, gate_scale], dim=1).contiguous(), + ) + + half = w13.shape[1] // 2 + return ( + torch.cat([w13[:, half:], w13[:, :half]], dim=1).contiguous(), + torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1).contiguous(), + ) + + def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( layer: "RoutedExperts", w13: torch.Tensor, @@ -114,9 +137,9 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts) if layer.activation.is_gated: - half = w13.shape[1] // 2 - w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1) - w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1) + w13, w13_scale = reorder_w13_to_w31_for_flashinfer_cutedsl( + layer.activation, w13, w13_scale + ) # Interleave up/gate rows for w13 weights and scales. w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)