[Perf][Quant][NVFP4] Prefer W4A4 linear kernels over weight-only ones on SM120/121 (#55170)

Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>

tests/kernels/quantization/test_nvfp4_kernel_selection.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for NVFP4 linear kernel selection order (CPU-only)."""
+
+import pytest
+
+from vllm.model_executor.kernels.linear import (
+    _POSSIBLE_NVFP4_KERNELS,
+    CutlassNvFp4LinearKernel,
+    FlashInferB12xNvFp4LinearKernel,
+    FlashInferCuteDslNvFp4W4A16LinearKernel,
+    FlashInferCutlassNvFp4LinearKernel,
+)
+from vllm.platforms.interface import PlatformEnum
+
+# W4A4 kernels that run on SM120/121, where the head of the list is gated to
+# sm_10x and selection falls through to whatever follows.
+W4A4_KERNELS_ON_SM12X = (
+    FlashInferCutlassNvFp4LinearKernel,
+    FlashInferB12xNvFp4LinearKernel,
+    CutlassNvFp4LinearKernel,
+)
+
+
+@pytest.mark.parametrize("w4a4_kernel", W4A4_KERNELS_ON_SM12X)
+def test_w4a16_kernel_does_not_precede_w4a4_kernels(w4a4_kernel):
+    candidates = _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA]
+    w4a16_index = candidates.index(FlashInferCuteDslNvFp4W4A16LinearKernel)
+    assert candidates.index(w4a4_kernel) < w4a16_index, (
+        f"{w4a4_kernel.__name__} must be preferred over "
+        f"{FlashInferCuteDslNvFp4W4A16LinearKernel.__name__}"
+    )

vllm/model_executor/kernels/linear/__init__.py

@@ -533,10 +533,11 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
 _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
     PlatformEnum.CUDA: [
         FlashInferCuteDslNvFp4LinearKernel,
-        FlashInferCuteDslNvFp4W4A16LinearKernel,
         FlashInferCutlassNvFp4LinearKernel,
         FlashInferB12xNvFp4LinearKernel,
         CutlassNvFp4LinearKernel,
+        # Weight-only: a fallback on a W4A4-capable checkpoint, not a preference.
+        FlashInferCuteDslNvFp4W4A16LinearKernel,
         MarlinNvFp4LinearKernel,
         FlashInferTrtllmNvFp4LinearKernel,
         FlashInferCudnnNvFp4LinearKernel,