diff --git a/tests/kernels/quantization/test_flashinfer_w4a16_nvfp4.py b/tests/kernels/quantization/test_flashinfer_w4a16_nvfp4.py new file mode 100644 index 000000000000..65af34302a6f --- /dev/null +++ b/tests/kernels/quantization/test_flashinfer_w4a16_nvfp4.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the FlashInfer W4A16 NVFP4 linear kernel (mm_bf16_fp4).""" + +from types import SimpleNamespace + +import pytest +import torch +from nvfp4_utils import ( + break_fp4_bytes, + convert_swizzled_to_linear, + quant_nvfp4_tensor, +) + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.model_executor.kernels.linear import ( + FlashInferW4A16NvFp4LinearKernel, + MarlinNvFp4LinearKernel, + NvFp4LinearLayerConfig, + init_nvfp4_linear_kernel, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +if not FlashInferW4A16NvFp4LinearKernel.is_supported()[0]: + pytest.skip( + reason="FlashInfer W4A16 NVFP4 requires SM100/SM110/SM12x and " + "FlashInfer with mm_bf16_fp4.", + allow_module_level=True, + ) + +# m, n, k +SHAPES = [ + (1, 128, 256), + (16, 256, 512), + (128, 512, 2048), + (4, 2048, 512), + # n not a multiple of the 64-wide repack tile (exercises N padding). + (8, 200, 512), + # k not a multiple of the 64-column tile (exercises K padding). + (8, 64, 1120), +] + +SEEDS = [42] +DEVICE = "cuda:0" +BLOCK_SIZE = 16 + + +def _assert_close_to_reference(out: torch.Tensor, ref: torch.Tensor) -> None: + # bf16-output rounding error grows with K, so compare relative L2 and + # cosine similarity instead of elementwise tolerances. + out_f = out.float().flatten() + ref_f = ref.float().flatten() + rel_l2 = (torch.linalg.vector_norm(out_f - ref_f) / ref_f.norm()).item() + cos = torch.nn.functional.cosine_similarity(out_f, ref_f, dim=0).item() + assert rel_l2 < 2e-2, f"relative L2 error {rel_l2:.4f} exceeds 2e-2" + assert cos > 0.999, f"cosine similarity {cos:.6f} below 0.999" + + +def _make_w4a16_layer(n: int, k: int, device: str): + """Build a layer holding checkpoint-format NVFP4 weights plus the + dequantized reference weight.""" + w = torch.randn((n, k), dtype=torch.bfloat16, device=device) + w_fp4, w_sf_swizzled, w_global_scale_encode = quant_nvfp4_tensor(w) + + # Checkpoints store linear (N, K // 16) fp8 block scales. The + # unswizzle helper keeps its column padding, so trim to K // 16. + sf_linear = ( + convert_swizzled_to_linear(w_sf_swizzled, n, k, BLOCK_SIZE)[ + :, : k // BLOCK_SIZE + ] + .view(torch.float8_e4m3fn) + .contiguous() + ) + + w_vals = break_fp4_bytes(w_fp4, torch.float32).reshape( + n, k // BLOCK_SIZE, BLOCK_SIZE + ) + sf = sf_linear.to(torch.float32) / w_global_scale_encode + w_ref = (w_vals * sf.unsqueeze(-1)).reshape(n, k).to(torch.bfloat16) + + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(w_fp4, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(sf_linear, requires_grad=False) + # Both quant methods store the global scale in its dequant form. + layer.weight_global_scale = torch.nn.Parameter( + 1.0 / w_global_scale_encode, requires_grad=False + ) + layer.input_size_per_partition = k + layer.output_size_per_partition = n + return layer, w_ref + + +@pytest.mark.parametrize("shape", SHAPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("use_bias", [False, True]) +@pytest.mark.parametrize("autotune", [False, True]) +@torch.inference_mode() +def test_flashinfer_w4a16_nvfp4_gemm( + shape: tuple[int, int, int], + seed: int, + use_bias: bool, + autotune: bool, +) -> None: + import flashinfer + + set_random_seed(seed) + m, n, k = shape + + layer, w_ref = _make_w4a16_layer(n, k, DEVICE) + kernel = FlashInferW4A16NvFp4LinearKernel(NvFp4LinearLayerConfig()) + kernel.process_weights_after_loading(layer) + + x = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + bias = torch.randn((n,), dtype=torch.bfloat16, device=DEVICE) if use_bias else None + + with flashinfer.autotune(autotune): + out = kernel.apply_weights(layer, x, bias) + + expected = x @ w_ref.t() + if use_bias: + expected = expected + bias + assert out.shape == (m, n) + assert out.dtype == torch.bfloat16 + _assert_close_to_reference(out, expected) + + +@torch.inference_mode() +def test_flashinfer_w4a16_nvfp4_gemm_3d_input() -> None: + set_random_seed(0) + n, k = 256, 512 + layer, w_ref = _make_w4a16_layer(n, k, DEVICE) + kernel = FlashInferW4A16NvFp4LinearKernel(NvFp4LinearLayerConfig()) + kernel.process_weights_after_loading(layer) + + x = torch.randn((2, 8, k), dtype=torch.bfloat16, device=DEVICE) + out = kernel.apply_weights(layer, x) + + assert out.shape == (2, 8, n) + _assert_close_to_reference(out, x @ w_ref.t()) + + +def test_w4a16_kernel_selection_auto(): + kernel = init_nvfp4_linear_kernel(use_a16=True) + if current_platform.is_device_capability(121): + assert isinstance(kernel, FlashInferW4A16NvFp4LinearKernel) + else: + assert isinstance(kernel, MarlinNvFp4LinearKernel) + + +def test_w4a16_kernel_selection_flashinfer_cutedsl(): + vllm_config = VllmConfig() + vllm_config.kernel_config.linear_backend = "flashinfer_cutedsl" + with set_current_vllm_config(vllm_config): + kernel = init_nvfp4_linear_kernel(use_a16=True) + assert isinstance(kernel, FlashInferW4A16NvFp4LinearKernel) + + +def test_w4a16_kernel_selection_marlin_backend(): + vllm_config = VllmConfig() + vllm_config.kernel_config.linear_backend = "marlin" + with set_current_vllm_config(vllm_config): + kernel = init_nvfp4_linear_kernel(use_a16=True) + assert isinstance(kernel, MarlinNvFp4LinearKernel) + + +def test_w4a16_kernel_selection_batch_invariant(monkeypatch): + """Batch-invariant mode must not route W4A16 onto a W4A4 kernel.""" + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True, raising=False) + kernel = init_nvfp4_linear_kernel(use_a16=True) + assert isinstance(kernel, MarlinNvFp4LinearKernel) + + +def test_w4a16_kernel_selection_rejects_w4a4_backend(): + """Backends with only W4A4 kernels cannot serve W4A16 layers.""" + vllm_config = VllmConfig() + vllm_config.kernel_config.linear_backend = "cutlass" + with ( + set_current_vllm_config(vllm_config), + pytest.raises(ValueError, match="W4A16"), + ): + init_nvfp4_linear_kernel(use_a16=True) + + +def test_w4a16_kernel_requires_bf16(monkeypatch): + """fp16 models must not select the FlashInfer W4A16 kernel.""" + import vllm.config as vllm_config_mod + + fake_config = SimpleNamespace(model_config=SimpleNamespace(dtype=torch.float16)) + monkeypatch.setattr( + vllm_config_mod, "get_current_vllm_config_or_none", lambda: fake_config + ) + ok, reason = FlashInferW4A16NvFp4LinearKernel.can_implement( + NvFp4LinearLayerConfig() + ) + assert not ok + assert "bfloat16" in reason diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 7287c1ce923f..29a05265f092 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -220,7 +220,8 @@ class KernelConfig: - "auto": Automatically select the best backend based on model and hardware - "cutlass": Use CUTLASS-based kernels - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels - - "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels (NVFP4, MXFP8) + - "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels (NVFP4, + MXFP8; for weight-only NVFP4 this selects the bf16 x fp4 W4A16 GEMM) - "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 b94c8f7d58d0..2423925765f5 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -150,6 +150,7 @@ FlashInferCuteDslNvFp4LinearKernel, FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, + FlashInferW4A16NvFp4LinearKernel, ) from vllm.model_executor.kernels.linear.nvfp4.humming import ( HummingNvFp4LinearKernel, @@ -265,6 +266,7 @@ def _get_linear_backend() -> str: "flashinfer_cutedsl": { FlashInferCuteDslNvFp4LinearKernel, FlashInferCutedslMxfp8LinearKernel, + FlashInferW4A16NvFp4LinearKernel, }, "flashinfer_trtllm": { FlashInferTrtllmNvFp4LinearKernel, @@ -1002,18 +1004,82 @@ def init_wfp8_a16_linear_kernel( ) +def _select_nvfp4_a16_kernel( + linear_backend: str, config: NvFp4LinearLayerConfig +) -> type[NvFp4LinearKernel]: + """Pick the kernel for weight-only (W4A16) NVFP4 layers. + + Only Marlin, Humming, and FlashInfer's bf16 x fp4 GEMM apply: the W4A4 + kernels quantize activations, and a W4A16 checkpoint has no input scale. + FlashInfer is the default on SM121, where it is tuned, and opt-in via + --linear-backend flashinfer_cutedsl elsewhere. + """ + fi_kernel = FlashInferW4A16NvFp4LinearKernel + marlin_kernel = MarlinNvFp4LinearKernel + humming_kernel = HummingNvFp4LinearKernel + + def usable(kernel: type[NvFp4LinearKernel]) -> tuple[bool, str | None]: + if kernel.__name__ in envs.VLLM_DISABLED_KERNELS: + return False, f"{kernel.__name__} is disabled via VLLM_DISABLED_KERNELS" + is_supported, reason = kernel.is_supported() + if not is_supported: + return False, reason + return kernel.can_implement(config) + + if linear_backend != "auto": + requested = _LINEAR_BACKEND_KERNEL_MAP.get(linear_backend, set()) + for kernel in (fi_kernel, marlin_kernel, humming_kernel): + if kernel not in requested: + continue + ok, reason = usable(kernel) + if not ok: + raise ValueError( + f"--linear-backend={linear_backend} cannot serve W4A16 " + f"NVFP4 layers: {reason}" + ) + return kernel + raise ValueError( + f"--linear-backend={linear_backend} has no kernel for W4A16 " + f"NVFP4 layers; use marlin, humming, or flashinfer_cutedsl." + ) + + # Auto selection. VLLM_BATCH_INVARIANT keeps the deterministic Marlin + # path; otherwise prefer FlashInfer where its kernel is tuned. + if ( + not envs.VLLM_BATCH_INVARIANT + and current_platform.is_device_capability(121) + and usable(fi_kernel)[0] + ): + return fi_kernel + if marlin_kernel.__name__ not in envs.VLLM_DISABLED_KERNELS: + return marlin_kernel + for kernel in (fi_kernel, humming_kernel): + ok, reason = usable(kernel) + if ok: + return kernel + raise ValueError(f"No W4A16 NVFP4 kernel is available: {reason}") + + 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 = ( + MarlinNvFp4LinearKernel, + HummingNvFp4LinearKernel, + FlashInferW4A16NvFp4LinearKernel, + ) # VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the # batch-invariant CUTLASS implementation when available, otherwise fall # back to emulation. It overrides --linear-backend. force_kernel: type[NvFp4LinearKernel] | None = None linear_backend = _get_linear_backend() - if envs.VLLM_BATCH_INVARIANT: + if use_a16: + # W4A16 has no input scale, so it must bypass the W4A4 paths + # below, including the batch-invariant CUTLASS branch. + force_kernel = _select_nvfp4_a16_kernel(linear_backend, config) + elif envs.VLLM_BATCH_INVARIANT: bi_supported, reason = CutlassNvFp4LinearKernel.is_supported() if bi_supported: if linear_backend not in ("auto", "cutlass"): @@ -1042,9 +1108,6 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: reason, ) force_kernel = EmulationNvFp4LinearKernel - elif linear_backend == "auto" and use_a16: - # Force a16 (Marlin) when running weight-only quantization. - force_kernel = MarlinNvFp4LinearKernel if force_kernel is not None: if use_a16 and force_kernel not in a16_kernels: @@ -1227,6 +1290,7 @@ def register_linear_kernel( "FlashInferCutlassNvFp4LinearKernel", "FlashInferTrtllmNvFp4LinearKernel", "FlashInferCudnnNvFp4LinearKernel", + "FlashInferW4A16NvFp4LinearKernel", "MarlinNvFp4LinearKernel", "_KernelT", "DeepGemmFp8BlockScaledMMKernel", diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index b721a135e267..0b011ce14934 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -20,10 +20,13 @@ ) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( + flashinfer_scaled_bf16_fp4_mm, flashinfer_scaled_fp4_mm, has_flashinfer, has_flashinfer_b12x_gemm, + has_flashinfer_bf16_fp4_gemm, ) +from vllm.utils.math_utils import round_up from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig @@ -305,6 +308,114 @@ def apply_weights( return out.view(*output_shape) +class FlashInferW4A16NvFp4LinearKernel(NvFp4LinearKernel): + """Weight-only NVFP4 GEMM via FlashInfer's mm_bf16_fp4: activations + stay in bf16 and the kernel dequantizes the FP4 weight.""" + + # Weights are zero-padded to the repack and kernel tile granularity. + N_ALIGN = 64 + K_ALIGN = 64 + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not ( + current_platform.is_device_capability_family(100) + or current_platform.is_device_capability_family(110) + or current_platform.is_device_capability_family(120) + ): + return False, "FlashInfer W4A16 NVFP4 requires SM100/SM110/SM12x" + if not has_flashinfer_bf16_fp4_gemm(): + return False, "FlashInfer with mm_bf16_fp4 (>= 0.6.14) required" + return True, None + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + # mm_bf16_fp4 is bf16-only. Without an active vllm config the dtype + # is unknown here; process_weights_after_loading re-checks it. + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + if ( + vllm_config is not None + and vllm_config.model_config is not None + and vllm_config.model_config.dtype != torch.bfloat16 + ): + return False, "FlashInfer W4A16 NVFP4 requires bfloat16 activations" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + from flashinfer import prepare_bf16_fp4_weights + + params_dtype = getattr(layer, "params_dtype", torch.bfloat16) + if params_dtype != torch.bfloat16: + raise ValueError( + f"FlashInfer W4A16 NVFP4 requires bfloat16 activations; " + f"this model runs {params_dtype}. Use --linear-backend " + f"marlin." + ) + + weight = layer.weight.data # (N, K // 2) packed uint8 + weight_scale = layer.weight_scale.data # (N, K // 16) fp8-e4m3 + n = weight.shape[0] + k = weight.shape[1] * 2 + + # Padding carries zero block scales, so it dequantizes to zero; + # apply_weights pads the activation K to match. + n_padded = round_up(n, self.N_ALIGN) + k_padded = round_up(k, self.K_ALIGN) + if n_padded != n or k_padded != k: + pad_rows = n_padded - n + pad_weight_cols = (k_padded - k) // 2 + pad_scale_cols = (k_padded - k) // 16 + weight = torch.nn.functional.pad(weight, (0, pad_weight_cols, 0, pad_rows)) + weight_scale = torch.nn.functional.pad( + weight_scale.view(torch.uint8), (0, pad_scale_cols, 0, pad_rows) + ).view(torch.float8_e4m3fn) + layer.w4a16_k_pad = k_padded - k + + alpha = layer.weight_global_scale.data.to( + device=weight.device, dtype=torch.float32 + ).reshape(1) + b_prepared, sf_prepared, alpha_prepared = prepare_bf16_fp4_weights( + weight, + swizzle_blockscale(weight_scale), + alpha, + backend="cute-dsl", + ) + layer.weight = torch.nn.Parameter(b_prepared, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(sf_prepared, requires_grad=False) + layer.alpha = torch.nn.Parameter(alpha_prepared, requires_grad=False) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + output_size = layer.output_size_per_partition + output_shape = [*x.shape[:-1], output_size] + + x2d = x.reshape(-1, x.shape[-1]) + k_pad = getattr(layer, "w4a16_k_pad", 0) + if k_pad: + x2d = torch.nn.functional.pad(x2d, (0, k_pad)) + + out = flashinfer_scaled_bf16_fp4_mm( + x2d, + layer.weight, + layer.weight_scale, + layer.alpha, + ) + + out = slice_nvfp4_output(out, output_size) + + if bias is not None: + out = out + bias + return out.view(*output_shape) + + class FlashInferB12xNvFp4LinearKernel(NvFp4LinearKernel): """NVFP4 GEMM via FlashInfer's b12x CuTe DSL warp-level MMA kernel (SM120+).""" diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e24c5656add3..9736d734751a 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1227,7 +1227,8 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase): """Linear method for ModelOpt NVFP4 W4A16. 4-bit NVFP4 weights, fp16/bf16 activations. Loads ModelOpt-style names - directly (no on-disk conversion) and dispatches to a W4A16 GEMM: + directly (no on-disk conversion) and dispatches to a W4A16 GEMM + (Marlin, or FlashInfer's bf16 x fp4 kernel where selected): weight uint8 packed NVFP4 (2 nibbles/byte along input dim) weight_scale fp8-e4m3 per 16-elem group along input dim diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 4f7d2eecd780..969d336f6c5e 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -345,6 +345,20 @@ def has_flashinfer_b12x_gemm() -> bool: ) +@functools.cache +def has_flashinfer_bf16_fp4_gemm() -> bool: + """Return True if the installed FlashInfer provides the bf16 x nvfp4 + weight-only GEMM (mm_bf16_fp4 and prepare_bf16_fp4_weights).""" + if not has_flashinfer_cutedsl(): + return False + mod = _get_submodule("flashinfer") + return ( + mod is not None + and hasattr(mod, "mm_bf16_fp4") + and hasattr(mod, "prepare_bf16_fp4_weights") + ) + + @functools.cache def has_flashinfer_b12x_moe() -> bool: """Return ``True`` if FlashInfer CuteDSL SM12x fused MoE is available.""" @@ -634,6 +648,33 @@ def flashinfer_mm_fp4_fake( ) -> 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, + alpha: torch.Tensor, + ) -> torch.Tensor: + from flashinfer import mm_bf16_fp4 as flashinfer_mm_bf16_fp4_ + + return flashinfer_mm_bf16_fp4_(A, B, B_scale, alpha, 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, + alpha: torch.Tensor, + ) -> torch.Tensor: + # Prepared weight is (K // 16, N * 2) int32. + return torch.empty(A.shape[0], B.shape[1] // 2, dtype=A.dtype, device=A.device) + @torch.library.custom_op( "vllm::flashinfer_mxfp4_quantize", mutates_args=[], @@ -884,6 +925,21 @@ def flashinfer_scaled_fp4_mm_out( return out +def flashinfer_scaled_bf16_fp4_mm( + a: torch.Tensor, + b: torch.Tensor, + b_descale: torch.Tensor, + alpha: torch.Tensor, +) -> torch.Tensor: + """W4A16 GEMM: bf16 activations x prepared nvfp4 weight. + + ``b``, ``b_descale``, and ``alpha`` must come from + ``flashinfer.prepare_bf16_fp4_weights(..., backend="cute-dsl")``. + """ + assert a.ndim == 2 and a.dtype == torch.bfloat16 + return flashinfer_mm_bf16_fp4(a, b, b_descale, alpha) + + def flashinfer_scaled_fp8_mm( a: torch.Tensor, b: torch.Tensor, @@ -1070,12 +1126,14 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: "has_flashinfer_cutedsl_moe_nvfp4", "has_flashinfer_b12x_moe", "has_flashinfer_b12x_gemm", + "has_flashinfer_bf16_fp4_gemm", "has_flashinfer_fp8_blockscale_gemm", "has_nvidia_artifactory", "supports_trtllm_attention", "can_use_trtllm_attention", "use_trtllm_attention", "flashinfer_mxfp4_quantize", + "flashinfer_scaled_bf16_fp4_mm", "flashinfer_scaled_fp4_mm", "flashinfer_scaled_fp4_mm_out", "flashinfer_scaled_fp8_mm",