| 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 | ) |