[Perf][Quant][NVFP4] Prefer W4A4 linear kernels over weight-only ones on SM120/121 (#55170)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
Stefano Castagnetta committed
Sep 8, 2026 at 13:06 UTC
13cf9e05c1eda0bfe5cbfb9344343ca2737d0723
2 files changed
+34
-1
tests/kernels/quantization/test_nvfp4_kernel_selection.py
new
+32
@@ -0,0 +1,32 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for NVFP4 linear kernel selection order (CPU-only)."""
4
+
5
+import pytest
6
+
7
+from vllm.model_executor.kernels.linear import (
8
+ _POSSIBLE_NVFP4_KERNELS,
9
+ CutlassNvFp4LinearKernel,
10
+ FlashInferB12xNvFp4LinearKernel,
11
+ FlashInferCuteDslNvFp4W4A16LinearKernel,
12
+ FlashInferCutlassNvFp4LinearKernel,
13
+)
14
+from vllm.platforms.interface import PlatformEnum
15
+
16
+# W4A4 kernels that run on SM120/121, where the head of the list is gated to
17
+# sm_10x and selection falls through to whatever follows.
18
+W4A4_KERNELS_ON_SM12X = (
19
+ FlashInferCutlassNvFp4LinearKernel,
20
+ FlashInferB12xNvFp4LinearKernel,
21
+ CutlassNvFp4LinearKernel,
22
+)
23
+
24
+
25
+@pytest.mark.parametrize("w4a4_kernel", W4A4_KERNELS_ON_SM12X)
26
+def test_w4a16_kernel_does_not_precede_w4a4_kernels(w4a4_kernel):
27
+ candidates = _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA]
28
+ w4a16_index = candidates.index(FlashInferCuteDslNvFp4W4A16LinearKernel)
29
+ assert candidates.index(w4a4_kernel) < w4a16_index, (
30
+ f"{w4a4_kernel.__name__} must be preferred over "
31
+ f"{FlashInferCuteDslNvFp4W4A16LinearKernel.__name__}"
32
+ )
vllm/model_executor/kernels/linear/__init__.py
+2
-1
@@ -533,10 +533,11 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
533
_POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
534
PlatformEnum.CUDA: [
535
FlashInferCuteDslNvFp4LinearKernel,
536
- FlashInferCuteDslNvFp4W4A16LinearKernel,
536
FlashInferCutlassNvFp4LinearKernel,
537
FlashInferB12xNvFp4LinearKernel,
538
CutlassNvFp4LinearKernel,
539
+ # Weight-only: a fallback on a W4A4-capable checkpoint, not a preference.
540
+ FlashInferCuteDslNvFp4W4A16LinearKernel,
541
MarlinNvFp4LinearKernel,
542
FlashInferTrtllmNvFp4LinearKernel,
543
FlashInferCudnnNvFp4LinearKernel,