| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """This module re-exports linear kernel implementations to provide a |
| 5 | stable import interface during an ongoing reorganization. Upcoming |
| 6 | PRs will remove the scaled_mm and mixed_precision subdirectories |
| 7 | and reorganize kernels by provider (aiter, cutlass, flashinfer, etc.) |
| 8 | rather than by precision type. By centralizing exports here, we |
| 9 | minimize the need to update imports across other modules when the |
| 10 | internal structure changes. If you are adding a new kernel selector |
| 11 | or kernel implementation, add it to this __init__.py to maintain |
| 12 | import stability. |
| 13 | """ |
| 14 | |
| 15 | from typing import TypeVar |
| 16 | |
| 17 | import torch |
| 18 | |
| 19 | import vllm.envs as envs |
| 20 | from vllm.logger import init_logger |
| 21 | from vllm.model_executor.kernels.linear.base import ( |
| 22 | MMLinearKernel, |
| 23 | MMLinearLayerConfig, |
| 24 | ) |
| 25 | from vllm.model_executor.kernels.linear.mixed_precision import ( |
| 26 | MPLinearKernel, |
| 27 | MPLinearLayerConfig, |
| 28 | ) |
| 29 | from vllm.model_executor.kernels.linear.mixed_precision.conch import ( |
| 30 | ConchLinearKernel, |
| 31 | ) |
| 32 | from vllm.model_executor.kernels.linear.mixed_precision.cpu import ( |
| 33 | CPUWNA16LinearKernel, |
| 34 | ) |
| 35 | from vllm.model_executor.kernels.linear.mixed_precision.cutlass import ( |
| 36 | CutlassW4A8LinearKernel, |
| 37 | ) |
| 38 | from vllm.model_executor.kernels.linear.mixed_precision.dynamic_4bit import ( |
| 39 | Dynamic4bitLinearKernel, |
| 40 | ) |
| 41 | from vllm.model_executor.kernels.linear.mixed_precision.exllama import ( |
| 42 | ExllamaLinearKernel, |
| 43 | ) |
| 44 | from vllm.model_executor.kernels.linear.mixed_precision.humming import ( |
| 45 | HummingLinearKernel, |
| 46 | ) |
| 47 | from vllm.model_executor.kernels.linear.mixed_precision.machete import ( |
| 48 | MacheteLinearKernel, |
| 49 | ) |
| 50 | from vllm.model_executor.kernels.linear.mixed_precision.marlin import ( |
| 51 | MarlinLinearKernel, |
| 52 | ) |
| 53 | from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( |
| 54 | RDNA3W4A16LinearKernel, |
| 55 | ) |
| 56 | from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import ( |
| 57 | RDNAHybridW4A16LinearKernel, |
| 58 | ) |
| 59 | from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import ( |
| 60 | TritonW4A16LinearKernel, |
| 61 | ) |
| 62 | from vllm.model_executor.kernels.linear.mixed_precision.xpu import ( |
| 63 | XPUW4A8IntLinearKernel, |
| 64 | XPUwNa16LinearKernel, |
| 65 | ) |
| 66 | from vllm.model_executor.kernels.linear.mixed_precision.zentorch import ( |
| 67 | ZentorchWNA16LinearKernel, |
| 68 | ) |
| 69 | from vllm.model_executor.kernels.linear.mxfp4 import ( |
| 70 | MxFp4LinearKernel, |
| 71 | MxFp4LinearLayerConfig, |
| 72 | ) |
| 73 | from vllm.model_executor.kernels.linear.mxfp4.aiter import ( |
| 74 | AiterMxfp4LinearKernel, |
| 75 | ) |
| 76 | from vllm.model_executor.kernels.linear.mxfp4.b12x import ( |
| 77 | B12xMxFp4LinearKernel, |
| 78 | ) |
| 79 | from vllm.model_executor.kernels.linear.mxfp4.emulation import ( |
| 80 | EmulationMxfp4LinearKernel, |
| 81 | ) |
| 82 | from vllm.model_executor.kernels.linear.mxfp4.flashinfer import ( |
| 83 | FlashInferMxFp4LinearKernel, |
| 84 | ) |
| 85 | from vllm.model_executor.kernels.linear.mxfp4.humming import ( |
| 86 | HummingMxFp4LinearKernel, |
| 87 | ) |
| 88 | from vllm.model_executor.kernels.linear.mxfp4.marlin import ( |
| 89 | MarlinMxFp4LinearKernel, |
| 90 | ) |
| 91 | from vllm.model_executor.kernels.linear.mxfp4.xpu import ( |
| 92 | XPUMxFp4LinearKernel, |
| 93 | ) |
| 94 | from vllm.model_executor.kernels.linear.mxfp6 import ( |
| 95 | MxFp6LinearKernel, |
| 96 | MxFp6LinearLayerConfig, |
| 97 | ) |
| 98 | from vllm.model_executor.kernels.linear.mxfp6.emulation import ( |
| 99 | EmulationMxfp6LinearKernel, |
| 100 | ) |
| 101 | from vllm.model_executor.kernels.linear.mxfp6.humming import HummingMxFp6LinearKernel |
| 102 | from vllm.model_executor.kernels.linear.mxfp8 import ( |
| 103 | Mxfp8LinearKernel, |
| 104 | Mxfp8LinearLayerConfig, |
| 105 | ) |
| 106 | from vllm.model_executor.kernels.linear.mxfp8.b12x import ( |
| 107 | B12xMxfp8LinearKernel, |
| 108 | ) |
| 109 | from vllm.model_executor.kernels.linear.mxfp8.deep_gemm import ( |
| 110 | DeepGemmMxfp8BmmLinearKernel, |
| 111 | ) |
| 112 | from vllm.model_executor.kernels.linear.mxfp8.emulation import ( |
| 113 | EmulationMxfp8LinearKernel, |
| 114 | ) |
| 115 | from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( |
| 116 | FlashInferCutedslMxfp8LinearKernel, |
| 117 | FlashInferCutlassMxfp8LinearKernel, |
| 118 | FlashInferTrtllmMxfp8LinearKernel, |
| 119 | ) |
| 120 | from vllm.model_executor.kernels.linear.mxfp8.humming import ( |
| 121 | HummingMxfp8LinearKernel, |
| 122 | ) |
| 123 | from vllm.model_executor.kernels.linear.mxfp8.marlin import ( |
| 124 | MarlinMxfp8LinearKernel, |
| 125 | ) |
| 126 | from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( |
| 127 | RocmDotScaledMxfp8LinearKernel, |
| 128 | ) |
| 129 | from vllm.model_executor.kernels.linear.mxfp8.xpu import ( |
| 130 | XPUMxFp8LinearKernel, |
| 131 | ) |
| 132 | from vllm.model_executor.kernels.linear.nvfp4 import ( |
| 133 | NvFp4LinearKernel, |
| 134 | NvFp4LinearLayerConfig, |
| 135 | ) |
| 136 | from vllm.model_executor.kernels.linear.nvfp4.b12x import ( |
| 137 | B12xNvFp4LinearKernel, |
| 138 | ) |
| 139 | from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( |
| 140 | CutlassNvFp4LinearKernel, |
| 141 | ) |
| 142 | from vllm.model_executor.kernels.linear.nvfp4.emulation import ( |
| 143 | EmulationNvFp4LinearKernel, |
| 144 | ) |
| 145 | from vllm.model_executor.kernels.linear.nvfp4.fbgemm import ( |
| 146 | FbgemmNvFp4LinearKernel, |
| 147 | ) |
| 148 | from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( |
| 149 | FlashInferB12xNvFp4LinearKernel, |
| 150 | FlashInferCudnnNvFp4LinearKernel, |
| 151 | FlashInferCuteDslNvFp4LinearKernel, |
| 152 | FlashInferCuteDslNvFp4W4A16LinearKernel, |
| 153 | FlashInferCutlassNvFp4LinearKernel, |
| 154 | FlashInferTrtllmNvFp4LinearKernel, |
| 155 | ) |
| 156 | from vllm.model_executor.kernels.linear.nvfp4.humming import ( |
| 157 | HummingNvFp4LinearKernel, |
| 158 | ) |
| 159 | from vllm.model_executor.kernels.linear.nvfp4.marlin import ( |
| 160 | MarlinNvFp4LinearKernel, |
| 161 | ) |
| 162 | from vllm.model_executor.kernels.linear.nvfp4.pytorch import ( |
| 163 | TorchNvFp4LinearKernel, |
| 164 | ) |
| 165 | from vllm.model_executor.kernels.linear.scaled_mm import ( |
| 166 | Fp8BlockScaledMMLinearKernel, |
| 167 | FP8ScaledMMLinearKernel, |
| 168 | FP8ScaledMMLinearLayerConfig, |
| 169 | Int8ScaledMMLinearKernel, |
| 170 | Int8ScaledMMLinearLayerConfig, |
| 171 | ScaledMMLinearKernel, |
| 172 | ) |
| 173 | from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( |
| 174 | AiterFp8BlockScaledMMKernel, |
| 175 | AiterHipbMMPerTokenFp8ScaledMMLinearKernel, |
| 176 | AiterInt8ScaledMMLinearKernel, |
| 177 | AiterPerTokenFp8ScaledMMLinearKernel, |
| 178 | AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, |
| 179 | ) |
| 180 | from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( |
| 181 | B12xFp8BlockScaledMMKernel, |
| 182 | B12xTensorFP8ScaledMMLinearKernel, |
| 183 | ) |
| 184 | from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( |
| 185 | CPUFp8BlockScaledMMKernel, |
| 186 | CPUFp8PerTensorScaledMMLinearKernel, |
| 187 | CPUFp8W8A8BlockScaledMMKernel, |
| 188 | CPUFP8W8A8ScaledMMLinearKernel, |
| 189 | CPUInt8ScaledMMLinearKernel, |
| 190 | ) |
| 191 | from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( |
| 192 | CutlassFp8BlockScaledMMKernel, |
| 193 | CutlassFP8ScaledMMLinearKernel, |
| 194 | CutlassInt8ScaledMMLinearKernel, |
| 195 | ) |
| 196 | from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import ( |
| 197 | DeepGemmFp8BlockScaledMMKernel, |
| 198 | ) |
| 199 | from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import ( |
| 200 | FlashInferFp8DeepGEMMDynamicBlockScaledKernel, |
| 201 | FlashInferFP8ScaledMMLinearKernel, |
| 202 | ) |
| 203 | from vllm.model_executor.kernels.linear.scaled_mm.humming import ( |
| 204 | HummingFP8ScaledMMLinearKernel, |
| 205 | HummingInt8ScaledMMLinearKernel, |
| 206 | ) |
| 207 | from vllm.model_executor.kernels.linear.scaled_mm.marlin import ( |
| 208 | MarlinFP8ScaledMMLinearKernel, |
| 209 | ) |
| 210 | from vllm.model_executor.kernels.linear.scaled_mm.pytorch import ( |
| 211 | BlockWiseTorchFP8ScaledMMLinearKernel, |
| 212 | ChannelWiseTorchFP8ScaledMMLinearKernel, |
| 213 | PerTensorTorchFP8ScaledMMLinearKernel, |
| 214 | RowWiseTorchFP8ScaledMMLinearKernel, |
| 215 | ) |
| 216 | from vllm.model_executor.kernels.linear.scaled_mm.rocm import ( |
| 217 | ROCmFP8ScaledMMLinearKernel, |
| 218 | ) |
| 219 | from vllm.model_executor.kernels.linear.scaled_mm.triton import ( |
| 220 | TritonFp8BlockScaledMMKernel, |
| 221 | TritonInt8ScaledMMLinearKernel, |
| 222 | ) |
| 223 | from vllm.model_executor.kernels.linear.scaled_mm.xpu import ( |
| 224 | XPUFp8BlockScaledMMKernel, |
| 225 | XPUW8A8FP8LinearKernel, |
| 226 | XPUW8A16FP8LinearKernel, |
| 227 | ) |
| 228 | from vllm.model_executor.kernels.linear.scaled_mm.zentorch import ( |
| 229 | ZentorchInt8ScaledMMLinearKernel, |
| 230 | ) |
| 231 | from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey |
| 232 | from vllm.platforms import PlatformEnum, current_platform |
| 233 | |
| 234 | logger = init_logger(__name__) |
| 235 | |
| 236 | |
| 237 | def _get_linear_backend(*, quantization: str) -> str: |
| 238 | """Get the linear_backend setting from the current vllm config.""" |
| 239 | from vllm.config import get_current_vllm_config_or_none |
| 240 | |
| 241 | if (config := get_current_vllm_config_or_none()) is None: |
| 242 | return "auto" |
| 243 | overrides = config.kernel_config.linear_backend_per_quant |
| 244 | if overrides is not None and (override := overrides.get(quantization)): |
| 245 | logger.info_once( |
| 246 | "Applied linear backend override for %r: %r", quantization, override |
| 247 | ) |
| 248 | return override |
| 249 | return config.kernel_config.linear_backend |
| 250 | |
| 251 | |
| 252 | # Kernel classes covered by each --linear-backend value. |
| 253 | _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { |
| 254 | "b12x": { |
| 255 | B12xFp8BlockScaledMMKernel, |
| 256 | B12xMxFp4LinearKernel, |
| 257 | B12xMxfp8LinearKernel, |
| 258 | B12xNvFp4LinearKernel, |
| 259 | B12xTensorFP8ScaledMMLinearKernel, |
| 260 | }, |
| 261 | "cutlass": { |
| 262 | CutlassInt8ScaledMMLinearKernel, |
| 263 | CutlassFP8ScaledMMLinearKernel, |
| 264 | CutlassFp8BlockScaledMMKernel, |
| 265 | CutlassW4A8LinearKernel, |
| 266 | CutlassNvFp4LinearKernel, |
| 267 | }, |
| 268 | "flashinfer_cutlass": { |
| 269 | FlashInferFP8ScaledMMLinearKernel, |
| 270 | FlashInferFp8DeepGEMMDynamicBlockScaledKernel, |
| 271 | FlashInferCutlassMxfp8LinearKernel, |
| 272 | FlashInferCutlassNvFp4LinearKernel, |
| 273 | FlashInferMxFp4LinearKernel, |
| 274 | }, |
| 275 | "flashinfer_cutedsl": { |
| 276 | FlashInferCuteDslNvFp4LinearKernel, |
| 277 | FlashInferCuteDslNvFp4W4A16LinearKernel, |
| 278 | FlashInferCutedslMxfp8LinearKernel, |
| 279 | }, |
| 280 | "flashinfer_trtllm": { |
| 281 | FlashInferTrtllmMxfp8LinearKernel, |
| 282 | FlashInferTrtllmNvFp4LinearKernel, |
| 283 | }, |
| 284 | "flashinfer_cudnn": { |
| 285 | FlashInferCudnnNvFp4LinearKernel, |
| 286 | }, |
| 287 | "flashinfer_b12x": { |
| 288 | FlashInferB12xNvFp4LinearKernel, |
| 289 | }, |
| 290 | "humming": { |
| 291 | HummingFP8ScaledMMLinearKernel, |
| 292 | HummingInt8ScaledMMLinearKernel, |
| 293 | HummingLinearKernel, |
| 294 | HummingMxfp8LinearKernel, |
| 295 | HummingMxFp4LinearKernel, |
| 296 | HummingMxFp6LinearKernel, |
| 297 | HummingNvFp4LinearKernel, |
| 298 | }, |
| 299 | "marlin": { |
| 300 | MarlinFP8ScaledMMLinearKernel, |
| 301 | MarlinLinearKernel, |
| 302 | MarlinMxfp8LinearKernel, |
| 303 | MarlinNvFp4LinearKernel, |
| 304 | MarlinMxFp4LinearKernel, |
| 305 | }, |
| 306 | "triton": { |
| 307 | TritonInt8ScaledMMLinearKernel, |
| 308 | TritonFp8BlockScaledMMKernel, |
| 309 | TritonW4A16LinearKernel, |
| 310 | }, |
| 311 | "deep_gemm": { |
| 312 | DeepGemmFp8BlockScaledMMKernel, |
| 313 | DeepGemmMxfp8BmmLinearKernel, |
| 314 | }, |
| 315 | "torch": { |
| 316 | PerTensorTorchFP8ScaledMMLinearKernel, |
| 317 | ChannelWiseTorchFP8ScaledMMLinearKernel, |
| 318 | RowWiseTorchFP8ScaledMMLinearKernel, |
| 319 | BlockWiseTorchFP8ScaledMMLinearKernel, |
| 320 | TorchNvFp4LinearKernel, |
| 321 | }, |
| 322 | "aiter": { |
| 323 | AiterInt8ScaledMMLinearKernel, |
| 324 | AiterFp8BlockScaledMMKernel, |
| 325 | AiterPerTokenFp8ScaledMMLinearKernel, |
| 326 | AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, |
| 327 | AiterMxfp4LinearKernel, |
| 328 | }, |
| 329 | "machete": { |
| 330 | MacheteLinearKernel, |
| 331 | }, |
| 332 | "fbgemm": { |
| 333 | FbgemmNvFp4LinearKernel, |
| 334 | }, |
| 335 | "conch": { |
| 336 | ConchLinearKernel, |
| 337 | }, |
| 338 | "exllama": { |
| 339 | ExllamaLinearKernel, |
| 340 | }, |
| 341 | "emulation": { |
| 342 | EmulationMxfp8LinearKernel, |
| 343 | EmulationNvFp4LinearKernel, |
| 344 | EmulationMxfp6LinearKernel, |
| 345 | EmulationMxfp4LinearKernel, |
| 346 | }, |
| 347 | "xpu": { |
| 348 | XPUW8A8FP8LinearKernel, |
| 349 | XPUFp8BlockScaledMMKernel, |
| 350 | }, |
| 351 | "xpu_woq": { |
| 352 | XPUW8A16FP8LinearKernel, |
| 353 | }, |
| 354 | } |
| 355 | |
| 356 | |
| 357 | def _filter_kernels_by_backend( |
| 358 | backend: str, |
| 359 | kernels: list[type], |
| 360 | ) -> list[type]: |
| 361 | """Prefer matching kernels, falling back to auto by layer type.""" |
| 362 | backend_kernels = _LINEAR_BACKEND_KERNEL_MAP.get(backend, set()) |
| 363 | filtered = [kernel for kernel in kernels if kernel in backend_kernels] |
| 364 | if not filtered: |
| 365 | logger.warning_once( |
| 366 | "--linear-backend=%s has no kernel for this linear layer type; " |
| 367 | "using automatic selection for that layer type.", |
| 368 | backend, |
| 369 | ) |
| 370 | return kernels |
| 371 | return filtered |
| 372 | |
| 373 | |
| 374 | def _resolve_backend_kernels( |
| 375 | kernels: list[type], |
| 376 | layer_desc: str, |
| 377 | *, |
| 378 | quantization: str, |
| 379 | ) -> list[type]: |
| 380 | """Apply --linear-backend filtering to one layer type's kernel list. |
| 381 | |
| 382 | When the requested backend has no kernel for this layer type, fall back |
| 383 | to the unfiltered list (with a WARNING log) instead of failing engine |
| 384 | startup: an explicitly requested backend usually provides kernels for a |
| 385 | single quantization scheme, while a model can contain several linear |
| 386 | layer types (e.g. NVFP4 MoE projections next to FP8 attention |
| 387 | projections). |
| 388 | """ |
| 389 | linear_backend = _get_linear_backend(quantization=quantization) |
| 390 | if linear_backend == "auto": |
| 391 | return kernels |
| 392 | filtered = _filter_kernels_by_backend(linear_backend, kernels) |
| 393 | if not filtered: |
| 394 | logger.warning_once( |
| 395 | "--linear-backend=%s was requested, but no %s kernel exists for " |
| 396 | "%s layers; falling back to normal kernel selection for this " |
| 397 | "layer.", |
| 398 | linear_backend, |
| 399 | linear_backend, |
| 400 | layer_desc, |
| 401 | scope="global", |
| 402 | ) |
| 403 | return kernels |
| 404 | return filtered |
| 405 | |
| 406 | |
| 407 | # in priority/performance order (when available) |
| 408 | _POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] = { |
| 409 | PlatformEnum.CPU: [ZentorchInt8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel], |
| 410 | PlatformEnum.CUDA: [ |
| 411 | CutlassInt8ScaledMMLinearKernel, |
| 412 | TritonInt8ScaledMMLinearKernel, |
| 413 | HummingInt8ScaledMMLinearKernel, |
| 414 | ], |
| 415 | PlatformEnum.ROCM: [AiterInt8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel], |
| 416 | PlatformEnum.XPU: [TritonInt8ScaledMMLinearKernel], |
| 417 | } |
| 418 | |
| 419 | # in priority/performance order (when available) |
| 420 | _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = { |
| 421 | PlatformEnum.CUDA: [ |
| 422 | FlashInferFP8ScaledMMLinearKernel, |
| 423 | CutlassFP8ScaledMMLinearKernel, |
| 424 | B12xTensorFP8ScaledMMLinearKernel, |
| 425 | PerTensorTorchFP8ScaledMMLinearKernel, |
| 426 | ChannelWiseTorchFP8ScaledMMLinearKernel, |
| 427 | MarlinFP8ScaledMMLinearKernel, |
| 428 | HummingFP8ScaledMMLinearKernel, |
| 429 | ], |
| 430 | PlatformEnum.ROCM: [ |
| 431 | AiterHipbMMPerTokenFp8ScaledMMLinearKernel, |
| 432 | AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, |
| 433 | AiterPerTokenFp8ScaledMMLinearKernel, |
| 434 | ROCmFP8ScaledMMLinearKernel, |
| 435 | PerTensorTorchFP8ScaledMMLinearKernel, |
| 436 | RowWiseTorchFP8ScaledMMLinearKernel, |
| 437 | ChannelWiseTorchFP8ScaledMMLinearKernel, |
| 438 | ], |
| 439 | PlatformEnum.CPU: [ |
| 440 | CPUFP8W8A8ScaledMMLinearKernel, |
| 441 | CPUFp8PerTensorScaledMMLinearKernel, |
| 442 | PerTensorTorchFP8ScaledMMLinearKernel, |
| 443 | ChannelWiseTorchFP8ScaledMMLinearKernel, |
| 444 | ], |
| 445 | PlatformEnum.XPU: [ |
| 446 | XPUW8A16FP8LinearKernel, |
| 447 | XPUW8A8FP8LinearKernel, |
| 448 | PerTensorTorchFP8ScaledMMLinearKernel, |
| 449 | ChannelWiseTorchFP8ScaledMMLinearKernel, |
| 450 | ], |
| 451 | } |
| 452 | |
| 453 | |
| 454 | # in priority/performance order (when available) |
| 455 | _POSSIBLE_FP8_BLOCK_KERNELS: dict[ |
| 456 | PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]] |
| 457 | ] = { |
| 458 | PlatformEnum.CUDA: [ |
| 459 | FlashInferFp8DeepGEMMDynamicBlockScaledKernel, |
| 460 | DeepGemmFp8BlockScaledMMKernel, |
| 461 | CutlassFp8BlockScaledMMKernel, |
| 462 | B12xFp8BlockScaledMMKernel, |
| 463 | MarlinFP8ScaledMMLinearKernel, |
| 464 | HummingFP8ScaledMMLinearKernel, |
| 465 | TritonFp8BlockScaledMMKernel, |
| 466 | BlockWiseTorchFP8ScaledMMLinearKernel, |
| 467 | ], |
| 468 | PlatformEnum.ROCM: [ |
| 469 | AiterFp8BlockScaledMMKernel, |
| 470 | TritonFp8BlockScaledMMKernel, |
| 471 | ], |
| 472 | PlatformEnum.CPU: [ |
| 473 | CPUFp8W8A8BlockScaledMMKernel, # W8A8 preferred; falls back to W8A16 below |
| 474 | CPUFp8BlockScaledMMKernel, |
| 475 | ], |
| 476 | PlatformEnum.XPU: [ |
| 477 | XPUFp8BlockScaledMMKernel, |
| 478 | TritonFp8BlockScaledMMKernel, |
| 479 | BlockWiseTorchFP8ScaledMMLinearKernel, |
| 480 | ], |
| 481 | } |
| 482 | |
| 483 | _POSSIBLE_WFP8A16_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = { |
| 484 | PlatformEnum.CUDA: [ |
| 485 | HummingFP8ScaledMMLinearKernel, |
| 486 | MarlinFP8ScaledMMLinearKernel, |
| 487 | ], |
| 488 | PlatformEnum.ROCM: [ |
| 489 | # To be added |
| 490 | ], |
| 491 | PlatformEnum.CPU: [ |
| 492 | # To be added |
| 493 | ], |
| 494 | PlatformEnum.XPU: [ |
| 495 | XPUW8A16FP8LinearKernel, |
| 496 | ], |
| 497 | } |
| 498 | |
| 499 | # in priority/performance order (when available) |
| 500 | _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = { |
| 501 | PlatformEnum.CUDA: [ |
| 502 | CutlassW4A8LinearKernel, |
| 503 | MacheteLinearKernel, |
| 504 | MarlinLinearKernel, |
| 505 | ConchLinearKernel, |
| 506 | ExllamaLinearKernel, |
| 507 | TritonW4A16LinearKernel, |
| 508 | HummingLinearKernel, |
| 509 | ], |
| 510 | PlatformEnum.ROCM: [ |
| 511 | RDNA3W4A16LinearKernel, |
| 512 | RDNAHybridW4A16LinearKernel, |
| 513 | TritonW4A16LinearKernel, |
| 514 | ConchLinearKernel, |
| 515 | ExllamaLinearKernel, |
| 516 | ], |
| 517 | PlatformEnum.XPU: [ |
| 518 | XPUW4A8IntLinearKernel, |
| 519 | XPUwNa16LinearKernel, |
| 520 | ], |
| 521 | PlatformEnum.CPU: [ |
| 522 | Dynamic4bitLinearKernel, |
| 523 | ZentorchWNA16LinearKernel, |
| 524 | CPUWNA16LinearKernel, |
| 525 | ], |
| 526 | } |
| 527 | |
| 528 | # in priority/performance order (when available) |
| 529 | _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = { |
| 530 | PlatformEnum.CUDA: [ |
| 531 | FlashInferCutedslMxfp8LinearKernel, |
| 532 | FlashInferCutlassMxfp8LinearKernel, |
| 533 | MarlinMxfp8LinearKernel, |
| 534 | B12xMxfp8LinearKernel, |
| 535 | EmulationMxfp8LinearKernel, |
| 536 | HummingMxfp8LinearKernel, |
| 537 | FlashInferTrtllmMxfp8LinearKernel, |
| 538 | ], |
| 539 | PlatformEnum.ROCM: [ |
| 540 | # Native CDNA4 (gfx950) MX linear; is_supported() gates to gfx95x and |
| 541 | # falls through to BF16 emulation (hipBLASLt) elsewhere / on regression. |
| 542 | RocmDotScaledMxfp8LinearKernel, |
| 543 | EmulationMxfp8LinearKernel, |
| 544 | ], |
| 545 | PlatformEnum.XPU: [ |
| 546 | XPUMxFp8LinearKernel, |
| 547 | EmulationMxfp8LinearKernel, |
| 548 | ], |
| 549 | } |
| 550 | |
| 551 | _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { |
| 552 | PlatformEnum.CUDA: [ |
| 553 | FlashInferCuteDslNvFp4LinearKernel, |
| 554 | FlashInferCutlassNvFp4LinearKernel, |
| 555 | FlashInferB12xNvFp4LinearKernel, |
| 556 | CutlassNvFp4LinearKernel, |
| 557 | # Weight-only: a fallback on a W4A4-capable checkpoint, not a preference. |
| 558 | FlashInferCuteDslNvFp4W4A16LinearKernel, |
| 559 | MarlinNvFp4LinearKernel, |
| 560 | FlashInferTrtllmNvFp4LinearKernel, |
| 561 | FlashInferCudnnNvFp4LinearKernel, |
| 562 | FbgemmNvFp4LinearKernel, |
| 563 | B12xNvFp4LinearKernel, |
| 564 | TorchNvFp4LinearKernel, |
| 565 | EmulationNvFp4LinearKernel, |
| 566 | HummingNvFp4LinearKernel, |
| 567 | ], |
| 568 | PlatformEnum.ROCM: [ |
| 569 | EmulationNvFp4LinearKernel, |
| 570 | ], |
| 571 | } |
| 572 | |
| 573 | _POSSIBLE_MXFP6_KERNELS: dict[PlatformEnum, list[type[MxFp6LinearKernel]]] = { |
| 574 | PlatformEnum.CUDA: [ |
| 575 | HummingMxFp6LinearKernel, |
| 576 | EmulationMxfp6LinearKernel, |
| 577 | ], |
| 578 | PlatformEnum.ROCM: [ |
| 579 | EmulationMxfp6LinearKernel, |
| 580 | ], |
| 581 | } |
| 582 | |
| 583 | _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = { |
| 584 | PlatformEnum.CUDA: [ |
| 585 | FlashInferMxFp4LinearKernel, |
| 586 | MarlinMxFp4LinearKernel, |
| 587 | HummingMxFp4LinearKernel, |
| 588 | B12xMxFp4LinearKernel, |
| 589 | EmulationMxfp4LinearKernel, |
| 590 | ], |
| 591 | PlatformEnum.ROCM: [ |
| 592 | AiterMxfp4LinearKernel, |
| 593 | EmulationMxfp4LinearKernel, |
| 594 | ], |
| 595 | PlatformEnum.XPU: [ |
| 596 | XPUMxFp4LinearKernel, |
| 597 | ], |
| 598 | } |
| 599 | |
| 600 | # TODO make all kernels inherit from MMLinearKernel |
| 601 | # then bound _KernelT only to MMLinearKernel |
| 602 | _KernelT = TypeVar("_KernelT", bound=ScaledMMLinearKernel | MMLinearKernel) |
| 603 | _KernelConfigT = TypeVar("_KernelConfigT", bound=MMLinearLayerConfig) |
| 604 | |
| 605 | |
| 606 | def is_supported_and_can_implement_kernel( |
| 607 | kernel: type[_KernelT], config: _KernelConfigT, compute_capability: int | None |
| 608 | ) -> tuple[bool, str]: |
| 609 | if kernel.__name__ in envs.VLLM_DISABLED_KERNELS: |
| 610 | return False, f" {kernel.__name__} is disabled by environment variable" |
| 611 | |
| 612 | if compute_capability is None: |
| 613 | _cc = current_platform.get_device_capability() |
| 614 | if _cc is not None: |
| 615 | compute_capability = _cc[0] * 10 + _cc[1] |
| 616 | |
| 617 | is_supported, failure_reason = kernel.is_supported(compute_capability) |
| 618 | if not is_supported: |
| 619 | return False, f"{kernel.__name__} {failure_reason}." |
| 620 | |
| 621 | can_implement, failure_reason = kernel.can_implement(config) |
| 622 | if not can_implement: |
| 623 | return ( |
| 624 | False, |
| 625 | f"{kernel.__name__} {failure_reason}.", |
| 626 | ) |
| 627 | |
| 628 | return True, "" |
| 629 | |
| 630 | |
| 631 | def choose_scaled_mm_linear_kernel( |
| 632 | config: _KernelConfigT, |
| 633 | possible_kernels: dict[PlatformEnum, list[type[_KernelT]]], |
| 634 | compute_capability: int | None = None, |
| 635 | force_kernel: type[_KernelT] | None = None, |
| 636 | *, |
| 637 | quantization: str, |
| 638 | ) -> type[_KernelT]: |
| 639 | """Choose a _KernelT that can implement the given config for the |
| 640 | given compute capability. Attempts to choose the best kernel in terms of |
| 641 | performance. |
| 642 | |
| 643 | Args: |
| 644 | config (_KernelConfigT): Description of the linear layer |
| 645 | to be implemented. |
| 646 | possible_kernels (dict[PlatformEnum, list[_KernelT]]): A |
| 647 | dictionary of platforms and their list of possible kernels. |
| 648 | compute_capability (Optional[int], optional): The compute capability of |
| 649 | the target device, if None uses `current_platform` to get the |
| 650 | compute capability. Defaults to None. |
| 651 | force_kernel (Optional[type[_KernelT]]): An Optional forced kernel to override |
| 652 | the possible_kernels if it can be implemented. If None, it will only try the |
| 653 | possible kernels. |
| 654 | quantization: Quantization scheme used to select a backend override. |
| 655 | |
| 656 | Raises: |
| 657 | ValueError: If no kernel can implement the given config. |
| 658 | |
| 659 | Returns: |
| 660 | _KernelT: Chosen kernel. |
| 661 | |
| 662 | """ |
| 663 | failure_reason_list = [] |
| 664 | |
| 665 | if force_kernel is not None: |
| 666 | can_implement, failure_reason = is_supported_and_can_implement_kernel( |
| 667 | force_kernel, config, compute_capability |
| 668 | ) |
| 669 | if can_implement: |
| 670 | return force_kernel |
| 671 | |
| 672 | logger.info_once( |
| 673 | "Tried to force %s, but the kernel couldn't be implemented", |
| 674 | force_kernel.__name__, |
| 675 | scope="global", |
| 676 | ) |
| 677 | |
| 678 | platform_kernels = possible_kernels.get(current_platform._enum, []) |
| 679 | |
| 680 | # Apply --linear-backend filtering when set. |
| 681 | platform_kernels = _resolve_backend_kernels( |
| 682 | platform_kernels, |
| 683 | "scaled-mm", |
| 684 | quantization=quantization, |
| 685 | ) |
| 686 | |
| 687 | for kernel in platform_kernels: |
| 688 | is_supported_and_can_implement, failure_reason = ( |
| 689 | is_supported_and_can_implement_kernel(kernel, config, compute_capability) |
| 690 | ) |
| 691 | if is_supported_and_can_implement: |
| 692 | return kernel |
| 693 | failure_reason_list.append(failure_reason) |
| 694 | |
| 695 | raise ValueError( |
| 696 | "Failed to find a kernel that can implement the " |
| 697 | "ScaledMM linear layer. Reasons: \n" + "\n".join(failure_reason_list) |
| 698 | ) |
| 699 | |
| 700 | |
| 701 | def init_fp8_linear_kernel( |
| 702 | activation_quant_key: QuantKey, |
| 703 | weight_quant_key: QuantKey, |
| 704 | input_dtype: torch.dtype, |
| 705 | out_dtype: torch.dtype, |
| 706 | weight_shape: tuple[int, int], |
| 707 | force_kernel: type[FP8ScaledMMLinearKernel] | None = None, |
| 708 | module_name: str | None = None, |
| 709 | ) -> FP8ScaledMMLinearKernel | Fp8BlockScaledMMLinearKernel: |
| 710 | scaled_mm_linear_kernel_config = FP8ScaledMMLinearLayerConfig( |
| 711 | weight_quant_key=weight_quant_key, |
| 712 | activation_quant_key=activation_quant_key, |
| 713 | input_dtype=input_dtype, |
| 714 | out_dtype=out_dtype, |
| 715 | weight_shape=weight_shape, |
| 716 | ) |
| 717 | |
| 718 | if activation_quant_key.scale.group_shape.is_per_group(): |
| 719 | kernel_type = choose_scaled_mm_linear_kernel( |
| 720 | config=scaled_mm_linear_kernel_config, |
| 721 | possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS, # type: ignore[misc] |
| 722 | quantization="fp8_block_w8a8", |
| 723 | force_kernel=force_kernel, |
| 724 | ) |
| 725 | if module_name: |
| 726 | logger.info_once( |
| 727 | "Selected %s for %s", |
| 728 | kernel_type.__name__, |
| 729 | module_name, |
| 730 | scope="global", |
| 731 | ) |
| 732 | |
| 733 | # TODO make scaled_mm kernels inherit from MMLinearKernel |
| 734 | # only MarlinFP8ScaledMMLinearKernel is a type of FP8ScaledMMLinearKernel. |
| 735 | if issubclass(kernel_type, FP8ScaledMMLinearKernel): |
| 736 | return kernel_type( |
| 737 | scaled_mm_linear_kernel_config, |
| 738 | layer_param_names=[ |
| 739 | "weight", |
| 740 | "weight_scale", |
| 741 | "input_scale", |
| 742 | "input_scale_ub", |
| 743 | ], |
| 744 | ) |
| 745 | |
| 746 | return kernel_type( |
| 747 | scaled_mm_linear_kernel_config, |
| 748 | ) |
| 749 | |
| 750 | else: |
| 751 | kernel_type = choose_scaled_mm_linear_kernel( |
| 752 | config=scaled_mm_linear_kernel_config, |
| 753 | possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type] |
| 754 | quantization="fp8_w8a8", |
| 755 | force_kernel=force_kernel, |
| 756 | ) |
| 757 | if module_name: |
| 758 | logger.info_once( |
| 759 | "Selected %s for %s", |
| 760 | kernel_type.__name__, |
| 761 | module_name, |
| 762 | scope="global", |
| 763 | ) |
| 764 | |
| 765 | return kernel_type( |
| 766 | scaled_mm_linear_kernel_config, |
| 767 | layer_param_names=[ |
| 768 | "weight", |
| 769 | "weight_scale", |
| 770 | "input_scale", |
| 771 | "input_scale_ub", |
| 772 | ], |
| 773 | ) |
| 774 | |
| 775 | |
| 776 | def init_int8_linear_kernel( |
| 777 | is_channelwise: bool, |
| 778 | is_static_input_scheme: bool, |
| 779 | input_symmetric: bool, |
| 780 | module_name: str, |
| 781 | ) -> Int8ScaledMMLinearKernel: |
| 782 | config = Int8ScaledMMLinearLayerConfig( |
| 783 | is_channelwise=is_channelwise, |
| 784 | is_static_input_scheme=is_static_input_scheme, |
| 785 | input_symmetric=input_symmetric, |
| 786 | ) |
| 787 | |
| 788 | kernel_type = choose_scaled_mm_linear_kernel( |
| 789 | config, |
| 790 | _POSSIBLE_INT8_KERNELS, |
| 791 | quantization="int8_w8a8", |
| 792 | ) |
| 793 | |
| 794 | logger.info_once( |
| 795 | "Selected %s for %s", |
| 796 | kernel_type.__name__, |
| 797 | module_name, |
| 798 | scope="global", |
| 799 | ) |
| 800 | |
| 801 | return kernel_type( |
| 802 | config, |
| 803 | layer_param_names=[ |
| 804 | "weight", |
| 805 | "weight_scale", |
| 806 | "input_scale", |
| 807 | "input_zero_point", |
| 808 | "azp_adj", |
| 809 | ], |
| 810 | ) |
| 811 | |
| 812 | |
| 813 | def choose_mp_linear_kernel( |
| 814 | config: MPLinearLayerConfig, compute_capability: int | None = None |
| 815 | ) -> type[MPLinearKernel]: |
| 816 | """Choose an MPLinearKernel that can implement the given config for the given |
| 817 | compute capability. Attempts to choose the best kernel in terms of |
| 818 | performance. |
| 819 | |
| 820 | Args: |
| 821 | config (MPLinearLayerConfig): Description of the linear layer to be |
| 822 | implemented. |
| 823 | compute_capability (Optional[int], optional): The compute capability of |
| 824 | the target device, if None uses `current_platform` to get |
| 825 | the compute capability. Defaults to None. |
| 826 | |
| 827 | Raises: |
| 828 | ValueError: If no kernel can implement the given config. |
| 829 | |
| 830 | Returns: |
| 831 | type[MPLinearKernel]: Chosen kernel. |
| 832 | |
| 833 | """ |
| 834 | if compute_capability is None: |
| 835 | if current_platform is None: |
| 836 | raise ValueError("Cannot determine compute capability") |
| 837 | _cc = current_platform.get_device_capability() |
| 838 | if _cc is not None: |
| 839 | compute_capability = _cc[0] * 10 + _cc[1] |
| 840 | |
| 841 | platform_kernels = _POSSIBLE_KERNELS.get(current_platform._enum, []) |
| 842 | |
| 843 | # Apply --linear-backend filtering when set. |
| 844 | platform_kernels = _resolve_backend_kernels( |
| 845 | platform_kernels, |
| 846 | "mixed-precision", |
| 847 | quantization="mixed_precision", |
| 848 | ) |
| 849 | |
| 850 | failure_reasons = [] |
| 851 | for kernel in platform_kernels: |
| 852 | if kernel.__name__ in envs.VLLM_DISABLED_KERNELS: |
| 853 | failure_reasons.append( |
| 854 | f" {kernel.__name__} disabled by environment variable" |
| 855 | ) |
| 856 | continue |
| 857 | if ( |
| 858 | compute_capability is not None |
| 859 | and kernel.get_min_capability() > compute_capability |
| 860 | ): |
| 861 | failure_reasons.append( |
| 862 | f"{kernel.__name__} requires capability " |
| 863 | f"{kernel.get_min_capability()}, current compute " |
| 864 | f" capability is {compute_capability}" |
| 865 | ) |
| 866 | continue |
| 867 | |
| 868 | can_implement, failure_reason = kernel.can_implement(config) |
| 869 | if can_implement: |
| 870 | logger.info_once("Using %s for mixed-precision linear", kernel.__name__) |
| 871 | return kernel |
| 872 | else: |
| 873 | failure_reasons.append( |
| 874 | f" {kernel.__name__} cannot implement due to: {failure_reason}" |
| 875 | ) |
| 876 | |
| 877 | raise ValueError( |
| 878 | "Failed to find a kernel that can implement the " |
| 879 | "WNA16 linear layer. Reasons: \n" + "\n".join(failure_reasons) |
| 880 | ) |
| 881 | |
| 882 | |
| 883 | def init_mxfp8_linear_kernel( |
| 884 | *, weight_shape: tuple[int, int], bmm_batch_size: int | None = None |
| 885 | ) -> Mxfp8LinearKernel: |
| 886 | """Select and instantiate the best MXFP8 linear kernel for the |
| 887 | current platform and `(N, K)` weight shape.""" |
| 888 | config = Mxfp8LinearLayerConfig( |
| 889 | weight_shape=weight_shape, |
| 890 | bmm_batch_size=bmm_batch_size, |
| 891 | ) |
| 892 | |
| 893 | platform = current_platform._enum |
| 894 | possible: list[type[Mxfp8LinearKernel]] |
| 895 | if bmm_batch_size is not None: |
| 896 | possible = ( |
| 897 | [DeepGemmMxfp8BmmLinearKernel, EmulationMxfp8LinearKernel] |
| 898 | if current_platform.is_cuda() |
| 899 | else [] |
| 900 | ) |
| 901 | else: |
| 902 | possible = list(_POSSIBLE_MXFP8_KERNELS.get(platform, [])) |
| 903 | |
| 904 | # Apply --linear-backend filtering when set. |
| 905 | possible = _resolve_backend_kernels( |
| 906 | possible, |
| 907 | "MXFP8", |
| 908 | quantization="mxfp8", |
| 909 | ) |
| 910 | |
| 911 | failure_reasons = [] |
| 912 | for kernel_cls in possible: |
| 913 | if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS: |
| 914 | failure_reasons.append( |
| 915 | f" {kernel_cls.__name__} disabled by environment variable" |
| 916 | ) |
| 917 | continue |
| 918 | |
| 919 | is_supported, reason = kernel_cls.is_supported() |
| 920 | if not is_supported: |
| 921 | failure_reasons.append(f"{kernel_cls.__name__}: {reason}") |
| 922 | continue |
| 923 | |
| 924 | can_implement, reason = kernel_cls.can_implement(config) |
| 925 | if not can_implement: |
| 926 | failure_reasons.append(f"{kernel_cls.__name__}: {reason}") |
| 927 | continue |
| 928 | |
| 929 | logger.info_once("Using %s for MXFP8 GEMM", kernel_cls.__name__) |
| 930 | return kernel_cls(config) |
| 931 | |
| 932 | raise ValueError( |
| 933 | "Failed to find a kernel that can implement the " |
| 934 | "MXFP8 linear layer. Reasons: \n" + "\n".join(failure_reasons) |
| 935 | ) |
| 936 | |
| 937 | |
| 938 | def init_mxfp4_linear_kernel( |
| 939 | activation_quant_key: QuantKey | None = None, |
| 940 | ) -> MxFp4LinearKernel: |
| 941 | """Select and instantiate the best MXFP4 linear kernel for the |
| 942 | current platform.""" |
| 943 | config = MxFp4LinearLayerConfig( |
| 944 | activation_quant_key=activation_quant_key, |
| 945 | ) |
| 946 | |
| 947 | platform = current_platform._enum |
| 948 | possible = list(_POSSIBLE_MXFP4_KERNELS.get(platform, [])) |
| 949 | |
| 950 | # Apply --linear-backend filtering when set. |
| 951 | possible = _resolve_backend_kernels( |
| 952 | possible, |
| 953 | "MXFP4", |
| 954 | quantization="mxfp4", |
| 955 | ) |
| 956 | |
| 957 | failure_reasons = [] |
| 958 | for kernel_cls in possible: |
| 959 | if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS: |
| 960 | failure_reasons.append( |
| 961 | f" {kernel_cls.__name__} disabled by environment variable" |
| 962 | ) |
| 963 | continue |
| 964 | |
| 965 | is_supported, reason = kernel_cls.is_supported() |
| 966 | if not is_supported: |
| 967 | failure_reasons.append(f"{kernel_cls.__name__}: {reason}") |
| 968 | continue |
| 969 | |
| 970 | can_implement, reason = kernel_cls.can_implement(config) |
| 971 | if not can_implement: |
| 972 | failure_reasons.append(f"{kernel_cls.__name__}: {reason}") |
| 973 | continue |
| 974 | |
| 975 | logger.info_once("Using %s for MXFP4 GEMM", kernel_cls.__name__) |
| 976 | return kernel_cls(config) |
| 977 | |
| 978 | raise ValueError( |
| 979 | "Failed to find a kernel that can implement the " |
| 980 | "MXFP4 linear layer. Reasons: \n" + "\n".join(failure_reasons) |
| 981 | ) |
| 982 | |
| 983 | |
| 984 | def init_mxfp6_linear_kernel( |
| 985 | weight_quant_key: QuantKey, |
| 986 | activation_quant_key: QuantKey | None = None, |
| 987 | ) -> MxFp6LinearKernel: |
| 988 | """Select and instantiate the best MXFP6 linear kernel for the |
| 989 | current platform.""" |
| 990 | config = MxFp6LinearLayerConfig( |
| 991 | weight_quant_key=weight_quant_key, |
| 992 | activation_quant_key=activation_quant_key, |
| 993 | ) |
| 994 | |
| 995 | linear_backend = _get_linear_backend(quantization="mxfp6") |
| 996 | |
| 997 | platform = current_platform._enum |
| 998 | possible = list(_POSSIBLE_MXFP6_KERNELS.get(platform, [])) |
| 999 | |
| 1000 | # Apply --linear-backend filtering when set. |
| 1001 | if linear_backend != "auto": |
| 1002 | filtered = _filter_kernels_by_backend(linear_backend, possible) |
| 1003 | if not filtered: |
| 1004 | raise ValueError( |
| 1005 | f"--linear-backend={linear_backend} was requested but no " |
| 1006 | f"'{linear_backend}' kernel exists for MXFP6 layers." |
| 1007 | ) |
| 1008 | possible = filtered |
| 1009 | |
| 1010 | failure_reasons = [] |
| 1011 | for kernel_cls in possible: |
| 1012 | if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS: |
| 1013 | failure_reasons.append( |
| 1014 | f" {kernel_cls.__name__} disabled by environment variable" |
| 1015 | ) |
| 1016 | continue |
| 1017 | |
| 1018 | is_supported, reason = kernel_cls.is_supported() |
| 1019 | if not is_supported: |
| 1020 | failure_reasons.append(f"{kernel_cls.__name__}: {reason}") |
| 1021 | continue |
| 1022 | |
| 1023 | can_implement, reason = kernel_cls.can_implement(config) |
| 1024 | if not can_implement: |
| 1025 | failure_reasons.append(f"{kernel_cls.__name__}: {reason}") |
| 1026 | continue |
| 1027 | |
| 1028 | logger.info_once("Using %s for MXFP6 GEMM", kernel_cls.__name__) |
| 1029 | return kernel_cls(config) |
| 1030 | |
| 1031 | raise ValueError( |
| 1032 | "Failed to find a kernel that can implement the " |
| 1033 | "MXFP6 linear layer. Reasons: \n" + "\n".join(failure_reasons) |
| 1034 | ) |
| 1035 | |
| 1036 | |
| 1037 | def init_wfp8_a16_linear_kernel( |
| 1038 | weight_quant_key: QuantKey, |
| 1039 | activation_quant_key: QuantKey, |
| 1040 | weight_shape: tuple[int, int], |
| 1041 | input_dtype: torch.dtype, |
| 1042 | out_dtype: torch.dtype, |
| 1043 | force_kernel: type[FP8ScaledMMLinearKernel] | None = None, |
| 1044 | module_name: str | None = None, |
| 1045 | ) -> FP8ScaledMMLinearKernel: |
| 1046 | config = FP8ScaledMMLinearLayerConfig( |
| 1047 | weight_quant_key=weight_quant_key, |
| 1048 | activation_quant_key=activation_quant_key, |
| 1049 | weight_shape=weight_shape, |
| 1050 | input_dtype=input_dtype, |
| 1051 | out_dtype=out_dtype, |
| 1052 | ) |
| 1053 | |
| 1054 | kernel_type = choose_scaled_mm_linear_kernel( |
| 1055 | config, |
| 1056 | _POSSIBLE_WFP8A16_KERNELS, |
| 1057 | quantization="w8a16_fp8", |
| 1058 | force_kernel=force_kernel, |
| 1059 | ) |
| 1060 | |
| 1061 | if module_name: |
| 1062 | logger.info_once( |
| 1063 | "Selected %s for %s", |
| 1064 | kernel_type.__name__, |
| 1065 | module_name, |
| 1066 | scope="global", |
| 1067 | ) |
| 1068 | |
| 1069 | return kernel_type( |
| 1070 | config, |
| 1071 | layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"], |
| 1072 | ) |
| 1073 | |
| 1074 | |
| 1075 | def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: |
| 1076 | """Select and instantiate the best NVFP4 linear kernel for the |
| 1077 | current platform.""" |
| 1078 | config = NvFp4LinearLayerConfig() |
| 1079 | a16_kernels = ( |
| 1080 | FlashInferCuteDslNvFp4W4A16LinearKernel, |
| 1081 | MarlinNvFp4LinearKernel, |
| 1082 | HummingNvFp4LinearKernel, |
| 1083 | ) |
| 1084 | |
| 1085 | # VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the |
| 1086 | # batch-invariant CUTLASS implementation when available, otherwise fall |
| 1087 | # back to emulation. It overrides --linear-backend. |
| 1088 | force_kernel: type[NvFp4LinearKernel] | None = None |
| 1089 | quantization = "nvfp4_w4a16" if use_a16 else "nvfp4_w4a4" |
| 1090 | linear_backend = _get_linear_backend(quantization=quantization) |
| 1091 | if envs.VLLM_BATCH_INVARIANT: |
| 1092 | bi_supported, reason = CutlassNvFp4LinearKernel.is_supported() |
| 1093 | if bi_supported: |
| 1094 | if linear_backend not in ("auto", "cutlass"): |
| 1095 | logger.warning_once( |
| 1096 | "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; " |
| 1097 | "using the CUTLASS backend for deterministic execution.", |
| 1098 | linear_backend, |
| 1099 | ) |
| 1100 | else: |
| 1101 | logger.info_once( |
| 1102 | "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the " |
| 1103 | "CUTLASS backend for deterministic execution." |
| 1104 | ) |
| 1105 | force_kernel = CutlassNvFp4LinearKernel |
| 1106 | else: |
| 1107 | if linear_backend not in ("auto", "emulation"): |
| 1108 | logger.warning_once( |
| 1109 | "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; " |
| 1110 | "using the emulation backend for deterministic execution.", |
| 1111 | linear_backend, |
| 1112 | ) |
| 1113 | logger.info_once( |
| 1114 | "VLLM_BATCH_INVARIANT is set but the batch-invariant NVFP4 " |
| 1115 | "kernel is not supported on this platform; falling back to " |
| 1116 | "emulation for deterministic execution. Reason: %s", |
| 1117 | reason, |
| 1118 | ) |
| 1119 | force_kernel = EmulationNvFp4LinearKernel |
| 1120 | elif linear_backend == "auto" and use_a16: |
| 1121 | _cc = current_platform.get_device_capability() |
| 1122 | compute_capability = _cc.to_int() if _cc is not None else None |
| 1123 | # Weight-only: prefer FlashInfer CuTe-DSL W4A16 on SM100/103, |
| 1124 | # otherwise Marlin. |
| 1125 | cutedsl_ok, _ = FlashInferCuteDslNvFp4W4A16LinearKernel.is_supported( |
| 1126 | compute_capability |
| 1127 | ) |
| 1128 | if compute_capability in (100, 103) and cutedsl_ok: |
| 1129 | force_kernel = FlashInferCuteDslNvFp4W4A16LinearKernel |
| 1130 | else: |
| 1131 | force_kernel = MarlinNvFp4LinearKernel |
| 1132 | |
| 1133 | if force_kernel is not None: |
| 1134 | if use_a16 and force_kernel not in a16_kernels: |
| 1135 | raise ValueError(f"{force_kernel.__name__} does not support W4A16") |
| 1136 | is_supported, reason = force_kernel.is_supported() |
| 1137 | if not is_supported: |
| 1138 | raise ValueError( |
| 1139 | f"Forced NVFP4 kernel {force_kernel.__name__} is not " |
| 1140 | f"supported: {reason}" |
| 1141 | ) |
| 1142 | logger.info_once("Using %s for NVFP4 GEMM", force_kernel.__name__) |
| 1143 | return force_kernel(config) |
| 1144 | |
| 1145 | # Auto-select from registry (or --linear-backend filtered). |
| 1146 | platform = current_platform._enum |
| 1147 | possible = list(_POSSIBLE_NVFP4_KERNELS.get(platform, [])) |
| 1148 | if use_a16: |
| 1149 | possible = [kernel for kernel in possible if kernel in a16_kernels] |
| 1150 | |
| 1151 | # Apply --linear-backend filtering when set. |
| 1152 | possible = _resolve_backend_kernels( |
| 1153 | possible, |
| 1154 | "NVFP4", |
| 1155 | quantization=quantization, |
| 1156 | ) |
| 1157 | |
| 1158 | failure_reasons = [] |
| 1159 | for kernel_cls in possible: |
| 1160 | if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS: |
| 1161 | failure_reasons.append( |
| 1162 | f" {kernel_cls.__name__} disabled by environment variable" |
| 1163 | ) |
| 1164 | continue |
| 1165 | |
| 1166 | is_supported, reason = kernel_cls.is_supported() |
| 1167 | if not is_supported: |
| 1168 | failure_reasons.append(f"{kernel_cls.__name__}: {reason}") |
| 1169 | continue |
| 1170 | |
| 1171 | can_implement, reason = kernel_cls.can_implement(config) |
| 1172 | if not can_implement: |
| 1173 | failure_reasons.append(f"{kernel_cls.__name__}: {reason}") |
| 1174 | continue |
| 1175 | |
| 1176 | if kernel_cls is EmulationNvFp4LinearKernel and failure_reasons: |
| 1177 | logger.warning_once( |
| 1178 | "NVFP4 linear falling back to the slow and unoptimized " |
| 1179 | "emulation backend as no optimized backend is available " |
| 1180 | "(unavailable reasons:\n - %s\n). " |
| 1181 | "In case you expect one of these backends to be used, " |
| 1182 | "please verify your environment.", |
| 1183 | "\n - ".join(failure_reasons), |
| 1184 | ) |
| 1185 | |
| 1186 | logger.info_once("Using %s for NVFP4 GEMM", kernel_cls.__name__) |
| 1187 | return kernel_cls(config) |
| 1188 | |
| 1189 | raise ValueError( |
| 1190 | "Failed to find a kernel that can implement the " |
| 1191 | "NVFP4 linear layer. Reasons: \n" + "\n".join(failure_reasons) |
| 1192 | ) |
| 1193 | |
| 1194 | |
| 1195 | def register_linear_kernel( |
| 1196 | kernel_class: type, |
| 1197 | platform: PlatformEnum, |
| 1198 | kernel_type: str = "mp", |
| 1199 | ) -> None: |
| 1200 | """Register a new linear kernel class to be considered in kernel selection. |
| 1201 | |
| 1202 | Args: |
| 1203 | kernel_class (type): The kernel class to register. |
| 1204 | platform (PlatformEnum): The platform for which this kernel is applicable. |
| 1205 | kernel_type (str): The type of the kernel, either "mp", "int8", or "fp8". |
| 1206 | Defaults to "mp". |
| 1207 | |
| 1208 | Raises: |
| 1209 | ValueError: If the kernel_type is not recognized. |
| 1210 | |
| 1211 | """ |
| 1212 | if kernel_type == "mp": |
| 1213 | if platform not in _POSSIBLE_KERNELS: |
| 1214 | _POSSIBLE_KERNELS[platform] = [] |
| 1215 | _POSSIBLE_KERNELS[platform].append(kernel_class) |
| 1216 | elif kernel_type == "int8": |
| 1217 | if platform not in _POSSIBLE_INT8_KERNELS: |
| 1218 | _POSSIBLE_INT8_KERNELS[platform] = [] |
| 1219 | _POSSIBLE_INT8_KERNELS[platform].append(kernel_class) |
| 1220 | elif kernel_type == "fp8": |
| 1221 | if platform not in _POSSIBLE_FP8_KERNELS: |
| 1222 | _POSSIBLE_FP8_KERNELS[platform] = [] |
| 1223 | _POSSIBLE_FP8_KERNELS[platform].append(kernel_class) |
| 1224 | elif kernel_type == "mxfp8": |
| 1225 | if platform not in _POSSIBLE_MXFP8_KERNELS: |
| 1226 | _POSSIBLE_MXFP8_KERNELS[platform] = [] |
| 1227 | _POSSIBLE_MXFP8_KERNELS[platform].append(kernel_class) |
| 1228 | elif kernel_type == "nvfp4": |
| 1229 | if platform not in _POSSIBLE_NVFP4_KERNELS: |
| 1230 | _POSSIBLE_NVFP4_KERNELS[platform] = [] |
| 1231 | _POSSIBLE_NVFP4_KERNELS[platform].append(kernel_class) |
| 1232 | elif kernel_type == "mxfp4": |
| 1233 | if platform not in _POSSIBLE_MXFP4_KERNELS: |
| 1234 | _POSSIBLE_MXFP4_KERNELS[platform] = [] |
| 1235 | _POSSIBLE_MXFP4_KERNELS[platform].append(kernel_class) |
| 1236 | elif kernel_type == "mxfp6": |
| 1237 | if platform not in _POSSIBLE_MXFP6_KERNELS: |
| 1238 | _POSSIBLE_MXFP6_KERNELS[platform] = [] |
| 1239 | _POSSIBLE_MXFP6_KERNELS[platform].append(kernel_class) |
| 1240 | else: |
| 1241 | raise ValueError(f"Unrecognized kernel type: {kernel_type}") |
| 1242 | |
| 1243 | |
| 1244 | __all__ = [ |
| 1245 | "init_fp8_linear_kernel", |
| 1246 | "init_int8_linear_kernel", |
| 1247 | "init_nvfp4_linear_kernel", |
| 1248 | "choose_mp_linear_kernel", |
| 1249 | "register_linear_kernel", |
| 1250 | "init_wfp8_a16_linear_kernel", |
| 1251 | "FP8ScaledMMLinearKernel", |
| 1252 | "Int8ScaledMMLinearKernel", |
| 1253 | "ScaledMMLinearKernel", |
| 1254 | "FP8ScaledMMLinearLayerConfig", |
| 1255 | "Int8ScaledMMLinearLayerConfig", |
| 1256 | "ScaledMMLinearLayerConfig", |
| 1257 | "AiterHipbMMPerTokenFp8ScaledMMLinearKernel", |
| 1258 | "AiterPreshuffledPerTokenFp8ScaledMMLinearKernel", |
| 1259 | "AiterPerTokenFp8ScaledMMLinearKernel", |
| 1260 | "NvFp4LinearKernel", |
| 1261 | "NvFp4LinearLayerConfig", |
| 1262 | "AiterInt8ScaledMMLinearKernel", |
| 1263 | "CPUInt8ScaledMMLinearKernel", |
| 1264 | "CutlassFP8ScaledMMLinearKernel", |
| 1265 | "CutlassInt8ScaledMMLinearKernel", |
| 1266 | "FlashInferFP8ScaledMMLinearKernel", |
| 1267 | "ChannelWiseTorchFP8ScaledMMLinearKernel", |
| 1268 | "PerTensorTorchFP8ScaledMMLinearKernel", |
| 1269 | "RowWiseTorchFP8ScaledMMLinearKernel", |
| 1270 | "ROCmFP8ScaledMMLinearKernel", |
| 1271 | "TritonInt8ScaledMMLinearKernel", |
| 1272 | "ZentorchInt8ScaledMMLinearKernel", |
| 1273 | "ZentorchWNA16LinearKernel", |
| 1274 | "MPLinearKernel", |
| 1275 | "MPLinearLayerConfig", |
| 1276 | "ConchLinearKernel", |
| 1277 | "CPUWNA16LinearKernel", |
| 1278 | "CutlassW4A8LinearKernel", |
| 1279 | "Dynamic4bitLinearKernel", |
| 1280 | "ExllamaLinearKernel", |
| 1281 | "RDNAHybridW4A16LinearKernel", |
| 1282 | "MacheteLinearKernel", |
| 1283 | "MarlinLinearKernel", |
| 1284 | "TritonW4A16LinearKernel", |
| 1285 | "XPUW4A8IntLinearKernel", |
| 1286 | "XPUwNa16LinearKernel", |
| 1287 | "init_mxfp8_linear_kernel", |
| 1288 | "Mxfp8LinearKernel", |
| 1289 | "Mxfp8LinearLayerConfig", |
| 1290 | "B12xMxfp8LinearKernel", |
| 1291 | "B12xMxFp4LinearKernel", |
| 1292 | "B12xNvFp4LinearKernel", |
| 1293 | "init_mxfp4_linear_kernel", |
| 1294 | "MxFp4LinearKernel", |
| 1295 | "MxFp4LinearLayerConfig", |
| 1296 | "MxFp6LinearKernel", |
| 1297 | "MxFp6LinearLayerConfig", |
| 1298 | "init_mxfp6_linear_kernel", |
| 1299 | "EmulationMxfp6LinearKernel", |
| 1300 | "HummingMxFp6LinearKernel", |
| 1301 | "AiterMxfp4LinearKernel", |
| 1302 | "EmulationMxfp4LinearKernel", |
| 1303 | "FlashInferMxFp4LinearKernel", |
| 1304 | "MarlinMxFp4LinearKernel", |
| 1305 | "FlashInferCutedslMxfp8LinearKernel", |
| 1306 | "FlashInferCutlassMxfp8LinearKernel", |
| 1307 | "FlashInferTrtllmMxfp8LinearKernel", |
| 1308 | "MarlinMxfp8LinearKernel", |
| 1309 | "XPUMxFp8LinearKernel", |
| 1310 | "EmulationMxfp8LinearKernel", |
| 1311 | "CutlassNvFp4LinearKernel", |
| 1312 | "EmulationNvFp4LinearKernel", |
| 1313 | "FbgemmNvFp4LinearKernel", |
| 1314 | "FlashInferCuteDslNvFp4LinearKernel", |
| 1315 | "FlashInferCuteDslNvFp4W4A16LinearKernel", |
| 1316 | "FlashInferB12xNvFp4LinearKernel", |
| 1317 | "FlashInferCutlassNvFp4LinearKernel", |
| 1318 | "FlashInferTrtllmNvFp4LinearKernel", |
| 1319 | "FlashInferCudnnNvFp4LinearKernel", |
| 1320 | "MarlinNvFp4LinearKernel", |
| 1321 | "TorchNvFp4LinearKernel", |
| 1322 | "_KernelT", |
| 1323 | "DeepGemmFp8BlockScaledMMKernel", |
| 1324 | "FlashInferFp8DeepGEMMDynamicBlockScaledKernel", |
| 1325 | "B12xFp8BlockScaledMMKernel", |
| 1326 | "B12xTensorFP8ScaledMMLinearKernel", |
| 1327 | ] |