wehub-resource-sync
1097 lines
36 KiB
Python
1097 lines
36 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
This module re-exports linear kernel implementations to provide a
|
|
stable import interface during an ongoing reorganization. Upcoming
|
|
PRs will remove the scaled_mm and mixed_precision subdirectories
|
|
and reorganize kernels by provider (aiter, cutlass, flashinfer, etc.)
|
|
rather than by precision type. By centralizing exports here, we
|
|
minimize the need to update imports across other modules when the
|
|
internal structure changes. If you are adding a new kernel selector
|
|
or kernel implementation, add it to this __init__.py to maintain
|
|
import stability.
|
|
"""
|
|
|
|
from typing import TypeVar
|
|
|
|
import torch
|
|
|
|
import vllm.envs as envs
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.kernels.linear.base import (
|
|
MMLinearKernel,
|
|
MMLinearLayerConfig,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision import (
|
|
MPLinearKernel,
|
|
MPLinearLayerConfig,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.allspark import (
|
|
AllSparkLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.conch import (
|
|
ConchLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.cpu import (
|
|
CPUWNA16LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.cutlass import (
|
|
CutlassW4A8LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.dynamic_4bit import (
|
|
Dynamic4bitLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.exllama import (
|
|
ExllamaLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.humming import (
|
|
HummingLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.machete import (
|
|
MacheteLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
|
|
MarlinLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import (
|
|
RDNA3W4A16LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import (
|
|
TritonW4A16LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
|
|
XPUW4A8IntLinearKernel,
|
|
XPUwNa16LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
|
|
ZentorchWNA16LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp4 import (
|
|
MxFp4LinearKernel,
|
|
MxFp4LinearLayerConfig,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp4.flashinfer import (
|
|
FlashInferMxFp4LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp4.humming import (
|
|
HummingMxFp4LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp4.marlin import (
|
|
MarlinMxFp4LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp4.xpu import (
|
|
XPUMxFp4LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp8 import (
|
|
Mxfp8LinearKernel,
|
|
Mxfp8LinearLayerConfig,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp8.emulation import (
|
|
EmulationMxfp8LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
|
|
FlashInferCutedslMxfp8LinearKernel,
|
|
FlashInferCutlassMxfp8LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp8.humming import (
|
|
HummingMxfp8LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp8.marlin import (
|
|
MarlinMxfp8LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
|
|
RocmDotScaledMxfp8LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mxfp8.xpu import (
|
|
XPUMxFp8LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.nvfp4 import (
|
|
NvFp4LinearKernel,
|
|
NvFp4LinearLayerConfig,
|
|
)
|
|
from vllm.model_executor.kernels.linear.nvfp4.cutlass import (
|
|
CutlassNvFp4LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.nvfp4.emulation import (
|
|
EmulationNvFp4LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.nvfp4.fbgemm import (
|
|
FbgemmNvFp4LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
|
|
FlashInferB12xNvFp4LinearKernel,
|
|
FlashInferCudnnNvFp4LinearKernel,
|
|
FlashInferCuteDslNvFp4LinearKernel,
|
|
FlashInferCutlassNvFp4LinearKernel,
|
|
FlashInferTrtllmNvFp4LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.nvfp4.humming import (
|
|
HummingNvFp4LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.nvfp4.marlin import (
|
|
MarlinNvFp4LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm import (
|
|
Fp8BlockScaledMMLinearKernel,
|
|
FP8ScaledMMLinearKernel,
|
|
FP8ScaledMMLinearLayerConfig,
|
|
Int8ScaledMMLinearKernel,
|
|
Int8ScaledMMLinearLayerConfig,
|
|
ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
|
|
AiterFp8BlockScaledMMKernel,
|
|
AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
|
|
AiterInt8ScaledMMLinearKernel,
|
|
AiterPerTokenFp8ScaledMMLinearKernel,
|
|
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
|
|
CPUFp8BlockScaledMMKernel,
|
|
CPUInt8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
|
|
CutlassFp8BlockScaledMMKernel,
|
|
CutlassFP8ScaledMMLinearKernel,
|
|
CutlassInt8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import (
|
|
DeepGemmFp8BlockScaledMMKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
|
|
FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
|
|
FlashInferFP8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.humming import (
|
|
HummingFP8ScaledMMLinearKernel,
|
|
HummingInt8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.marlin import (
|
|
MarlinFP8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.pytorch import (
|
|
ChannelWiseTorchFP8ScaledMMLinearKernel,
|
|
PerTensorTorchFP8ScaledMMLinearKernel,
|
|
RowWiseTorchFP8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.rocm import (
|
|
ROCmFP8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.triton import (
|
|
TritonFp8BlockScaledMMKernel,
|
|
TritonInt8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
|
|
XPUFp8BlockScaledMMKernel,
|
|
XPUW8A8FP8LinearKernel,
|
|
XPUW8A16FP8LinearKernel,
|
|
)
|
|
from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
|
|
ZentorchInt8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
|
|
from vllm.platforms import PlatformEnum, current_platform
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _get_linear_backend() -> str:
|
|
"""Get the linear_backend setting from the current vllm config."""
|
|
from vllm.config import get_current_vllm_config_or_none
|
|
|
|
config = get_current_vllm_config_or_none()
|
|
if config is not None:
|
|
return config.kernel_config.linear_backend
|
|
return "auto"
|
|
|
|
|
|
# Mapping from linear_backend name to the set of kernel classes it covers.
|
|
# When a user sets --linear-backend <name>, only kernels in the corresponding
|
|
# set are considered candidates. If none can implement the layer config,
|
|
# an error is raised to respect the user's explicit intent.
|
|
_LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
|
|
"cutlass": {
|
|
CutlassInt8ScaledMMLinearKernel,
|
|
CutlassFP8ScaledMMLinearKernel,
|
|
CutlassFp8BlockScaledMMKernel,
|
|
CutlassW4A8LinearKernel,
|
|
CutlassNvFp4LinearKernel,
|
|
},
|
|
"flashinfer_cutlass": {
|
|
FlashInferFP8ScaledMMLinearKernel,
|
|
FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
|
|
FlashInferCutlassMxfp8LinearKernel,
|
|
FlashInferCutlassNvFp4LinearKernel,
|
|
FlashInferMxFp4LinearKernel,
|
|
},
|
|
"flashinfer_cutedsl": {
|
|
FlashInferCuteDslNvFp4LinearKernel,
|
|
FlashInferCutedslMxfp8LinearKernel,
|
|
},
|
|
"flashinfer_trtllm": {
|
|
FlashInferTrtllmNvFp4LinearKernel,
|
|
},
|
|
"flashinfer_cudnn": {
|
|
FlashInferCudnnNvFp4LinearKernel,
|
|
},
|
|
"flashinfer_b12x": {
|
|
FlashInferB12xNvFp4LinearKernel,
|
|
},
|
|
"humming": {
|
|
HummingFP8ScaledMMLinearKernel,
|
|
HummingInt8ScaledMMLinearKernel,
|
|
HummingLinearKernel,
|
|
HummingMxfp8LinearKernel,
|
|
HummingMxFp4LinearKernel,
|
|
HummingNvFp4LinearKernel,
|
|
},
|
|
"marlin": {
|
|
MarlinFP8ScaledMMLinearKernel,
|
|
MarlinLinearKernel,
|
|
MarlinMxfp8LinearKernel,
|
|
MarlinNvFp4LinearKernel,
|
|
MarlinMxFp4LinearKernel,
|
|
},
|
|
"triton": {
|
|
TritonInt8ScaledMMLinearKernel,
|
|
TritonFp8BlockScaledMMKernel,
|
|
TritonW4A16LinearKernel,
|
|
},
|
|
"deep_gemm": {
|
|
DeepGemmFp8BlockScaledMMKernel,
|
|
},
|
|
"torch": {
|
|
PerTensorTorchFP8ScaledMMLinearKernel,
|
|
ChannelWiseTorchFP8ScaledMMLinearKernel,
|
|
RowWiseTorchFP8ScaledMMLinearKernel,
|
|
},
|
|
"aiter": {
|
|
AiterInt8ScaledMMLinearKernel,
|
|
AiterFp8BlockScaledMMKernel,
|
|
AiterPerTokenFp8ScaledMMLinearKernel,
|
|
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
|
|
},
|
|
"machete": {
|
|
MacheteLinearKernel,
|
|
},
|
|
"fbgemm": {
|
|
FbgemmNvFp4LinearKernel,
|
|
},
|
|
"conch": {
|
|
ConchLinearKernel,
|
|
},
|
|
"exllama": {
|
|
ExllamaLinearKernel,
|
|
},
|
|
"emulation": {
|
|
EmulationMxfp8LinearKernel,
|
|
EmulationNvFp4LinearKernel,
|
|
},
|
|
"xpu": {
|
|
XPUW8A8FP8LinearKernel,
|
|
XPUFp8BlockScaledMMKernel,
|
|
},
|
|
"xpu_woq": {
|
|
XPUW8A16FP8LinearKernel,
|
|
},
|
|
}
|
|
|
|
|
|
def _filter_kernels_by_backend(
|
|
backend: str,
|
|
kernels: list[type],
|
|
) -> list[type]:
|
|
"""Filter a kernel priority list to only those matching the backend."""
|
|
backend_kernels = _LINEAR_BACKEND_KERNEL_MAP.get(backend, set())
|
|
return [k for k in kernels if k in backend_kernels]
|
|
|
|
|
|
# in priority/performance order (when available)
|
|
_POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] = {
|
|
PlatformEnum.CPU: [ZentorchInt8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel],
|
|
PlatformEnum.CUDA: [
|
|
CutlassInt8ScaledMMLinearKernel,
|
|
TritonInt8ScaledMMLinearKernel,
|
|
HummingInt8ScaledMMLinearKernel,
|
|
],
|
|
PlatformEnum.ROCM: [AiterInt8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel],
|
|
}
|
|
|
|
# in priority/performance order (when available)
|
|
_POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = {
|
|
PlatformEnum.CUDA: [
|
|
MarlinFP8ScaledMMLinearKernel,
|
|
FlashInferFP8ScaledMMLinearKernel,
|
|
CutlassFP8ScaledMMLinearKernel,
|
|
PerTensorTorchFP8ScaledMMLinearKernel,
|
|
ChannelWiseTorchFP8ScaledMMLinearKernel,
|
|
HummingFP8ScaledMMLinearKernel,
|
|
],
|
|
PlatformEnum.ROCM: [
|
|
AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
|
|
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
|
|
AiterPerTokenFp8ScaledMMLinearKernel,
|
|
ROCmFP8ScaledMMLinearKernel,
|
|
PerTensorTorchFP8ScaledMMLinearKernel,
|
|
RowWiseTorchFP8ScaledMMLinearKernel,
|
|
ChannelWiseTorchFP8ScaledMMLinearKernel,
|
|
],
|
|
PlatformEnum.CPU: [
|
|
PerTensorTorchFP8ScaledMMLinearKernel,
|
|
ChannelWiseTorchFP8ScaledMMLinearKernel,
|
|
],
|
|
PlatformEnum.XPU: [
|
|
XPUW8A16FP8LinearKernel,
|
|
XPUW8A8FP8LinearKernel,
|
|
],
|
|
}
|
|
|
|
|
|
# in priority/performance order (when available)
|
|
_POSSIBLE_FP8_BLOCK_KERNELS: dict[
|
|
PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]]
|
|
] = {
|
|
PlatformEnum.CUDA: [
|
|
FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
|
|
DeepGemmFp8BlockScaledMMKernel,
|
|
CutlassFp8BlockScaledMMKernel,
|
|
MarlinFP8ScaledMMLinearKernel,
|
|
TritonFp8BlockScaledMMKernel,
|
|
HummingFP8ScaledMMLinearKernel,
|
|
],
|
|
PlatformEnum.ROCM: [
|
|
AiterFp8BlockScaledMMKernel,
|
|
TritonFp8BlockScaledMMKernel,
|
|
],
|
|
PlatformEnum.CPU: [
|
|
CPUFp8BlockScaledMMKernel,
|
|
],
|
|
PlatformEnum.XPU: [
|
|
XPUFp8BlockScaledMMKernel,
|
|
TritonFp8BlockScaledMMKernel,
|
|
],
|
|
}
|
|
|
|
_POSSIBLE_WFP8A16_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = {
|
|
PlatformEnum.CUDA: [
|
|
HummingFP8ScaledMMLinearKernel,
|
|
MarlinFP8ScaledMMLinearKernel,
|
|
],
|
|
PlatformEnum.ROCM: [
|
|
# To be added
|
|
],
|
|
PlatformEnum.CPU: [
|
|
# To be added
|
|
],
|
|
PlatformEnum.XPU: [
|
|
XPUW8A16FP8LinearKernel,
|
|
],
|
|
}
|
|
|
|
# in priority/performance order (when available)
|
|
_POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
|
|
PlatformEnum.CUDA: [
|
|
CutlassW4A8LinearKernel,
|
|
MacheteLinearKernel,
|
|
AllSparkLinearKernel,
|
|
MarlinLinearKernel,
|
|
ConchLinearKernel,
|
|
ExllamaLinearKernel,
|
|
TritonW4A16LinearKernel,
|
|
HummingLinearKernel,
|
|
],
|
|
PlatformEnum.ROCM: [
|
|
RDNA3W4A16LinearKernel,
|
|
TritonW4A16LinearKernel,
|
|
ConchLinearKernel,
|
|
ExllamaLinearKernel,
|
|
],
|
|
PlatformEnum.XPU: [
|
|
XPUW4A8IntLinearKernel,
|
|
XPUwNa16LinearKernel,
|
|
],
|
|
PlatformEnum.CPU: [
|
|
Dynamic4bitLinearKernel,
|
|
ZentorchWNA16LinearKernel,
|
|
CPUWNA16LinearKernel,
|
|
],
|
|
}
|
|
|
|
# in priority/performance order (when available)
|
|
_POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
|
|
PlatformEnum.CUDA: [
|
|
FlashInferCutedslMxfp8LinearKernel,
|
|
FlashInferCutlassMxfp8LinearKernel,
|
|
MarlinMxfp8LinearKernel,
|
|
EmulationMxfp8LinearKernel,
|
|
HummingMxfp8LinearKernel,
|
|
],
|
|
PlatformEnum.ROCM: [
|
|
# Native CDNA4 (gfx950) MX linear; is_supported() gates to gfx95x and
|
|
# falls through to BF16 emulation (hipBLASLt) elsewhere / on regression.
|
|
RocmDotScaledMxfp8LinearKernel,
|
|
EmulationMxfp8LinearKernel,
|
|
],
|
|
PlatformEnum.XPU: [
|
|
XPUMxFp8LinearKernel,
|
|
EmulationMxfp8LinearKernel,
|
|
],
|
|
}
|
|
|
|
_POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
|
|
PlatformEnum.CUDA: [
|
|
FlashInferCuteDslNvFp4LinearKernel,
|
|
# FlashInferB12xNvFp4LinearKernel excluded from auto-selection until
|
|
# upstream CUTLASS SM121 MMA op guard is resolved; use
|
|
# --linear-backend flashinfer_b12x to opt in explicitly.
|
|
FlashInferCutlassNvFp4LinearKernel,
|
|
CutlassNvFp4LinearKernel,
|
|
MarlinNvFp4LinearKernel,
|
|
FlashInferTrtllmNvFp4LinearKernel,
|
|
FlashInferCudnnNvFp4LinearKernel,
|
|
FbgemmNvFp4LinearKernel,
|
|
EmulationNvFp4LinearKernel,
|
|
HummingNvFp4LinearKernel,
|
|
],
|
|
PlatformEnum.ROCM: [
|
|
EmulationNvFp4LinearKernel,
|
|
],
|
|
}
|
|
|
|
_POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = {
|
|
PlatformEnum.CUDA: [
|
|
FlashInferMxFp4LinearKernel,
|
|
MarlinMxFp4LinearKernel,
|
|
HummingMxFp4LinearKernel,
|
|
],
|
|
PlatformEnum.XPU: [
|
|
XPUMxFp4LinearKernel,
|
|
],
|
|
}
|
|
|
|
# TODO make all kernels inherit from MMLinearKernel
|
|
# then bound _KernelT only to MMLinearKernel
|
|
_KernelT = TypeVar("_KernelT", bound=ScaledMMLinearKernel | MMLinearKernel)
|
|
_KernelConfigT = TypeVar("_KernelConfigT", bound=MMLinearLayerConfig)
|
|
|
|
|
|
def is_supported_and_can_implement_kernel(
|
|
kernel: type[_KernelT], config: _KernelConfigT, compute_capability: int | None
|
|
) -> tuple[bool, str]:
|
|
if kernel.__name__ in envs.VLLM_DISABLED_KERNELS:
|
|
return False, f" {kernel.__name__} is disabled by environment variable"
|
|
|
|
if compute_capability is None:
|
|
_cc = current_platform.get_device_capability()
|
|
if _cc is not None:
|
|
compute_capability = _cc[0] * 10 + _cc[1]
|
|
|
|
is_supported, failure_reason = kernel.is_supported(compute_capability)
|
|
if not is_supported:
|
|
return False, f"{kernel.__name__} {failure_reason}."
|
|
|
|
can_implement, failure_reason = kernel.can_implement(config)
|
|
if not can_implement:
|
|
return (
|
|
False,
|
|
f"{kernel.__name__} {failure_reason}.",
|
|
)
|
|
|
|
return True, ""
|
|
|
|
|
|
def choose_scaled_mm_linear_kernel(
|
|
config: _KernelConfigT,
|
|
possible_kernels: dict[PlatformEnum, list[type[_KernelT]]],
|
|
compute_capability: int | None = None,
|
|
force_kernel: type[_KernelT] | None = None,
|
|
) -> type[_KernelT]:
|
|
"""
|
|
Choose a _KernelT that can implement the given config for the
|
|
given compute capability. Attempts to choose the best kernel in terms of
|
|
performance.
|
|
|
|
Args:
|
|
config (_KernelConfigT): Description of the linear layer
|
|
to be implemented.
|
|
possible_kernels (dict[PlatformEnum, list[_KernelT]]): A
|
|
dictionary of platforms and their list of possible kernels.
|
|
compute_capability (Optional[int], optional): The compute capability of
|
|
the target device, if None uses `current_platform` to get the
|
|
compute capability. Defaults to None.
|
|
force_kernel (Optional[type[_KernelT]]): An Optional forced kernel to override
|
|
the possible_kernels if it can be implemented. If None, it will only try the
|
|
possible kernels.
|
|
|
|
Raises:
|
|
ValueError: If no kernel can implement the given config.
|
|
|
|
Returns:
|
|
_KernelT: Chosen kernel.
|
|
"""
|
|
|
|
failure_reason_list = []
|
|
|
|
if force_kernel is not None:
|
|
can_implement, failure_reason = is_supported_and_can_implement_kernel(
|
|
force_kernel, config, compute_capability
|
|
)
|
|
if can_implement:
|
|
return force_kernel
|
|
|
|
logger.info_once(
|
|
"Tried to force %s, but the kernel couldn't be implemented",
|
|
force_kernel.__name__,
|
|
scope="global",
|
|
)
|
|
|
|
platform_kernels = possible_kernels[current_platform._enum]
|
|
|
|
# Apply --linear-backend filtering when set.
|
|
linear_backend = _get_linear_backend()
|
|
if linear_backend != "auto":
|
|
filtered = _filter_kernels_by_backend(linear_backend, platform_kernels)
|
|
if not filtered:
|
|
raise ValueError(
|
|
f"--linear-backend={linear_backend} was requested but no "
|
|
f"'{linear_backend}' kernel exists for this layer type."
|
|
)
|
|
platform_kernels = filtered
|
|
|
|
for kernel in platform_kernels:
|
|
is_supported_and_can_implement, failure_reason = (
|
|
is_supported_and_can_implement_kernel(kernel, config, compute_capability)
|
|
)
|
|
if is_supported_and_can_implement:
|
|
return kernel
|
|
failure_reason_list.append(failure_reason)
|
|
|
|
raise ValueError(
|
|
"Failed to find a kernel that can implement the "
|
|
"ScaledMM linear layer. Reasons: \n" + "\n".join(failure_reason_list)
|
|
)
|
|
|
|
|
|
def init_fp8_linear_kernel(
|
|
activation_quant_key: QuantKey,
|
|
weight_quant_key: QuantKey,
|
|
input_dtype: torch.dtype,
|
|
out_dtype: torch.dtype,
|
|
weight_shape: tuple[int, int],
|
|
force_kernel: type[FP8ScaledMMLinearKernel] | None = None,
|
|
module_name: str | None = None,
|
|
) -> FP8ScaledMMLinearKernel | Fp8BlockScaledMMLinearKernel:
|
|
scaled_mm_linear_kernel_config = FP8ScaledMMLinearLayerConfig(
|
|
weight_quant_key=weight_quant_key,
|
|
activation_quant_key=activation_quant_key,
|
|
input_dtype=input_dtype,
|
|
out_dtype=out_dtype,
|
|
weight_shape=weight_shape,
|
|
)
|
|
|
|
if activation_quant_key.scale.group_shape.is_per_group():
|
|
kernel_type = choose_scaled_mm_linear_kernel(
|
|
config=scaled_mm_linear_kernel_config,
|
|
possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS, # type: ignore[misc]
|
|
force_kernel=force_kernel,
|
|
)
|
|
if module_name:
|
|
logger.info_once(
|
|
"Selected %s for %s",
|
|
kernel_type.__name__,
|
|
module_name,
|
|
scope="global",
|
|
)
|
|
|
|
# TODO make scaled_mm kernels inherit from MMLinearKernel
|
|
# only MarlinFP8ScaledMMLinearKernel is a type of FP8ScaledMMLinearKernel.
|
|
if issubclass(kernel_type, FP8ScaledMMLinearKernel):
|
|
return kernel_type(
|
|
scaled_mm_linear_kernel_config,
|
|
layer_param_names=[
|
|
"weight",
|
|
"weight_scale",
|
|
"input_scale",
|
|
"input_scale_ub",
|
|
],
|
|
)
|
|
|
|
return kernel_type(
|
|
scaled_mm_linear_kernel_config,
|
|
)
|
|
|
|
else:
|
|
kernel_type = choose_scaled_mm_linear_kernel(
|
|
config=scaled_mm_linear_kernel_config,
|
|
possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type]
|
|
force_kernel=force_kernel,
|
|
)
|
|
if module_name:
|
|
logger.info_once(
|
|
"Selected %s for %s",
|
|
kernel_type.__name__,
|
|
module_name,
|
|
scope="global",
|
|
)
|
|
|
|
return kernel_type(
|
|
scaled_mm_linear_kernel_config,
|
|
layer_param_names=[
|
|
"weight",
|
|
"weight_scale",
|
|
"input_scale",
|
|
"input_scale_ub",
|
|
],
|
|
)
|
|
|
|
|
|
def init_int8_linear_kernel(
|
|
is_channelwise: bool,
|
|
is_static_input_scheme: bool,
|
|
input_symmetric: bool,
|
|
module_name: str,
|
|
) -> Int8ScaledMMLinearKernel:
|
|
config = Int8ScaledMMLinearLayerConfig(
|
|
is_channelwise=is_channelwise,
|
|
is_static_input_scheme=is_static_input_scheme,
|
|
input_symmetric=input_symmetric,
|
|
)
|
|
|
|
kernel_type = choose_scaled_mm_linear_kernel(
|
|
config,
|
|
_POSSIBLE_INT8_KERNELS,
|
|
)
|
|
|
|
logger.info_once(
|
|
"Selected %s for %s",
|
|
kernel_type.__name__,
|
|
module_name,
|
|
scope="global",
|
|
)
|
|
|
|
return kernel_type(
|
|
config,
|
|
layer_param_names=[
|
|
"weight",
|
|
"weight_scale",
|
|
"input_scale",
|
|
"input_zero_point",
|
|
"azp_adj",
|
|
],
|
|
)
|
|
|
|
|
|
def choose_mp_linear_kernel(
|
|
config: MPLinearLayerConfig, compute_capability: int | None = None
|
|
) -> type[MPLinearKernel]:
|
|
"""
|
|
Choose an MPLinearKernel that can implement the given config for the given
|
|
compute capability. Attempts to choose the best kernel in terms of
|
|
performance.
|
|
|
|
Args:
|
|
config (MPLinearLayerConfig): Description of the linear layer to be
|
|
implemented.
|
|
compute_capability (Optional[int], optional): The compute capability of
|
|
the target device, if None uses `current_platform` to get
|
|
the compute capability. Defaults to None.
|
|
|
|
Raises:
|
|
ValueError: If no kernel can implement the given config.
|
|
|
|
Returns:
|
|
type[MPLinearKernel]: Chosen kernel.
|
|
"""
|
|
if compute_capability is None:
|
|
if current_platform is None:
|
|
raise ValueError("Cannot determine compute capability")
|
|
_cc = current_platform.get_device_capability()
|
|
if _cc is not None:
|
|
compute_capability = _cc[0] * 10 + _cc[1]
|
|
|
|
platform_kernels = _POSSIBLE_KERNELS[current_platform._enum]
|
|
|
|
# Apply --linear-backend filtering when set.
|
|
linear_backend = _get_linear_backend()
|
|
if linear_backend != "auto":
|
|
filtered = _filter_kernels_by_backend(linear_backend, platform_kernels)
|
|
if not filtered:
|
|
raise ValueError(
|
|
f"--linear-backend={linear_backend} was requested but no "
|
|
f"'{linear_backend}' kernel exists for mixed-precision layers."
|
|
)
|
|
platform_kernels = filtered
|
|
|
|
failure_reasons = []
|
|
for kernel in platform_kernels:
|
|
if kernel.__name__ in envs.VLLM_DISABLED_KERNELS:
|
|
failure_reasons.append(
|
|
f" {kernel.__name__} disabled by environment variable"
|
|
)
|
|
continue
|
|
if (
|
|
compute_capability is not None
|
|
and kernel.get_min_capability() > compute_capability
|
|
):
|
|
failure_reasons.append(
|
|
f"{kernel.__name__} requires capability "
|
|
f"{kernel.get_min_capability()}, current compute "
|
|
f" capability is {compute_capability}"
|
|
)
|
|
continue
|
|
|
|
can_implement, failure_reason = kernel.can_implement(config)
|
|
if can_implement:
|
|
return kernel
|
|
else:
|
|
failure_reasons.append(
|
|
f" {kernel.__name__} cannot implement due to: {failure_reason}"
|
|
)
|
|
|
|
raise ValueError(
|
|
"Failed to find a kernel that can implement the "
|
|
"WNA16 linear layer. Reasons: \n" + "\n".join(failure_reasons)
|
|
)
|
|
|
|
|
|
def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel:
|
|
"""Select and instantiate the best MXFP8 linear kernel for the
|
|
current platform."""
|
|
config = Mxfp8LinearLayerConfig()
|
|
|
|
platform = current_platform._enum
|
|
possible = list(_POSSIBLE_MXFP8_KERNELS.get(platform, []))
|
|
|
|
# Apply --linear-backend filtering when set.
|
|
linear_backend = _get_linear_backend()
|
|
if linear_backend != "auto":
|
|
filtered = _filter_kernels_by_backend(linear_backend, possible)
|
|
if not filtered:
|
|
raise ValueError(
|
|
f"--linear-backend={linear_backend} was requested but no "
|
|
f"'{linear_backend}' kernel exists for MXFP8 layers."
|
|
)
|
|
possible = filtered
|
|
|
|
failure_reasons = []
|
|
for kernel_cls in possible:
|
|
if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS:
|
|
failure_reasons.append(
|
|
f" {kernel_cls.__name__} disabled by environment variable"
|
|
)
|
|
continue
|
|
|
|
is_supported, reason = kernel_cls.is_supported()
|
|
if not is_supported:
|
|
failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
|
|
continue
|
|
|
|
can_implement, reason = kernel_cls.can_implement(config)
|
|
if not can_implement:
|
|
failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
|
|
continue
|
|
|
|
logger.info_once("Using %s for MXFP8 GEMM", kernel_cls.__name__)
|
|
return kernel_cls(config)
|
|
|
|
raise ValueError(
|
|
"Failed to find a kernel that can implement the "
|
|
"MXFP8 linear layer. Reasons: \n" + "\n".join(failure_reasons)
|
|
)
|
|
|
|
|
|
def init_mxfp4_linear_kernel() -> MxFp4LinearKernel:
|
|
"""Select and instantiate the best MXFP4 linear kernel for the
|
|
current platform."""
|
|
linear_backend = _get_linear_backend()
|
|
|
|
platform = current_platform._enum
|
|
possible = list(_POSSIBLE_MXFP4_KERNELS.get(platform, []))
|
|
|
|
# Apply --linear-backend filtering when set.
|
|
if linear_backend != "auto":
|
|
filtered = _filter_kernels_by_backend(linear_backend, possible)
|
|
if not filtered:
|
|
raise ValueError(
|
|
f"--linear-backend={linear_backend} was requested but no "
|
|
f"'{linear_backend}' kernel exists for MXFP4 layers."
|
|
)
|
|
possible = filtered
|
|
|
|
failure_reasons = []
|
|
for kernel_cls in possible:
|
|
if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS:
|
|
failure_reasons.append(
|
|
f" {kernel_cls.__name__} disabled by environment variable"
|
|
)
|
|
continue
|
|
|
|
is_supported, reason = kernel_cls.is_supported()
|
|
if not is_supported:
|
|
failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
|
|
continue
|
|
|
|
logger.info_once("Using %s for MXFP4 GEMM", kernel_cls.__name__)
|
|
return kernel_cls(MxFp4LinearLayerConfig())
|
|
|
|
raise ValueError(
|
|
"Failed to find a kernel that can implement the "
|
|
"MXFP4 linear layer. Reasons: \n" + "\n".join(failure_reasons)
|
|
)
|
|
|
|
|
|
def init_wfp8_a16_linear_kernel(
|
|
weight_quant_key: QuantKey,
|
|
activation_quant_key: QuantKey,
|
|
weight_shape: tuple[int, int],
|
|
input_dtype: torch.dtype,
|
|
out_dtype: torch.dtype,
|
|
force_kernel: type[FP8ScaledMMLinearKernel] | None = None,
|
|
module_name: str | None = None,
|
|
) -> FP8ScaledMMLinearKernel:
|
|
config = FP8ScaledMMLinearLayerConfig(
|
|
weight_quant_key=weight_quant_key,
|
|
activation_quant_key=activation_quant_key,
|
|
weight_shape=weight_shape,
|
|
input_dtype=input_dtype,
|
|
out_dtype=out_dtype,
|
|
)
|
|
|
|
kernel_type = choose_scaled_mm_linear_kernel(
|
|
config, _POSSIBLE_WFP8A16_KERNELS, force_kernel=force_kernel
|
|
)
|
|
|
|
if module_name:
|
|
logger.info_once(
|
|
"Selected %s for %s",
|
|
kernel_type.__name__,
|
|
module_name,
|
|
scope="global",
|
|
)
|
|
|
|
return kernel_type(
|
|
config,
|
|
layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"],
|
|
)
|
|
|
|
|
|
def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:
|
|
"""Select and instantiate the best NVFP4 linear kernel for the
|
|
current platform."""
|
|
config = NvFp4LinearLayerConfig()
|
|
|
|
# VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the
|
|
# batch-invariant CUTLASS implementation when available, otherwise fall
|
|
# back to emulation. It overrides --linear-backend.
|
|
force_kernel: type[NvFp4LinearKernel] | None = None
|
|
linear_backend = _get_linear_backend()
|
|
if envs.VLLM_BATCH_INVARIANT:
|
|
bi_supported, reason = CutlassNvFp4LinearKernel.is_supported()
|
|
if bi_supported:
|
|
if linear_backend not in ("auto", "cutlass"):
|
|
logger.warning_once(
|
|
"VLLM_BATCH_INVARIANT overrides --linear-backend=%s; "
|
|
"using the CUTLASS backend for deterministic execution.",
|
|
linear_backend,
|
|
)
|
|
else:
|
|
logger.info_once(
|
|
"VLLM_BATCH_INVARIANT forces NVFP4 linear to use the "
|
|
"CUTLASS backend for deterministic execution."
|
|
)
|
|
force_kernel = CutlassNvFp4LinearKernel
|
|
else:
|
|
if linear_backend not in ("auto", "emulation"):
|
|
logger.warning_once(
|
|
"VLLM_BATCH_INVARIANT overrides --linear-backend=%s; "
|
|
"using the emulation backend for deterministic execution.",
|
|
linear_backend,
|
|
)
|
|
logger.info_once(
|
|
"VLLM_BATCH_INVARIANT is set but the batch-invariant NVFP4 "
|
|
"kernel is not supported on this platform; falling back to "
|
|
"emulation for deterministic execution. Reason: %s",
|
|
reason,
|
|
)
|
|
force_kernel = EmulationNvFp4LinearKernel
|
|
elif linear_backend == "auto" and use_a16:
|
|
# Force a16 (Marlin) when running weight-only quantization.
|
|
force_kernel = MarlinNvFp4LinearKernel
|
|
|
|
if force_kernel is not None:
|
|
is_supported, reason = force_kernel.is_supported()
|
|
if not is_supported:
|
|
raise ValueError(
|
|
f"Forced NVFP4 kernel {force_kernel.__name__} is not "
|
|
f"supported: {reason}"
|
|
)
|
|
logger.info_once("Using %s for NVFP4 GEMM", force_kernel.__name__)
|
|
return force_kernel(config)
|
|
|
|
# Auto-select from registry (or --linear-backend filtered).
|
|
platform = current_platform._enum
|
|
possible = list(_POSSIBLE_NVFP4_KERNELS.get(platform, []))
|
|
|
|
# Apply --linear-backend filtering when set.
|
|
if linear_backend != "auto":
|
|
filtered = _filter_kernels_by_backend(linear_backend, possible)
|
|
if not filtered:
|
|
raise ValueError(
|
|
f"--linear-backend={linear_backend} was requested but no "
|
|
f"'{linear_backend}' kernel exists for NVFP4 layers."
|
|
)
|
|
possible = filtered
|
|
|
|
failure_reasons = []
|
|
for kernel_cls in possible:
|
|
if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS:
|
|
failure_reasons.append(
|
|
f" {kernel_cls.__name__} disabled by environment variable"
|
|
)
|
|
continue
|
|
|
|
is_supported, reason = kernel_cls.is_supported()
|
|
if not is_supported:
|
|
failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
|
|
continue
|
|
|
|
can_implement, reason = kernel_cls.can_implement(config)
|
|
if not can_implement:
|
|
failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
|
|
continue
|
|
|
|
if kernel_cls is EmulationNvFp4LinearKernel and failure_reasons:
|
|
logger.warning_once(
|
|
"NVFP4 linear falling back to the slow and unoptimized "
|
|
"emulation backend as no optimized backend is available "
|
|
"(unavailable reasons:\n - %s\n). "
|
|
"In case you expect one of these backends to be used, "
|
|
"please verify your environment.",
|
|
"\n - ".join(failure_reasons),
|
|
)
|
|
|
|
logger.info_once("Using %s for NVFP4 GEMM", kernel_cls.__name__)
|
|
return kernel_cls(config)
|
|
|
|
raise ValueError(
|
|
"Failed to find a kernel that can implement the "
|
|
"NVFP4 linear layer. Reasons: \n" + "\n".join(failure_reasons)
|
|
)
|
|
|
|
|
|
def register_linear_kernel(
|
|
kernel_class: type,
|
|
platform: PlatformEnum,
|
|
kernel_type: str = "mp",
|
|
) -> None:
|
|
"""
|
|
Register a new linear kernel class to be considered in kernel selection.
|
|
|
|
Args:
|
|
kernel_class (type): The kernel class to register.
|
|
platform (PlatformEnum): The platform for which this kernel is applicable.
|
|
kernel_type (str): The type of the kernel, either "mp", "int8", or "fp8".
|
|
Defaults to "mp".
|
|
|
|
Raises:
|
|
ValueError: If the kernel_type is not recognized.
|
|
"""
|
|
if kernel_type == "mp":
|
|
if platform not in _POSSIBLE_KERNELS:
|
|
_POSSIBLE_KERNELS[platform] = []
|
|
_POSSIBLE_KERNELS[platform].append(kernel_class)
|
|
elif kernel_type == "int8":
|
|
if platform not in _POSSIBLE_INT8_KERNELS:
|
|
_POSSIBLE_INT8_KERNELS[platform] = []
|
|
_POSSIBLE_INT8_KERNELS[platform].append(kernel_class)
|
|
elif kernel_type == "fp8":
|
|
if platform not in _POSSIBLE_FP8_KERNELS:
|
|
_POSSIBLE_FP8_KERNELS[platform] = []
|
|
_POSSIBLE_FP8_KERNELS[platform].append(kernel_class)
|
|
elif kernel_type == "mxfp8":
|
|
if platform not in _POSSIBLE_MXFP8_KERNELS:
|
|
_POSSIBLE_MXFP8_KERNELS[platform] = []
|
|
_POSSIBLE_MXFP8_KERNELS[platform].append(kernel_class)
|
|
elif kernel_type == "nvfp4":
|
|
if platform not in _POSSIBLE_NVFP4_KERNELS:
|
|
_POSSIBLE_NVFP4_KERNELS[platform] = []
|
|
_POSSIBLE_NVFP4_KERNELS[platform].append(kernel_class)
|
|
elif kernel_type == "mxfp4":
|
|
if platform not in _POSSIBLE_MXFP4_KERNELS:
|
|
_POSSIBLE_MXFP4_KERNELS[platform] = []
|
|
_POSSIBLE_MXFP4_KERNELS[platform].append(kernel_class)
|
|
else:
|
|
raise ValueError(f"Unrecognized kernel type: {kernel_type}")
|
|
|
|
|
|
__all__ = [
|
|
"init_fp8_linear_kernel",
|
|
"init_int8_linear_kernel",
|
|
"init_nvfp4_linear_kernel",
|
|
"choose_mp_linear_kernel",
|
|
"register_linear_kernel",
|
|
"init_wfp8_a16_linear_kernel",
|
|
"FP8ScaledMMLinearKernel",
|
|
"Int8ScaledMMLinearKernel",
|
|
"ScaledMMLinearKernel",
|
|
"FP8ScaledMMLinearLayerConfig",
|
|
"Int8ScaledMMLinearLayerConfig",
|
|
"ScaledMMLinearLayerConfig",
|
|
"AiterHipbMMPerTokenFp8ScaledMMLinearKernel",
|
|
"AiterPreshuffledPerTokenFp8ScaledMMLinearKernel",
|
|
"AiterPerTokenFp8ScaledMMLinearKernel",
|
|
"NvFp4LinearKernel",
|
|
"NvFp4LinearLayerConfig",
|
|
"AiterInt8ScaledMMLinearKernel",
|
|
"CPUInt8ScaledMMLinearKernel",
|
|
"CutlassFP8ScaledMMLinearKernel",
|
|
"CutlassInt8ScaledMMLinearKernel",
|
|
"FlashInferFP8ScaledMMLinearKernel",
|
|
"ChannelWiseTorchFP8ScaledMMLinearKernel",
|
|
"PerTensorTorchFP8ScaledMMLinearKernel",
|
|
"RowWiseTorchFP8ScaledMMLinearKernel",
|
|
"ROCmFP8ScaledMMLinearKernel",
|
|
"TritonInt8ScaledMMLinearKernel",
|
|
"ZentorchInt8ScaledMMLinearKernel",
|
|
"ZentorchWNA16LinearKernel",
|
|
"MPLinearKernel",
|
|
"MPLinearLayerConfig",
|
|
"AllSparkLinearKernel",
|
|
"ConchLinearKernel",
|
|
"CPUWNA16LinearKernel",
|
|
"CutlassW4A8LinearKernel",
|
|
"Dynamic4bitLinearKernel",
|
|
"ExllamaLinearKernel",
|
|
"MacheteLinearKernel",
|
|
"MarlinLinearKernel",
|
|
"TritonW4A16LinearKernel",
|
|
"XPUW4A8IntLinearKernel",
|
|
"XPUwNa16LinearKernel",
|
|
"init_mxfp8_linear_kernel",
|
|
"Mxfp8LinearKernel",
|
|
"Mxfp8LinearLayerConfig",
|
|
"init_mxfp4_linear_kernel",
|
|
"MxFp4LinearKernel",
|
|
"MxFp4LinearLayerConfig",
|
|
"FlashInferMxFp4LinearKernel",
|
|
"MarlinMxFp4LinearKernel",
|
|
"FlashInferCutedslMxfp8LinearKernel",
|
|
"FlashInferCutlassMxfp8LinearKernel",
|
|
"MarlinMxfp8LinearKernel",
|
|
"XPUMxFp8LinearKernel",
|
|
"EmulationMxfp8LinearKernel",
|
|
"CutlassNvFp4LinearKernel",
|
|
"EmulationNvFp4LinearKernel",
|
|
"FbgemmNvFp4LinearKernel",
|
|
"FlashInferCuteDslNvFp4LinearKernel",
|
|
"FlashInferB12xNvFp4LinearKernel",
|
|
"FlashInferCutlassNvFp4LinearKernel",
|
|
"FlashInferTrtllmNvFp4LinearKernel",
|
|
"FlashInferCudnnNvFp4LinearKernel",
|
|
"MarlinNvFp4LinearKernel",
|
|
"_KernelT",
|
|
"DeepGemmFp8BlockScaledMMKernel",
|
|
"FlashInferFp8DeepGEMMDynamicBlockScaledKernel",
|
|
]
|