From 882ca8d6964458ace0379e0eaa8aac9039b1bffc Mon Sep 17 00:00:00 2001 From: Lain Date: Mon, 31 Aug 2026 19:49:14 -0700 Subject: [PATCH] [Kernel] add Flashinfer cutedsl w4a16 linear (#53014) Signed-off-by: Siyuan Fu --- docs/features/quantization/modelopt.md | 10 ++- .../test_flashinfer_nvfp4_scaled_mm.py | 60 ++++++++++++++ tests/quantization/test_modelopt.py | 12 ++- vllm/config/kernel.py | 2 +- .../model_executor/kernels/linear/__init__.py | 23 +++++- .../kernels/linear/nvfp4/flashinfer.py | 80 +++++++++++++++++++ .../layers/quantization/modelopt.py | 4 +- vllm/utils/flashinfer.py | 49 ++++++++++++ 8 files changed, 230 insertions(+), 10 deletions(-) diff --git a/docs/features/quantization/modelopt.md b/docs/features/quantization/modelopt.md index 139497abab..c93e7e15df 100644 --- a/docs/features/quantization/modelopt.md +++ b/docs/features/quantization/modelopt.md @@ -17,6 +17,8 @@ following `quantization.quant_algo` values: - `FP8_PER_CHANNEL_PER_TOKEN`: per-channel weight scale and dynamic per-token activation quantization. - `FP8_PB_WO` (ModelOpt may emit `fp8_pb_wo`): block-scaled FP8 weight-only (typically 128×128 blocks). - `NVFP4`: ModelOpt NVFP4 checkpoints (use `quantization="modelopt_fp4"`). +- `W4A16_NVFP4`: weight-only ModelOpt NVFP4 checkpoints with 16-bit + activations (use `quantization="modelopt_fp4"`). - `MXFP8`: ModelOpt MXFP8 checkpoints (use `quantization="modelopt_mxfp8"`). !!! note @@ -28,10 +30,12 @@ following `quantization.quant_algo` values: workloads. Use `--linear-backend` to override the automatic selection (this replaces the deprecated `VLLM_NVFP4_GEMM_BACKEND` environment variable). Values relevant to NVFP4 include `cutlass`, - `flashinfer_cutlass`, `flashinfer_trtllm`, `flashinfer_cudnn`, and - `marlin`; the full list is documented under `KernelConfig` on the + `flashinfer_cutlass`, `flashinfer_cutedsl`, `flashinfer_trtllm`, + `flashinfer_cudnn`, and `marlin`; see `KernelConfig` on the [Engine Arguments](../../configuration/engine_args.md) page and shown by - `vllm serve --help=KernelConfig`. + `vllm serve --help=KernelConfig`. For `W4A16_NVFP4`, `auto` currently + selects Marlin. Models with BF16 activations can explicitly select the + FlashInfer CuTe-DSL backend with `--linear-backend flashinfer_cutedsl`. !!! note For models quantized to MXFP8 with BF16 activations on SM100-family GPUs, diff --git a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py index a1e76d73a1..54dcaa0b97 100644 --- a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py +++ b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch +import torch.nn.functional as F from nvfp4_utils import ( FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX, @@ -10,6 +11,13 @@ from nvfp4_utils import ( ) from vllm import _custom_ops as ops +from vllm.model_executor.kernels.linear.nvfp4 import NvFp4LinearLayerConfig +from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4W4A16LinearKernel, +) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + dequantize_to_dtype, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( flashinfer_scaled_fp4_mm, @@ -166,3 +174,55 @@ def test_flashinfer_nvfp4_gemm( ) torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1) + + +@pytest.mark.parametrize("shape", SHAPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_flashinfer_cutedsl_nvfp4_w4a16_linear( + shape: tuple[int, int, int], + seed: int, + device: str, +) -> None: + supported, reason = FlashInferCuteDslNvFp4W4A16LinearKernel.is_supported() + if not supported: + pytest.skip(reason) + + set_random_seed(seed) + m, n, packed_k = shape + k = packed_k * 2 + x = torch.randn((m, k), dtype=torch.bfloat16, device=device) + weight = torch.randn((n, k), dtype=torch.bfloat16, device=device) + weight_quant_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / weight.abs().max()).to( + torch.float32 + ) + weight_global_scale = weight_quant_scale.reciprocal() + weight_fp4, weight_scale = ops.scaled_fp4_quant( + weight, + weight_quant_scale, + is_sf_swizzled_layout=False, + ) + weight_ref = dequantize_to_dtype( + weight_fp4, + weight_scale, + weight_global_scale, + dtype=torch.bfloat16, + swizzle=False, + ) + + layer = torch.nn.Module() + layer.output_size_per_partition = n + layer.weight = torch.nn.Parameter(weight_fp4, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + layer.weight_global_scale = torch.nn.Parameter( + weight_global_scale, requires_grad=False + ) + kernel = FlashInferCuteDslNvFp4W4A16LinearKernel(NvFp4LinearLayerConfig()) + + kernel.process_weights_after_loading(layer) + output = kernel.apply_weights(layer, x) + + expected = F.linear(x, weight_ref) + assert output.shape == (m, n) + torch.testing.assert_close(output, expected, atol=1e-1, rtol=1e-1) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 052462f7f3..c4b549a91a 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -16,6 +16,7 @@ from tests.quantization.utils import is_quant_method_supported from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.model import ModelConfig from vllm.model_executor.kernels.linear import ( + FlashInferCuteDslNvFp4W4A16LinearKernel, HummingNvFp4LinearKernel, MarlinNvFp4LinearKernel, ) @@ -584,10 +585,19 @@ def test_modelopt_nvfp4_config_dispatches_w4a16_method(): @pytest.mark.parametrize( ("linear_backend", "kernel_cls"), - [("auto", MarlinNvFp4LinearKernel), ("humming", HummingNvFp4LinearKernel)], + [ + ("auto", MarlinNvFp4LinearKernel), + ("humming", HummingNvFp4LinearKernel), + ("flashinfer_cutedsl", FlashInferCuteDslNvFp4W4A16LinearKernel), + ], ) @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls): + if linear_backend != "auto": + is_supported, reason = kernel_cls.is_supported() + if not is_supported: + pytest.skip(reason) + vllm_config = VllmConfig() vllm_config.kernel_config.linear_backend = linear_backend with set_current_vllm_config(vllm_config): diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 56b3546a2f..0b5c368dcc 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -276,7 +276,7 @@ class KernelConfig: - "cutlass": Use CUTLASS-based kernels - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels - "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels - (BF16, NVFP4, MXFP8) + (BF16, NVFP4, MXFP8, W4A16_NVFP4) - "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels - "flashinfer_cudnn": Use FlashInfer with cuDNN kernels - "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index e02334b84c..fb0df9f063 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -149,6 +149,7 @@ from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( FlashInferB12xNvFp4LinearKernel, FlashInferCudnnNvFp4LinearKernel, FlashInferCuteDslNvFp4LinearKernel, + FlashInferCuteDslNvFp4W4A16LinearKernel, FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, ) @@ -262,6 +263,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { }, "flashinfer_cutedsl": { FlashInferCuteDslNvFp4LinearKernel, + FlashInferCuteDslNvFp4W4A16LinearKernel, FlashInferCutedslMxfp8LinearKernel, }, "flashinfer_trtllm": { @@ -531,6 +533,7 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = { _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { PlatformEnum.CUDA: [ FlashInferCuteDslNvFp4LinearKernel, + FlashInferCuteDslNvFp4W4A16LinearKernel, FlashInferCutlassNvFp4LinearKernel, FlashInferB12xNvFp4LinearKernel, CutlassNvFp4LinearKernel, @@ -1014,7 +1017,11 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: """Select and instantiate the best NVFP4 linear kernel for the current platform.""" config = NvFp4LinearLayerConfig() - a16_kernels = (MarlinNvFp4LinearKernel, HummingNvFp4LinearKernel) + a16_kernels = ( + FlashInferCuteDslNvFp4W4A16LinearKernel, + MarlinNvFp4LinearKernel, + HummingNvFp4LinearKernel, + ) # VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the # batch-invariant CUTLASS implementation when available, otherwise fall @@ -1051,8 +1058,17 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: ) force_kernel = EmulationNvFp4LinearKernel elif linear_backend == "auto" and use_a16: - # Force a16 (Marlin) when running weight-only quantization. - force_kernel = MarlinNvFp4LinearKernel + _cc = current_platform.get_device_capability() + compute_capability = _cc.to_int() if _cc is not None else None + # Weight-only: prefer FlashInfer CuTe-DSL W4A16 on SM100/103, + # otherwise Marlin. + cutedsl_ok, _ = FlashInferCuteDslNvFp4W4A16LinearKernel.is_supported( + compute_capability + ) + if compute_capability in (100, 103) and cutedsl_ok: + force_kernel = FlashInferCuteDslNvFp4W4A16LinearKernel + else: + force_kernel = MarlinNvFp4LinearKernel if force_kernel is not None: if use_a16 and force_kernel not in a16_kernels: @@ -1232,6 +1248,7 @@ __all__ = [ "EmulationNvFp4LinearKernel", "FbgemmNvFp4LinearKernel", "FlashInferCuteDslNvFp4LinearKernel", + "FlashInferCuteDslNvFp4W4A16LinearKernel", "FlashInferB12xNvFp4LinearKernel", "FlashInferCutlassNvFp4LinearKernel", "FlashInferTrtllmNvFp4LinearKernel", diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index b721a135e2..09fc3a4811 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -20,14 +20,94 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( + flashinfer_prepare_bf16_fp4_weights, flashinfer_scaled_fp4_mm, has_flashinfer, has_flashinfer_b12x_gemm, + has_flashinfer_bf16_fp4, ) from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig +class FlashInferCuteDslNvFp4W4A16LinearKernel(NvFp4LinearKernel): + """BF16 x NVFP4 GEMM via FlashInfer's CuTe-DSL backend.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if compute_capability is None: + if not current_platform.is_cuda(): + return False, "FlashInfer CuTe-DSL W4A16 requires CUDA" + capability = current_platform.get_device_capability() + if capability is None: + return False, "CUDA compute capability is unavailable" + compute_capability = capability.to_int() + + if compute_capability not in (100, 103) and not ( + compute_capability >= 120 and compute_capability < 130 + ): + return False, "FlashInfer CuTe-DSL W4A16 requires sm_100 or sm_12x" + if not has_flashinfer_bf16_fp4(): + return False, "FlashInfer CuTe-DSL BF16 x FP4 GEMM is unavailable" + return True, None + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + padded_weight, weights_padding_bytes = pad_nvfp4_weight_for_cutlass( + layer.weight.data, alignment=64 + ) + swizzled_scale = swizzle_blockscale(layer.weight_scale.data) + weight, weight_scale, global_scale = flashinfer_prepare_bf16_fp4_weights( + padded_weight, + swizzled_scale, + layer.weight_global_scale.data.reshape(1), + backend="cute-dsl", + ) + assert global_scale is not None + + layer.weight = torch.nn.Parameter(weight, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + layer.weight_global_scale = torch.nn.Parameter( + global_scale, requires_grad=False + ) + layer.weights_padding_cols = weights_padding_bytes + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if x.dtype != torch.bfloat16: + raise ValueError( + f"FlashInfer CuTe-DSL W4A16 requires BF16 input, got {x.dtype}" + ) + + output_size = layer.output_size_per_partition + output_shape = [*x.shape[:-1], output_size] + x_2d = x.reshape(-1, x.shape[-1]) + weights_padding_bytes = getattr(layer, "weights_padding_cols", 0) + if weights_padding_bytes: + x_2d = torch.nn.functional.pad(x_2d, (0, weights_padding_bytes * 2)) + x_2d = x_2d.contiguous() + + out = torch.ops.vllm.flashinfer_mm_bf16_fp4( + x_2d, + layer.weight, + layer.weight_scale, + layer.weight_global_scale, + ) + out = slice_nvfp4_output(out, output_size) + if bias is not None: + out = out + bias + return out.view(*output_shape) + + class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel): """NVFP4 GEMM via FlashInfer's cutedsl backend.""" diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 575c122d6b..aa11f9269e 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1038,7 +1038,7 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase): # Select LinearMethod implementation based on quant_algo (FP8 pattern). # NVFP4 -> W4A4: cutlass NVFP4 GEMM with input quantization - # W4A16_NVFP4 -> W4A16: FP4 Marlin GEMM with bf16/fp16 activations + # W4A16_NVFP4 -> W4A16: flashinfer cute-dsl or Marlin with 16-bit inputs if quant_method == "NVFP4": self.LinearMethodCls = ModelOptNvFp4LinearMethod elif quant_method == "W4A16_NVFP4": @@ -1283,7 +1283,7 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase): self.quant_config = quant_config self.marlin_input_dtype = None # `init_nvfp4_linear_kernel(use_a16=True)` is best of both worlds: - # 1. `use_a16=True` forces `Marlin`: https://github.com/vllm-project/vllm/commit/e68988a#diff-7135ab92aa94dfacb1ad3c77fc13f9c4ffe0b977f8eac5d86c2afe243e5f92a6R842-R889 + # 1. `use_a16=True` forces flashinfer cutedsl for sm100 and `Marlin` for others. # for `--linear-backend=auto`, avoiding a W4A4 kernel that requires input_scale. # 2. Specifying e.g. `--linear-backend=humming` will override. self.kernel = init_nvfp4_linear_kernel(use_a16=True) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index ed9cfbd720..e20fac6a6c 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -197,6 +197,9 @@ flashinfer_cutlass_fused_moe = _lazy_import_wrapper( flashinfer_cutedsl_grouped_gemm_nt_masked = _lazy_import_wrapper( "flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked" ) +flashinfer_prepare_bf16_fp4_weights = _lazy_import_wrapper( + "flashinfer.gemm", "prepare_bf16_fp4_weights" +) flashinfer_fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize") flashinfer_mxfp4_quantize = _lazy_import_wrapper("flashinfer", "mxfp4_quantize") nvfp4_batched_quantize = _lazy_import_wrapper("flashinfer", "nvfp4_batched_quantize") @@ -333,6 +336,17 @@ def has_flashinfer_cutedsl() -> bool: ) +@functools.cache +def has_flashinfer_bf16_fp4() -> bool: + """Return ``True`` if FlashInfer's CuTe-DSL W4A16 GEMM is available.""" + if not has_flashinfer_cutedsl(): + return False + mod = _get_submodule("flashinfer.gemm") + return mod is not None and all( + hasattr(mod, name) for name in ("mm_bf16_fp4", "prepare_bf16_fp4_weights") + ) + + @functools.cache def has_flashinfer_trtllm_fused_moe() -> bool: """Return `True` if FlashInfer TRTLLM fused MoE is available.""" @@ -706,6 +720,39 @@ if has_flashinfer(): ) -> torch.Tensor: return torch.empty(A.shape[0], B.shape[1], dtype=dtype, device=A.device) + @torch.library.custom_op( + "vllm::flashinfer_mm_bf16_fp4", + mutates_args=[], + device_types="cuda", + ) + def flashinfer_mm_bf16_fp4( + A: torch.Tensor, + B: torch.Tensor, + B_scale: torch.Tensor, + global_scale: torch.Tensor, + ) -> torch.Tensor: + from flashinfer.gemm import mm_bf16_fp4 + + return mm_bf16_fp4( + A, + B, + B_scale, + global_scale, + backend="cute-dsl", + ) + + @torch.library.register_fake( + "vllm::flashinfer_mm_bf16_fp4", + ) + def flashinfer_mm_bf16_fp4_fake( + A: torch.Tensor, + B: torch.Tensor, + B_scale: torch.Tensor, + global_scale: torch.Tensor, + ) -> torch.Tensor: + output_size = B.shape[0] if B.dtype == torch.uint8 else B.shape[1] // 2 + return torch.empty(A.shape[0], output_size, dtype=A.dtype, device=A.device) + @torch.library.custom_op( "vllm::flashinfer_mxfp4_quantize", mutates_args=[], @@ -1158,6 +1205,7 @@ __all__ = [ "flashinfer_trtllm_fp8_block_scale_moe", "flashinfer_cutlass_fused_moe", "flashinfer_cutedsl_grouped_gemm_nt_masked", + "flashinfer_prepare_bf16_fp4_weights", "flashinfer_fp4_quantize", "silu_and_mul_scaled_nvfp4_experts_quantize", "scaled_fp4_grouped_quantize", @@ -1177,6 +1225,7 @@ __all__ = [ "has_flashinfer_cutlass_fused_moe", "has_flashinfer_cutedsl_grouped_gemm_nt_masked", "has_flashinfer_cutedsl_moe_nvfp4", + "has_flashinfer_bf16_fp4", "has_flashinfer_b12x_moe", "has_flashinfer_b12x_gemm", "has_flashinfer_fp8_blockscale_gemm",