| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import ABC, abstractmethod |
| 5 | from dataclasses import dataclass |
| 6 | from typing import Any, ClassVar, Generic, TypeVar |
| 7 | |
| 8 | import torch |
| 9 | from typing_extensions import Self |
| 10 | |
| 11 | from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey |
| 12 | |
| 13 | |
| 14 | @dataclass |
| 15 | class MMLinearLayerConfig: ... |
| 16 | |
| 17 | |
| 18 | @dataclass |
| 19 | class Params: |
| 20 | """Base class for quantized layer parameters. |
| 21 | |
| 22 | This class provides a typed interface for accessing quantized weights and scales |
| 23 | from layer modules. It serves as a parameter container that can be extracted from |
| 24 | layers and passed to kernel implementations. |
| 25 | |
| 26 | Attributes: |
| 27 | weight: The quantized weight tensor |
| 28 | weight_scale: weight scaling factors |
| 29 | input_scale: Optional input scaling factors |
| 30 | |
| 31 | Class Variables: |
| 32 | WEIGHT: Attribute name for weight tensor on the layer module |
| 33 | WEIGHT_SCALE: Attribute name for weight scale tensor on the layer module |
| 34 | INPUT_SCALE: Attribute name for input scale tensor on the layer module |
| 35 | |
| 36 | Important: |
| 37 | The string values of WEIGHT, WEIGHT_SCALE, and INPUT_SCALE class variables |
| 38 | MUST match the attribute names used in the corresponding quantization method's |
| 39 | create_weights() implementation. |
| 40 | For example, if FP8LinearMethod.create_weights() |
| 41 | sets layer.weight and layer.weight_scale, |
| 42 | then WEIGHT="weight" and |
| 43 | WEIGHT_SCALE="weight_scale" must be used here. |
| 44 | |
| 45 | Usage: |
| 46 | ```python |
| 47 | # Extract parameters from a quantized layer |
| 48 | params = Params.from_layer(layer) |
| 49 | |
| 50 | # Access typed parameters |
| 51 | output = func(input, params.weight, params.weight_scale) |
| 52 | ``` |
| 53 | |
| 54 | """ |
| 55 | |
| 56 | weight: torch.Tensor |
| 57 | weight_scale: torch.Tensor |
| 58 | input_scale: torch.Tensor | None |
| 59 | |
| 60 | # Attribute names on the layer |
| 61 | WEIGHT: ClassVar[str] = "weight" |
| 62 | WEIGHT_SCALE: ClassVar[str] = "weight_scale" |
| 63 | INPUT_SCALE: ClassVar[str] = "input_scale" |
| 64 | |
| 65 | @classmethod |
| 66 | def from_layer(cls, layer: torch.nn.Module) -> Self: |
| 67 | return cls( |
| 68 | weight=getattr(layer, cls.WEIGHT), |
| 69 | weight_scale=getattr(layer, cls.WEIGHT_SCALE), |
| 70 | input_scale=getattr(layer, cls.INPUT_SCALE, None), |
| 71 | ) |
| 72 | |
| 73 | |
| 74 | @dataclass |
| 75 | class FP8Params(Params): |
| 76 | """FP8 layer parameters with typed fields.""" |
| 77 | |
| 78 | input_scale_ub: torch.Tensor | None |
| 79 | |
| 80 | INPUT_SCALE_UB: ClassVar[str] = "input_scale_ub" |
| 81 | |
| 82 | @classmethod |
| 83 | def from_layer(cls, layer: torch.nn.Module) -> "FP8Params": |
| 84 | """Extract parameters from layer.""" |
| 85 | return cls( |
| 86 | weight=getattr(layer, cls.WEIGHT), |
| 87 | weight_scale=getattr(layer, cls.WEIGHT_SCALE), |
| 88 | input_scale=getattr(layer, cls.INPUT_SCALE, None), |
| 89 | input_scale_ub=getattr(layer, cls.INPUT_SCALE_UB, None), |
| 90 | ) |
| 91 | |
| 92 | |
| 93 | _ParamsT = TypeVar("_ParamsT", bound=Params) |
| 94 | _ConfigT = TypeVar("_ConfigT", bound=MMLinearLayerConfig) |
| 95 | |
| 96 | |
| 97 | class MMLinearKernel(ABC, Generic[_ConfigT, _ParamsT]): |
| 98 | """Abstract base class for quantized matrix multiplication kernels. |
| 99 | |
| 100 | This class provides the interface for implementing custom quantized linear layer |
| 101 | kernels in vLLM. Subclasses should implement specific quantization strategies |
| 102 | (e.g., FP8, INT8) and their corresponding compute kernels. |
| 103 | |
| 104 | Generic Type Parameters: |
| 105 | _ConfigT: Configuration type for the kernel (subclass of MMLinearLayerConfig). |
| 106 | Contains kernel-specific settings like quantization keys, dtypes, etc. |
| 107 | _ParamsT: Parameter type for the kernel (subclass of Params). |
| 108 | Defines the quantized weights and scales needed by the kernel. |
| 109 | |
| 110 | Typical Usage: |
| 111 | 1. Define a config dataclass inheriting from MMLinearLayerConfig |
| 112 | 2. Define a params dataclass inheriting from Params |
| 113 | 3. Subclass MMLinearKernel with your config and params types |
| 114 | 4. Implement all abstract methods |
| 115 | 5. Register the kernel with the quantization method |
| 116 | |
| 117 | Example: |
| 118 | ```python |
| 119 | @dataclass |
| 120 | class MyKernelConfig(MMLinearLayerConfig): |
| 121 | static: bool |
| 122 | output_dtype: torch.dtype |
| 123 | |
| 124 | |
| 125 | @dataclass |
| 126 | class MyKernelParams(FP8Params): |
| 127 | custom_scale: torch.Tensor |
| 128 | CUSTOM_SCALE: ClassVar[str] = "custom_scale" |
| 129 | |
| 130 | |
| 131 | class MyKernel(MMLinearKernel[MyKernelConfig, MyKernelParams]): |
| 132 | @classmethod |
| 133 | def is_supported(cls, compute_capability=None): |
| 134 | if compute_capability and compute_capability < 90: |
| 135 | return False, "Requires compute capability >= 9.0" |
| 136 | return True, None |
| 137 | |
| 138 | @classmethod |
| 139 | def can_implement(cls, config): |
| 140 | if not config.static: |
| 141 | return False, "Only static quantization supported" |
| 142 | return True, None |
| 143 | |
| 144 | def process_weights_after_loading(self, layer): |
| 145 | # Preprocess weights for the kernel |
| 146 | params = self._get_layer_params(layer) |
| 147 | processed = preprocess_weights(params.weight) |
| 148 | replace_parameter(layer, params.WEIGHT, processed) |
| 149 | |
| 150 | def _get_layer_params(self, layer, **kwargs): |
| 151 | return MyKernelParams.from_layer(layer) |
| 152 | |
| 153 | def apply_weights(self, layer, x, bias=None, **kwargs): |
| 154 | params = self._get_layer_params(layer) |
| 155 | # Call your custom kernel |
| 156 | output = my_custom_kernel(x, params.weight, params.weight_scale) |
| 157 | if bias is not None: |
| 158 | output += bias |
| 159 | return output |
| 160 | ``` |
| 161 | |
| 162 | Lifecycle: |
| 163 | 1. Kernel selection: is_supported() and can_implement() check compatibility |
| 164 | 2. Initialization: __init__() creates kernel instance with config |
| 165 | 3. Weight loading: process_weights_after_loading() preprocesses weights |
| 166 | 4. Inference: apply_weights() executes the quantized matmul |
| 167 | |
| 168 | """ |
| 169 | |
| 170 | @classmethod |
| 171 | @abstractmethod |
| 172 | def is_supported( |
| 173 | cls, compute_capability: int | None = None |
| 174 | ) -> tuple[bool, str | None]: |
| 175 | """Check if this kernel is supported on the current hardware. |
| 176 | |
| 177 | This method checks hardware-level compatibility (e.g., GPU architecture, |
| 178 | compute capability, available instructions). It's called during kernel |
| 179 | selection to filter out kernels that cannot run on the current device. |
| 180 | |
| 181 | Args: |
| 182 | compute_capability: GPU compute capability (e.g., 80 for A100, 90 for H100). |
| 183 | If None, should check the current device. |
| 184 | |
| 185 | Returns: |
| 186 | A tuple of (is_supported, reason): |
| 187 | - is_supported: True if the kernel can run on this hardware |
| 188 | - reason: If not supported, a string explaining why; otherwise None |
| 189 | |
| 190 | """ |
| 191 | raise NotImplementedError |
| 192 | |
| 193 | @classmethod |
| 194 | @abstractmethod |
| 195 | def can_implement(cls, config: _ConfigT) -> tuple[bool, str | None]: |
| 196 | """Check if this kernel can implement the given configuration. |
| 197 | |
| 198 | This method checks configuration-level compatibility (e.g., quantization |
| 199 | scheme, group sizes, static vs dynamic quantization). It's called after |
| 200 | is_supported() to determine if this kernel can handle the specific |
| 201 | quantization configuration. |
| 202 | |
| 203 | Args: |
| 204 | config: The kernel configuration to check |
| 205 | |
| 206 | Returns: |
| 207 | A tuple of (can_implement, reason): |
| 208 | - can_implement: True if this kernel supports the config |
| 209 | - reason: If not supported, a string explaining why; otherwise None |
| 210 | |
| 211 | """ |
| 212 | raise NotImplementedError |
| 213 | |
| 214 | def __init__(self, config: _ConfigT) -> None: |
| 215 | """Initialize the kernel with the given configuration. |
| 216 | |
| 217 | Args: |
| 218 | config: Kernel-specific configuration containing settings like |
| 219 | quantization keys, output dtypes, etc. |
| 220 | |
| 221 | """ |
| 222 | self.config = config |
| 223 | |
| 224 | def input_quant_key(self) -> QuantKey | None: |
| 225 | """Return the input quantization key supported by this kernel. If the kernel |
| 226 | does not support input quantization outside of the kernel, return None. |
| 227 | """ |
| 228 | return None |
| 229 | |
| 230 | @abstractmethod |
| 231 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 232 | """Process and transform weights after loading from checkpoint. |
| 233 | |
| 234 | This method is called once after weights are loaded but before inference. |
| 235 | Use it to preprocess weights into the format required by your kernel |
| 236 | (e.g., reordering, padding, format conversion). |
| 237 | |
| 238 | Modifications should be done in-place using replace_parameter() to ensure |
| 239 | the layer's parameters are properly updated. |
| 240 | |
| 241 | Args: |
| 242 | layer: The layer module containing the weights to process |
| 243 | |
| 244 | Example: |
| 245 | ```python |
| 246 | def process_weights_after_loading(self, layer): |
| 247 | params = self._get_layer_params(layer) |
| 248 | # Reorder weights for better memory access |
| 249 | weight_reordered = reorder_weights(params.weight) |
| 250 | replace_parameter(layer, params.WEIGHT, weight_reordered) |
| 251 | ``` |
| 252 | |
| 253 | """ |
| 254 | raise NotImplementedError |
| 255 | |
| 256 | # return a covariant type in the subclass |
| 257 | @abstractmethod |
| 258 | def _get_layer_params(self, layer: torch.nn.Module, **kwargs: Any) -> _ParamsT: |
| 259 | """Extract typed parameters from the layer module. |
| 260 | |
| 261 | This internal method retrieves the quantized weights and scales from |
| 262 | the layer as a typed parameter object. Subclasses should typically |
| 263 | delegate to ParamsClass.from_layer(). |
| 264 | |
| 265 | Args: |
| 266 | layer: The layer module containing the parameters |
| 267 | **kwargs: Additional arguments |
| 268 | |
| 269 | Returns: |
| 270 | A typed parameter object containing weights, scales, and other |
| 271 | quantization parameters |
| 272 | |
| 273 | Example: |
| 274 | ```python |
| 275 | def _get_layer_params(self, layer, **kwargs): |
| 276 | return MyKernelParams.from_layer(layer) |
| 277 | ``` |
| 278 | |
| 279 | """ |
| 280 | raise NotImplementedError |
| 281 | |
| 282 | def get_output_padding(self) -> int | None: |
| 283 | """Get the number of output tokens to pad for this kernel. |
| 284 | |
| 285 | Some kernels require input padding for optimal performance. |
| 286 | Override this method to specify padding requirements. |
| 287 | |
| 288 | Returns: |
| 289 | Number of tokens to pad, or None for no padding (default) |
| 290 | |
| 291 | """ |
| 292 | return None |
| 293 | |
| 294 | @abstractmethod |
| 295 | def apply_weights( |
| 296 | self, |
| 297 | layer: torch.nn.Module, |
| 298 | x: torch.Tensor, |
| 299 | bias: torch.Tensor | None = None, |
| 300 | **kwargs: Any, |
| 301 | ) -> torch.Tensor: |
| 302 | """Apply the quantized weights to the input tensor. |
| 303 | |
| 304 | This is the main inference method that performs the quantized matrix |
| 305 | multiplication. It should handle input quantization (if needed), call |
| 306 | the underlying kernel, and apply bias. |
| 307 | |
| 308 | Args: |
| 309 | layer: The layer module containing the quantized weights |
| 310 | x: Input tensor of shape [..., in_features] |
| 311 | bias: Optional bias tensor of shape [out_features] |
| 312 | **kwargs: Additional kernel-specific arguments |
| 313 | |
| 314 | Returns: |
| 315 | Output tensor of shape [..., out_features] |
| 316 | |
| 317 | """ |
| 318 | raise NotImplementedError |