From 858d7ceb25ef99bb17dc88da9ed64f3b864f2109 Mon Sep 17 00:00:00 2001 From: Yicheng Jin Date: Mon, 20 Jul 2026 22:39:16 -0700 Subject: [PATCH 1/3] [Kernel] Add FlashInfer bf16 x nvfp4 W4A16 GEMM kernel for NVFP4 weight-only linear layers Wire FlashInfer's mm_bf16_fp4 (flashinfer-ai/flashinfer#3597, in FlashInfer v0.6.14, vLLM's current pin) into the NVFP4 linear kernel registry as FlashInferW4A16NvFp4LinearKernel. Like Marlin, the weight is repacked once at load time via prepare_bf16_fp4_weights; at run time the kernel dequantizes the FP4 weight and runs the matrix multiply in bf16, so activations are never quantized. The kernel is the default on SM121 (DGX Spark), where it is tuned, and opt-in via --linear-backend flashinfer_cutedsl on SM100/SM110/SM120. Marlin remains the default elsewhere and for fp16-activation models. The ModelOpt W4A16 method routes through init_nvfp4_linear_kernel(use_a16=True) so both it and the compressed-tensors W4A16 scheme share the selection. The GEMM is exposed as torch custom op vllm::flashinfer_mm_bf16_fp4 so it works under torch.compile and CUDA graph capture, and the existing FlashInfer autotune warmup tunes it during engine startup. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 --- .../test_flashinfer_w4a16_nvfp4.py | 200 ++++++++++++++++++ vllm/config/kernel.py | 3 +- .../model_executor/kernels/linear/__init__.py | 70 +++++- .../kernels/linear/nvfp4/flashinfer.py | 120 +++++++++++ .../layers/quantization/modelopt.py | 3 +- vllm/utils/flashinfer.py | 59 ++++++ 6 files changed, 449 insertions(+), 6 deletions(-) create mode 100644 tests/kernels/quantization/test_flashinfer_w4a16_nvfp4.py 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..5dfc7eab8910 --- /dev/null +++ b/tests/kernels/quantization/test_flashinfer_w4a16_nvfp4.py @@ -0,0 +1,200 @@ +# 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: + # Elementwise tolerances do not fit a bf16-output GEMM, whose rounding + # error grows with K, so compare relative L2 error and cosine + # similarity against the dequantized-weight reference. + 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..2f9f4c2e5298 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,6 +1004,63 @@ 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. + + W4A16 layers can never use the W4A4 priority list, because those + kernels quantize activations and a weight-only checkpoint has no input + scale. The choice is between Marlin and FlashInfer's bf16 x fp4 GEMM: + the FlashInfer kernel is the default on SM121 (DGX Spark), where it is + tuned, and opt-in via --linear-backend flashinfer_cutedsl elsewhere. + Explicit --linear-backend values that only contain W4A4 kernels are + rejected. The caller validates the returned kernel with is_supported(). + """ + fi_kernel = FlashInferW4A16NvFp4LinearKernel + marlin_kernel = MarlinNvFp4LinearKernel + + 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): + 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 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 + ok, reason = usable(fi_kernel) + if ok: + return fi_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.""" @@ -1013,7 +1072,12 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: # 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: + # Weight-only layers must never reach the W4A4 kernels (including + # the batch-invariant CUTLASS path below): those quantize + # activations, and a W4A16 checkpoint has no input scale. + 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 +1106,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 +1288,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..605f01d84b23 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,123 @@ def apply_weights( return out.view(*output_shape) +class FlashInferW4A16NvFp4LinearKernel(NvFp4LinearKernel): + """Weight-only NVFP4 GEMM (bf16 activations) via FlashInfer's mm_bf16_fp4. + + Unlike the other FlashInfer kernels in this file, activations are not + quantized: the kernel dequantizes the FP4 weight and runs the matrix + multiply in bf16. + """ + + # FlashInfer repacks the weight in 64-row tiles of N; K is padded to + # the same granularity to satisfy the kernel's K tiling. + 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 only accepts bfloat16 activations; fp16 models must + # fall back to Marlin. Without an active vLLM config the dtype is + # unknown here; process_weights_after_loading checks the layer's + # params_dtype as a backstop. + 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 + + # Zero-pad N and K to the repack tile sizes. The padded rows and + # columns have zero block scales, so they dequantize to zero; + # apply_weights pads the activation K to match and slices the + # extra output columns off. + 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..19ad1e3feefe 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,34 @@ 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: + # The prepared weight is (K // 16, N * 2) int32, so N is + # B.shape[1] // 2. + 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 +926,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 +1127,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", From 5e822ad889274d64ae2c5833c156d27c0972e50f Mon Sep 17 00:00:00 2001 From: Yicheng Jin Date: Wed, 22 Jul 2026 10:49:38 -0700 Subject: [PATCH 2/3] Trim comments to key design notes for an experienced-reviewer audience Comment-only change. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 --- .../test_flashinfer_w4a16_nvfp4.py | 5 ++-- .../model_executor/kernels/linear/__init__.py | 16 +++++-------- .../kernels/linear/nvfp4/flashinfer.py | 23 ++++++------------- vllm/utils/flashinfer.py | 3 +-- 4 files changed, 16 insertions(+), 31 deletions(-) diff --git a/tests/kernels/quantization/test_flashinfer_w4a16_nvfp4.py b/tests/kernels/quantization/test_flashinfer_w4a16_nvfp4.py index 5dfc7eab8910..65af34302a6f 100644 --- a/tests/kernels/quantization/test_flashinfer_w4a16_nvfp4.py +++ b/tests/kernels/quantization/test_flashinfer_w4a16_nvfp4.py @@ -47,9 +47,8 @@ def _assert_close_to_reference(out: torch.Tensor, ref: torch.Tensor) -> None: - # Elementwise tolerances do not fit a bf16-output GEMM, whose rounding - # error grows with K, so compare relative L2 error and cosine - # similarity against the dequantized-weight reference. + # 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() diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 2f9f4c2e5298..994827579e7d 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -1009,13 +1009,10 @@ def _select_nvfp4_a16_kernel( ) -> type[NvFp4LinearKernel]: """Pick the kernel for weight-only (W4A16) NVFP4 layers. - W4A16 layers can never use the W4A4 priority list, because those - kernels quantize activations and a weight-only checkpoint has no input - scale. The choice is between Marlin and FlashInfer's bf16 x fp4 GEMM: - the FlashInfer kernel is the default on SM121 (DGX Spark), where it is - tuned, and opt-in via --linear-backend flashinfer_cutedsl elsewhere. - Explicit --linear-backend values that only contain W4A4 kernels are - rejected. The caller validates the returned kernel with is_supported(). + Only Marlin 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 @@ -1073,9 +1070,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: force_kernel: type[NvFp4LinearKernel] | None = None linear_backend = _get_linear_backend() if use_a16: - # Weight-only layers must never reach the W4A4 kernels (including - # the batch-invariant CUTLASS path below): those quantize - # activations, and a W4A16 checkpoint has no input scale. + # 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() diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index 605f01d84b23..0b011ce14934 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -309,15 +309,10 @@ def apply_weights( class FlashInferW4A16NvFp4LinearKernel(NvFp4LinearKernel): - """Weight-only NVFP4 GEMM (bf16 activations) via FlashInfer's mm_bf16_fp4. + """Weight-only NVFP4 GEMM via FlashInfer's mm_bf16_fp4: activations + stay in bf16 and the kernel dequantizes the FP4 weight.""" - Unlike the other FlashInfer kernels in this file, activations are not - quantized: the kernel dequantizes the FP4 weight and runs the matrix - multiply in bf16. - """ - - # FlashInfer repacks the weight in 64-row tiles of N; K is padded to - # the same granularity to satisfy the kernel's K tiling. + # Weights are zero-padded to the repack and kernel tile granularity. N_ALIGN = 64 K_ALIGN = 64 @@ -337,10 +332,8 @@ def is_supported( @classmethod def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: - # mm_bf16_fp4 only accepts bfloat16 activations; fp16 models must - # fall back to Marlin. Without an active vLLM config the dtype is - # unknown here; process_weights_after_loading checks the layer's - # params_dtype as a backstop. + # 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() @@ -368,10 +361,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: n = weight.shape[0] k = weight.shape[1] * 2 - # Zero-pad N and K to the repack tile sizes. The padded rows and - # columns have zero block scales, so they dequantize to zero; - # apply_weights pads the activation K to match and slices the - # extra output columns off. + # 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: diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 19ad1e3feefe..969d336f6c5e 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -672,8 +672,7 @@ def flashinfer_mm_bf16_fp4_fake( B_scale: torch.Tensor, alpha: torch.Tensor, ) -> torch.Tensor: - # The prepared weight is (K // 16, N * 2) int32, so N is - # B.shape[1] // 2. + # 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( From 26a4951f2a0daab95a4a978fcbedd5ce01dac4cb Mon Sep 17 00:00:00 2001 From: Yicheng Jin Date: Thu, 20 Aug 2026 01:00:24 +0000 Subject: [PATCH 3/3] Keep Humming usable for W4A16 and register the FlashInfer kernel in a16_kernels Upstream added HummingNvFp4LinearKernel as a W4A16-capable kernel after this PR was written; fold it into _select_nvfp4_a16_kernel so an explicit --linear-backend=humming still serves W4A16 layers, and add the FlashInfer kernel to a16_kernels so the force-kernel guard accepts it. Co-Authored-By: Claude Fable 5 --- .../model_executor/kernels/linear/__init__.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 994827579e7d..2423925765f5 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -1009,13 +1009,14 @@ def _select_nvfp4_a16_kernel( ) -> type[NvFp4LinearKernel]: """Pick the kernel for weight-only (W4A16) NVFP4 layers. - Only Marlin and FlashInfer's bf16 x fp4 GEMM apply: the W4A4 kernels - quantize activations, and a W4A16 checkpoint has no input scale. + 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: @@ -1027,7 +1028,7 @@ def usable(kernel: type[NvFp4LinearKernel]) -> tuple[bool, str | None]: if linear_backend != "auto": requested = _LINEAR_BACKEND_KERNEL_MAP.get(linear_backend, set()) - for kernel in (fi_kernel, marlin_kernel): + for kernel in (fi_kernel, marlin_kernel, humming_kernel): if kernel not in requested: continue ok, reason = usable(kernel) @@ -1039,7 +1040,7 @@ def usable(kernel: type[NvFp4LinearKernel]) -> tuple[bool, str | None]: return kernel raise ValueError( f"--linear-backend={linear_backend} has no kernel for W4A16 " - f"NVFP4 layers; use marlin or flashinfer_cutedsl." + f"NVFP4 layers; use marlin, humming, or flashinfer_cutedsl." ) # Auto selection. VLLM_BATCH_INVARIANT keeps the deterministic Marlin @@ -1052,9 +1053,10 @@ def usable(kernel: type[NvFp4LinearKernel]) -> tuple[bool, str | None]: return fi_kernel if marlin_kernel.__name__ not in envs.VLLM_DISABLED_KERNELS: return marlin_kernel - ok, reason = usable(fi_kernel) - if ok: - return fi_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}") @@ -1062,7 +1064,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 = ( + MarlinNvFp4LinearKernel, + HummingNvFp4LinearKernel, + FlashInferW4A16NvFp4LinearKernel, + ) # VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the # batch-invariant CUTLASS implementation when available, otherwise fall