| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for ScaledMM kernel selection logic (CPU-only). |
| 4 | |
| 5 | Run `pytest tests/kernels/quantization/test_scaled_mm_kernel_selection.py`. |
| 6 | """ |
| 7 | |
| 8 | import inspect |
| 9 | from abc import ABC |
| 10 | from unittest.mock import patch |
| 11 | |
| 12 | import pytest |
| 13 | import torch |
| 14 | |
| 15 | from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config |
| 16 | from vllm.model_executor.kernels.linear import ( |
| 17 | AiterInt8ScaledMMLinearKernel, |
| 18 | CPUInt8ScaledMMLinearKernel, |
| 19 | HummingFP8ScaledMMLinearKernel, |
| 20 | Int8ScaledMMLinearKernel, |
| 21 | Int8ScaledMMLinearLayerConfig, |
| 22 | ScaledMMLinearKernel, |
| 23 | _get_linear_backend, |
| 24 | _resolve_backend_kernels, |
| 25 | init_fp8_linear_kernel, |
| 26 | init_int8_linear_kernel, |
| 27 | register_linear_kernel, |
| 28 | ) |
| 29 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 30 | kFp8StaticTensorSym, |
| 31 | ) |
| 32 | from vllm.platforms import PlatformEnum |
| 33 | |
| 34 | pytestmark = pytest.mark.cpu_test |
| 35 | |
| 36 | |
| 37 | def test_auto_linear_backend_is_unchanged_without_overrides(): |
| 38 | kernels: list[type] = [object] |
| 39 | quantizations = ( |
| 40 | "fp8_block_w8a8", |
| 41 | "fp8_w8a8", |
| 42 | "int8_w8a8", |
| 43 | "mixed_precision", |
| 44 | "mxfp8", |
| 45 | "mxfp4", |
| 46 | "mxfp6", |
| 47 | "w8a16_fp8", |
| 48 | "nvfp4_w4a4", |
| 49 | "nvfp4_w4a16", |
| 50 | ) |
| 51 | config = VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) |
| 52 | |
| 53 | with set_current_vllm_config(config): |
| 54 | for quantization in quantizations: |
| 55 | assert _get_linear_backend(quantization=quantization) == "auto" |
| 56 | assert ( |
| 57 | _resolve_backend_kernels(kernels, "test", quantization=quantization) |
| 58 | is kernels |
| 59 | ) |
| 60 | |
| 61 | |
| 62 | def test_linear_backend_override_is_quantization_specific(): |
| 63 | config = VllmConfig( |
| 64 | kernel_config=KernelConfig( |
| 65 | linear_backend="cutlass", |
| 66 | linear_backend_per_quant={"nvfp4_w4a16": "humming"}, |
| 67 | ) |
| 68 | ) |
| 69 | |
| 70 | with set_current_vllm_config(config): |
| 71 | assert _get_linear_backend(quantization="nvfp4_w4a16") == "humming" |
| 72 | assert _get_linear_backend(quantization="nvfp4_w4a4") == "cutlass" |
| 73 | assert _get_linear_backend(quantization="fp8_w8a8") == "cutlass" |
| 74 | |
| 75 | |
| 76 | @patch.object(HummingFP8ScaledMMLinearKernel, "is_supported", return_value=(True, None)) |
| 77 | @patch("vllm.model_executor.kernels.linear.current_platform") |
| 78 | def test_fp8_linear_backend_override(platform_mock, _): |
| 79 | platform_mock._enum = PlatformEnum.CUDA |
| 80 | config = VllmConfig( |
| 81 | kernel_config=KernelConfig( |
| 82 | linear_backend="cutlass", |
| 83 | linear_backend_per_quant={"fp8_w8a8": "humming"}, |
| 84 | ) |
| 85 | ) |
| 86 | |
| 87 | with set_current_vllm_config(config): |
| 88 | kernel = init_fp8_linear_kernel( |
| 89 | activation_quant_key=kFp8StaticTensorSym, |
| 90 | weight_quant_key=kFp8StaticTensorSym, |
| 91 | input_dtype=torch.bfloat16, |
| 92 | out_dtype=torch.bfloat16, |
| 93 | weight_shape=(128, 128), |
| 94 | ) |
| 95 | |
| 96 | assert type(kernel) is HummingFP8ScaledMMLinearKernel |
| 97 | |
| 98 | |
| 99 | def test_is_supported_is_abstract(): |
| 100 | """Test that is_supported() is properly defined as abstract.""" |
| 101 | assert issubclass(ScaledMMLinearKernel, ABC) |
| 102 | assert hasattr(ScaledMMLinearKernel, "is_supported") |
| 103 | |
| 104 | |
| 105 | def test_cpu_kernel_implements_is_supported(): |
| 106 | """Test that CPUInt8ScaledMMLinearKernel implements is_supported() method.""" |
| 107 | assert hasattr(CPUInt8ScaledMMLinearKernel, "is_supported"), ( |
| 108 | "CPUInt8ScaledMMLinearKernel missing is_supported() method" |
| 109 | ) |
| 110 | # Verify it's a classmethod by checking if it can be called with the class |
| 111 | # and by checking the method type |
| 112 | assert inspect.ismethod( |
| 113 | CPUInt8ScaledMMLinearKernel.is_supported |
| 114 | ) or inspect.isfunction(CPUInt8ScaledMMLinearKernel.is_supported), ( |
| 115 | "CPUInt8ScaledMMLinearKernel.is_supported() should be a classmethod" |
| 116 | ) |
| 117 | # Verify it can be called as a classmethod |
| 118 | result, reason = CPUInt8ScaledMMLinearKernel.is_supported() |
| 119 | assert isinstance(result, bool), "is_supported() should return a bool" |
| 120 | assert reason is None or isinstance(reason, str), "reason should be str or None" |
| 121 | |
| 122 | |
| 123 | def test_aiter_kernel_implements_is_supported(): |
| 124 | """Test that AiterInt8ScaledMMLinearKernel implements is_supported() method.""" |
| 125 | assert hasattr(AiterInt8ScaledMMLinearKernel, "is_supported"), ( |
| 126 | "AiterInt8ScaledMMLinearKernel missing is_supported() method" |
| 127 | ) |
| 128 | # Verify it's a classmethod by checking if it can be called with the class |
| 129 | # and by checking the method type |
| 130 | assert inspect.ismethod( |
| 131 | AiterInt8ScaledMMLinearKernel.is_supported |
| 132 | ) or inspect.isfunction(AiterInt8ScaledMMLinearKernel.is_supported), ( |
| 133 | "AiterInt8ScaledMMLinearKernel.is_supported() should be a classmethod" |
| 134 | ) |
| 135 | # Verify it can be called as a classmethod |
| 136 | # (will return False on CPU, which is expected) |
| 137 | result, reason = AiterInt8ScaledMMLinearKernel.is_supported() |
| 138 | assert isinstance(result, bool), "is_supported() should return a bool" |
| 139 | assert reason is None or isinstance(reason, str), "reason should be str or None" |
| 140 | # On CPU, it should return False with a reason about requiring ROCm |
| 141 | # This validates the method works correctly even on non-ROCm platforms |
| 142 | |
| 143 | |
| 144 | def test_cpu_kernel_accepts_all_configs(): |
| 145 | """Test that CPUInt8ScaledMMLinearKernel accepts all config combinations.""" |
| 146 | configs = [ |
| 147 | Int8ScaledMMLinearLayerConfig( |
| 148 | is_channelwise=False, |
| 149 | is_static_input_scheme=True, |
| 150 | input_symmetric=True, |
| 151 | ), |
| 152 | Int8ScaledMMLinearLayerConfig( |
| 153 | is_channelwise=True, |
| 154 | is_static_input_scheme=False, |
| 155 | input_symmetric=False, |
| 156 | ), |
| 157 | ] |
| 158 | |
| 159 | for config in configs: |
| 160 | can_impl, reason = CPUInt8ScaledMMLinearKernel.can_implement(config) |
| 161 | assert can_impl, ( |
| 162 | f"CPUInt8ScaledMMLinearKernel should accept config {config}: {reason}" |
| 163 | ) |
| 164 | |
| 165 | |
| 166 | class OOTInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): |
| 167 | @classmethod |
| 168 | def is_supported( |
| 169 | cls, compute_capability: int | None = None |
| 170 | ) -> tuple[bool, str | None]: |
| 171 | return True, None |
| 172 | |
| 173 | @classmethod |
| 174 | def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 175 | return True, None |
| 176 | |
| 177 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 178 | pass |
| 179 | |
| 180 | def apply_weights( |
| 181 | self, |
| 182 | layer: torch.nn.Module, |
| 183 | x: torch.Tensor, |
| 184 | bias: torch.Tensor | None = None, |
| 185 | ) -> torch.Tensor: |
| 186 | pass |
| 187 | |
| 188 | |
| 189 | @patch("vllm.model_executor.kernels.linear.current_platform") |
| 190 | def test_register_oot_linear_kernel(platform_mock): |
| 191 | """Test that the linear kernel registration works correctly.""" |
| 192 | platform_mock._enum = PlatformEnum.OOT |
| 193 | register_linear_kernel(OOTInt8ScaledMMLinearKernel, PlatformEnum.OOT, "int8") |
| 194 | |
| 195 | kernel = init_int8_linear_kernel(True, True, True, "module") |
| 196 | |
| 197 | assert isinstance(kernel, OOTInt8ScaledMMLinearKernel), ( |
| 198 | "init_int8_linear_kernel should return an instance of the registered kernel" |
| 199 | ) |