iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kernel] Support Nvfp4 Cutedsl Moe Swiglu-oai and Relu2(non-gated) Activation (#47106)

Browse Source 
Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
This commit is contained in:
Bi Tiekai
2026-08-06 17:42:45 -07:00
committed by GitHub
co-authored by Claude Opus 5 Yongye Zhu
parent b1e12d142d
commit e08111211b
9 changed files with 216 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_flashinfer_cutedsl_layout.py
+43
View File
@@ -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)
tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py
+101 -8
View File
@@ -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"
vllm/model_executor/layers/fused_moe/config.py
+4
View File
@@ -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,
)
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
+26 -1
View File
@@ -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,
)
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+9 -2
View File
@@ -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,
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+2
View File
@@ -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,
)
vllm/model_executor/layers/quantization/modelopt.py
+2
View File
@@ -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,
)
vllm/model_executor/layers/quantization/quark/quark_moe.py
+3
View File
@@ -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,
)
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+26 -3
View File
@@ -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)