chore: import upstream snapshot with attribution
pre-commit / pre-run-check (push) Has been cancelled
pre-commit / pre-commit (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 12:55:37 +08:00
commit 7ce4c8e27e
5900 changed files with 1668062 additions and 0 deletions
File diff suppressed because it is too large Load Diff
+332
View File
@@ -0,0 +1,332 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, ClassVar, Generic, TypeVar
import torch
from typing_extensions import Self
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
@dataclass
class MMLinearLayerConfig: ...
@dataclass
class Params:
"""Base class for quantized layer parameters.
This class provides a typed interface for accessing quantized weights and scales
from layer modules. It serves as a parameter container that can be extracted from
layers and passed to kernel implementations.
Attributes:
weight: The quantized weight tensor
weight_scale: weight scaling factors
input_scale: Optional input scaling factors
Class Variables:
WEIGHT: Attribute name for weight tensor on the layer module
WEIGHT_SCALE: Attribute name for weight scale tensor on the layer module
INPUT_SCALE: Attribute name for input scale tensor on the layer module
Important:
The string values of WEIGHT, WEIGHT_SCALE, and INPUT_SCALE class variables
MUST match the attribute names used in the corresponding quantization method's
create_weights() implementation.
For example, if FP8LinearMethod.create_weights()
sets layer.weight and layer.weight_scale,
then WEIGHT="weight" and
WEIGHT_SCALE="weight_scale" must be used here.
Usage:
```python
# Extract parameters from a quantized layer
params = Params.from_layer(layer)
# Access typed parameters
output = func(input, params.weight, params.weight_scale)
```
"""
weight: torch.Tensor
weight_scale: torch.Tensor
input_scale: torch.Tensor | None
# Attribute names on the layer
WEIGHT: ClassVar[str] = "weight"
WEIGHT_SCALE: ClassVar[str] = "weight_scale"
INPUT_SCALE: ClassVar[str] = "input_scale"
@classmethod
def from_layer(cls, layer: torch.nn.Module) -> Self:
return cls(
weight=getattr(layer, cls.WEIGHT),
weight_scale=getattr(layer, cls.WEIGHT_SCALE),
input_scale=getattr(layer, cls.INPUT_SCALE, None),
)
@dataclass
class FP8Params(Params):
"""FP8 layer parameters with typed fields"""
input_scale_ub: torch.Tensor | None
INPUT_SCALE_UB: ClassVar[str] = "input_scale_ub"
@classmethod
def from_layer(cls, layer: torch.nn.Module) -> "FP8Params":
"""Extract parameters from layer"""
return cls(
weight=getattr(layer, cls.WEIGHT),
weight_scale=getattr(layer, cls.WEIGHT_SCALE),
input_scale=getattr(layer, cls.INPUT_SCALE, None),
input_scale_ub=getattr(layer, cls.INPUT_SCALE_UB, None),
)
@dataclass
class Int8Params(Params):
"""Int8 layer parameters with typed fields"""
input_zero_point: torch.Tensor | None
azp_adj: torch.Tensor | None
INPUT_ZERO_POINT: ClassVar[str] = "input_zero_point"
AZP_ADJ: ClassVar[str] = "azp_adj"
@classmethod
def from_layer(cls, layer: torch.nn.Module) -> "Int8Params":
"""Extract parameters from layer"""
return cls(
weight=getattr(layer, cls.WEIGHT),
weight_scale=getattr(layer, cls.WEIGHT_SCALE),
input_scale=getattr(layer, cls.INPUT_SCALE, None),
input_zero_point=getattr(layer, cls.INPUT_ZERO_POINT, None),
azp_adj=getattr(layer, cls.AZP_ADJ, None),
)
_ParamsT = TypeVar("_ParamsT", bound=Params)
_ConfigT = TypeVar("_ConfigT", bound=MMLinearLayerConfig)
class MMLinearKernel(ABC, Generic[_ConfigT, _ParamsT]):
"""Abstract base class for quantized matrix multiplication kernels.
This class provides the interface for implementing custom quantized linear layer
kernels in vLLM. Subclasses should implement specific quantization strategies
(e.g., FP8, INT8) and their corresponding compute kernels.
Generic Type Parameters:
_ConfigT: Configuration type for the kernel (subclass of MMLinearLayerConfig).
Contains kernel-specific settings like quantization keys, dtypes, etc.
_ParamsT: Parameter type for the kernel (subclass of Params).
Defines the quantized weights and scales needed by the kernel.
Typical Usage:
1. Define a config dataclass inheriting from MMLinearLayerConfig
2. Define a params dataclass inheriting from Params (or FP8Params/Int8Params)
3. Subclass MMLinearKernel with your config and params types
4. Implement all abstract methods
5. Register the kernel with the quantization method
Example:
```python
@dataclass
class MyKernelConfig(MMLinearLayerConfig):
static: bool
output_dtype: torch.dtype
@dataclass
class MyKernelParams(FP8Params):
custom_scale: torch.Tensor
CUSTOM_SCALE: ClassVar[str] = "custom_scale"
class MyKernel(MMLinearKernel[MyKernelConfig, MyKernelParams]):
@classmethod
def is_supported(cls, compute_capability=None):
if compute_capability and compute_capability < 90:
return False, "Requires compute capability >= 9.0"
return True, None
@classmethod
def can_implement(cls, config):
if not config.static:
return False, "Only static quantization supported"
return True, None
def process_weights_after_loading(self, layer):
# Preprocess weights for the kernel
params = self._get_layer_params(layer)
processed = preprocess_weights(params.weight)
replace_parameter(layer, params.WEIGHT, processed)
def _get_layer_params(self, layer, **kwargs):
return MyKernelParams.from_layer(layer)
def apply_weights(self, layer, x, bias=None, **kwargs):
params = self._get_layer_params(layer)
# Call your custom kernel
output = my_custom_kernel(x, params.weight, params.weight_scale)
if bias is not None:
output += bias
return output
```
Lifecycle:
1. Kernel selection: is_supported() and can_implement() check compatibility
2. Initialization: __init__() creates kernel instance with config
3. Weight loading: process_weights_after_loading() preprocesses weights
4. Inference: apply_weights() executes the quantized matmul
"""
@classmethod
@abstractmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
"""Check if this kernel is supported on the current hardware.
This method checks hardware-level compatibility (e.g., GPU architecture,
compute capability, available instructions). It's called during kernel
selection to filter out kernels that cannot run on the current device.
Args:
compute_capability: GPU compute capability (e.g., 80 for A100, 90 for H100).
If None, should check the current device.
Returns:
A tuple of (is_supported, reason):
- is_supported: True if the kernel can run on this hardware
- reason: If not supported, a string explaining why; otherwise None
"""
raise NotImplementedError
@classmethod
@abstractmethod
def can_implement(cls, config: _ConfigT) -> tuple[bool, str | None]:
"""Check if this kernel can implement the given configuration.
This method checks configuration-level compatibility (e.g., quantization
scheme, group sizes, static vs dynamic quantization). It's called after
is_supported() to determine if this kernel can handle the specific
quantization configuration.
Args:
config: The kernel configuration to check
Returns:
A tuple of (can_implement, reason):
- can_implement: True if this kernel supports the config
- reason: If not supported, a string explaining why; otherwise None
```
"""
raise NotImplementedError
def __init__(self, config: _ConfigT) -> None:
"""Initialize the kernel with the given configuration.
Args:
config: Kernel-specific configuration containing settings like
quantization keys, output dtypes, etc.
"""
self.config = config
def input_quant_key(self) -> QuantKey | None:
"""Return the input quantization key supported by this kernel. If the kernel
does not support input quantization outside of the kernel, return None.
"""
return None
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""Process and transform weights after loading from checkpoint.
This method is called once after weights are loaded but before inference.
Use it to preprocess weights into the format required by your kernel
(e.g., reordering, padding, format conversion).
Modifications should be done in-place using replace_parameter() to ensure
the layer's parameters are properly updated.
Args:
layer: The layer module containing the weights to process
Example:
```python
def process_weights_after_loading(self, layer):
params = self._get_layer_params(layer)
# Reorder weights for better memory access
weight_reordered = reorder_weights(params.weight)
replace_parameter(layer, params.WEIGHT, weight_reordered)
```
"""
raise NotImplementedError
# return a covariant type in the subclass
@abstractmethod
def _get_layer_params(self, layer: torch.nn.Module, **kwargs: Any) -> _ParamsT:
"""Extract typed parameters from the layer module.
This internal method retrieves the quantized weights and scales from
the layer as a typed parameter object. Subclasses should typically
delegate to ParamsClass.from_layer().
Args:
layer: The layer module containing the parameters
**kwargs: Additional arguments
Returns:
A typed parameter object containing weights, scales, and other
quantization parameters
Example:
```python
def _get_layer_params(self, layer, **kwargs):
return MyKernelParams.from_layer(layer)
```
"""
raise NotImplementedError
def get_output_padding(self) -> int | None:
"""Get the number of output tokens to pad for this kernel.
Some kernels require input padding for optimal performance.
Override this method to specify padding requirements.
Returns:
Number of tokens to pad, or None for no padding (default)
"""
return None
@abstractmethod
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
**kwargs: Any,
) -> torch.Tensor:
"""Apply the quantized weights to the input tensor.
This is the main inference method that performs the quantized matrix
multiplication. It should handle input quantization (if needed), call
the underlying kernel, and apply bias.
Args:
layer: The layer module containing the quantized weights
x: Input tensor of shape [..., in_features]
bias: Optional bias tensor of shape [out_features]
**kwargs: Additional kernel-specific arguments
Returns:
Output tensor of shape [..., out_features]
"""
raise NotImplementedError
@@ -0,0 +1,94 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import ABC, abstractmethod
from collections.abc import Callable
from dataclasses import dataclass
import torch
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.scalar_type import ScalarType
@dataclass
class MPLinearLayerConfig:
full_weight_shape: tuple[int, int] # [in, out]
partition_weight_shape: tuple[int, int]
weight_type: ScalarType
act_type: torch.dtype
group_size: int
zero_points: bool
has_g_idx: bool
out_type: torch.dtype | None = None
class MPLinearKernel(ABC):
@classmethod
@abstractmethod
def get_min_capability(cls) -> int:
raise NotImplementedError
@classmethod
@abstractmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
raise NotImplementedError
def __init__(
self,
c: MPLinearLayerConfig,
w_q_param_name: str,
w_s_param_name: str,
w_zp_param_name: str | None = None,
w_gidx_param_name: str | None = None,
) -> None:
assert self.can_implement(c)
self.config = c
self.w_q_name = w_q_param_name
self.w_s_name = w_s_param_name
if c.zero_points:
assert w_zp_param_name is not None
if c.has_g_idx:
assert w_gidx_param_name is not None
self.w_zp_name: str | None = w_zp_param_name
self.w_gidx_name: str | None = w_gidx_param_name
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
raise NotImplementedError
@abstractmethod
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
raise NotImplementedError
def _transform_param(
self, layer: torch.nn.Module, name: str | None, fn: Callable
) -> None:
if name is not None and getattr(layer, name, None) is not None:
old_param = getattr(layer, name)
new_param = fn(old_param)
# replace the parameter with torch.nn.Parameter for TorchDynamo
# compatibility
replace_parameter(
layer, name, torch.nn.Parameter(new_param.data, requires_grad=False)
)
def _get_weight_params(
self, layer: torch.nn.Module
) -> tuple[
torch.Tensor, # w_q
torch.Tensor, # w_s
torch.Tensor | None, # w_zp,
torch.Tensor | None, # w_gidx
]:
return (
getattr(layer, self.w_q_name),
getattr(layer, self.w_s_name),
getattr(layer, self.w_zp_name or "", None),
getattr(layer, self.w_gidx_name or "", None),
)
@@ -0,0 +1,66 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
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.MPLinearKernel import (
MPLinearKernel,
MPLinearLayerConfig,
)
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,
)
__all__ = [
"MPLinearKernel",
"MPLinearLayerConfig",
"AllSparkLinearKernel",
"ConchLinearKernel",
"CPUWNA16LinearKernel",
"CutlassW4A8LinearKernel",
"Dynamic4bitLinearKernel",
"ExllamaLinearKernel",
"HummingLinearKernel",
"MacheteLinearKernel",
"MarlinLinearKernel",
"RDNA3W4A16LinearKernel",
"TritonW4A16LinearKernel",
"XPUW4A8IntLinearKernel",
"XPUwNa16LinearKernel",
"ZentorchWNA16LinearKernel",
]
@@ -0,0 +1,116 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.model_executor.layers.quantization.utils.allspark_utils import (
ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
check_allspark_supported_dtype_shape,
)
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
from vllm.utils.platform_utils import num_compute_units
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
class AllSparkLinearKernel(MPLinearKernel):
@classmethod
def get_min_capability(cls) -> int:
return 80
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if c.has_g_idx:
return False, "Act reordering currently not supported by AllSpark"
if c.zero_points:
return False, "Zero points currently not supported by AllSpark"
return check_allspark_supported_dtype_shape(
c.partition_weight_shape[0], # in_features
c.partition_weight_shape[1], # out_features
c.group_size,
c.weight_type,
c.act_type,
)
# note assumes that
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
# `weight_scale` is: {input_dim = 0, output_dim = 1}
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
device = getattr(layer, self.w_q_name).device
c = self.config
# prepare the parameters required for the kernel
properties = torch.cuda.get_device_properties(device.index)
sm_count = num_compute_units(device.index)
sm_version = properties.major * 10 + properties.minor
gemm_args = {}
gemm_args["sm_count"] = sm_count
gemm_args["sm_version"] = sm_version
self.gemm_args = gemm_args
# transform param weight, scale
old_weight_param = getattr(layer, self.w_q_name)
old_scale_param = getattr(layer, self.w_s_name)
assert isinstance(old_weight_param, BasevLLMParameter)
permute_param_layout_(old_weight_param, input_dim=0, output_dim=1, packed_dim=0)
assert isinstance(old_scale_param, BasevLLMParameter)
permute_param_layout_(old_scale_param, input_dim=0, output_dim=1)
# unpack weight from K / 4 x N int32 to K x N uint8
new_weight_param = torch.nn.Parameter(
old_weight_param.data, requires_grad=False
)
new_weight_param.data = (
new_weight_param.data.t().contiguous().view(dtype=torch.uint8)
)
new_weight_param.data = new_weight_param.data.t().contiguous()
new_scale_param = torch.nn.Parameter(old_scale_param.data, requires_grad=False)
# reorder K x N weight as N32K16 format for Ampere W8A16
new_weight_param.data, new_scale_param.data, _ = ops.allspark_repack_weight(
new_weight_param.data, new_scale_param.data, None, c.zero_points
)
replace_parameter(layer, self.w_q_name, new_weight_param.data)
replace_parameter(layer, self.w_s_name, new_scale_param.data)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
c = self.config
gemm_args = self.gemm_args
w_q, w_s, _, _ = self._get_weight_params(layer)
reshaped_x = x.reshape(-1, x.shape[-1])
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
output = ops.allspark_w8a16_gemm(
a=reshaped_x,
b_qweight=w_q,
b_scales=w_s,
b_qzeros=None,
n=c.partition_weight_shape[1],
group_size=c.group_size,
sm_count=gemm_args["sm_count"],
sm_version=gemm_args["sm_version"],
CUBLAS_M_THRESHOLD=ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
has_zp=c.zero_points,
n32k16_reorder=True,
)
if bias is not None:
output.add_(bias) # In-place add
return output.reshape(out_shape)
@@ -0,0 +1,159 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from importlib.util import find_spec
from typing import Final
import torch
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
from vllm.scalar_type import scalar_types
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
_CONCH_SUPPORTED_WEIGHT_TYPES: Final = [
scalar_types.uint4,
scalar_types.uint8,
scalar_types.uint4b8,
scalar_types.uint8b128,
]
_CONCH_SUPPORTED_GROUP_SIZES: Final = [-1, 128]
class ConchLinearKernel(MPLinearKernel):
@classmethod
def get_min_capability(cls) -> int:
return 80
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if c.weight_type not in _CONCH_SUPPORTED_WEIGHT_TYPES:
error_msg = (
f"Weight type ({c.weight_type}) not supported by "
"ConchLinearKernel, supported types are: "
f"{_CONCH_SUPPORTED_WEIGHT_TYPES}"
)
return False, error_msg
if c.group_size not in _CONCH_SUPPORTED_GROUP_SIZES:
error_msg = (
f"Group size ({c.group_size}) not supported by "
"ConchLinearKernel, supported group sizes are: "
f"{_CONCH_SUPPORTED_GROUP_SIZES}"
)
return False, error_msg
if c.has_g_idx:
return (
False,
"Activation reordering (g_idx) is not supported by ConchLinearKernel",
)
if find_spec("conch") is None:
error_msg = (
"conch-triton-kernels is not installed, please "
"install it via `pip install conch-triton-kernels` "
"and try again!"
)
return False, error_msg
return True, None
# note assumes that
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
# `weight_scale` is: {input_dim = 0, output_dim = 1}
# `weight_zero_point` is: {input_dim = 1, output_dim = 0, packed_dim = 0}
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
def transform_w_q(x):
assert isinstance(x, BasevLLMParameter)
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
x.data = x.data.contiguous()
return x
def transform_w_s(x):
assert isinstance(x, BasevLLMParameter)
permute_param_layout_(x, input_dim=0, output_dim=1)
x.data = x.data.contiguous()
return x
def transform_w_zp(x):
# Zero points are stored PACKED as [N//pack_factor, K//G]
# The Conch kernel expects UNPACKED zeros: [K//G, N]
# We need to unpack and reorder
assert isinstance(x, BasevLLMParameter)
packed = x.data # shape: [N//pack_factor, K//G], dtype: int32
# Determine packing based on weight bit width
size_bits = self.config.weight_type.size_bits
pack_factor = 32 // size_bits # 8 for 4-bit, 4 for 8-bit
mask = (1 << size_bits) - 1 # 0xF for 4-bit, 0xFF for 8-bit
n_packed, k_groups = packed.shape
n_full = n_packed * pack_factor
# Unpack using vectorized bitwise ops
# shifts = [0, size_bits, 2*size_bits, ...] for each packed position
shifts = torch.arange(
0, 32, size_bits, dtype=torch.int32, device=packed.device
)
# packed: [N//pack_factor, K//G] -> [N//pack_factor, K//G, 1]
# shifts: [pack_factor] -> [1, 1, pack_factor]
# Result: [N//pack_factor, K//G, pack_factor]
unpacked = (packed.unsqueeze(-1) >> shifts) & mask
# Permute to [K//G, N//pack_factor, pack_factor] then reshape to [K//G, N]
unpacked = unpacked.permute(1, 0, 2).reshape(k_groups, n_full)
x.data = unpacked.to(torch.uint8).contiguous()
# Update metadata - zeros are no longer packed
if hasattr(x, "_input_dim"):
x._input_dim = 0
if hasattr(x, "_output_dim"):
x._output_dim = 1
if hasattr(x, "_packed_factor"):
x._packed_factor = 1
return x
self._transform_param(layer, self.w_q_name, transform_w_q)
self._transform_param(layer, self.w_s_name, transform_w_s)
if self.config.zero_points:
self._transform_param(layer, self.w_zp_name, transform_w_zp)
elif self.w_zp_name is not None:
layer.register_parameter(self.w_zp_name, None)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
from conch.ops.quantization.gemm import mixed_precision_gemm
w_q, w_s, w_zp, _ = self._get_weight_params(layer)
# Map channelwise group_size=-1 to the actual input dimension K.
# The conch kernel computes stride_mul = block_k / group_size;
# passing -1 produces a negative stride that reads out-of-bounds
# scale values for all K-blocks after the first.
group_size = self.config.group_size
if group_size == -1:
group_size = x.shape[-1]
x_2d = x.reshape(-1, x.shape[-1])
out_shape = x.shape[:-1] + (self.config.partition_weight_shape[1],)
output = mixed_precision_gemm(
x=x_2d,
w_q_packed=w_q.data,
w_s=w_s.data,
w_zp=w_zp.data if w_zp is not None else None,
weight_size_bits=self.config.weight_type.size_bits,
weight_bias=self.config.weight_type.bias,
group_size=group_size,
)
if bias is not None:
output.add_(bias) # In-place add
return output.reshape(out_shape)
@@ -0,0 +1,222 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm import _custom_ops as ops
from vllm import envs
from vllm.model_executor.layers.quantization.utils.quant_utils import (
pack_quantized_values_into_int32,
unpack_quantized_values_into_int32,
)
from vllm.platforms import CpuArchEnum, current_platform
from vllm.scalar_type import scalar_types
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
_CPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8)
class CPUWNA16LinearKernel(MPLinearKernel):
@classmethod
def get_min_capability(cls) -> int:
return -1
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if not current_platform.is_cpu():
return False, "CPUWNA16 only supported on CPU"
if c.weight_type not in _CPUWNA16_SUPPORTED_QUANT_TYPES:
return (
False,
f"Quant type ({c.weight_type}) not supported by "
"CPUWNA16, supported types are: "
f"{_CPUWNA16_SUPPORTED_QUANT_TYPES}",
)
if c.group_size != -1 and c.group_size % 2 != 0:
return (
False,
f"Group size ({c.group_size}) not supported by "
"CPUWNA16, supported group sizes are multiples of 2",
)
if c.partition_weight_shape[0] % 32 != 0:
return (
False,
f"Input size ({c.partition_weight_shape[0]}) not supported by "
"CPUWNA16, supported sizes are multiples of 32",
)
if c.partition_weight_shape[1] % 32 != 0:
return (
False,
f"Output size ({c.partition_weight_shape[1]}) not supported by "
"CPUWNA16, supported sizes are multiples of 32",
)
return True, None
# note assumes that
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
# `weight_scale` is: {input_dim = 0, output_dim = 1}
# `weight_zp` is: {input_dim = 0, output_dim = 1, packed_dim = 1}
def _process_gptq_weights_w4a16(self, layer: torch.nn.Module):
packed_weight = getattr(layer, self.w_q_name)
bits = self.config.weight_type.mantissa
pack_factor = 32 // bits
p_w_k, _ = packed_weight.size()
input_size = p_w_k * pack_factor
isa_hint = _get_isa_hint(getattr(layer, self.w_s_name).dtype)
layer.isa_hint = isa_hint
# convert input dim packed to output dim packed
weight = unpack_quantized_values_into_int32(
packed_weight, self.config.weight_type, 0
)
weight = pack_quantized_values_into_int32(weight, self.config.weight_type, 1)
# make 16 output channel as a block and transpose to the make
# the block contiguous
weight = (
weight.view(input_size, -1, 16 // pack_factor)
.permute(1, 0, 2)
.reshape(-1, input_size * 16 // pack_factor)
.contiguous()
)
getattr(layer, self.w_q_name).data = weight
# note assumes that
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
# `weight_scale` is: {input_dim = 0, output_dim = 1}
# `weight_zp` is: {input_dim = 0, output_dim = 1, packed_dim = 1}
def _process_gptq_weights_w4a8(self, layer: torch.nn.Module):
packed_weight = getattr(layer, self.w_q_name)
scales = getattr(layer, self.w_s_name)
group_num = scales.data.size(0)
zp_output_size = scales.data.size(1) // 8
if self.config.zero_points:
assert self.w_zp_name
packed_zp = getattr(layer, self.w_zp_name)
else:
# w4a8 kernel always requires zp, allocate a fake zp
assert self.w_zp_name
packed_zp = torch.nn.Parameter(
torch.ones(group_num, zp_output_size, dtype=torch.int32) * -2004318072,
requires_grad=False,
)
setattr(layer, self.w_zp_name, packed_zp)
# FIXME: some bugs in convert_weight_packed_scale_zp with GPTQ format,
# repack to AWQ weight
weight = unpack_quantized_values_into_int32(
packed_weight, self.config.weight_type, 0
)
input_size, output_size = weight.size()
weight = weight.view(input_size, output_size // 8, 8)
weight = weight[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(input_size, output_size)
weight = pack_quantized_values_into_int32(
weight, self.config.weight_type, 1
).contiguous()
zp = unpack_quantized_values_into_int32(packed_zp, self.config.weight_type, 1)
zp = zp.view(group_num, output_size // 8, 8)
zp = zp[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(group_num, output_size)
zp = pack_quantized_values_into_int32(
zp, self.config.weight_type, 1
).contiguous()
blocked_w, blocked_zp, blocked_s = ops.convert_weight_packed_scale_zp(
weight,
zp,
scales.data,
ops.CPUQuantAlgo.AWQ,
)
if layer.bias is not None:
layer.bias.data = layer.bias.float()
packed_weight.data = blocked_w
scales.data = blocked_s
packed_zp.data = blocked_zp
def process_weights_after_loading(self, layer: torch.nn.Module):
if (not self.config.zero_points) and (self.w_zp_name is not None):
setattr(layer, self.w_zp_name, None)
if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
setattr(layer, self.w_gidx_name, None)
weights = getattr(layer, self.w_q_name)
# Require GPTQ pack format
assert weights.input_dim == weights.packed_dim
# Weights in CT format is [output_size, input_size]
if weights.input_dim == 1:
weights.data = weights.t()
# Scales in CT format is [output_size, group_num]
scales = getattr(layer, self.w_s_name)
if scales.output_dim == 0:
scales.data = scales.t().contiguous()
# Zero points in CT format is [output_size, group_num]
# Zero points in awq_marlin format is [output_size, group_num]
if self.config.zero_points:
assert self.w_zp_name
zp = getattr(layer, self.w_zp_name)
if zp.output_dim == 0:
zp.data = zp.t().contiguous()
supports_amx = torch.cpu._is_amx_tile_supported()
supports_riscv = current_platform.get_cpu_architecture() == CpuArchEnum.RISCV
layer.use_w4a8 = (
envs.VLLM_CPU_INT4_W4A8
and not self.config.has_g_idx
and self.config.act_type == torch.bfloat16
and (supports_amx or supports_riscv)
)
# layer.use_w4a8 = False
# AWQ format will be converted to GPTQ format in `AutoAWQMarlinLinearMethod`
if layer.use_w4a8:
self._process_gptq_weights_w4a8(layer)
else:
self._process_gptq_weights_w4a16(layer)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
if layer.use_w4a8:
x = ops.int4_scaled_mm_cpu(
x=x,
w=w_q,
w_zeros=w_zp,
w_scales=w_s,
bias=bias,
)
else:
x = ops.cpu_gemm_wna16(
input=x,
q_weight=w_q,
scales=w_s,
zeros=w_zp,
g_idx=w_gidx,
bias=bias,
pack_factor=8, # 32 // 4
isa_hint=layer.isa_hint,
)
return x
def _get_isa_hint(dtype: torch.dtype) -> str:
supports_amx = torch.cpu._is_amx_tile_supported()
if supports_amx and dtype in (torch.bfloat16,):
return "amx"
elif current_platform.get_cpu_architecture() == CpuArchEnum.RISCV:
return "rvv"
else:
return "vec"
@@ -0,0 +1,131 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
from vllm.model_executor.layers.quantization.utils.quant_utils import (
GroupShape,
convert_bf16_scales_to_fp8,
convert_packed_uint4b8_to_signed_int4_inplace,
)
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
class CutlassW4A8LinearKernel(MPLinearKernel):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# dynamic per-tok fp8 activation quantization
self.quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN)
@classmethod
def get_min_capability(cls) -> int:
return 90
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "CUTLASS only supported on CUDA"
if not current_platform.is_device_capability(90):
return False, "CUTLASS W4A8 requires compute capability of 90 (Hopper)"
if c.act_type != torch.float8_e4m3fn:
return False, "CUTLASS W4A8 only supports FP8 (e4m3) activations"
if c.has_g_idx:
return False, "Act reordering not supported by CUTLASS W4A8"
if c.zero_points:
return False, "Zero points not supported by CUTLASS W4A8"
if c.weight_type != scalar_types.int4:
return (
False,
f"Quant type ({c.weight_type}) not supported by "
"CUTLASS W4A8, only supported int4",
)
if c.group_size != 128:
return False, "Only group_size 128 is supported"
in_features, out_features = c.partition_weight_shape
if in_features % 128 or out_features % 128:
return (
False,
f"K and N must be divisible by 128, got {c.partition_weight_shape}",
)
if c.out_type != torch.bfloat16:
return (
False,
f"Only bfloat16 output type currently supportedgot {c.out_type=}",
)
return True, None
# note assumes that
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
# `weight_scale` is: {input_dim = 0, output_dim = 1}
def process_weights_after_loading(self, layer: torch.nn.Module):
def transform_w_q(x):
assert isinstance(x, BasevLLMParameter)
convert_packed_uint4b8_to_signed_int4_inplace(x.data)
torch.accelerator.synchronize()
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
x.data = ops.cutlass_encode_and_reorder_int4b(x.data.t().contiguous().t())
return x
def transform_w_s(x):
assert isinstance(x, BasevLLMParameter)
permute_param_layout_(x, input_dim=0, output_dim=1)
x.data = x.data.contiguous().to(torch.float8_e4m3fn)
x.data = ops.cutlass_pack_scale_fp8(x.data)
return x
w_s = getattr(layer, self.w_s_name)
fp8_scales, chan_scales = convert_bf16_scales_to_fp8(self.quant_fp8, w_s.data)
w_s.data = fp8_scales
# register per-channel scales
layer.register_parameter(
"weight_chan_scale", torch.nn.Parameter(chan_scales, requires_grad=False)
)
# Encode/reorder weights and pack scales
self._transform_param(layer, self.w_q_name, transform_w_q)
self._transform_param(layer, self.w_s_name, transform_w_s)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
c = self.config
w_q, w_s, _, _ = self._get_weight_params(layer)
w_ch_s = layer.weight_chan_scale
x_2d = x.reshape(-1, x.shape[-1])
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
x_2d, act_scales = self.quant_fp8(x_2d)
output = ops.cutlass_w4a8_mm(
a=x_2d,
b_q=w_q,
b_group_scales=w_s,
b_group_size=c.group_size,
a_token_scales=act_scales,
b_channel_scales=w_ch_s,
)
if bias is not None:
output.add_(bias) # In-place add
return output.reshape(out_shape)
@@ -0,0 +1,159 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.platforms import CpuArchEnum, current_platform
from vllm.scalar_type import scalar_types
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
# This implementation is for the KleidiAI-accelerated w4a8int quantization
# scheme on Arm CPUs:
# torch.ops.aten._dyn_quant_matmul_4bit performs dynamic quantized matmul
# it takes:
# - int4 weights packed along with bias/scales by
# torch.ops.aten._dyn_quant_pack_4bit_weight
# - float32/bfloat16 activations
# then it leverages KleidiAI ukernels that:
# - dynamically quantize the activations to int8
# - unpack the int4 weights to int8
# - perform int8 x int8 -> int32 matmul
# - dequantize the int32 output to float32/bfloat16 outputs
class Dynamic4bitLinearKernel(MPLinearKernel):
SUPPORTED_QUANT_TYPES = [scalar_types.int4]
@classmethod
def get_min_capability(cls) -> int:
return 1
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if not current_platform.is_cpu():
return False, "Only CPU is supported"
if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
return False, f"Unsupported quant type {c.weight_type}"
if (
current_platform.get_cpu_architecture() == CpuArchEnum.ARM
and c.act_type
not in [
torch.float32,
torch.bfloat16,
torch.float16,
]
):
return (
False,
"Dynamic4bitLinearKernel on Arm requires Float32 or"
" BFloat16 or Float16 activations",
)
if c.full_weight_shape[0] % c.group_size != 0:
return (
False,
f"Group size ({c.group_size}) does not evenly divide"
" the number of input features "
f"({c.full_weight_shape[0]})",
)
if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
try:
# Attempt to retrieve the operation
_ = torch.ops.aten._dyn_quant_matmul_4bit
except AttributeError:
return (
False,
f"PyTorch {torch.__version__} does not support"
" _dyn_quant_matmul_4bit. Install a newer version",
)
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module):
c = self.config
packed_weight = getattr(layer, self.w_q_name)
packed_weight = packed_weight.add(8)
uint8_packed = (packed_weight[::, 1::2] << 4 | packed_weight[::, ::2]).to(
torch.uint8
)
scales = getattr(layer, self.w_s_name)
block_size = c.group_size
# Handle scaling factors for partitioned weights
if block_size == c.partition_weight_shape[0]:
scales = scales.to(
torch.float32
) # Float32 & Bfloat16 variants requires float32 scales
scales = scales.view(-1, 1) # Channel-wise scales
if layer.bias is not None:
# Float32 & Bfloat16 variants requires float32 bias
replace_parameter(
layer,
"bias",
torch.nn.Parameter(
layer.bias.to(torch.float32), requires_grad=False
),
)
else:
# KleidiAI kernel requires bfloat16 scales with groupwise scheme
scales = scales.to(torch.bfloat16)
# Repack weights as per kernel requirement
w = torch.ops.aten._dyn_quant_pack_4bit_weight(
uint8_packed,
scales,
layer.bias,
block_size,
c.partition_weight_shape[0],
c.partition_weight_shape[1],
)
replace_parameter(
layer, self.w_q_name, torch.nn.Parameter(w, requires_grad=False)
)
setattr(layer, self.w_s_name, None)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
# PyTorch / KleidiAI kernels natively support the following configs:
# - channelwise with bfloat16 / float32 activations
# - groupwise with float32 activations
# To support:
# - groupwise with bfloat16/float16 activations: we need to upcast
# activations to float32 before matmul and downcast back to bfloat16/float16
# - channelwise with float16 activations, we need to upcast activations to
# float32 before matmul and downcast back to float16
# Note: these activations will be dynamically quantized to int8 by the kernel.
c = self.config
is_groupwise = c.group_size != c.partition_weight_shape[0]
# dtype of activations before they get dynamically quantized to int8
original_pre_quant_act_dtype = x.dtype
pre_quant_act_dtype = original_pre_quant_act_dtype
if (
is_groupwise and pre_quant_act_dtype == torch.bfloat16
) or pre_quant_act_dtype == torch.float16:
pre_quant_act_dtype = torch.float32
x_2d = x.reshape(-1, x.shape[-1])
if pre_quant_act_dtype != original_pre_quant_act_dtype:
x_2d = x_2d.to(pre_quant_act_dtype)
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
w_q = getattr(layer, self.w_q_name)
output = torch.ops.aten._dyn_quant_matmul_4bit(
x_2d,
w_q,
c.group_size,
c.partition_weight_shape[0],
c.partition_weight_shape[1],
).reshape(out_shape)
if pre_quant_act_dtype != original_pre_quant_act_dtype:
output = output.to(original_pre_quant_act_dtype)
return output
@@ -0,0 +1,175 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils.quant_utils import (
pack_quantized_values_into_int32,
)
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
class ExllamaLinearKernel(MPLinearKernel):
SUPPORTED_QUANT_TYPES = [scalar_types.uint4b8, scalar_types.uint8b128]
# In theory supports `scalar_types.uint2b2, scalar_types.uint3b4` too but
# currently untested so not added to the list
@classmethod
def get_min_capability(cls) -> int:
return 60
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if not current_platform.is_cuda_alike():
return (
False,
"Exllama is only supported on CUDA and ROCm",
)
if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:
return (
False,
"Act reordering currently not supported by Exllama, "
"when the input features are partitioned across "
"devices",
)
if c.partition_weight_shape[1] % (32 // c.weight_type.size_bits) != 0:
return (
False,
"Output features must be a multiple of the pack "
"factor (32 / num_bits) so that we can correctly "
"pack the zero points",
)
if c.act_type != torch.float16:
return False, "Exllama only supports float16 activations"
if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
return (
False,
f"Quant type ({c.weight_type}) not supported by "
"Exllama, supported types are: "
f"{cls.SUPPORTED_QUANT_TYPES}",
)
if c.group_size <= 0:
return (
False,
f"Group size ({c.group_size}) must be positive, "
"Exllama does not support channelwise quantization",
)
if c.full_weight_shape[0] % c.group_size != 0:
return (
False,
f"Group size ({c.group_size}) does not evenly divide"
" the number of input features "
f"({c.full_weight_shape[0]})",
)
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module):
c = self.config
# For Exllama, we need to set a zero-point tensor if there is not one
if not c.zero_points:
self.w_zp_name = "qzeros"
device = getattr(layer, self.w_q_name).device
groups = c.partition_weight_shape[0] // c.group_size
out_features = c.partition_weight_shape[1]
if c.weight_type.has_bias():
# if the type has a bias we have to create a zeros tensor that
# contains the bias values repeated for each group (-1 due to
# a bug in the original GPTQ checkpoint format leading to
# exllama kernel adding 1 to the zero points during inference)
# Documentation of the bug can be found here:
# https://garden.danieldk.eu/GPTQ-Checkpoint-Format
zeros = torch.full(
(groups, out_features),
c.weight_type.bias - 1,
dtype=torch.int32,
device=device,
)
else:
raise NotImplementedError(
"A 0 zero-point is not supported by Exllama due to "
"a bug in the original GPTQ checkpoint format leading to "
"exllama kernel adding 1 to the zero points during "
"inference"
)
zeros = pack_quantized_values_into_int32(zeros, c.weight_type, packed_dim=1)
setattr(
layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False)
)
if c.has_g_idx:
def transform_w_g_idx(x):
# Exllama wants the permutation array instead of the group
# indices
return torch.argsort(x).to(torch.int)
self._transform_param(layer, self.w_gidx_name, transform_w_g_idx) # type: ignore
else:
self.w_gidx_name = "g_idx"
empty_g_idx = torch.nn.Parameter(
torch.empty((0,), dtype=torch.int, device=device), requires_grad=False
)
setattr(layer, self.w_gidx_name, empty_g_idx)
def transform_w_q(x):
assert isinstance(x, BasevLLMParameter)
assert self.w_gidx_name is not None
g_idx = getattr(layer, self.w_gidx_name)
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
x_cont = x.data.contiguous()
ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits)
return x_cont
def transform_w_s(x):
assert isinstance(x, BasevLLMParameter)
permute_param_layout_(x, input_dim=0, output_dim=1)
x.data = x.data.contiguous()
return x.to(dtype=c.act_type)
# Repack weights and scales for Machete
self._transform_param(layer, self.w_q_name, transform_w_q)
self._transform_param(layer, self.w_s_name, transform_w_s)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
c = self.config
x_2d = x.reshape(-1, x.shape[-1])
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer)
# gptq_gemm supports GPTQv2 format by passing use_v2_format=True.
# However, the MPLinearLayerConfig doesn't contain format info.
# So hardcode GPTQv1 format here, to keep its behavior unchanged.
use_v2_format = False
assert w_zp is not None, "Zero points are required by Exllama"
assert w_g_idx is not None, "Group index is required by Exllama"
output = ops.gptq_gemm(
x_2d, w_q, w_zp, w_s, w_g_idx, True, use_v2_format, c.weight_type.size_bits
)
if bias is not None:
output.add_(bias)
return output.reshape(out_shape)
@@ -0,0 +1,64 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Humming GEMM as a mixed-precision WNA16Int linear kernel."""
import torch
from vllm.platforms import current_platform
from vllm.utils.import_utils import has_humming
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
class HummingLinearKernel(MPLinearKernel):
@classmethod
def get_min_capability(cls) -> int:
return 75
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "Humming is only supported on CUDA"
if not has_humming():
return False, "Humming is not installed"
if c.has_g_idx:
return False, "Humming does not support act-order (g_idx)"
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
from vllm.model_executor.layers.quantization.utils.humming_utils import (
convert_linear_layer_to_humming_standard,
prepare_humming_layer,
)
name_map = {"weight": self.w_q_name, "weight_scale": self.w_s_name}
group_size = self.config.group_size
quant_config = {
"quant_method": "humming",
"dtype": "int" + str(self.config.weight_type.size_bits),
"group_size": 0 if group_size == -1 else group_size,
}
if self.config.zero_points:
assert self.w_zp_name is not None
name_map["zero_point"] = self.w_zp_name
quant_config["has_zero_point"] = True
convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
prepare_humming_layer(layer, quant_config)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
from vllm.utils.humming import HummingMethod
flatten_inputs = x.view(-1, x.size(-1))
output = HummingMethod.forward_layer(
layer=layer,
inputs=flatten_inputs,
compute_config=layer.compute_config,
)
return output.view(*x.shape[:-1], output.size(-1))
@@ -0,0 +1,159 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from functools import partial
import torch
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils.machete_utils import (
check_machete_supports_shape,
query_machete_supported_group_sizes,
query_machete_supported_quant_types,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
pack_quantized_values_into_int32,
unpack_quantized_values_into_int32,
)
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
from vllm.platforms import current_platform
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
class MacheteLinearKernel(MPLinearKernel):
@classmethod
def get_min_capability(cls) -> int:
return 90
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
# Machete uses CUTLASS, so it can only be compatible with Nvidia
if not current_platform.is_cuda():
return False, "Machete only supported on CUDA"
if not current_platform.is_device_capability(90):
return False, "Machete requires compute capability of 90 (Hopper)"
if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:
return (
False,
"Act reordering currently not supported by Machete, "
"when the input features are partitioned across "
"devices",
)
if c.weight_type not in query_machete_supported_quant_types(c.zero_points):
return (
False,
f"Quant type ({c.weight_type}) not supported by "
"Machete, supported types are: "
f"{query_machete_supported_quant_types(c.zero_points)}",
)
if c.group_size not in query_machete_supported_group_sizes(c.act_type):
return (
False,
f"Group size ({c.group_size}) not supported by "
"Machete, supported group sizes are: "
f"{query_machete_supported_group_sizes(c.act_type)}",
)
return check_machete_supports_shape(
c.partition_weight_shape[0], c.partition_weight_shape[1]
)
# note assumes that
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
# `weight_scale` is: {input_dim = 0, output_dim = 1}
# `weight_zp` is: {input_dim = 0, output_dim = 1, packed_dim = 1}
def process_weights_after_loading(self, layer: torch.nn.Module):
c = self.config
if c.has_g_idx:
assert self.w_gidx_name is not None
perm = torch.argsort(getattr(layer, self.w_gidx_name)).to(torch.int)
self.act_perm = lambda x: x[:, perm]
# use `ops.permute_cols` if possible
if (
c.act_type in [torch.float16, torch.bfloat16]
and c.partition_weight_shape[0] % 8 == 0
):
self.act_perm = partial(ops.permute_cols, perm=perm)
def transform_w_q(x):
assert isinstance(x, BasevLLMParameter)
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
if c.has_g_idx:
x_unpacked = unpack_quantized_values_into_int32(
x.data, c.weight_type, packed_dim=0
)
x_perm = x_unpacked[perm, :]
x.data = pack_quantized_values_into_int32(
x_perm, c.weight_type, packed_dim=0
)
x.data = ops.machete_prepack_B(
x.data.t().contiguous().t(),
a_type=c.act_type,
b_type=c.weight_type,
group_scales_type=c.act_type,
)
return x
def transform_w_s(x):
assert isinstance(x, BasevLLMParameter)
permute_param_layout_(x, input_dim=0, output_dim=1)
x.data = x.data.contiguous()
return x
def transform_w_zp(x):
assert isinstance(x, BasevLLMParameter)
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=1)
x_unpacked = unpack_quantized_values_into_int32(
x.data, c.weight_type, packed_dim=1
)
w_s = getattr(layer, self.w_s_name).data
# pre-apply scales to zero-points
x.data = (-1.0 * w_s * (x_unpacked.to(w_s.dtype))).contiguous()
return x
# Repack weights and scales for Machete
self._transform_param(layer, self.w_q_name, transform_w_q)
self._transform_param(layer, self.w_s_name, transform_w_s)
if c.zero_points:
self._transform_param(layer, self.w_zp_name, transform_w_zp)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
c = self.config
w_q, w_s, w_zp, _ = self._get_weight_params(layer)
x_2d = x.reshape(-1, x.shape[-1])
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
if c.has_g_idx:
x_2d = self.act_perm(x_2d)
if c.zero_points:
assert w_zp is not None
else:
w_zp = None
output = ops.machete_mm(
a=x_2d,
b_q=w_q,
b_type=c.weight_type,
b_group_zeros=w_zp,
b_group_scales=w_s,
b_group_size=c.group_size,
)
if bias is not None:
output.add_(bias) # In-place add
return output.reshape(out_shape)
@@ -0,0 +1,245 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
MARLIN_SUPPORTED_GROUP_SIZES,
apply_gptq_marlin_linear,
check_marlin_supports_shape,
marlin_act_int8_process_scales,
marlin_is_k_full,
marlin_make_empty_g_idx,
marlin_make_workspace_new,
marlin_pad_dim,
marlin_pad_qweight,
marlin_pad_scales,
marlin_padded_nk,
marlin_permute_bias,
marlin_permute_scales,
marlin_sort_g_idx,
marlin_zero_points,
query_marlin_supported_quant_types,
unpack_cols,
)
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
class MarlinLinearKernel(MPLinearKernel):
@classmethod
def get_min_capability(cls) -> int:
return 75
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
# Marlin uses inline PTX, so it can only be compatible with Nvidia
if not current_platform.is_cuda():
return False, "Marlin only supported on CUDA"
quant_types = query_marlin_supported_quant_types(c.zero_points)
if c.weight_type not in quant_types:
return (
False,
f"Quant type ({c.weight_type}) not supported by"
f" Marlin, supported types are: {quant_types}",
)
if c.group_size not in MARLIN_SUPPORTED_GROUP_SIZES:
return (
False,
f"Group size ({c.group_size}) not supported by "
"Marlin, supported group sizes are: "
f"{MARLIN_SUPPORTED_GROUP_SIZES}",
)
if c.has_g_idx:
# Act-order couples K to the full-model group layout, so tile
# padding is not supported; keep the strict shape check.
return check_marlin_supports_shape(
c.partition_weight_shape[1], # out_features
c.partition_weight_shape[0], # in_features
c.full_weight_shape[0], # in_features
c.group_size,
)
# A group straddling TP ranks cannot be fixed by padding.
if (
c.group_size != -1
and c.group_size < c.full_weight_shape[0]
and c.partition_weight_shape[0] % c.group_size != 0
):
return False, (
f"in_features per partition {c.partition_weight_shape[0]} is "
f"not divisible by group_size = {c.group_size}."
)
# Tile misalignment is fixed by zero-padding at weight prep.
return True, None
# note assumes that
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
# `weight_scale` is: {input_dim = 0, output_dim = 1}
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
device = getattr(layer, self.w_q_name).device
c = self.config
is_a_8bit = c.act_type is not None and c.act_type.itemsize == 1
if is_a_8bit:
assert c.weight_type == scalar_types.uint4b8, (
"W8A8 is not supported by marlin kernel."
)
if c.act_type == torch.float8_e4m3fn:
ops.marlin_int4_fp8_preprocess(getattr(layer, self.w_q_name), inplace=True)
getattr(layer, self.w_s_name).data = (
getattr(layer, self.w_s_name).data * 512
)
row_parallel = c.partition_weight_shape[0] != c.full_weight_shape[0]
self.is_k_full = marlin_is_k_full(c.has_g_idx, row_parallel)
size_k, size_n = c.partition_weight_shape
if c.has_g_idx:
# Act-order shapes were strictly validated in can_implement.
padded_n, padded_k = size_n, size_k
else:
padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size)
# Allocate marlin workspace.
self.workspace = marlin_make_workspace_new(device)
# Default names since marlin requires empty parameters for these,
# TODO: remove this requirement from marlin (allow optional tensors)
if self.w_gidx_name is None:
self.w_gidx_name = "g_idx"
if self.w_zp_name is None:
self.w_zp_name = "w_zp"
def transform_w_q(x):
assert isinstance(x, BasevLLMParameter)
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
x.data = ops.gptq_marlin_repack(
marlin_pad_qweight(
x.data.contiguous(), size_n, size_k, padded_n, padded_k
),
perm=layer.g_idx_sort_indices,
size_k=padded_k,
size_n=padded_n,
num_bits=c.weight_type.size_bits,
is_a_8bit=is_a_8bit,
)
return x
def transform_w_s(x):
assert isinstance(x, BasevLLMParameter)
permute_param_layout_(x, input_dim=0, output_dim=1)
x.data = marlin_permute_scales(
marlin_pad_scales(
x.data.contiguous(),
size_n,
size_k,
padded_n,
padded_k,
c.group_size,
),
size_k=padded_k,
size_n=padded_n,
group_size=c.group_size,
is_a_8bit=is_a_8bit,
)
if c.group_size == -1:
num_groups = 1
else:
num_groups = c.partition_weight_shape[0] // c.group_size
if c.act_type == torch.int8 and num_groups > 1:
x.data, input_global_scale = marlin_act_int8_process_scales(x.data)
layer.register_parameter(
"input_global_scale",
torch.nn.Parameter(input_global_scale, requires_grad=False),
)
else:
layer.input_global_scale = None
return x
if c.has_g_idx:
g_idx, g_idx_sort_indices = marlin_sort_g_idx(
getattr(layer, self.w_gidx_name)
)
self._transform_param(layer, self.w_gidx_name, lambda _: g_idx)
layer.g_idx_sort_indices = g_idx_sort_indices
else:
setattr(layer, self.w_gidx_name, marlin_make_empty_g_idx(device))
layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
if c.zero_points:
grouped_k = size_k // c.group_size if c.group_size != -1 else 1
padded_grouped_k = padded_k // c.group_size if c.group_size != -1 else 1
self._transform_param(
layer,
self.w_zp_name,
lambda x: marlin_zero_points(
marlin_pad_scales(
unpack_cols(
x.t(),
c.weight_type.size_bits,
grouped_k,
size_n,
),
size_n,
size_k,
padded_n,
padded_k,
c.group_size,
),
size_k=padded_grouped_k,
size_n=padded_n,
num_bits=c.weight_type.size_bits,
is_a_8bit=is_a_8bit,
),
)
else:
setattr(layer, self.w_zp_name, marlin_make_empty_g_idx(device))
self._transform_param(layer, self.w_q_name, transform_w_q)
self._transform_param(layer, self.w_s_name, transform_w_s)
if hasattr(layer, "bias") and layer.bias is not None:
layer.bias.data = marlin_permute_bias(
marlin_pad_dim(layer.bias, size_n, padded_n)
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
c = self.config
w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
# `process_weights_after_loading` will ensure w_zp and w_gidx are not
# None for marlin
return apply_gptq_marlin_linear(
input=x,
weight=w_q,
weight_scale=w_s,
weight_zp=w_zp, # type: ignore
g_idx=w_gidx, # type: ignore
g_idx_sort_indices=layer.g_idx_sort_indices,
workspace=self.workspace,
wtype=c.weight_type,
input_size_per_partition=c.partition_weight_shape[0],
output_size_per_partition=c.partition_weight_shape[1],
is_k_full=self.is_k_full,
input_global_scale=getattr(layer, "input_global_scale", None),
bias=bias,
input_dtype=c.act_type,
)
@@ -0,0 +1,193 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""W4A16 GPTQ kernel for AMD RDNA3 (gfx1100) — fp16 + bf16.
Drop-in replacement for ExllamaLinearKernel on RDNA3 that adds native bf16
support. The HIP kernel lives in ``csrc/rocm/q_gemm_rdna3.cu``
and is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3``.
Registered ahead of TritonW4A16LinearKernel for the ROCm-RDNA3 path; falls
through to the Triton kernel on non-RDNA3 ROCm devices (e.g. CDNA/MI300).
"""
import torch
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils.quant_utils import (
pack_quantized_values_into_int32,
)
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
class RDNA3W4A16LinearKernel(MPLinearKernel):
SUPPORTED_QUANT_TYPES = [scalar_types.uint4b8]
@classmethod
def get_min_capability(cls) -> int:
# ROCm gates via on_gfx1100() in can_implement.
return 60
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if not current_platform.is_rocm():
return False, "RDNA3 W4A16 kernel is ROCm-only"
from vllm.platforms.rocm import on_gfx1100
if not on_gfx1100():
return False, "RDNA3 W4A16 kernel requires gfx1100"
# The HIP op is registered by the C++ extension; if a user is running
# against a vLLM build that doesn't include it (e.g. partial rebuild),
# fall through gracefully to the next kernel in the registry.
if not (
hasattr(torch.ops, "_rocm_C")
and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
):
return (
False,
"torch.ops._rocm_C.gptq_gemm_rdna3 missing — rebuild C++ extension",
)
if c.act_type not in (torch.float16, torch.bfloat16):
return False, "RDNA3 W4A16 kernel only supports fp16 and bf16"
if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
return (
False,
f"Quant type ({c.weight_type}) not supported by "
f"RDNA3 W4A16 kernel; supported: {cls.SUPPORTED_QUANT_TYPES}",
)
if c.group_size <= 0:
return (
False,
"RDNA3 W4A16 kernel does not support channelwise quantization",
)
if c.full_weight_shape[0] % c.group_size != 0:
return (
False,
f"Group size ({c.group_size}) does not evenly divide K "
f"({c.full_weight_shape[0]})",
)
# Output features must be a multiple of the pack factor (8 nibbles per
# int32) and of 8 so that qzeros (packed 4-bit per col) align cleanly
# against the BLOCK_KN_SIZE*4 = 512 N-stride and per-thread 4 columns.
if c.partition_weight_shape[1] % 8 != 0:
return (
False,
"Output features must be a multiple of 8 for the RDNA3 "
"W4A16 kernel (qzeros packing)",
)
if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:
return (
False,
"Act-order with TP-partitioned input features is not "
"supported by the RDNA3 W4A16 kernel",
)
return True, None
# ----- Weight prep (identical layout/shuffle as ExllamaLinearKernel) -----
def process_weights_after_loading(self, layer: torch.nn.Module):
c = self.config
device = getattr(layer, self.w_q_name).device
# Synthesize zero points if the checkpoint doesn't carry them.
if not c.zero_points:
self.w_zp_name = "qzeros"
groups = c.partition_weight_shape[0] // c.group_size
out_features = c.partition_weight_shape[1]
if c.weight_type.has_bias():
# GPTQv1 quirk: the kernel adds 1 to the stored zero, so we
# encode (bias - 1) here. See exllama.py for the link to the
# documentation of this checkpoint-format wart.
zeros = torch.full(
(groups, out_features),
c.weight_type.bias - 1,
dtype=torch.int32,
device=device,
)
else:
raise NotImplementedError(
"RDNA3 W4A16 kernel: zero-bias 4-bit quant requires "
"explicit zero points (GPTQv1 +1 quirk)."
)
zeros = pack_quantized_values_into_int32(zeros, c.weight_type, packed_dim=1)
setattr(
layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False)
)
# Act-order: convert g_idx to the inverse permutation array exllama
# expects (kernel reads a[perm[k]] instead of using groups indirected
# by g_idx[k]).
if c.has_g_idx:
def transform_w_g_idx(x):
return torch.argsort(x).to(torch.int)
self._transform_param(layer, self.w_gidx_name, transform_w_g_idx) # type: ignore
else:
self.w_gidx_name = "g_idx"
empty_g_idx = torch.nn.Parameter(
torch.empty((0,), dtype=torch.int, device=device),
requires_grad=False,
)
setattr(layer, self.w_gidx_name, empty_g_idx)
def transform_w_q(x):
assert isinstance(x, BasevLLMParameter)
assert self.w_gidx_name is not None
g_idx = getattr(layer, self.w_gidx_name)
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
x_cont = x.data.contiguous()
# Same 4-bit shuffle as exllama. The RDNA3 kernel reads weights in
# the same shuffled int32 layout and uses the (qa & 0x000F000F)
# bit-trick on top.
ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits)
return x_cont
def transform_w_s(x):
assert isinstance(x, BasevLLMParameter)
permute_param_layout_(x, input_dim=0, output_dim=1)
x.data = x.data.contiguous()
# Keep scales in the activation dtype (fp16 OR bf16) — the kernel
# branches on dtype internally.
return x.to(dtype=c.act_type)
self._transform_param(layer, self.w_q_name, transform_w_q)
self._transform_param(layer, self.w_s_name, transform_w_s)
# ----- Forward --------------------------------------------------------
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
c = self.config
x_2d = x.reshape(-1, x.shape[-1])
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer)
assert w_zp is not None, "Zero points are required by RDNA3 W4A16"
assert w_g_idx is not None, "g_idx tensor (possibly empty) required"
output = ops.gptq_gemm_rdna3(x_2d, w_q, w_zp, w_s, w_g_idx, False)
if bias is not None:
output.add_(bias)
return output.reshape(out_shape)
@@ -0,0 +1,438 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Triton-based W4A16 GEMM kernel for ROCm MI300.
Implements fused int4-weight dequantization + fp16 GEMM in a single kernel,
using GPTQ sequential packing (8 int4 values per int32, shifts [0,4,...,28]).
Plugs into the MPLinearKernel selection system and is preferred over
MarlinLinearKernel/ExllamaLinearKernel on ROCm.
Weight layout expected by this kernel (post-process_weights_after_loading):
qweight: [K, N//8] int32 — rows=K (input), cols=N//8 (N is packed)
scales: [K//G, N] fp16/bf16
qzeros: [K//G, N//8] int32 (optional; None for symmetric uint4b8)
Checkpoint layout from compressed_tensors_wNa16 create_weights:
weight_packed: [N, K//8] int32 (output_dim=0, input_dim=1, packed_dim=1)
weight_scale: [N, K//G] fp16 (output_dim=0, input_dim=1)
weight_zero_point: [N//8, K//G] int32 (output_dim=0, packed_dim=0)
"""
import torch
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types
from vllm.triton_utils import tl, triton
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
TRITON_W4A16_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128, 256]
TRITON_W4A16_SUPPORTED_QUANT_TYPES = [
scalar_types.uint4b8, # symmetric GPTQ (bias=8)
scalar_types.uint4, # asymmetric with explicit zeros
]
@triton.jit
def triton_w4a16_gemm_kernel(
# Pointers
a_ptr, # [M, K] fp16/bf16 activations
b_ptr, # [K, N//8] int32 packed 4-bit weights (N is the packed dim)
scales_ptr, # [K//G, N] fp16/bf16 scales
zeros_ptr, # [K//G, N//8] int32 packed zeros (unused when HAS_ZP=False)
c_ptr, # [M, N] fp16/bf16 output
# Dimensions
M,
N,
K,
# Strides
stride_am,
stride_ak,
stride_bk,
stride_bn, # stride in b along the packed N//8 dim
stride_cm,
stride_cn,
# Quantization parameters
group_size,
# Whether explicit zero points are provided
HAS_ZP: tl.constexpr,
# Zero bias used when HAS_ZP is False (e.g. 8 for uint4b8)
ZP_BIAS: tl.constexpr,
# Block sizes (tuned for MI300 wavefront=64)
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_K: tl.constexpr,
):
"""
Fused W4A16 GEMM: C[M,N] = A[M,K] @ dequant(B)[K,N]
B is stored as [K, N//8] int32 using GPTQ sequential packing:
each int32 packs 8 consecutive N-values at bit offsets [0,4,8,12,16,20,24,28].
Dequant: w_fp = (w_int4 - zero) * scale
HAS_ZP=True: zero is loaded from zeros_ptr and unpacked
HAS_ZP=False: zero = ZP_BIAS constant (e.g. 8 for uint4b8 symmetric)
"""
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
# Row/col offsets for this tile
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
# b/zeros are stored with N packed: N//8 int32 columns per K row
offs_bn = pid_n * (BLOCK_N // 8) + tl.arange(0, BLOCK_N // 8)
# GPTQ sequential shifts tiled across BLOCK_N:
# [0,4,8,...,28] repeating for every group of 8 N-values.
# Build 1D shifts_1d of length BLOCK_N: column j gets shift (j % 8) * 4.
shifts_row = tl.arange(0, 8) * 4 # [8]
shifts_1d_2d = tl.broadcast_to(shifts_row[None, :], (BLOCK_N // 8, 8))
shifts_1d = tl.reshape(shifts_1d_2d, (BLOCK_N,)) # [BLOCK_N]
# Broadcast to [BLOCK_K, BLOCK_N] for weight unpacking
shifts = tl.broadcast_to(shifts_1d[None, :], (BLOCK_K, BLOCK_N))
# Scales column offsets: full N-width (one scale per output neuron)
offs_sn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k_start in range(0, tl.cdiv(K, BLOCK_K)):
offs_k = k_start * BLOCK_K + tl.arange(0, BLOCK_K)
mask_k = offs_k < K
# ---- Load activations A: [BLOCK_M, BLOCK_K] ----
a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
mask_a = (offs_m[:, None] < M) & mask_k[None, :]
a = tl.load(a_ptrs, mask=mask_a, other=0.0)
# ---- Load packed weights B: [BLOCK_K, BLOCK_N//8] int32 ----
b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
mask_b = mask_k[:, None] & (offs_bn[None, :] < N // 8)
b_packed = tl.load(b_ptrs, mask=mask_b, other=0)
# ---- Unpack int4 weights → [BLOCK_K, BLOCK_N] ----
# tl.interleave(x, x) doubles the last dim by interleaving.
# Starting from [BLOCK_K, BLOCK_N//8], three interleaves give
# [BLOCK_K, BLOCK_N], where each int32 is replicated 8 times.
b = tl.interleave(b_packed, b_packed)
b = tl.interleave(b, b)
b = tl.interleave(b, b)
# Extract the correct 4-bit nibble for each output column
b = (b >> shifts) & 0xF
# ---- Compute scale/zero group row index ----
g_idx = (k_start * BLOCK_K) // group_size
# ---- Load scales: [BLOCK_N] → broadcast to [BLOCK_K, BLOCK_N] ----
scale_offset = g_idx * N + offs_sn
scale_mask = offs_sn < N
scales = tl.load(scales_ptr + scale_offset, mask=scale_mask, other=1.0)
scales = tl.broadcast_to(scales[None, :], (BLOCK_K, BLOCK_N))
# ---- Load / compute zeros ----
if HAS_ZP:
# Load packed zeros row: [BLOCK_N//8] int32
zero_offset = g_idx * (N // 8) + offs_bn
zero_mask = offs_bn < N // 8
z_packed = tl.load(zeros_ptr + zero_offset, mask=zero_mask, other=0)
# Unpack to [BLOCK_N] using same interleave+shift pattern
z = tl.interleave(z_packed, z_packed)
z = tl.interleave(z, z)
z = tl.interleave(z, z)
z = (z >> shifts_1d) & 0xF
z = tl.broadcast_to(z[None, :], (BLOCK_K, BLOCK_N))
else:
z = tl.full((BLOCK_K, BLOCK_N), ZP_BIAS, dtype=tl.int32)
# ---- Dequantize: (w - zero) * scale ----
b_fp = (b - z).to(a.dtype) * scales
# ---- Accumulate ----
accumulator += tl.dot(a, b_fp, out_dtype=tl.float32)
# ---- Store output C: [BLOCK_M, BLOCK_N] ----
c = accumulator.to(c_ptr.type.element_ty)
c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
mask_c = (offs_m[:, None] < M) & (offs_n[None, :] < N)
tl.store(c_ptrs, c, mask=mask_c)
def triton_w4a16_gemm(
a: torch.Tensor, # [M, K] fp16/bf16
b_q: torch.Tensor, # [K, N//8] int32
scales: torch.Tensor, # [K//G, N] fp16/bf16
qzeros: torch.Tensor | None, # [K//G, N//8] int32, or None
group_size: int,
zp_bias: int = 8, # bias for uint4b8 when qzeros is None
) -> torch.Tensor:
"""
Fused W4A16 GEMM using GPTQ-packed int4 weights.
Args:
a: Activation matrix [M, K], float16 or bfloat16.
b_q: Packed weight matrix [K, N//8], int32 (GPTQ sequential).
scales: Per-group scales [K//G, N], same dtype as a.
qzeros: Per-group packed zero points [K//G, N//8] int32, or None
for symmetric quantization (uses zp_bias instead).
group_size: Quantization group size (resolved from -1 to K by caller).
zp_bias: Constant zero used when qzeros is None (default 8 for uint4b8).
Returns:
Output matrix [M, N], same dtype as a.
"""
assert a.is_contiguous(), "Activation matrix must be contiguous"
assert b_q.is_contiguous(), "Weight matrix must be contiguous"
assert scales.is_contiguous(), "Scales must be contiguous"
M, K = a.shape
N = b_q.shape[1] * 8
assert b_q.shape == (K, N // 8), (
f"b_q shape mismatch: {b_q.shape} vs ({K}, {N // 8})"
)
assert scales.shape == (K // group_size, N), (
f"scales shape mismatch: {scales.shape} vs ({K // group_size}, {N})"
)
if qzeros is not None:
assert qzeros.shape == (K // group_size, N // 8), (
f"qzeros shape mismatch: {qzeros.shape}"
)
c = torch.empty((M, N), dtype=a.dtype, device=a.device)
has_zp = qzeros is not None
# Provide a dummy pointer when HAS_ZP=False (Triton requires a valid ptr)
zeros_ptr = qzeros if has_zp else b_q
if current_platform.is_rocm():
from vllm.platforms.rocm import on_gfx1x
if on_gfx1x():
# Tuned for RDNA 3.5 (gfx1151, 40 CUs, 32-wide wavefronts).
if M <= 32:
BLOCK_M, BLOCK_N, BLOCK_K = 32, 32, 64
elif M <= 64:
BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
else:
BLOCK_M, BLOCK_N, BLOCK_K = 128, 32, 64
else:
# Tuned for MI300 (gfx942, 304 CUs, 64-wide wavefronts).
if M <= 32:
BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32
elif M <= 64:
BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
else:
BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32
else:
if M <= 32:
BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32
elif M <= 64:
BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
else:
BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32
# The kernel loads scales/zeros for a single group per BLOCK_K tile
# (one g_idx per iteration). If BLOCK_K > group_size, rows at the tail
# of the tile dequantize with the wrong group's scales, silently
# corrupting the output. Clamp BLOCK_K to group_size to keep one
# scale group per tile.
if group_size < BLOCK_K:
BLOCK_K = group_size
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
triton_w4a16_gemm_kernel[grid](
a,
b_q,
scales,
zeros_ptr,
c,
M,
N,
K,
a.stride(0),
a.stride(1),
b_q.stride(0),
b_q.stride(1),
c.stride(0),
c.stride(1),
group_size=group_size,
HAS_ZP=has_zp,
ZP_BIAS=zp_bias,
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
BLOCK_K=BLOCK_K,
)
return c
class TritonW4A16LinearKernel(MPLinearKernel):
"""
Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer).
Supports GPTQ-format int4 weights (uint4b8 symmetric, uint4 asymmetric)
with grouped quantization. Weight tensors are transposed from the
compressed-tensors checkpoint layout to the kernel's [K, N//8] layout.
"""
SUPPORTED_QUANT_TYPES = TRITON_W4A16_SUPPORTED_QUANT_TYPES
@classmethod
def get_min_capability(cls) -> int:
# Triton handles capability checks itself
return 0
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if not (current_platform.is_rocm() or current_platform.is_cuda()):
return False, "TritonW4A16LinearKernel requires CUDA or ROCm"
if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
return (
False,
f"Quant type {c.weight_type} not supported; "
f"supported: {cls.SUPPORTED_QUANT_TYPES}",
)
if c.act_type not in (torch.float16, torch.bfloat16):
return False, "Only float16/bfloat16 activations are supported"
N = c.partition_weight_shape[1]
if N % 8 != 0:
return (
False,
f"Output features ({N}) must be divisible by 8 "
"(8 int4 values packed per int32)",
)
if c.has_g_idx:
return (
False,
"Activation reordering (g_idx) is not supported by "
"TritonW4A16LinearKernel",
)
gs = c.group_size
if (
gs not in TRITON_W4A16_SUPPORTED_GROUP_SIZES
and gs != c.full_weight_shape[0]
):
return (
False,
f"Group size {gs} not supported; "
f"supported: {TRITON_W4A16_SUPPORTED_GROUP_SIZES} "
f"or full K ({c.full_weight_shape[0]})",
)
K = c.partition_weight_shape[0]
eff_gs = gs if gs != -1 else K
if K % eff_gs != 0:
return (False, f"Input features {K} not divisible by group size {eff_gs}")
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""
Convert compressed-tensors checkpoint layout to kernel layout.
Checkpoint (from compressed_tensors_wNa16.create_weights):
weight_packed: [N, K//8] int32 input_dim=1, output_dim=0, packed_dim=1
weight_scale: [N, K//G] fp16 input_dim=1, output_dim=0
weight_zero_point: [N//8, K//G] int32 output_dim=0, packed_dim=0
Kernel needs:
qweight: [K, N//8] int32 (transpose weight_packed)
scales: [K//G, N] fp16 (transpose weight_scale)
qzeros: [K//G, N//8] int32 (transpose weight_zero_point)
"""
# ---- Transform qweight: [N, K//8] → [K//8, N] → back to [K, N//8] ----
# permute_param_layout_(x, input_dim=0, output_dim=1) rearranges so that
# the input(K) dimension is at physical dim 0 and output(N) at dim 1.
# Checkpoint has input_dim=1, output_dim=0, packed_dim=1 (K is packed).
# After permute we get [K//8, N] (K packed at dim 0, N at dim 1).
# The kernel wants [K, N//8] (K at dim 0, N packed at dim 1), so we
# then transpose: [K//8, N].T = [N, K//8] — that's not right.
#
# Actually we need to change WHAT is packed:
# Original packing: K packed into K//8 (8 K-values per int32)
# Kernel packing: N packed into N//8 (8 N-values per int32)
# These require a full repack, not just a transpose.
#
# Simple approach: unpack → transpose the full [N, K] → repack as [K, N//8].
# This is done CPU-side at load time (one-time cost).
def repack_w_q(x: BasevLLMParameter) -> BasevLLMParameter:
# x.data is [N, K//8] int32, K packed (GPTQ checkpoint format)
# Step 1: bring to [N, K//8] with output(N) at dim 0
permute_param_layout_(x, input_dim=1, output_dim=0, packed_dim=1)
w = x.data # [N, K//8] int32
N_dim, K8 = w.shape
K_dim = K8 * 8
# Step 2: unpack to [N, K] int32 (vectorized)
shifts = torch.arange(8, device=w.device, dtype=torch.int32) * 4
w_unpacked = ((w.unsqueeze(-1) >> shifts) & 0xF).reshape(N_dim, K_dim)
# Step 3: transpose to [K, N] int32
w_KN = w_unpacked.t().contiguous()
# Step 4: repack N into N//8 int32 values → [K, N//8] (vectorized)
N8 = N_dim // 8
w_repacked = torch.sum(
(w_KN.view(K_dim, N8, 8) & 0xF) << shifts,
dim=2,
dtype=torch.int32,
)
x.data = w_repacked.contiguous()
return x
def repack_w_s(x: BasevLLMParameter) -> BasevLLMParameter:
# x.data is [N, K//G] fp16, bring to [K//G, N]
permute_param_layout_(x, input_dim=1, output_dim=0)
x.data = x.data.t().contiguous()
return x
self._transform_param(layer, self.w_q_name, repack_w_q)
self._transform_param(layer, self.w_s_name, repack_w_s)
if self.w_zp_name is not None:
zp = getattr(layer, self.w_zp_name, None)
if zp is not None:
# Checkpoint: [N//8, K//G] int32 (N packed at dim 0, K//G at dim 1)
# Kernel needs: [K//G, N//8] — just transpose
replace_parameter(
layer,
self.w_zp_name,
torch.nn.Parameter(zp.data.t().contiguous(), requires_grad=False),
)
def apply_weights(
self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None
) -> torch.Tensor:
c = self.config
w_q, w_s, w_zp, _ = self._get_weight_params(layer)
x_2d = x.reshape(-1, x.shape[-1]).contiguous()
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
K = c.partition_weight_shape[0]
group_size = c.group_size if c.group_size != -1 else K
# For symmetric types (uint4b8), use the scalar bias; no zeros tensor
zp_bias = c.weight_type.bias if c.weight_type.has_bias() else 0
output = triton_w4a16_gemm(
a=x_2d,
b_q=w_q,
scales=w_s,
qzeros=w_zp,
group_size=group_size,
zp_bias=zp_bias,
)
if bias is not None:
output.add_(bias)
return output.reshape(out_shape)
@@ -0,0 +1,231 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from torch.nn.parameter import Parameter
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
_XPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8)
logger = init_logger(__name__)
class XPUwNa16LinearKernel(MPLinearKernel):
@classmethod
def get_min_capability(cls) -> int:
return -1
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if not current_platform.is_xpu():
return False, "XPUwNa16 only supported on XPU"
if c.act_type != torch.bfloat16 and c.act_type != torch.float16:
return False, "XPUwNa16 only supports BF16/FP16 activations"
if c.weight_type not in _XPUWNA16_SUPPORTED_QUANT_TYPES:
return (
False,
f"Quant type ({c.weight_type}) not supported by "
"XPUwNa16, supported types are: "
f"{_XPUWNA16_SUPPORTED_QUANT_TYPES}",
)
if c.group_size != -1 and c.group_size % 32 != 0:
return (
False,
f"Group size ({c.group_size}) not supported by "
"XPUwNa16, supported group sizes are multiples of 32",
)
if c.partition_weight_shape[0] % 32 != 0:
return (
False,
f"Input size ({c.partition_weight_shape[0]}) not supported by "
"XPUwNa16, supported sizes are multiples of 32",
)
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module):
# Default names since marlin requires empty parameters for these,
# TODO: remove this requirement from marlin (allow optional tensors)
if self.w_gidx_name is None:
self.w_gidx_name = "g_idx"
if self.w_zp_name is None:
self.w_zp_name = "w_zp"
need_transpose = False
qweight_shape = getattr(layer, self.w_q_name).shape
scale_shape = getattr(layer, self.w_s_name).shape
# gptq marlin and compressed tensors wna16 expect different default
# layouts for weight and scale, so we check the shapes to determine
# if we need to transpose
if qweight_shape[0] != scale_shape[0]:
need_transpose = True
if need_transpose:
getattr(layer, self.w_q_name).data = (
getattr(layer, self.w_q_name).data.t().contiguous()
)
getattr(layer, self.w_s_name).data = getattr(layer, self.w_s_name).data
else:
getattr(layer, self.w_s_name).data = (
getattr(layer, self.w_s_name).data.t().contiguous()
)
if self.config.zero_points:
# (FIXME): maybe zero points should also be transposed.
getattr(layer, self.w_zp_name).data = (
getattr(layer, self.w_zp_name).data.t().contiguous()
)
else:
weight_zero_point = torch.Tensor([8]).to(torch.int8).to("xpu")
setattr(
layer, self.w_zp_name, Parameter(weight_zero_point, requires_grad=False)
)
if self.config.has_g_idx:
setattr(
layer,
self.w_gidx_name,
Parameter(
getattr(layer, self.w_gidx_name).data.t().contiguous(),
requires_grad=False,
),
)
else:
setattr(layer, self.w_gidx_name, None)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
reshaped_x = x.reshape(-1, x.shape[-1])
w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
out = torch.ops._xpu_C.int4_gemm_w4a16(
reshaped_x,
w_q.t(),
bias if bias is not None else None,
w_s,
w_zp,
self.config.group_size,
w_gidx,
)
return out
class XPUW4A8IntLinearKernel(MPLinearKernel):
"""XPU kernel for W4A8 integer quantization using oneDNN int4_gemm_w4a8.
Weights are symmetric group-quantized int4 packed as uint4.
Activations are dynamically quantized per-token to symmetric int8.
"""
@classmethod
def get_min_capability(cls) -> int:
return -1
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
if not current_platform.is_xpu():
return False, "XPUW4A8Int only supported on XPU"
if c.act_type not in (torch.bfloat16, torch.float16):
return False, "XPUW4A8Int requires BF16/FP16 activations"
if c.weight_type != scalar_types.int4:
return (
False,
f"XPUW4A8Int requires int4 weights, got {c.weight_type}",
)
if c.zero_points:
return False, "XPUW4A8Int only supports symmetric weight quantization"
if c.group_size != -1 and c.group_size % 32 != 0:
return (
False,
f"Group size ({c.group_size}) not supported by XPUW4A8Int, "
"must be a multiple of 32",
)
in_size, out_size = c.partition_weight_shape
if in_size % 8 != 0 or out_size % 8 != 0:
return (
False,
f"in/out sizes ({in_size}, {out_size}) must be multiples of 8",
)
if c.act_type != torch.float16:
logger.warning_once(
"XPUW4A8IntLinearKernel is running with model dtype %s, "
"but int4_gemm_w4a8 produces float16 output. Recommend "
"setting --dtype float16 for best performance.",
c.act_type,
)
return True, None
def _pack_int4_weight(self, w: torch.Tensor) -> torch.Tensor:
# w is [N, K] int8 with values in [-8, 7]
w_u4 = w.to(torch.int32) + 8 # shift to [0, 15]
w_u4 = w_u4.reshape(w.shape[0], w.shape[1] // 8, 8) # [N, K/8, 8]
shifts = torch.arange(0, 32, 4, dtype=torch.int32, device=w.device)
packed = ((w_u4 & 0xF) << shifts[None, None, :]).sum(dim=2).to(torch.int32)
return packed
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
layer.weight_scale.data = layer.weight_scale.data.t().contiguous()
device = layer.weight_packed.device
# TODO: support asymmetric quantization
weight_zero_point = torch.tensor([8], dtype=torch.int8, device=device)
layer.weight_zero_point = Parameter(weight_zero_point, requires_grad=False)
# weight_packed is [out, in] int8, signed int4 values in [-8, 7]
w = layer.weight_packed.data # [out, in]
# TODO: implement asym case
packed = self._pack_int4_weight(w) # [out, in/8] packed uint4
replace_parameter(
layer,
self.w_q_name,
torch.nn.Parameter(packed, requires_grad=False),
)
# Free the original unpacked int8 weight (still registered as "weight")
# to avoid double-storing both int8 [N, K] and int32 [N, K/8] in memory.
layer.register_parameter("weight", None)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
reshaped_x = x.reshape(-1, x.shape[-1]) # [M, K]
from vllm._xpu_ops import xpu_ops as ops
# TODO: static and asymmetric quantization case
# Common code for CompressedTensorsW4A8Int does not read act symmetry data
quant_x, x_scale, x_zero = ops.dynamic_per_token_int8_quant_ref(
reshaped_x, True, 8
)
out = torch.ops._xpu_C.int4_gemm_w4a8(
quant_x,
x_scale,
x_zero,
layer.weight_packed.t(),
layer.weight_scale,
layer.weight_zero_point,
self.config.group_size,
None, # g_idx not currently supported
bias,
)
return out.to(x.dtype)
@@ -0,0 +1,211 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Zentorch W4A16 GPTQ weight-only-quantized linear kernel for AMD Zen CPUs.
Selected by ``choose_mp_linear_kernel`` ahead of the generic oneDNN-backed
``CPUWNA16LinearKernel``. When ``can_implement`` rejects a layer, the selector
falls through to the next kernel in ``_POSSIBLE_KERNELS[PlatformEnum.CPU]``.
"""
import torch
from vllm.logger import init_logger
from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types
from .cpu import CPUWNA16LinearKernel
from .MPLinearKernel import MPLinearLayerConfig
logger = init_logger(__name__)
def _import_unpack_from_int32():
"""Import compressed-tensors' ``unpack_from_int32`` across versions."""
try:
from compressed_tensors.compressors.pack_quantized.helpers import (
unpack_from_int32,
)
except ImportError:
from compressed_tensors.compressors.quantized_compressors.pack_quantized import ( # type: ignore[import-not-found] # noqa: E501
unpack_from_int32,
)
return unpack_from_int32
class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
"""W4A16 GPTQ kernel backed by ``torch.ops.zentorch.zentorch_woq_linear``."""
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
ok, reason = super().can_implement(c)
if not ok:
return ok, reason
if not current_platform.is_zen_cpu():
return False, "ZentorchWNA16 requires an AMD Zen CPU."
if not has_zentorch_op(["zentorch_woq_repack_weight", "zentorch_woq_linear"]):
return (
False,
"torch.ops.zentorch.{zentorch_woq_repack_weight, "
"zentorch_woq_linear} are not registered.",
)
if c.has_g_idx:
return False, "ZentorchWNA16 does not support activation re-ordering."
return True, None
def _zentorch_woq_eligible(self, layer: torch.nn.Module) -> bool:
"""Eligibility predicate for the zentorch W4A16 GPTQ fast path.
Constraints (any failure -> ``cpu_gemm_wna16`` path via ``super()``
with ``layer`` untouched).
"""
if (
self.w_gidx_name is not None
and getattr(layer, self.w_gidx_name, None) is not None
) or (getattr(self.config, "has_g_idx", False)):
return False
weight_packed = getattr(layer, self.w_q_name, None)
weight_scale = getattr(layer, self.w_s_name, None)
if weight_packed is None or weight_scale is None:
return False
bits = self.config.weight_type.mantissa
pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
# 4-bit -> 8 values per int32;
if pack_factor != 8:
return False
# GPTQ-only. AWQ packs along the output dim instead.
in_dim = getattr(weight_packed, "input_dim", None)
pk_dim = getattr(weight_packed, "packed_dim", None)
if in_dim is None or pk_dim is None or in_dim != pk_dim:
return False
is_ct_format = in_dim == pk_dim == 1
if not is_ct_format:
return False
if weight_packed.dim() != 2 or weight_scale.dim() != 2:
return False
# 4-bit -> 8 values per int32; in_features must be divisible by num_groups.
in_features = weight_packed.shape[1] * 8
num_groups = weight_scale.shape[1]
return num_groups > 0 and in_features % num_groups == 0
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""Repack CT GPTQ weights into the zentorch WOQ layout.
Falls back to ``CPUWNA16LinearKernel.process_weights_after_loading``
via ``super()`` when the layer doesn't satisfy
``_zentorch_woq_eligible``.
On success, ``layer._zentorch_processed_weights`` is set to ``True``
"""
if getattr(layer, "_zentorch_processed_weights", False):
return
if not self._zentorch_woq_eligible(layer):
logger.info_once(
"[zen_cpu] ZentorchWNA16 fast path not eligible for this "
"layer (AWQ pack layout, g_idx, or non-int32 storage); "
"falling back to CPUWNA16LinearKernel (cpu_gemm_wna16)."
)
super().process_weights_after_loading(layer)
return
if (not self.config.zero_points) and (self.w_zp_name is not None):
setattr(layer, self.w_zp_name, None)
if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
setattr(layer, self.w_gidx_name, None)
weight_q = getattr(layer, self.w_q_name)
weight_s = getattr(layer, self.w_s_name)
weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q
weight_scale = weight_s.data if hasattr(weight_s, "data") else weight_s
bits = self.config.weight_type.mantissa
pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
out_features, num_groups = weight_scale.shape[0], weight_scale.shape[1]
in_features = weight_packed.shape[1] * pack_factor
original_shape = torch.Size([out_features, in_features])
unpack_from_int32 = _import_unpack_from_int32()
repack_op = torch.ops.zentorch.zentorch_woq_repack_weight.default
weight_unpacked = unpack_from_int32(
weight_packed,
bits,
original_shape,
packed_dim=weight_q.packed_dim,
)
zp_param = (
getattr(layer, self.w_zp_name, None) if self.w_zp_name is not None else None
)
needs_unsigned_offset = self.config.weight_type == scalar_types.uint4
if needs_unsigned_offset:
weight_unpacked = (weight_unpacked.to(torch.int32) + 8).clamp(0, 15)
repacked = repack_op(weight_unpacked.to(torch.int8).contiguous())
if zp_param is None:
zp_tc = None
else:
zp_tensor = zp_param.data if hasattr(zp_param, "data") else zp_param
zp = unpack_from_int32(
zp_tensor,
bits,
(out_features, num_groups),
packed_dim=zp_param.packed_dim,
)
if needs_unsigned_offset:
zp = (zp.to(torch.int32) + 8).clamp(0, 15)
zp_tc = zp.to(torch.int8).t().contiguous()
layer._zentorch_woq_packed = repacked.t()
layer._zentorch_woq_scale = weight_scale.t().contiguous()
layer._zentorch_woq_zero_point = zp_tc
for param_name in (self.w_q_name, self.w_s_name, self.w_zp_name):
if param_name is None:
continue
param = getattr(layer, param_name, None)
if param is None:
continue
if hasattr(param, "data"):
param.data = torch.empty(0)
else:
setattr(layer, param_name, torch.empty(0))
layer._zentorch_kind = "compressed_tensors_w4a16_gptq"
layer._zentorch_processed_weights = True
logger.info_once(
"[zen_cpu] Using zentorch_woq_linear for W4A16 GPTQ "
"(weight_type=%s, has_zp=%s)",
self.config.weight_type,
zp_tc is not None,
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
if getattr(layer, "_zentorch_processed_weights", False):
return torch.ops.zentorch.zentorch_woq_linear.default(
x,
layer._zentorch_woq_packed,
layer._zentorch_woq_scale,
layer._zentorch_woq_zero_point,
bias,
)
return super().apply_weights(layer, x, bias)
__all__ = ["ZentorchWNA16LinearKernel"]
@@ -0,0 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.model_executor.kernels.linear.mxfp4.base import (
MxFp4LinearKernel,
MxFp4LinearLayerConfig,
)
__all__ = [
"MxFp4LinearKernel",
"MxFp4LinearLayerConfig",
]
@@ -0,0 +1,67 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import ABC, abstractmethod
from dataclasses import dataclass
import torch
@dataclass
class MxFp4LinearLayerConfig:
"""Configuration for an MXFP4 linear layer.
All MXFP4 layers share the same structure: packed uint8 weights (2 FP4 values per
byte) and per-block weight scales (group size 32).
"""
pass
class MxFp4LinearKernel(ABC):
"""Base class for MXFP4 quantized linear kernels.
Each subclass implements a specific GEMM backend (CUTLASS, Marlin, etc).
The kernel selection mechanism iterates over registered subclasses in
priority order,calling ``is_supported`` and ``can_implement`` to find the best
match for the current hardware.
"""
def __init__(self, config: MxFp4LinearLayerConfig) -> None:
assert self.can_implement(config)[0]
assert self.is_supported()[0]
self.config = config
@classmethod
@abstractmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
"""Return whether this kernel can run on the current platform."""
raise NotImplementedError
@classmethod
@abstractmethod
def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
"""Return whether this kernel can handle *config*."""
raise NotImplementedError
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""Transform weights into the format required by this kernel.
Called once after checkpoint weights have been loaded onto the
device. Implementations should repack / swizzle / pad weights
and scales in-place on *layer*.
"""
raise NotImplementedError
@abstractmethod
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
"""Run the quantized GEMM."""
raise NotImplementedError
@@ -0,0 +1,74 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from torch.nn.parameter import Parameter
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
swizzle_mxfp4_scales,
)
from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer_cutedsl
from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
_MXFP4_GROUP_SIZE = 32
class FlashInferMxFp4LinearKernel(MxFp4LinearKernel):
"""MXFP4 W4A4 GEMM via FlashInfer CUTLASS (SM100+)."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if current_platform.has_device_capability(100) and has_flashinfer_cutedsl():
return True, None
return False, "FlashInfer + >=sm_100 (Blackwell) required"
@classmethod
def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
N, scale_K = layer.weight_scale.shape
K = scale_K * _MXFP4_GROUP_SIZE
# swizzle pads N to the next multiple of 128 for CUTLASS tiling
padded_N = ((N + 127) // 128) * 128
layer.weight_scale = Parameter(
swizzle_mxfp4_scales(layer.weight_scale.data, N, K).reshape(padded_N, -1),
requires_grad=False,
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
from vllm.utils.flashinfer import (
flashinfer_mxfp4_quantize,
flashinfer_scaled_fp4_mm,
)
weight = layer.weight
out_shape = x.shape[:-1] + (layer.output_size_per_partition,)
x_2d = x.reshape(-1, x.shape[-1])
x_fp4, x_scale = flashinfer_mxfp4_quantize(x_2d.contiguous())
out = flashinfer_scaled_fp4_mm(
x_fp4,
weight,
x_scale,
layer.weight_scale,
alpha=None,
out_dtype=x.dtype,
backend="cute-dsl",
block_size=_MXFP4_GROUP_SIZE,
use_nvfp4=False,
)
if bias is not None:
out = out + bias
return out.view(out_shape)
@@ -0,0 +1,63 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.model_executor.layers.quantization.utils.humming_utils import (
convert_linear_layer_to_humming_standard,
prepare_humming_layer,
)
from vllm.platforms import current_platform
from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
class HummingMxFp4LinearKernel(MxFp4LinearKernel):
"""Humming GEMM Kernel for MXFP4."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "Humming only supported on CUDA"
if not current_platform.has_device_capability(75):
return False, "Humming only supported on SM75+"
return True, None
@classmethod
def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu)
name_map = {"weight": "weight", "weight_scale": "weight_scale"}
quant_config = {
"quant_method": "humming",
"dtype": "float4e2m1",
"scale_dtype": "float8e8m0",
"group_size": 32,
"weight_scale_type": "group",
}
convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
prepare_humming_layer(layer, quant_config)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
from vllm.utils.humming import HummingMethod
flatten_inputs = x.view(-1, x.size(-1))
output = HummingMethod.forward_layer(
layer=layer,
inputs=flatten_inputs,
compute_config=layer.compute_config,
)
return output.view(*x.shape[:-1], output.size(-1))
@@ -0,0 +1,52 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
class MarlinMxFp4LinearKernel(MxFp4LinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
is_fp4_marlin_supported,
)
if is_fp4_marlin_supported():
return True, None
return False, "Marlin FP4 not available"
@classmethod
def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
prepare_fp4_layer_for_marlin,
)
prepare_fp4_layer_for_marlin(layer)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
apply_fp4_marlin_linear,
)
return apply_fp4_marlin_linear(
input=x,
weight=layer.weight,
weight_scale=layer.weight_scale,
weight_global_scale=None,
workspace=layer.workspace,
size_n=layer.output_size_per_partition,
size_k=layer.input_size_per_partition,
bias=bias,
)
@@ -0,0 +1,53 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
xpu_mxfp4_quantize as quant_mxfp4,
)
from vllm.model_executor.utils import replace_parameter
from vllm.platforms import current_platform
from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
class XPUMxFp4LinearKernel(MxFp4LinearKernel):
"""MXFP4 W4A4 GEMM on XPU."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_xpu():
return False, "XPUMxFp4 only support on XPU"
return True, None
@classmethod
def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight = layer.weight.view(torch.float4_e2m1fn_x2)
replace_parameter(layer, "weight", weight.data.t())
weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu)
weight_scale = weight_scale.t().contiguous()
replace_parameter(layer, "weight_scale", weight_scale.data)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
out_dtype = x.dtype
x_fp4, x_blockscale = quant_mxfp4(x)
return torch.ops._xpu_C.fp4_gemm(
x_fp4,
layer.weight,
x_blockscale,
layer.weight_scale,
out_dtype,
bias,
)
@@ -0,0 +1,56 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import ABC, abstractmethod
from dataclasses import dataclass
import torch
@dataclass
class Mxfp8LinearLayerConfig:
"""Configuration for an MXFP8 linear layer.
All MXFP8 layers share the same structure: FP8-E4M3 weights with
uint8 (E8M0) per-block scales at block size 32.
"""
pass
class Mxfp8LinearKernel(ABC):
"""Base class for MXFP8 quantized linear kernels.
Each subclass implements a specific GEMM backend (FlashInfer CUTLASS,
Marlin, emulation).
"""
def __init__(self, c: Mxfp8LinearLayerConfig) -> None:
assert self.can_implement(c)[0]
assert self.is_supported()[0]
self.config = c
@classmethod
@abstractmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
raise NotImplementedError
@classmethod
@abstractmethod
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
raise NotImplementedError
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
raise NotImplementedError
@abstractmethod
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
raise NotImplementedError
@@ -0,0 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
Mxfp8LinearKernel,
Mxfp8LinearLayerConfig,
)
__all__ = [
"Mxfp8LinearKernel",
"Mxfp8LinearLayerConfig",
]
@@ -0,0 +1,84 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from torch.nn.parameter import Parameter
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
MXFP8_BLOCK_SIZE,
MXFP8_SCALE_DTYPE,
dequant_mxfp8_to_bf16,
)
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
"""Software emulation fallback for MXFP8 (dequant to BF16)."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
return True, None
@classmethod
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight = layer.weight.data # [N, K]
N, K = weight.shape
scale_k = K // MXFP8_BLOCK_SIZE
weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
# Dequantize MXFP8 -> BF16 ONCE here, at load time, so apply_weights runs
# a plain BF16 linear with no per-step dequant -- i.e. run as if from a
# BF16 checkpoint. The 1-byte MXFP8 weight is replaced by BF16 (2x its
# size, but linear weights are small vs the MoE experts); the tiny E8M0
# scale is kept for the dtype/ndim asserts but is otherwise unused.
# Opt out (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0) to keep the MXFP8
# weight and dequant per-step in apply_weights instead.
import vllm.envs as envs
if envs.VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD:
weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale)
layer.weight = Parameter(weight.contiguous(), requires_grad=False)
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
weight = layer.weight
# Load-time dequant path: weights are already BF16/FP16 (>= 2-byte), so
# run a plain linear -- no per-step dequant. (MXFP8 weights are 1-byte.)
if weight.element_size() >= 2:
# F.linear requires x and weight share a dtype; .to() is a no-op when
# they already match (e.g. both BF16).
output = torch.nn.functional.linear(x, weight.to(x.dtype), bias)
return output.to(x.dtype)
# Fallback: weights still in MXFP8 -- dequant on the fly (other archs /
# if a future caller skips the load-time conversion above).
weight_scale = layer.weight_scale
if weight_scale.dtype != MXFP8_SCALE_DTYPE:
raise ValueError(
f"Emulation backend requires {MXFP8_SCALE_DTYPE} "
f"weight_scale dtype, got {weight_scale.dtype}."
)
if weight_scale.ndim != 2:
raise ValueError(
f"Emulation backend requires 2D weight_scale, "
f"got {weight_scale.ndim}D. "
f"Ensure process_weights_after_loading was called."
)
# Cast to x's dtype: dequant yields BF16, but F.linear needs both operands
# to match (e.g. an FP16 model). No-op when x is already BF16.
weight_bf16 = dequant_mxfp8_to_bf16(weight, weight_scale).to(x.dtype)
output = torch.nn.functional.linear(x, weight_bf16, bias)
return output.to(x.dtype)
@@ -0,0 +1,175 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from torch.nn.parameter import Parameter
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
MXFP8_BLOCK_SIZE,
mxfp8_e4m3_quantize,
swizzle_mxfp8_scale,
)
from vllm.platforms import current_platform
from vllm.utils import flashinfer as vllm_flashinfer
from vllm.utils.flashinfer import has_flashinfer_cutedsl
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
"""MXFP8 W8A8 GEMM via FlashInfer CUTLASS (SM100+)."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if current_platform.has_device_capability(100):
return True, None
return False, "requires >=sm_100 (Blackwell)"
@classmethod
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight = layer.weight.data # [N, K]
N, K = weight.shape
scale_k = K // MXFP8_BLOCK_SIZE
weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous()
weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K)
layer.weight = Parameter(weight.contiguous(), requires_grad=False)
layer.weight_scale = Parameter(
weight_scale_swizzled.contiguous(), requires_grad=False
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
weight = layer.weight
weight_scale = layer.weight_scale
out_dtype = x.dtype
N, K = weight.shape
input_shape = x.shape
input_2d = x.view(-1, K)
min_dim = 128
assert min_dim <= K, (
f"mm_mxfp8 requires K >= {min_dim}, got K={K}. "
f"in_features is too small for mm_mxfp8."
)
assert K % MXFP8_BLOCK_SIZE == 0, (
f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}."
)
assert min_dim <= N, (
f"mm_mxfp8 requires N >= {min_dim}, got N={N}. "
f"out_features is too small for mm_mxfp8."
)
input_mxfp8, input_scale = mxfp8_e4m3_quantize(
input_2d, is_sf_swizzled_layout=True
)
if not weight.is_contiguous():
weight = weight.contiguous()
output = vllm_flashinfer.mm_mxfp8(
input_mxfp8,
weight.t(),
input_scale,
weight_scale,
out_dtype=out_dtype,
backend="cutlass",
)
if bias is not None:
output = output + bias
output_shape = (*input_shape[:-1], N)
return output.view(output_shape)
class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
"""MXFP8 W8A8 GEMM via FlashInfer CuTe-DSL (SM100/SM103)."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not (
current_platform.is_cuda()
and current_platform.is_device_capability_family(100)
):
return False, "requires sm_100/sm_103 (Blackwell)"
if not has_flashinfer_cutedsl():
return False, "requires FlashInfer CuTe-DSL module"
return True, None
@classmethod
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight = layer.weight.data # [N, K]
N, K = weight.shape
scale_k = K // MXFP8_BLOCK_SIZE
weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous()
weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K)
# Store weight column-major [K, N] as mm_mxfp8 expects for operand B.
layer.weight = Parameter(weight.contiguous().t(), requires_grad=False)
layer.weight_scale = Parameter(
weight_scale_swizzled.contiguous(), requires_grad=False
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
weight = layer.weight # [K, N], column-major
weight_scale = layer.weight_scale
out_dtype = x.dtype
K, N = weight.shape
input_shape = x.shape
input_2d = x.view(-1, K)
min_dim = 128
assert min_dim <= K, (
f"mm_mxfp8 requires K >= {min_dim}, got K={K}. "
f"in_features is too small for mm_mxfp8."
)
assert K % MXFP8_BLOCK_SIZE == 0, (
f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}."
)
assert min_dim <= N, (
f"mm_mxfp8 requires N >= {min_dim}, got N={N}. "
f"out_features is too small for mm_mxfp8."
)
input_mxfp8, input_scale = mxfp8_e4m3_quantize(
input_2d, is_sf_swizzled_layout=True
)
output = vllm_flashinfer.mm_mxfp8(
input_mxfp8,
weight,
input_scale,
weight_scale,
out_dtype=out_dtype,
backend="cute-dsl",
)
if bias is not None:
output = output + bias
output_shape = (*input_shape[:-1], N)
return output.view(output_shape)
@@ -0,0 +1,63 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.model_executor.layers.quantization.utils.humming_utils import (
convert_linear_layer_to_humming_standard,
prepare_humming_layer,
)
from vllm.platforms import current_platform
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
class HummingMxfp8LinearKernel(Mxfp8LinearKernel):
"""Humming GEMM Kernel for MXFP8."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "Humming only supported on CUDA"
if not current_platform.has_device_capability(75):
return False, "Humming only supported on SM75+"
return True, None
@classmethod
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu)
name_map = {"weight": "weight", "weight_scale": "weight_scale"}
quant_config = {
"quant_method": "humming",
"dtype": "float8e4m3",
"scale_dtype": "float8e8m0",
"group_size": 32,
"weight_scale_type": "group",
}
convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
prepare_humming_layer(layer, quant_config)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
from vllm.utils.humming import HummingMethod
flatten_inputs = x.view(-1, x.size(-1))
output = HummingMethod.forward_layer(
layer=layer,
inputs=flatten_inputs,
compute_config=layer.compute_config,
)
return output.view(*x.shape[:-1], output.size(-1))
@@ -0,0 +1,53 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
class MarlinMxfp8LinearKernel(Mxfp8LinearKernel):
"""MXFP8 W8A16 GEMM via Marlin (SM80+)."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
is_fp8_marlin_supported,
)
if is_fp8_marlin_supported():
return True, None
return False, "Marlin FP8 not available"
@classmethod
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
prepare_mxfp8_layer_for_marlin,
)
prepare_mxfp8_layer_for_marlin(layer)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
apply_mxfp8_marlin_linear,
)
return apply_mxfp8_marlin_linear(
input=x,
weight=layer.weight,
weight_scale=layer.weight_scale,
workspace=layer.workspace,
size_n=layer.output_size_per_partition,
size_k=layer.input_size_per_partition,
bias=bias,
)
@@ -0,0 +1,247 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Native MXFP8 linear GEMM for AMD CDNA4 (gfx950) via Triton ``tl.dot_scaled``.
Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16);
activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling
matrix cores. Falls back (via the kernel selector) to the BF16
``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with
``K % 128 != 0``.
"""
import torch
from torch.nn.parameter import Parameter
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
MXFP8_BLOCK_SIZE,
MXFP8_SCALE_DTYPE,
dequant_mxfp8_to_bf16,
mxfp8_e4m3_quantize,
)
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
@triton.jit
def _mxfp8_linear_kernel(
x_ptr,
xs_ptr,
w_ptr,
ws_ptr,
out_ptr,
M,
N,
K,
stride_xm,
stride_xk,
stride_xsm,
stride_xsk,
stride_wn,
stride_wk,
stride_wsn,
stride_wsk,
stride_om,
stride_on,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_K: tl.constexpr,
):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, BLOCK_K)
offs_sk = tl.arange(0, BLOCK_K // 32)
m_mask = offs_m < M
n_mask = offs_n < N
x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk
xs_ptrs = xs_ptr + offs_m[:, None] * stride_xsm + offs_sk[None, :] * stride_xsk
w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk
ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :] * stride_wsk
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for _ in range(0, tl.cdiv(K, BLOCK_K)):
x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0)
ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
acc += tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3")
x_ptrs += BLOCK_K * stride_xk
w_ptrs += BLOCK_K * stride_wk
xs_ptrs += (BLOCK_K // 32) * stride_xsk
ws_ptrs += (BLOCK_K // 32) * stride_wsk
o_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
tl.store(
o_ptrs, acc.to(out_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :]
)
def _mxfp8_dot_scaled_linear(
x: torch.Tensor, # [M, K] bf16/fp16
w: torch.Tensor, # [N, K] fp8 e4m3
w_scale: torch.Tensor, # [N, K//32] uint8 (E8M0)
) -> torch.Tensor:
M, K = x.shape
N = w.shape[0]
x_q, x_scale = mxfp8_e4m3_quantize(x)
out = torch.empty((M, N), dtype=x.dtype, device=x.device)
BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _select_cfg(M, N, K)
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
_mxfp8_linear_kernel[grid](
x_q,
x_scale,
w,
w_scale,
out,
M,
N,
K,
x_q.stride(0),
x_q.stride(1),
x_scale.stride(0),
x_scale.stride(1),
w.stride(0),
w.stride(1),
w_scale.stride(0),
w_scale.stride(1),
out.stride(0),
out.stride(1),
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
BLOCK_K=BLOCK_K,
num_warps=num_warps,
num_stages=num_stages,
)
return out
def _select_cfg(M, N, K):
"""(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) — graph-tuned on gfx950.
The M-bucketed, shape-adaptive tile selection here is the speedup over the
upstream 2-bucket launcher. Tiles are pipelined (num_stages>=2, larger BLOCK_K)
and occupancy- and shape-aware: keyed on the LOCAL (M, N, K), so it adapts to the
TP-sharded shapes (e.g. MiniMax-M3 TP=4 vs TP=8, where local N and K differ) —
large-K prefill uses BLOCK_K=256; short-K (K=768) widens N. BLOCK_K must divide K
(the K-loop is unmasked), so every BLOCK_K below is guarded to be K-divisible
(served K: 384/768/1024/2048/6144).
"""
if M <= 64:
# decode (M in {1,32,64}): tiny-M GEMV is weight-BW + GPU-OCCUPANCY bound. The
# lever is NARROW BLOCK_N=16 (maximize N-tiles so more CUs stream the weight in
# parallel) + LARGE BLOCK_K (fewer K-iters, bigger coalesced weight loads).
# Tuned by CUDA-graph replay latency. Optimal at both TP=4 and TP=8.
if K % 1024 == 0: # K=2048, 6144 -> graph-best 16x16x1024 (all M)
return 16, 16, 1024, 2, 2
if K % 512 == 0:
return 16, 16, 512, 2, 3
if K % 256 == 0: # K=768 (shared_down) -> graph-best 16x32x256
return 16, 32, 256, 4, 3
return 16, 32, 128, 4, 3
# mid-M (65..256) on SMALL local-N: still occupancy-bound (a 64x64 tile makes too
# few N-tiles), so the narrow-BLOCK_N decode-style tile fills the CUs better.
# N<=1536 covers the real fused-qkv local N at TP=8: q heads shard but the GQA KV
# (4) + sparse-indexer (4) heads are < TP=8, so vLLM replicates them to 1/rank ->
# N = 1024 + 4*128 = 1536 (not 2560/2=1280). For the wider 1280<N<=1536 band the
# narrow tile only wins up to M=128 (at M=256 the 64x64 tile is better), so cap it
# there; N<=1280 keeps the narrow tile through M=256. TP=4 qkv N=2560 is unchanged.
if (M <= 256 and N <= 1280) or (M <= 128 and N <= 1536):
if K % 1024 == 0:
return 16, 16, 1024, 2, 2
if K % 512 == 0:
return 16, 16, 512, 2, 3
if K % 256 == 0:
return 16, 16, 256, 2, 3
return 16, 16, 128, 2, 3
# right-sized launch grid (host-side), used to gate the tall 256-BLOCK_M tile.
occ = triton.cdiv(M, 256) * triton.cdiv(N, 128)
if K <= 1024: # short-K (shared_down K=384/768; TP=8 o_proj K=1024)
if M <= 256:
return (64, 64, 256, 8, 2) if K % 256 == 0 else (64, 64, 128, 8, 2)
# large prefill. (The former 256x128x256 tall tile was faster only on triton
# 3.6; on triton 3.7 its large BLOCK_M register/LDS footprint spills or hits
# "out of resources", so use 128x128x256 -- within the known-good footprint.)
if M >= 4096 and K >= 1024 and K % 256 == 0 and occ >= 256:
return 128, 128, 256, 8, 3
return 128, 128, 128, 8, 3
# large-K (K >= 2048). BLOCK_K is K-divisibility-guarded (the K-loop is unmasked):
# served large-K is 2048/6144 (%256==0), but fall back to 128 (always divides, since
# the entry requires K%128==0) for any other K to stay correct.
if M <= 256: # conc~128 decode + small prefill chunk: occupancy tile
if K % 512 == 0:
return 64, 64, 512, 8, 2
return (64, 64, 256, 8, 2) if K % 256 == 0 else (64, 64, 128, 8, 2)
if M <= 1024: # medium prefill chunk: BN=64 keeps small-N occupied
return (128, 64, 256, 8, 3) if K % 256 == 0 else (128, 64, 128, 8, 3)
# large prefill (M > 1024). The previously graph-tuned tall 256x128x256 and deep
# 128x128x512 tiles won only on triton 3.6; on triton 3.7 their larger BLOCK_M /
# BLOCK_K register+LDS footprint spills (or hits "out of resources" on stricter
# ROCm/triton builds). The 128x128x256 tile is equal-or-faster on triton 3.7 (the
# M=4096,N=2048,K=6144 shape: ~104us vs the 256x128x256 tile's ~184us), within ~5%
# on 3.6, and inside the footprint of the tiles used elsewhere in this selector.
# Covers the qkv-class local N=1536 (TP=8 qkv / TP=4 shared_gate_up) and the deep-K
# / very-large-M shapes.
if K % 256 == 0 and (1280 < N <= 1536 or (occ >= 128 and (K >= 4096 or M >= 4096))):
return 128, 128, 256, 8, 3
# small local-N (e.g. TP=8 shared_gate_up N=768): a 64-wide BLOCK_N doubles the
# N-tile count -> better CU fill than 128x128 at this mid-large M (~1.4x there).
if N <= 1024 and K % 256 == 0:
return 128, 64, 256, 8, 3
return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3)
class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
"""Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_rocm():
return False, "not ROCm"
# supports_mx() == gfx95x (CDNA4 native microscaling hardware). On other
# archs dot_scaled would upcast to BF16, so the kernel selector falls
# through to the BF16 emulation (hipBLASLt) path instead.
if not current_platform.supports_mx():
return False, "native MX requires CDNA4 (gfx95x)"
return True, None
@classmethod
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight = layer.weight.data # [N, K] fp8
N, K = weight.shape
scale_k = K // MXFP8_BLOCK_SIZE
weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
layer.weight = Parameter(weight.contiguous(), requires_grad=False)
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE:
raise ValueError(
f"Expected {MXFP8_SCALE_DTYPE} weight_scale, got "
f"{layer.weight_scale.dtype}."
)
out_shape = (*x.shape[:-1], layer.weight.shape[0])
x2d = x.reshape(-1, x.shape[-1])
if x2d.shape[-1] % 128 == 0:
out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
else:
# dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise.
w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale)
out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype)
out = out.reshape(out_shape)
if bias is not None:
out = out + bias
return out
@@ -0,0 +1,51 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
xpu_mxfp8_quantize as quant_mxfp8,
)
from vllm.model_executor.utils import replace_parameter
from vllm.platforms import current_platform
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
class XPUMxFp8LinearKernel(Mxfp8LinearKernel):
"""MXFP8 W8A8 GEMM on XPU."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_xpu():
return False, "XPUMxFp8 only support on XPU"
return True, None
@classmethod
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu)
weight_scale = weight_scale.t().contiguous()
replace_parameter(layer, "weight", layer.weight.t())
replace_parameter(layer, "weight_scale", weight_scale.data)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
out_dtype = x.dtype
x_fp8, x_scale = quant_mxfp8(x)
return torch.ops._xpu_C.fp8_gemm(
x_fp8,
layer.weight,
out_dtype,
x_scale,
layer.weight_scale,
bias,
)
@@ -0,0 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.model_executor.kernels.linear.nvfp4.base import (
NvFp4LinearKernel,
NvFp4LinearLayerConfig,
)
__all__ = [
"NvFp4LinearKernel",
"NvFp4LinearLayerConfig",
]
@@ -0,0 +1,76 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import ABC, abstractmethod
from dataclasses import dataclass
import torch
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
@dataclass
class NvFp4LinearLayerConfig:
"""Configuration for an NVFP4 linear layer.
All NVFP4 layers share the same structure: packed uint8 weights (2 FP4 values per
byte), FP8-E4M3 per-block weight scales (group size 16), and scalar global
scales for both weights and activations.
"""
pass
class NvFp4LinearKernel(ABC):
"""Base class for NVFP4 quantized linear kernels.
Each subclass implements a specific GEMM backend (CUTLASS, Marlin, etc).
The kernel selection mechanism iterates over registered subclasses in
priority order,calling ``is_supported`` and ``can_implement`` to find the best
match for the current hardware.
"""
def __init__(self, config: NvFp4LinearLayerConfig) -> None:
assert self.can_implement(config)[0]
assert self.is_supported()[0]
self.config = config
def input_quant_key(self) -> QuantKey | None:
"""Return the input quantization key supported by this kernel. If the kernel
does not support input quantization outside of the kernel, return None.
"""
return None
@classmethod
@abstractmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
"""Return whether this kernel can run on the current platform."""
raise NotImplementedError
@classmethod
@abstractmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
"""Return whether this kernel can handle *config*."""
raise NotImplementedError
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""Transform weights into the format required by this kernel.
Called once after checkpoint weights have been loaded onto the
device. Implementations should repack / swizzle / pad weights
and scales in-place on *layer*.
"""
raise NotImplementedError
@abstractmethod
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
"""Run the quantized GEMM."""
raise NotImplementedError
@@ -0,0 +1,77 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm._custom_ops import (
cutlass_scaled_fp4_mm,
scaled_fp4_quant,
)
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
cutlass_fp4_supported,
pad_nvfp4_weight_for_cutlass,
slice_nvfp4_output,
swizzle_blockscale,
)
from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
class CutlassNvFp4LinearKernel(NvFp4LinearKernel):
"""NVFP4 GEMM via the vLLM CUTLASS kernel."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if cutlass_fp4_supported():
return True, None
return False, "CUTLASS FP4 kernels not available"
@classmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
layer.weight_scale = torch.nn.Parameter(
swizzle_blockscale(layer.weight_scale.data), requires_grad=False
)
padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
layer.weight.data
)
layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
layer.weights_padding_cols = weights_padding_cols
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
output_size = layer.output_size_per_partition
output_dtype = x.dtype
output_shape = [*x.shape[:-1], output_size]
weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)
x_fp4, x_blockscale = scaled_fp4_quant(
x,
layer.input_global_scale_inv,
is_sf_swizzled_layout=True,
backend="cutlass",
padded_n=x.shape[-1] + weights_padding_bytes * 2,
)
out = cutlass_scaled_fp4_mm(
x_fp4,
layer.weight,
x_blockscale,
layer.weight_scale,
layer.alpha,
output_dtype,
)
out = slice_nvfp4_output(out, output_size)
if bias is not None:
out = out + bias
return out.view(*output_shape)
@@ -0,0 +1,49 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
kE2M1ToFloat_handle,
run_nvfp4_emulations,
)
from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
class EmulationNvFp4LinearKernel(NvFp4LinearKernel):
"""Software emulation fallback for NVFP4 (dequant → BF16 matmul)."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
# Always available as a last-resort fallback.
return True, None
@classmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# Move the E2M1 lookup table to the device now, because
# `.to(device)` is not allowed during CUDA graph capture.
kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
out = run_nvfp4_emulations(
x=x,
input_global_scale=layer.input_global_scale_inv,
weight=layer.weight,
weight_scale_swizzled=layer.weight_scale,
weight_global_scale=layer.weight_global_scale,
swizzle=False,
)
if bias is not None:
out = out + bias
return out
@@ -0,0 +1,69 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm._custom_ops import scaled_fp4_quant
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
slice_nvfp4_output,
swizzle_blockscale,
)
from vllm.utils.import_utils import has_fbgemm_gpu
from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
class FbgemmNvFp4LinearKernel(NvFp4LinearKernel):
"""NVFP4 GEMM via FBGEMM."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if has_fbgemm_gpu():
return True, None
return False, "fbgemm_gpu required"
@classmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
swizzled = swizzle_blockscale(layer.weight_scale.data)
layer.weight_scale = torch.nn.Parameter(
swizzled.view(-1).view(torch.uint8), requires_grad=False
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
import fbgemm_gpu # noqa: F401 - registers torch.ops.fbgemm.*
output_size = layer.output_size_per_partition
output_dtype = x.dtype
output_shape = [*x.shape[:-1], output_size]
x_fp4, x_blockscale = scaled_fp4_quant(
x,
layer.input_global_scale_inv,
is_sf_swizzled_layout=True,
backend="fbgemm",
)
out = torch.ops.fbgemm.f4f4bf16(
x_fp4,
layer.weight,
x_blockscale.view(-1).view(torch.uint8),
layer.weight_scale,
layer.alpha,
use_mx=False,
).to(output_dtype)
out = slice_nvfp4_output(out, output_size)
if bias is not None:
out = out + bias
return out.view(*output_shape)
@@ -0,0 +1,372 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm._custom_ops import scaled_fp4_quant
from vllm.model_executor.layers.fusion.quant_activation import (
QuantizedActivation,
as_quantized_activation,
)
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
pad_nvfp4_activation_for_cutlass,
pad_nvfp4_weight_for_cutlass,
slice_nvfp4_output,
swizzle_blockscale,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
QuantKey,
kNvfp4Dynamic,
)
from vllm.platforms import current_platform
from vllm.utils.flashinfer import (
flashinfer_scaled_fp4_mm,
has_flashinfer,
has_flashinfer_b12x_gemm,
)
from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel):
"""NVFP4 GEMM via FlashInfer's cutedsl backend."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_device_capability_family(100):
return False, "FlashInfer cutedsl requires sm_10x"
if not has_flashinfer():
return False, "FlashInfer required"
return True, None
@classmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# cutedsl uses the same swizzled + padded layout as cutlass.
layer.weight_scale = torch.nn.Parameter(
swizzle_blockscale(layer.weight_scale.data), requires_grad=False
)
padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
layer.weight.data
)
layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
layer.weights_padding_cols = weights_padding_cols
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
output_size = layer.output_size_per_partition
output_dtype = x.dtype
output_shape = [*x.shape[:-1], output_size]
x_fp4, x_blockscale = scaled_fp4_quant(
x,
layer.input_global_scale_inv,
is_sf_swizzled_layout=True,
backend="flashinfer-cutedsl",
)
x_fp4 = pad_nvfp4_activation_for_cutlass(
x_fp4, getattr(layer, "weights_padding_cols", 0)
)
out = flashinfer_scaled_fp4_mm(
x_fp4,
layer.weight,
x_blockscale,
layer.weight_scale,
layer.alpha,
output_dtype,
backend="cute-dsl",
)
out = slice_nvfp4_output(out, output_size)
if bias is not None:
out = out + bias
return out.view(*output_shape)
class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):
"""NVFP4 GEMM via FlashInfer's CUTLASS wrapper."""
def input_quant_key(self) -> QuantKey | None:
"""This kernel supports dynamic quantization of the input. By
convention, pre-quantized blockscales must use the swizzled layout."""
return kNvfp4Dynamic
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
cutlass_fp4_supported,
)
if (
cutlass_fp4_supported()
and current_platform.has_device_capability(100)
and has_flashinfer()
):
return True, None
return False, "FlashInfer + >=sm_100 required"
@classmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
layer.weight_scale = torch.nn.Parameter(
swizzle_blockscale(layer.weight_scale.data), requires_grad=False
)
padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
layer.weight.data
)
layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
layer.weights_padding_cols = weights_padding_cols
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor | QuantizedActivation,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
output_size = layer.output_size_per_partition
weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)
qa = as_quantized_activation(x, self.input_quant_key())
if qa is not None:
x_fp4, x_blockscale = qa.data, qa.scale
x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_bytes)
output_dtype = qa.orig_dtype
output_shape = [*qa.orig_shape[:-1], output_size]
else:
assert isinstance(x, torch.Tensor)
output_dtype = x.dtype
output_shape = [*x.shape[:-1], output_size]
x_fp4, x_blockscale = scaled_fp4_quant(
x,
layer.input_global_scale_inv,
is_sf_swizzled_layout=True,
backend="flashinfer-cutlass",
padded_n=x.shape[-1] + weights_padding_bytes * 2,
)
out = flashinfer_scaled_fp4_mm(
x_fp4,
layer.weight,
x_blockscale,
layer.weight_scale,
layer.alpha,
output_dtype,
backend="cutlass",
)
out = slice_nvfp4_output(out, output_size)
if bias is not None:
out = out + bias
return out.view(*output_shape)
class FlashInferTrtllmNvFp4LinearKernel(NvFp4LinearKernel):
"""NVFP4 GEMM via FlashInfer's TensorRT-LLM wrapper."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if has_flashinfer():
return True, None
return False, "FlashInfer required"
@classmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
weight = layer.weight.data
weight_scale = layer.weight_scale.data
epilogue_tile_m = 128
layer.weight = torch.nn.Parameter(
shuffle_matrix_a(weight.view(torch.uint8), epilogue_tile_m),
requires_grad=False,
)
layer.weight_scale = torch.nn.Parameter(
shuffle_matrix_sf_a(weight_scale.view(torch.uint8), epilogue_tile_m)
.reshape(weight_scale.shape)
.view(torch.float8_e4m3fn),
requires_grad=False,
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
output_size = layer.output_size_per_partition
output_dtype = x.dtype
output_shape = [*x.shape[:-1], output_size]
x_fp4, x_blockscale = scaled_fp4_quant(
x,
layer.input_global_scale_inv,
is_sf_swizzled_layout=True,
backend="flashinfer-trtllm",
)
out = flashinfer_scaled_fp4_mm(
x_fp4,
layer.weight,
x_blockscale,
layer.weight_scale,
layer.alpha,
output_dtype,
backend="trtllm",
)
out = slice_nvfp4_output(out, output_size)
if bias is not None:
out = out + bias
return out.view(*output_shape)
class FlashInferCudnnNvFp4LinearKernel(NvFp4LinearKernel):
"""NVFP4 GEMM via FlashInfer's cuDNN wrapper."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if has_flashinfer():
return True, None
return False, "FlashInfer required"
@classmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# cuDNN uses the same swizzled + padded layout as CUTLASS
layer.weight_scale = torch.nn.Parameter(
swizzle_blockscale(layer.weight_scale.data), requires_grad=False
)
padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
layer.weight.data
)
layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
layer.weights_padding_cols = weights_padding_cols
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
output_size = layer.output_size_per_partition
output_dtype = x.dtype
output_shape = [*x.shape[:-1], output_size]
weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)
x_fp4, x_blockscale = scaled_fp4_quant(
x,
layer.input_global_scale_inv,
is_sf_swizzled_layout=True,
backend="flashinfer-cudnn",
padded_n=x.shape[-1] + weights_padding_bytes * 2,
)
out = flashinfer_scaled_fp4_mm(
x_fp4,
layer.weight,
x_blockscale,
layer.weight_scale,
layer.alpha,
output_dtype,
backend="cudnn",
)
out = slice_nvfp4_output(out, output_size)
if bias is not None:
out = out + bias
return out.view(*output_shape)
class FlashInferB12xNvFp4LinearKernel(NvFp4LinearKernel):
"""NVFP4 GEMM via FlashInfer's b12x CuTe DSL warp-level MMA kernel (SM120+)."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if current_platform.has_device_capability(120) and has_flashinfer_b12x_gemm():
return True, None
return (
False,
"FlashInfer b12x requires SM120+ and FlashInfer "
"with Sm120BlockScaledDenseGemmKernel",
)
@classmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
layer.weight_scale = torch.nn.Parameter(
swizzle_blockscale(layer.weight_scale.data), requires_grad=False
)
padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
layer.weight.data
)
layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
layer.weights_padding_cols = weights_padding_cols
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
output_size = layer.output_size_per_partition
output_dtype = x.dtype
output_shape = [*x.shape[:-1], output_size]
x_fp4, x_blockscale = scaled_fp4_quant(
x,
layer.input_global_scale_inv,
is_sf_swizzled_layout=True,
backend="b12x",
)
x_fp4 = pad_nvfp4_activation_for_cutlass(
x_fp4, getattr(layer, "weights_padding_cols", 0)
)
out = flashinfer_scaled_fp4_mm(
x_fp4,
layer.weight,
x_blockscale,
layer.weight_scale,
layer.alpha,
output_dtype,
backend="b12x",
)
out = slice_nvfp4_output(out, output_size)
if bias is not None:
out = out + bias
return out.view(*output_shape)
@@ -0,0 +1,73 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils.humming_utils import (
prepare_humming_layer,
)
from vllm.platforms import current_platform
from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
logger = init_logger(__name__)
class HummingNvFp4LinearKernel(NvFp4LinearKernel):
"""Humming GEMM Kernel for NVFP4."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "Humming only supported on CUDA"
if not current_platform.has_device_capability(75):
return False, "Humming only supported on SM75+"
return True, None
@classmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# Route through humming's compressed-tensors nvfp4 loader (same path as
# the MoE oracle); the native group_tensor schema mishandles a scalar
# global scale.
quant_config = {
"quant_method": "compressed-tensors",
"format": "nvfp4-pack-quantized",
"type": "float",
"num_bits": 4,
"strategy": "group",
"group_size": 16,
}
# CT pack-quantized reads `weight_packed`; the scheme renamed it to `weight`.
if not hasattr(layer, "weight_packed"):
layer.weight_packed = layer.weight
del layer.weight
# The CT linear scheme inverts the global scale (1/scale) for
# marlin/cutlass; humming wants the original.
layer.weight_global_scale = torch.nn.Parameter(
1.0 / layer.weight_global_scale, requires_grad=False
)
prepare_humming_layer(layer, quant_config)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
from vllm.utils.humming import HummingMethod
flatten_inputs = x.view(-1, x.size(-1))
output = HummingMethod.forward_layer(
layer=layer,
inputs=flatten_inputs,
compute_config=layer.compute_config,
)
return output.view(*x.shape[:-1], output.size(-1))
@@ -0,0 +1,57 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
apply_fp4_marlin_linear,
is_fp4_marlin_supported,
prepare_fp4_layer_for_marlin,
)
from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
logger = init_logger(__name__)
class MarlinNvFp4LinearKernel(NvFp4LinearKernel):
"""NVFP4 weight-only GEMM via Marlin (W4A16)."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if is_fp4_marlin_supported():
return True, None
return False, "Marlin FP4 not available"
@classmethod
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
logger.warning_once(
"Your GPU does not have native support for FP4 computation but "
"FP4 quantization is being used. Weight-only FP4 compression "
"will be used leveraging the Marlin kernel. This may degrade "
"performance for compute-heavy workloads."
)
prepare_fp4_layer_for_marlin(layer)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
return apply_fp4_marlin_linear(
input=x,
weight=layer.weight,
weight_scale=layer.weight_scale,
weight_global_scale=layer.weight_global_scale,
workspace=layer.workspace,
size_n=layer.output_size_per_partition,
size_k=layer.input_size_per_partition,
bias=bias,
)
@@ -0,0 +1,209 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import ClassVar
import torch
from typing_extensions import Self
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
process_fp8_weight_block_strategy,
)
from vllm.model_executor.utils import replace_parameter
from ..base import (
FP8Params,
MMLinearKernel,
)
from .ScaledMMLinearKernel import FP8ScaledMMLinearLayerConfig
@dataclass
class FP8BlockParams(FP8Params):
weight_scale_inv: torch.Tensor | None
weight_scale: torch.Tensor | None
WEIGHT_SCALE_INV: ClassVar[str] = "weight_scale_inv"
@classmethod
def from_layer(cls, layer: torch.nn.Module) -> Self:
return cls(
weight=getattr(layer, cls.WEIGHT),
weight_scale_inv=getattr(layer, cls.WEIGHT_SCALE_INV, None),
weight_scale=getattr(layer, cls.WEIGHT_SCALE, None),
input_scale=getattr(layer, cls.INPUT_SCALE, None),
input_scale_ub=getattr(layer, cls.INPUT_SCALE_UB, None),
)
class Fp8BlockScaledMMLinearKernel(
MMLinearKernel[FP8ScaledMMLinearLayerConfig, FP8BlockParams], ABC
):
# Set to False in subclasses that accept BF16 input directly (e.g. FlashInfer)
# and therefore do not need the input quantization step in apply_weights.
apply_input_quant: ClassVar[bool] = True
def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
super().__init__(config)
act_scale_descriptor = config.activation_quant_key.scale
self.weight_group_shape = config.weight_quant_key.scale.group_shape
self.quant_fp8 = QuantFP8(
static=act_scale_descriptor.static,
group_shape=act_scale_descriptor.group_shape,
num_token_padding=self.get_output_padding(),
use_ue8m0=False,
)
self.use_triton = False
@classmethod
def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
act_quant_key = config.activation_quant_key
if act_quant_key.scale.static:
return (
False,
"Only dynamic per token group activation quantization is supported.",
)
return True, None
def _get_layer_params(self, layer: torch.nn.Module, **kwargs) -> FP8BlockParams:
return FP8BlockParams.from_layer(layer)
def process_weights_after_loading(self, layer: torch.nn.Module):
params = self._get_layer_params(layer)
# Fp8LinearMethod registered weight scale
# buffer as weight_scale_inv unlike compressed tensors.
weight_scale = (
params.weight_scale
if params.weight_scale_inv is None
else params.weight_scale_inv
)
scale_attr_name = (
params.WEIGHT_SCALE
if params.weight_scale_inv is None
else params.WEIGHT_SCALE_INV
)
new_weight, new_weight_scale = process_fp8_weight_block_strategy(
params.weight,
weight_scale,
)
replace_parameter(layer, params.WEIGHT, new_weight.data)
replace_parameter(layer, scale_attr_name, new_weight_scale.data)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
out_dtype = self.config.out_dtype
params = self._get_layer_params(layer)
weight = params.weight
weight_scale = (
params.weight_scale
if params.weight_scale_inv is None
else params.weight_scale_inv
)
input_scale = params.input_scale
scale_up = params.input_scale_ub
# View input as 2D matrix for fp8 methods
input_2d = x.view(-1, x.shape[-1])
output_shape = [*x.shape[:-1], weight.shape[0]]
if self.apply_input_quant:
q_input, input_scale = self.quant_fp8(
input_2d, input_scale, scale_up, use_triton=self.use_triton
)
else:
q_input = input_2d
# Provide a concrete placeholder so apply_block_scaled_mm args are
# always Tensors. Subclasses with apply_input_quant=False must not
# use As in apply_block_scaled_mm.
input_scale = (
input_scale if input_scale is not None else input_2d.new_empty(1)
)
output = self.apply_block_scaled_mm(
A=q_input,
B=weight,
As=input_scale,
Bs=weight_scale,
)
if bias is not None:
output = output + bias
return output.to(dtype=out_dtype).view(*output_shape)
@abstractmethod
def apply_block_scaled_mm(
self,
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
) -> torch.Tensor:
raise NotImplementedError
class Fp8BlockScaledDynamicMMLinearKernel(Fp8BlockScaledMMLinearKernel, ABC):
"""Dynamic FP8 block-scaled kernel that dispatches at runtime.
Extends Fp8BlockScaledMMLinearKernel to inherit apply_weights and overrides
apply_block_scaled_mm to dispatch between two sub-kernels using torch.cond.
Subclasses must define:
base_type: The primary kernel class.
fallback_type: The fallback kernel class.
"""
base_type: ClassVar[type[Fp8BlockScaledMMLinearKernel]]
fallback_type: ClassVar[type[Fp8BlockScaledMMLinearKernel]]
def __init__(self, config: "FP8ScaledMMLinearLayerConfig") -> None:
super().__init__(config)
self.base = self.base_type(config)
self.fallback = self.fallback_type(config)
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
is_base_supported, reason_1 = cls.base_type.is_supported(compute_capability)
is_fallback_supported, reason_2 = cls.fallback_type.is_supported(
compute_capability
)
if is_base_supported and is_fallback_supported:
return True, None
if not is_base_supported and not is_fallback_supported:
return (
False,
f"base is not supported due to {reason_1}; "
f"fallback is not supported due to {reason_2}",
)
if not is_base_supported:
return False, f"base is not supported due to {reason_1}"
return False, f"fallback is not supported due to {reason_2}"
@classmethod
def can_implement(
cls, config: "FP8ScaledMMLinearLayerConfig"
) -> tuple[bool, str | None]:
can_implement_base, reason_1 = cls.base_type.can_implement(config)
can_implement_fallback, reason_2 = cls.fallback_type.can_implement(config)
if can_implement_base and can_implement_fallback:
return True, None
if not can_implement_base and not can_implement_fallback:
return (
False,
f"base cannot implement due to {reason_1}; "
f"fallback cannot implement due to {reason_2}",
)
if not can_implement_base:
return False, f"base cannot implement due to {reason_1}"
return False, f"fallback cannot implement due to {reason_2}"
@@ -0,0 +1,201 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import ABC, abstractmethod
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Generic, TypeVar
import torch
from vllm.model_executor.layers.fusion.quant_activation import (
QuantizedActivation,
as_quantized_activation,
)
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
from vllm.model_executor.layers.quantization.utils.quant_utils import (
QuantKey,
)
from vllm.platforms import current_platform
from ..base import MMLinearLayerConfig
@dataclass
class Int8ScaledMMLinearLayerConfig(MMLinearLayerConfig):
# TODO: Change to QuantKey like FP8ScaledMMLinearLayerConfig
is_static_input_scheme: bool
is_channelwise: bool
input_symmetric: bool
@dataclass
class FP8ScaledMMLinearLayerConfig(MMLinearLayerConfig):
weight_quant_key: QuantKey
activation_quant_key: QuantKey
weight_shape: tuple[int, int]
input_dtype: torch.dtype
out_dtype: torch.dtype
_FP8ParamsT = tuple[
torch.Tensor, # weight
torch.Tensor, # weight_scale
torch.Tensor | None, # input_scale,
torch.Tensor | None, # input_scale_ub,
]
_Int8ParamsT = tuple[
torch.Tensor, # weight
torch.Tensor, # weight_scale
torch.Tensor | None, # input_scale,
torch.Tensor | None, # input_zp
torch.Tensor | None, # azp_adj
]
_ParamsT = TypeVar("_ParamsT", _Int8ParamsT, _FP8ParamsT)
_ConfigT = TypeVar("_ConfigT", bound=MMLinearLayerConfig)
class ScaledMMLinearKernel(Generic[_ConfigT, _ParamsT], ABC):
@classmethod
@abstractmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
raise NotImplementedError
@classmethod
@abstractmethod
def can_implement(cls, c: _ConfigT) -> tuple[bool, str | None]:
raise NotImplementedError
def __init__(self, c: _ConfigT, layer_param_names: Sequence[str]) -> None:
assert self.can_implement(c)[0]
assert self.is_supported()[0]
self.config = c
self.layer_param_names = layer_param_names
def input_quant_key(self) -> QuantKey | None:
"""The activation quant key this kernel can consume pre-quantized.
Manual fusion uses this to decide whether to hoist activation
quantization out of apply_weights into an upstream fused kernel.
Return None when the kernel needs in-kernel quantization (custom
padding or swizzling, dynamic scales, etc.). Kernels that return a
key must consume the activation via as_quantized_activation.
"""
return None
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
raise NotImplementedError
@abstractmethod
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
raise NotImplementedError
# return a covariant type in the subclass
@abstractmethod
def _get_layer_params(self, layer) -> _ParamsT:
raise NotImplementedError
class FP8ScaledMMLinearKernel(
ScaledMMLinearKernel[FP8ScaledMMLinearLayerConfig, _FP8ParamsT], ABC
):
def __init__(
self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
) -> None:
act_scale_descriptor = c.activation_quant_key.scale
self.quant_fp8 = QuantFP8(
static=act_scale_descriptor.static,
group_shape=act_scale_descriptor.group_shape,
num_token_padding=self.get_output_padding(),
)
self.fp8_dtype = current_platform.fp8_dtype()
super().__init__(c, layer_param_names)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
pass
def _get_layer_params(self, layer) -> _FP8ParamsT:
w, w_s, x_s, x_s_ub = self.layer_param_names
return (
getattr(layer, w),
getattr(layer, w_s),
getattr(layer, x_s, None),
getattr(layer, x_s_ub, None),
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor | QuantizedActivation,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
fp8_dtype = self.fp8_dtype
maybe_out_dtype = self.config.out_dtype
w, w_s, x_s, x_s_ub = self._get_layer_params(layer)
qa = as_quantized_activation(x, self.input_quant_key())
if qa is not None:
x_data, x_s = qa.data, qa.scale
orig_shape, orig_dtype = qa.orig_shape, qa.orig_dtype
assert x_data.dtype == fp8_dtype
else:
assert isinstance(x, torch.Tensor)
x_data = x
orig_shape, orig_dtype = x.shape, x.dtype
x_2d = x_data.view(-1, x_data.shape[-1])
output_shape = [*orig_shape[:-1], w.shape[1]]
out_dtype = orig_dtype if maybe_out_dtype is None else maybe_out_dtype
x_2d_q = x_2d
if qa is None:
x_2d_q, x_s = self.quant_fp8(x_2d, x_s, x_s_ub)
return self.apply_scaled_mm(
A=x_2d_q,
B=w,
out_dtype=out_dtype,
As=x_s,
Bs=w_s,
bias=bias,
output_shape=output_shape,
)
@abstractmethod
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
raise NotImplementedError
def get_output_padding(self) -> int | None:
return None
class Int8ScaledMMLinearKernel(
ScaledMMLinearKernel[Int8ScaledMMLinearLayerConfig, _Int8ParamsT], ABC
):
def _get_layer_params(self, layer) -> _Int8ParamsT:
w_q, w_s, i_s, i_zp, azp_adj = self.layer_param_names
return (
getattr(layer, w_q),
getattr(layer, w_s),
getattr(layer, i_s, None),
getattr(layer, i_zp, None),
getattr(layer, azp_adj, None),
)
@@ -0,0 +1,71 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
AiterInt8ScaledMMLinearKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel import (
Fp8BlockScaledMMLinearKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
CPUFp8BlockScaledMMKernel,
CPUInt8ScaledMMLinearKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
CutlassFP8ScaledMMLinearKernel,
CutlassInt8ScaledMMLinearKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
FlashInferFP8ScaledMMLinearKernel,
)
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.ScaledMMLinearKernel import (
FP8ScaledMMLinearKernel,
FP8ScaledMMLinearLayerConfig,
Int8ScaledMMLinearKernel,
Int8ScaledMMLinearLayerConfig,
ScaledMMLinearKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.triton import (
TritonInt8ScaledMMLinearKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
XPUFp8BlockScaledMMKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
ZentorchInt8ScaledMMLinearKernel,
)
__all__ = [
"FP8ScaledMMLinearKernel",
"FP8ScaledMMLinearLayerConfig",
"Int8ScaledMMLinearKernel",
"Int8ScaledMMLinearLayerConfig",
"ScaledMMLinearKernel",
"ScaledMMLinearLayerConfig",
"AiterInt8ScaledMMLinearKernel",
"CPUInt8ScaledMMLinearKernel",
"CutlassFP8ScaledMMLinearKernel",
"CutlassInt8ScaledMMLinearKernel",
"FlashInferFP8ScaledMMLinearKernel",
"MarlinFP8ScaledMMLinearKernel",
"ChannelWiseTorchFP8ScaledMMLinearKernel",
"PerTensorTorchFP8ScaledMMLinearKernel",
"RowWiseTorchFP8ScaledMMLinearKernel",
"ROCmFP8ScaledMMLinearKernel",
"TritonInt8ScaledMMLinearKernel",
"ZentorchInt8ScaledMMLinearKernel",
"Fp8BlockScaledMMLinearKernel",
"CPUFp8BlockScaledMMKernel",
"XPUFp8BlockScaledMMKernel",
]
@@ -0,0 +1,431 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm import _custom_ops as ops
from vllm._aiter_ops import (
rocm_aiter_ops,
)
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils.quant_utils import (
GroupShape,
)
from vllm.model_executor.utils import replace_parameter
from vllm.platforms import current_platform
from .BlockScaledMMLinearKernel import (
Fp8BlockScaledMMLinearKernel,
)
from .cutlass import CutlassInt8ScaledMMLinearKernel
from .ScaledMMLinearKernel import (
FP8ScaledMMLinearKernel,
FP8ScaledMMLinearLayerConfig,
Int8ScaledMMLinearLayerConfig,
)
logger = init_logger(__name__)
class AiterInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_rocm():
return False, "Requires ROCm."
if compute_capability is not None and compute_capability < 90:
return False, "requires compute capability 90 and above."
try:
import aiter # noqa: F401 # deliberately attempt to import aiter
except Exception:
return False, "requires `aiter` to be installed."
if not rocm_aiter_ops.is_linear_enabled():
return (
False,
"requires setting `VLLM_ROCM_USE_AITER=1` "
"and `VLLM_ROCM_USE_AITER_LINEAR=1`. "
"`VLLM_ROCM_USE_AITER_LINEAR` default is True.",
)
return True, None
@classmethod
def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
if not c.input_symmetric:
return False, "supports symmetric quantization only."
return True, None
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
"""
`AiterInt8ScaledMMLinearKernel` implements a fused version of
`output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
where scale_a * a and scale_b * b are implemented using numpy-style
broadcasting.
Currently only support per-tensor-per-tensor GEMM
and per-token-per-channel GEMM through AITER
w8a8 scaled gemm. `AiterInt8ScaledMMLinearKernel` also does not support
ATIER block scaled GEMM and mix-precision GEMM.
"""
w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
# ops.scaled_int8_quant supports both dynamic and static quant:
# * dynamic, i_s is None and x_s computed from x.
# * static, i_s is scalar and x_s is i_s.
symmetric = azp_adj is None
assert symmetric, (
"AiterInt8ScaledMMLinearKernel only supports symmetric quantization."
)
x_q, x_s, x_zp = ops.scaled_int8_quant(x, i_s, i_zp, symmetric=symmetric)
assert x_zp is None, (
"AiterInt8ScaledMMLinearKernel only supports symmetric quantization."
)
out_dtype = x.dtype
assert w_q.shape[0] % 16 == 0 and w_q.shape[1] % 16 == 0
assert out_dtype is torch.bfloat16 or out_dtype is torch.float16
assert bias is None or bias.shape[0] == w_q.shape[1] and bias.dtype == out_dtype
m = x_q.shape[0] # a
n = w_q.shape[1] # b
per_tensor_scale_a = x_s.numel() == 1
per_tensor_scale_b = w_s.numel() == 1
per_token_scale_a = x_s.numel() == m
per_channel_scale_b = w_s.numel() == n
# @TODO:
# Maybe broadcast the per-tensor-scale into per-channel-scale
# if one of the scale is a per-channel-scale.
# For now, it only supports:
# - per-tensor-per-tensor a8w8 scaled GEMM, and
# - per-token-per-channel a8w8 scaled GEMM
assert (per_tensor_scale_a and per_tensor_scale_b) or (
per_token_scale_a and per_channel_scale_b
), (
"Currently only support per-tensor-per-tensor GEMM "
" and per-token-per-channel GEMM through AITER"
" w8a8 scaled gemm. `AiterInt8ScaledMMLinearKernel` "
"does not support AITER block scaled GEMM."
)
# gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
# a to be [M, K]
# b to be [N, K]
# CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format
return rocm_aiter_ops.w8a8_gemm(x_q, w_q.t(), x_s, w_s, bias, out_dtype)
class AiterPreshuffledPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_rocm():
return False, "requires ROCm."
if not rocm_aiter_ops.is_linear_fp8_enabled():
return (
False,
"requires setting `VLLM_ROCM_USE_AITER=1` "
"and `VLLM_ROCM_USE_AITER_LINEAR=1`. "
"`VLLM_ROCM_USE_AITER_LINEAR` default is True.",
)
try:
import aiter # noqa: F401
except Exception:
return False, "requires aiter library to be installed."
return True, None
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
is_ptpc = (
c.activation_quant_key.scale.group_shape.is_per_token()
and c.weight_quant_key.scale.group_shape.is_per_channel()
)
if c.weight_shape is None:
return False, "weight_shape is required for Aiter kernels"
N, K = c.weight_shape
fp8_dtype = current_platform.fp8_dtype()
if c.out_dtype is not torch.bfloat16:
return False, "requires bfloat16 output dtype."
if not is_ptpc:
return (
False,
"requires per token activation scales and per channel weight scales.",
)
if not (N % 16 == 0 and K % 16 == 0):
return (
False,
f"requires N and K dimensions divisible by 16, received "
f"N={N} and K={K}.",
)
# Aiter's shuffled per-token Gemm performs better than torch only when its
# tuned.
if not rocm_aiter_ops.is_shuffled_per_token_w8a8_gemm_tuned(N, K, fp8_dtype):
return (
False,
f"requires a tuned configuration for N: {N} and K: {K} "
f"and fp8 dtype {fp8_dtype}.",
)
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
w_name, *_ = self.layer_param_names
w, *_ = self._get_layer_params(layer)
replace_parameter(
layer,
w_name,
torch.nn.Parameter(
rocm_aiter_ops.shuffle_weight(w.t().contiguous()).data,
requires_grad=False,
),
)
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
return rocm_aiter_ops.preshuffled_per_token_w8a8_gemm(
A, B, As, Bs, bias, out_dtype
)
class AiterHipbMMPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_rocm():
return False, "requires ROCm."
if not rocm_aiter_ops.is_linear_hipbmm_enabled():
return (
False,
"requires setting `VLLM_ROCM_USE_AITER=1`, "
"`VLLM_ROCM_USE_AITER_LINEAR=1`, "
"and `VLLM_ROCM_USE_AITER_LINEAR_HIPBMM=1`.",
)
try:
import aiter # noqa: F401
except Exception:
return False, "requires aiter library to be installed."
if not hasattr(aiter, "hipb_mm"):
return False, "requires aiter hipb_mm support."
return True, None
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
is_ptpc = (
c.activation_quant_key.scale.group_shape.is_per_token()
and c.weight_quant_key.scale.group_shape.is_per_channel()
)
if c.weight_shape is None:
return False, "weight_shape is required for Aiter kernels"
N, K = c.weight_shape
if c.out_dtype is not torch.bfloat16:
return False, "requires bfloat16 output dtype."
if not is_ptpc:
return (
False,
"requires per token activation scales and per channel weight scales.",
)
if not (N >= 16 and N % 16 == 0 and K % 16 == 0):
return (
False,
"requires N >= 16 and both N and K divisible by 16, "
f"received N={N} and K={K}.",
)
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
w_name, w_s_name, *_ = self.layer_param_names
w, w_s, *_ = self._get_layer_params(layer)
# Pre-apply the transposes that used to live in
# _rocm_aiter_hipb_mm_fp8_impl so the kernel can consume B/Bs directly.
# The `.t()` on the shuffled weight is kept as a non-contiguous view —
# materializing it with `.contiguous()` would re-arrange the bytes and
# break the `bpreshuffle` layout.
shuffled_w = rocm_aiter_ops.shuffle_weight(w.t().contiguous())
replace_parameter(
layer,
w_name,
torch.nn.Parameter(shuffled_w.t(), requires_grad=False),
)
if w_s.ndim > 1:
replace_parameter(
layer,
w_s_name,
torch.nn.Parameter(w_s.t().contiguous(), requires_grad=False),
)
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
output_shape[-1] = B.shape[1]
return rocm_aiter_ops.hipb_mm_fp8(A, B, As, Bs, bias, out_dtype).view(
*output_shape
)
class AiterPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported(
compute_capability
)
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
is_ptpc = (
c.activation_quant_key.scale.group_shape.is_per_token()
and c.weight_quant_key.scale.group_shape.is_per_channel()
)
if c.weight_shape is None:
return False, "weight_shape is required for Aiter kernels"
N, K = c.weight_shape
fp8_dtype = current_platform.fp8_dtype()
if not is_ptpc:
return (
False,
"requires per token activation scales and per channel weight scales.",
)
# Aiter's per-token Gemm performs better than torch oonly when its
# tuned.
if not rocm_aiter_ops.is_per_token_w8a8_gemm_tuned(N, K, fp8_dtype):
return (
False,
f"requires a tuned configuration for N: {N} and K: {K} "
f"and fp8 dtype {fp8_dtype}.",
)
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
w_name, *_ = self.layer_param_names
w, *_ = self._get_layer_params(layer)
replace_parameter(
layer,
w_name,
torch.nn.Parameter(w.t(), requires_grad=False),
)
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
return rocm_aiter_ops.w8a8_gemm(A, B, As, Bs, bias, out_dtype)
class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
def __init__(self, config: FP8ScaledMMLinearLayerConfig):
super().__init__(config)
n, k = config.weight_shape
self.use_triton = (
not current_platform.is_fp8_fnuz()
and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k)
)
@classmethod
def is_supported(cls, compute_capability=None):
return (
rocm_aiter_ops.is_linear_enabled(),
"Only supported on ROCm platform \
with aiter package installed.",
)
@classmethod
def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
can_implement_base, reason = super().can_implement(config)
if not can_implement_base:
return can_implement_base, reason
act_quant_desc = config.activation_quant_key.scale
if act_quant_desc.group_shape != GroupShape(1, 128):
return (
False,
"Supports only dynamic per token group activation "
"quantization with group_shape=(1,128).",
)
return True, None
def apply_block_scaled_mm(
self,
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
) -> torch.Tensor:
if As.dtype != Bs.dtype:
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
_upcast_e8m0_to_fp32,
)
if As.dtype == torch.float8_e8m0fnu:
As = _upcast_e8m0_to_fp32(As).contiguous()
else:
As = As.to(torch.float32)
if Bs.dtype == torch.float8_e8m0fnu:
Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
else:
Bs = Bs.to(torch.float32)
out_dtype = self.config.out_dtype
if self.use_triton:
gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale
else:
gemm_a8w8_blockscale_op = rocm_aiter_ops.gemm_a8w8_blockscale
return gemm_a8w8_blockscale_op(
A, B, As, Bs, list(self.weight_group_shape), output_dtype=out_dtype
)
@@ -0,0 +1,327 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm import _custom_ops as ops
from vllm import envs
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
convert_to_channelwise,
)
from vllm.model_executor.layers.utils import check_cpu_sgl_kernel
from vllm.platforms import current_platform
from vllm.platforms.interface import CpuArchEnum
from .BlockScaledMMLinearKernel import (
Fp8BlockScaledMMLinearKernel,
FP8ScaledMMLinearLayerConfig,
)
from .ScaledMMLinearKernel import (
Int8ScaledMMLinearKernel,
Int8ScaledMMLinearLayerConfig,
)
class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cpu():
return False, "requires CPU."
return True, None
@classmethod
def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
w_q_name, _, _, _, _ = self.layer_param_names
weight = getattr(layer, w_q_name)
dtype = weight.dtype
N, K = weight.size()
if (
current_platform.get_cpu_architecture() == CpuArchEnum.X86
and envs.VLLM_CPU_SGL_KERNEL
and self.config.input_symmetric
and check_cpu_sgl_kernel(N, K, dtype)
):
self.linear_method = self._apply_weights_sgl
self.process_weights_for_sgl(layer)
else:
self.linear_method = self._apply_weights_onednn
self.process_weights_for_onednn(layer)
def process_weights_for_onednn(self, layer: torch.nn.Module) -> None:
# WEIGHT
# Transpose to [K, N] for convenience
w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names
weight = getattr(layer, w_q_name)
replace_parameter(
layer,
w_q_name,
torch.nn.Parameter(weight.t().data, requires_grad=False),
)
# WEIGHT SCALE
# oneDNN kernels support only per-tensor and per-channel.
# If we have a fused module (QKV, MLP) with per tensor scales (thus N
# scales being passed to the kernel), convert to the per-channel case.
is_fused_module = len(layer.logical_widths) > 1
weight_scale = getattr(layer, w_s_name)
if is_fused_module and not self.config.is_channelwise:
weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
replace_parameter(
layer,
w_s_name,
torch.nn.Parameter(weight_scale.data, requires_grad=False),
)
# INPUT SCALE
if self.config.is_static_input_scheme:
input_scale = getattr(layer, i_s_name)
if self.config.input_symmetric:
replace_parameter(
layer,
i_s_name,
torch.nn.Parameter(input_scale.max(), requires_grad=False),
)
else:
input_zero_point = getattr(layer, i_zp_name)
# reconstruct the ranges
int8_traits = torch.iinfo(torch.int8)
azps = input_zero_point.to(dtype=torch.int32)
range_max = (input_scale * (int8_traits.max - azps)).max()
range_min = (input_scale * (int8_traits.min - azps)).min()
scale = (range_max - range_min) / (int8_traits.max - int8_traits.min)
replace_parameter(
layer, i_s_name, torch.nn.Parameter(scale, requires_grad=False)
)
azp = (
(int8_traits.min - range_min / scale).round().to(dtype=torch.int32)
)
replace_parameter(
layer, i_zp_name, torch.nn.Parameter(azp, requires_grad=False)
)
# Different from cutlass, oneDNN kernels only need the AZP adjustment
# term for dynamic quantization. And s_b should be folded into the
# term. Such as:
# s_a * s_b * [(A - zp_a)B] + bias =
# s_a * (s_b * AB) - s_a * s_b * zp_a * B + bias =
# s_a * GEMM_output - s_a * zp_a * adj + bias
if not (self.config.input_symmetric and self.config.is_static_input_scheme):
weight = getattr(layer, w_q_name)
weight_scale = getattr(layer, w_s_name)
azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.float32)
azp_adj = azp_adj * weight_scale.squeeze()
setattr(
layer,
azp_adj_name,
torch.nn.Parameter(azp_adj, requires_grad=False),
)
weight = getattr(layer, w_q_name)
self.dnnl_handler = ops.create_onednn_scaled_mm(
weight,
getattr(layer, w_s_name),
torch.get_default_dtype(),
getattr(layer, i_s_name) is None,
not self.config.input_symmetric,
32,
)
# weight is prepacked and maintained by the dnnl_handler,
# release the original weight
setattr(layer, w_q_name, None)
del weight
def process_weights_for_sgl(self, layer: torch.nn.Module) -> None:
w_q_name, w_s_name, _, _, _ = self.layer_param_names
# WEIGHT
weight = getattr(layer, w_q_name)
packed_weight = torch.ops._C.convert_weight_packed(weight)
replace_parameter(
layer, w_q_name, torch.nn.Parameter(packed_weight, requires_grad=False)
)
if layer.bias is not None:
bias = layer.bias
layer.register_parameter(
"bias_fp32", torch.nn.Parameter(bias.float().data, requires_grad=False)
)
# WEIGHT SCALE
# CPU SGL kernels only support per-channel.
# For per-tensor quant, convert to the per-channel case.
weight_scale = getattr(layer, w_s_name)
if not self.config.is_channelwise:
weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
replace_parameter(
layer,
w_s_name,
torch.nn.Parameter(weight_scale.data, requires_grad=False),
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
return self.linear_method(
layer,
x,
bias,
)
def _apply_weights_onednn(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
x_shape = x.shape
x = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x
w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
# ops.scaled_int8_quant supports both dynamic and static quant:
# * dynamic, i_s is None and x_s computed from x.
# * static, i_s is scalar and x_s is i_s.
x_q, x_s, x_zp = ops.onednn_scaled_int8_quant(
x, i_s, i_zp, self.config.input_symmetric
)
m = x.size(0)
n = self.dnnl_handler.n
out = torch.empty((m, n), dtype=x.dtype)
ops.onednn_scaled_mm(self.dnnl_handler, x_q, out, x_s, x_zp, azp_adj, bias)
out = out.reshape(x_shape[:-1] + (n,)) if len(x_shape) > 2 else out
return out
def _apply_weights_sgl(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
w_q, w_s, _, _, _ = self._get_layer_params(layer)
return torch.ops._C.int8_scaled_mm_with_quant(
x,
w_q,
w_s,
layer.bias_fp32 if bias is not None else None,
x.dtype,
True,
)
class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
"""FP8 W8A16 block-quantized GEMM via AMX BRGEMM on CPU."""
# Input stays BF16 — no FP8 activation quantization.
apply_input_quant = False
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cpu():
return False, "requires CPU platform."
if not torch.cpu._is_amx_tile_supported():
return False, "requires AMX tile support (Sapphire Rapids or newer)."
if not ops._supports_cpu_fp8_w8a16:
return False, "fp8_scaled_mm_cpu op not available."
return True, None
@classmethod
def can_implement(
cls, config: FP8ScaledMMLinearLayerConfig
) -> tuple[bool, str | None]:
# Validate weight block shape
weight_gs = config.weight_quant_key.scale.group_shape
if weight_gs.col <= 0 or weight_gs.col != 128:
return False, (
"CPU FP8 kernel requires K-dimension block size of 128, "
f"got {weight_gs.col}."
)
if weight_gs.row <= 0 or weight_gs.row % 32 != 0:
return False, (
"CPU FP8 kernel requires N-dimension block size to be "
f"a positive multiple of 32, got {weight_gs.row}."
)
if config.out_dtype not in (torch.bfloat16, torch.float32):
return False, "Only bfloat16/float32 output dtype supported."
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# Skip the base class process (FP8 padding / fnuz normalization)
# which is GPU-oriented. Instead, VNNI-prepack weights for AMX.
params = self._get_layer_params(layer)
packed_weight = torch.ops._C.convert_weight_packed(params.weight)
replace_parameter(
layer,
params.WEIGHT,
torch.nn.Parameter(packed_weight, requires_grad=False),
)
# Re-wrap scale as a plain Parameter so the kernel can read it
# without weight-loader metadata interfering.
scale_attr = (
params.WEIGHT_SCALE_INV
if params.weight_scale_inv is not None
else params.WEIGHT_SCALE
)
weight_scale = (
params.weight_scale_inv
if params.weight_scale_inv is not None
else params.weight_scale
)
assert weight_scale is not None
replace_parameter(
layer,
scale_attr,
torch.nn.Parameter(weight_scale.data, requires_grad=False),
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
params = self._get_layer_params(layer)
weight_scale = (
params.weight_scale_inv
if params.weight_scale_inv is not None
else params.weight_scale
)
x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x
out = torch.ops._C.fp8_scaled_mm_cpu(
x_2d,
params.weight,
weight_scale,
list(self.weight_group_shape),
bias,
x.dtype,
True, # is_vnni (weight already prepacked)
)
return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
def apply_block_scaled_mm(
self,
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
) -> torch.Tensor:
raise NotImplementedError(
"CPUFp8BlockScaledMMKernel overrides apply_weights directly."
)
@@ -0,0 +1,343 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Sequence
import torch
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.model_executor.layers.quantization.utils.quant_utils import (
GroupShape,
QuantKey,
kFp8StaticTensorSym,
)
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
CUTLASS_BLOCK_FP8_SUPPORTED,
convert_to_channelwise,
)
from vllm.model_executor.utils import set_weight_attrs
from vllm.platforms import current_platform
from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
from .ScaledMMLinearKernel import (
FP8ScaledMMLinearKernel,
FP8ScaledMMLinearLayerConfig,
Int8ScaledMMLinearKernel,
Int8ScaledMMLinearLayerConfig,
)
class CutlassInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "requires CUDA."
return True, None
@classmethod
def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names
config = self.config
# WEIGHT
# Cutlass kernels need transposed weight.
weight = getattr(layer, w_q_name)
replace_parameter(
layer,
w_q_name,
torch.nn.Parameter(weight.t().data, requires_grad=False),
)
# WEIGHT SCALE
# Cutlass kernels support only per-tensor and per-channel.
# If we have a fused module (QKV, MLP) with per tensor scales (thus N
# scales being passed to the kernel), convert to the per-channel case.
is_fused_module = len(layer.logical_widths) > 1
weight_scale = getattr(layer, w_s_name)
if is_fused_module and not config.is_channelwise:
weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
replace_parameter(
layer,
w_s_name,
torch.nn.Parameter(weight_scale.data, requires_grad=False),
)
# INPUT SCALE
if config.is_static_input_scheme:
input_scale = getattr(layer, i_s_name)
if config.input_symmetric:
replace_parameter(
layer,
i_s_name,
torch.nn.Parameter(input_scale.max(), requires_grad=False),
)
setattr(layer, i_zp_name, None)
else:
input_zero_point = getattr(layer, i_zp_name)
# reconstruct the ranges
int8_traits = torch.iinfo(torch.int8)
azps = input_zero_point.to(dtype=torch.int32)
range_max = (input_scale * (int8_traits.max - azps)).max()
range_min = (input_scale * (int8_traits.min - azps)).min()
scale = (range_max - range_min) / (int8_traits.max - int8_traits.min)
replace_parameter(
layer, i_s_name, torch.nn.Parameter(scale, requires_grad=False)
)
# AZP loaded as int8 but used as int32
azp = (int8_traits.min - range_min / scale).to(dtype=torch.int32)
replace_parameter(
layer, i_zp_name, torch.nn.Parameter(azp, requires_grad=False)
)
# azp_adj is the AZP adjustment term, used to account for weights.
# It does not depend on scales or azp, so it is the same for
# static and dynamic quantization.
# For more details, see csrc/quantization/w8a8/cutlass/Epilogues.md
# https://github.com/vllm-project/vllm/blob/main/csrc/quantization/w8a8/cutlass/Epilogues.md
if not config.input_symmetric:
weight = getattr(layer, w_q_name)
azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32)
if config.is_static_input_scheme:
# cutlass_w8a8 requires azp to be folded into azp_adj
# in the per-tensor case
azp_adj = getattr(layer, i_zp_name) * azp_adj
setattr(
layer,
azp_adj_name,
torch.nn.Parameter(azp_adj, requires_grad=False),
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
# ops.scaled_int8_quant supports both dynamic and static quant:
# * dynamic, i_s is None and x_s computed from x.
# * static, i_s is scalar and x_s is i_s.
symmetric = azp_adj is None
x_q, x_s, x_zp = ops.scaled_int8_quant(
x.contiguous(), i_s, i_zp, symmetric=symmetric
)
if x_zp is not None:
# Currently, static is always per-tensor and dynamic is per-token
static = i_zp is not None
azp = None if static else x_zp
return ops.cutlass_scaled_mm_azp(
x_q,
w_q,
scale_a=x_s,
scale_b=w_s,
out_dtype=x.dtype,
azp_adj=azp_adj,
azp=azp,
bias=bias,
)
return ops.cutlass_scaled_mm(
x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias
)
class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
def __init__(
self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
) -> None:
self.logical_output_size: int | None = None
super().__init__(c, layer_param_names)
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "requires CUDA."
return True, None
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def input_quant_key(self) -> QuantKey | None:
"""Only static per-tensor activation quantization is supported for external
quantization."""
if self.config.activation_quant_key == kFp8StaticTensorSym:
return kFp8StaticTensorSym
return None
@staticmethod
def _pad_to_alignment(
x: torch.Tensor, dim: int, alignment: int, value: float = 0.0
) -> torch.Tensor:
"""Pad tensor ``x`` along ``dim`` to the next multiple of
``alignment``."""
remainder = x.shape[dim] % alignment
if remainder == 0:
return x
pad_size = alignment - remainder
pad_spec = [0] * (2 * x.dim())
pad_spec[-(2 * dim + 1)] = pad_size
return torch.nn.functional.pad(x, pad_spec, value=value)
@staticmethod
def padded_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
if loaded_weight.shape != param.shape:
slices = tuple(slice(0, s) for s in loaded_weight.shape)
param.data[slices].copy_(loaded_weight)
else:
param.data.copy_(loaded_weight.view(param.shape))
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight_name, weight_scale_name, _, _ = self.layer_param_names
weight = getattr(layer, weight_name)
# keep the logical output width so runtime can slice away static padding.
self.logical_output_size = weight.shape[1]
pad_k = (16 - weight.shape[0] % 16) % 16
pad_n = (16 - weight.shape[1] % 16) % 16
if pad_k == 0 and pad_n == 0:
return
# B is column-major [K, N]
padded_weight = torch.nn.functional.pad(
weight.t().contiguous(),
(0, pad_k, 0, pad_n),
).t()
replace_parameter(layer, weight_name, padded_weight.data)
set_weight_attrs(
getattr(layer, weight_name),
{
"weight_loader": self.padded_weight_loader,
},
)
weight_scale = getattr(layer, weight_scale_name, None)
if weight_scale is not None and pad_n > 0 and weight_scale.numel() > 1:
flat_scale = weight_scale.reshape(-1)
padded_scale = self._pad_to_alignment(
flat_scale, dim=0, alignment=16, value=1.0
).view(-1, *weight_scale.shape[1:])
replace_parameter(layer, weight_scale_name, padded_scale.data)
set_weight_attrs(
getattr(layer, weight_name),
{
"weight_loader": self.padded_weight_loader,
},
)
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
padded_k, padded_n = B.shape
output_size = self.logical_output_size
assert output_size is not None
pad_k = padded_k - A.shape[1]
pad_n = padded_n - output_size
if pad_k > 0:
A = self._pad_to_alignment(A, dim=1, alignment=16)
if pad_n > 0 and bias is not None:
bias = self._pad_to_alignment(bias, dim=0, alignment=16)
output = ops.cutlass_scaled_mm(
A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
)
if pad_n > 0:
output = output[..., :output_size].contiguous()
return output.view(*output_shape[:-1], output_size)
class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
super().__init__(config)
act_scale_descriptor = config.activation_quant_key.scale
self.quant_fp8 = QuantFP8(
static=act_scale_descriptor.static,
group_shape=act_scale_descriptor.group_shape,
num_token_padding=self.get_output_padding(),
use_ue8m0=False,
column_major_scales=True,
)
@classmethod
def is_supported(cls, compute_capability=None):
if not CUTLASS_BLOCK_FP8_SUPPORTED:
return (
False,
"The device compute capability of"
f"{compute_capability} is not supported.",
)
return True, None
@classmethod
def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
can_implement_base, reason = super().can_implement(config)
if not can_implement_base:
return can_implement_base, reason
act_quant_desc = config.activation_quant_key.scale
if act_quant_desc.group_shape != GroupShape(1, 128):
return (
False,
"Supports only dynamic per token group activation "
"quantization with group_shape=(1,128).",
)
return True, None
def apply_block_scaled_mm(
self,
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
) -> torch.Tensor:
out_dtype = self.config.out_dtype
return ops.cutlass_scaled_mm(
A,
B.T,
out_dtype=out_dtype,
scale_a=As,
scale_b=Bs.T,
)
def cutlass_scaled_mm(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
block_size: list[int],
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
return ops.cutlass_scaled_mm(
A,
B.T,
out_dtype=output_dtype,
scale_a=As,
scale_b=Bs.T,
)
@@ -0,0 +1,158 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import vllm.envs as envs
from vllm.config import get_current_vllm_config
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
deepgemm_post_process_fp8_weight_block,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
GroupShape,
)
from vllm.model_executor.utils import replace_parameter
from vllm.platforms import current_platform
from vllm.utils.deep_gemm import (
fp8_gemm_nt,
is_deep_gemm_e8m0_used,
is_deep_gemm_supported,
should_auto_disable_deep_gemm,
should_use_deepgemm_for_fp8_linear,
)
from vllm.utils.torch_utils import direct_register_custom_op
from .BlockScaledMMLinearKernel import (
Fp8BlockScaledMMLinearKernel,
FP8ScaledMMLinearLayerConfig,
)
class DeepGemmFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
def __init__(self, config: FP8ScaledMMLinearLayerConfig):
super().__init__(config)
self.use_deep_gemm_e8m0 = is_deep_gemm_e8m0_used()
act_scale_descriptor = config.activation_quant_key.scale
self.is_deep_gemm_supported = is_deep_gemm_supported()
self.quant_fp8 = QuantFP8(
static=False,
group_shape=act_scale_descriptor.group_shape,
use_ue8m0=self.use_deep_gemm_e8m0,
tma_aligned_scales=envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES,
column_major_scales=True,
)
@classmethod
def is_supported(cls, compute_capability=None):
if not current_platform.is_cuda():
return False, "DeepGEMM is only supported on cuda platform"
if not is_deep_gemm_supported():
return False, "Currently, only Hopper and Blackwell GPUs are supported."
return True, None
@classmethod
def can_implement(cls, config):
can_implement_base, reason = super().can_implement(config)
if not can_implement_base:
return can_implement_base, reason
if config.out_dtype != torch.bfloat16:
return (False, "Supports only output dtype of bfloat16")
act_quant_desc = config.activation_quant_key.scale
if act_quant_desc.group_shape != GroupShape(1, 128):
return (
False,
"Supports only dynamic per token group activation "
"quantization with group_shape=(1,128).",
)
model_config = get_current_vllm_config().model_config
if model_config is None:
return False, "Model configuration is required."
model_type = getattr(model_config.hf_text_config, "model_type", None)
if should_auto_disable_deep_gemm(model_type):
return False, f"Should not use deepgemm for model {model_type}"
if not should_use_deepgemm_for_fp8_linear(
config.out_dtype, config.weight_shape
):
return False, "The provided metadata is not supported."
return True, None
def process_weights_after_loading(self, layer):
super().process_weights_after_loading(layer)
params = self._get_layer_params(layer)
assert layer.weight_block_size is not None
if self.is_deep_gemm_supported:
weight_scale_invs = params.weight_scale_inv
scale_attr = (
params.WEIGHT_SCALE_INV
if weight_scale_invs is not None
else params.WEIGHT_SCALE
)
dg_weight, dg_weight_scale = deepgemm_post_process_fp8_weight_block(
wq=params.weight,
ws=weight_scale_invs
if weight_scale_invs is not None
else params.weight_scale,
quant_block_shape=tuple(layer.weight_block_size),
use_e8m0=self.use_deep_gemm_e8m0,
is_bmm=getattr(layer, "is_bmm", False),
bmm_batch_size=getattr(layer, "bmm_batch_size", 0),
)
replace_parameter(layer, params.WEIGHT, dg_weight)
replace_parameter(layer, scale_attr, dg_weight_scale)
def apply_block_scaled_mm(
self,
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
) -> torch.Tensor:
out_dtype = self.config.out_dtype
output = torch.empty(
(A.shape[0], B.shape[0]),
dtype=out_dtype,
device=A.device,
)
torch.ops.vllm.fp8_gemm_nt_op(A, As, B, Bs, output, self.use_deep_gemm_e8m0)
return output
def _fp8_gemm_nt_op(
q_input: torch.Tensor,
input_scale: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
output: torch.Tensor,
use_deep_gemm_e8m0: bool,
) -> None:
fp8_gemm_nt(
(q_input, input_scale),
(weight, weight_scale),
output,
is_deep_gemm_e8m0_used=use_deep_gemm_e8m0,
)
def _fp8_gemm_nt_op_fake(
q_input: torch.Tensor,
input_scale: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
output: torch.Tensor,
use_deep_gemm_e8m0: bool,
) -> None:
return None
direct_register_custom_op(
"fp8_gemm_nt_op",
_fp8_gemm_nt_op,
mutates_args=["output"],
fake_impl=_fp8_gemm_nt_op_fake,
)
@@ -0,0 +1,338 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import ClassVar
import torch
import vllm.envs as envs
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
per_token_group_quant_fp8,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
GroupShape,
QuantKey,
kFp8StaticTensorSym,
)
from vllm.platforms import current_platform
from vllm.utils.flashinfer import (
flashinfer_fp8_blockscale_gemm,
flashinfer_scaled_fp8_mm,
has_flashinfer,
is_flashinfer_fp8_blockscale_gemm_supported,
should_use_flashinfer_for_blockscale_fp8_gemm,
)
from vllm.utils.torch_utils import direct_register_custom_op
from .BlockScaledMMLinearKernel import (
Fp8BlockScaledDynamicMMLinearKernel,
Fp8BlockScaledMMLinearKernel,
)
from .deep_gemm import DeepGemmFp8BlockScaledMMKernel, fp8_gemm_nt
from .ScaledMMLinearKernel import (
FP8ScaledMMLinearKernel,
FP8ScaledMMLinearLayerConfig,
)
class FlashInferFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "requires CUDA."
if not has_flashinfer():
return False, "requires FlashInfer to be installed."
if compute_capability is not None and compute_capability < 100:
return False, "requires compute capability 100 and above."
return True, None
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
per_tensor_activation_scales = (
c.activation_quant_key.scale.group_shape.is_per_tensor()
)
per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
if not (per_tensor_activation_scales and per_tensor_weight_scales):
return False, "requires per tensor activation and weight scales."
return True, None
def input_quant_key(self) -> QuantKey | None:
if self.config.activation_quant_key == kFp8StaticTensorSym:
return kFp8StaticTensorSym
return None
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
return flashinfer_scaled_fp8_mm(
A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
)
class FlashInferFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
# FlashInfer accepts BF16 input and handles FP8 conversion internally.
apply_input_quant: ClassVar[bool] = False
def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
super().__init__(config)
@classmethod
def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
can_implement_base, reason = super().can_implement(config)
if not can_implement_base:
return can_implement_base, reason
act_quant_desc = config.activation_quant_key.scale
if act_quant_desc.group_shape != GroupShape(1, 128):
return (
False,
"Supports only dynamic per token group activation "
"quantization with group_shape=(1,128).",
)
if not should_use_flashinfer_for_blockscale_fp8_gemm(
is_flashinfer_fp8_blockscale_gemm_supported(),
config.out_dtype,
config.input_dtype,
config.weight_quant_key.dtype,
config.weight_shape,
):
return (
False,
"The provided metadata is not supported.",
)
return True, None
@classmethod
def is_supported(cls, compute_capability=None):
if not current_platform.is_cuda():
return False, "only cuda devices are supported."
if not is_flashinfer_fp8_blockscale_gemm_supported():
return False, "FlashInfer block-scale FP8 GEMM is not available."
return True, None
def apply_block_scaled_mm(
self,
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
) -> torch.Tensor:
# A is BF16 — FlashInfer handles FP8 conversion internally.
# As is a placeholder (apply_input_quant=False) and is not used here.
return torch.ops.vllm.flashinfer_fp8_blockscale_gemm(
A, # BF16 input
B, # FP8 weight
Bs, # Weight scales
)
class FlashInferFp8DeepGEMMDynamicBlockScaledKernel(
Fp8BlockScaledDynamicMMLinearKernel
):
"""
Conditional FlashInfer / DeepGEMM FP8 block-scaled GEMM.
Dispatches between two kernels based on input batch size:
- Small batches (M < 32): FlashInfer's swapAB trick for better utilisation.
- Large batches (M >= 32): DeepGEMM for peak throughput.
apply_input_quant is False because FlashInfer accepts BF16 input and
handles FP8 conversion internally. The DeepGEMM branch therefore
quantises BF16→FP8 inside apply_mm via a closure before dispatching to
the DeepGEMM kernel — keeping both branches compatible with the single
BF16 tensor operand list passed by torch.cond.
"""
base_type: ClassVar[type[FlashInferFp8BlockScaledMMKernel]] = (
FlashInferFp8BlockScaledMMKernel
)
fallback_type: ClassVar[type[DeepGemmFp8BlockScaledMMKernel]] = (
DeepGemmFp8BlockScaledMMKernel
)
apply_input_quant: ClassVar[bool] = False
def __init__(self, config: FP8ScaledMMLinearLayerConfig):
super().__init__(config)
self.base: FlashInferFp8BlockScaledMMKernel
self.fallback: DeepGemmFp8BlockScaledMMKernel
def process_weights_after_loading(self, layer: torch.nn.Module):
# DeepGEMM need post-processing; both kernels share the same
# parameter tensor layout so processing once is sufficient.
self.fallback.process_weights_after_loading(layer)
def apply_block_scaled_mm(
self,
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
) -> torch.Tensor:
group_size = self.weight_group_shape.col
use_deep_gemm_e8m0 = self.fallback.use_deep_gemm_e8m0
return torch.ops.vllm.dynamic_flashinfer_deepgemm_blockscale_gemm(
A, B, Bs, group_size, use_deep_gemm_e8m0
)
def _flashinfer_fp8_blockscale_gemm_impl(
input: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
) -> torch.Tensor:
return flashinfer_fp8_blockscale_gemm(
input=input,
weight=weight,
weight_scale=weight_scale,
out_dtype=torch.bfloat16,
)
def _flashinfer_fp8_blockscale_gemm_fake(
input: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
) -> torch.Tensor:
"""
Required fake/meta implementation for torch.compile graph tracing.
"""
return torch.empty(
input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device
)
direct_register_custom_op(
"flashinfer_fp8_blockscale_gemm",
_flashinfer_fp8_blockscale_gemm_impl,
fake_impl=_flashinfer_fp8_blockscale_gemm_fake,
)
def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl(
input: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
group_size: int,
use_deep_gemm_e8m0: bool,
) -> torch.Tensor:
"""
Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection.
This function switches between two optimized kernels based on the input batch size:
- For small batches (M < 32): Uses FlashInfer's DeepGEMM swapAB optimization.
- For larger batches (M >= 32): Uses the official DeepGEMM kernel.
The conditional logic must use torch.cond() instead of a simple if-else statement
to maintain compatibility with torch.compile graph compilation.
This batch-size-dependent selection is essential for maintaining model accuracy.
Benchmarks on GSM8K show a significant accuracy gap (88% vs 95%) for DeepSeek-V3.1
when using FlashInfer's DeepGEMM on M>=32. The M < 32 strategy fixes the accuracy
drop.
Args:
input: Input tensor of shape (batch_size, input_dim) in FP8 format
weight: Weight tensor of shape (output_dim, input_dim) in FP8 format
weight_scale: Scale factors for weight quantization (per-group)
group_size: Quantization group size for the weight tensor
use_deep_gemm_e8m0: Whether to use the E8M0 format in DeepGEMM quantization
Returns:
Output tensor of shape (batch_size, output_dim) in bfloat16 format
"""
def run_flashinfer_deepgemm_swapAB(
input: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
) -> torch.Tensor:
return flashinfer_fp8_blockscale_gemm(
input=input,
weight=weight,
weight_scale=weight_scale,
out_dtype=torch.bfloat16,
)
def run_deepgemm(
input: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
) -> torch.Tensor:
q_input, input_scale = per_token_group_quant_fp8(
input,
group_size=group_size,
column_major_scales=True,
use_ue8m0=use_deep_gemm_e8m0,
)
output = torch.empty(
(q_input.shape[0], weight.shape[0]),
dtype=torch.bfloat16,
device=q_input.device,
)
fp8_gemm_nt(
(q_input, input_scale),
(weight, weight_scale),
output,
is_deep_gemm_e8m0_used=use_deep_gemm_e8m0,
)
return output
if envs.VLLM_BATCH_INVARIANT:
return run_deepgemm(input, weight, weight_scale)
condition = input.shape[0] < 32
# PyTorch's torch.compile cannot handle input-dependent control flow in standard
# Python conditionals. torch.cond() explicitly registers both code paths in the
# computation graph, allowing torch.compile to capture both branches.
# without torch.cond, the M < 32 condition won't be able to be captured by torch
# compile
return torch.cond(
condition,
run_flashinfer_deepgemm_swapAB,
run_deepgemm,
(input, weight, weight_scale),
)
def _dynamic_flashinfer_deepgemm_blockscale_gemm_fake(
input: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
group_size: int,
use_deep_gemm_e8m0: bool,
) -> torch.Tensor:
"""
Required fake/meta implementation for torch.compile graph tracing.
"""
return torch.empty(
input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device
)
direct_register_custom_op(
"dynamic_flashinfer_deepgemm_blockscale_gemm",
_dynamic_flashinfer_deepgemm_blockscale_gemm_impl,
fake_impl=_dynamic_flashinfer_deepgemm_blockscale_gemm_fake,
)
@@ -0,0 +1,156 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils.humming_utils import (
convert_linear_layer_to_humming_standard,
prepare_humming_layer,
)
from vllm.platforms import current_platform
from .ScaledMMLinearKernel import (
FP8ScaledMMLinearKernel,
FP8ScaledMMLinearLayerConfig,
Int8ScaledMMLinearKernel,
Int8ScaledMMLinearLayerConfig,
)
logger = init_logger(__name__)
class HummingFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
"""Humming GEMM Kernel for FP8."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "Humming only supported on CUDA"
if not current_platform.has_device_capability(75):
return False, "Humming only supported on SM75+"
return True, None
@classmethod
def can_implement(
cls, config: FP8ScaledMMLinearLayerConfig
) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
from vllm.utils.humming import dtypes
name_map = {"weight": "weight", "weight_scale": "weight_scale"}
scale_torch_dtype = self.config.weight_quant_key.scale.dtype
scale_dtype = dtypes.DataType.from_torch_dtype(scale_torch_dtype)
quant_config = {
"quant_method": "humming",
"dtype": "float8e4m3",
"scale_dtype": scale_dtype,
}
assert self.config.weight_quant_key.scale2 is None
scale_group_shape = self.config.weight_quant_key.scale.group_shape
if scale_group_shape.is_per_tensor():
quant_config["weight_scale_type"] = "tensor"
if not hasattr(layer, "global_scale") and hasattr(layer, "weight_scale"):
del name_map["weight_scale"]
name_map["global_scale"] = "weight_scale"
elif scale_group_shape.is_per_channel():
quant_config["weight_scale_type"] = "channel"
elif scale_group_shape.is_per_group():
quant_config["weight_scale_type"] = "group"
quant_config["group_size"] = scale_group_shape.col
else:
assert scale_group_shape.row > 0 and scale_group_shape.col > 0
quant_config["weight_scale_type"] = "block"
quant_config["weight_scale_group_size_n"] = scale_group_shape.row
quant_config["weight_scale_group_size"] = scale_group_shape.col
if hasattr(layer, "weight_scale_inv"):
name_map["weight_scale"] = "weight_scale_inv"
convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
prepare_humming_layer(layer, quant_config)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
from vllm.utils.humming import HummingMethod
flatten_inputs = x.view(-1, x.size(-1))
output = HummingMethod.forward_layer(
layer=layer,
inputs=flatten_inputs,
compute_config=layer.compute_config,
)
return output.view(*x.shape[:-1], output.size(-1))
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
pass
class HummingInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
"""Humming GEMM Kernel for INT8."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "Humming only supported on CUDA"
if not current_platform.has_device_capability(75):
return False, "Humming only supported on SM75+"
return True, None
@classmethod
def can_implement(
cls, config: Int8ScaledMMLinearLayerConfig
) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight_name, weight_scale_name, *_ = self.layer_param_names
name_map = {"weight": weight_name, "weight_scale": weight_scale_name}
quant_config = {"quant_method": "humming", "dtype": "int8"}
weight = getattr(layer, weight_name)
weight.data = weight.data + 128
convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
prepare_humming_layer(layer, quant_config)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
from vllm.utils.humming import HummingMethod
flatten_inputs = x.view(-1, x.size(-1))
output = HummingMethod.forward_layer(
layer=layer,
inputs=flatten_inputs,
compute_config=layer.compute_config,
)
return output.view(*x.shape[:-1], output.size(-1))
@@ -0,0 +1,118 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Sequence
import torch
import vllm.envs as envs
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
process_fp8_weight_block_strategy,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
apply_fp8_marlin_linear,
is_fp8_marlin_supported,
prepare_fp8_layer_for_marlin,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
kFp8Static128BlockSym,
)
from vllm.model_executor.utils import replace_parameter
from vllm.platforms import current_platform
from .ScaledMMLinearKernel import (
FP8ScaledMMLinearKernel,
FP8ScaledMMLinearLayerConfig,
)
class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
"""
FP8 Marlin kernel for GPUs that lack FP8 hardware support.
Leverages the Marlin kernel for fast weight-only FP8 quantization.
"""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cuda():
return False, "requires CUDA."
# Check if platform supports FP8 Marlin
if not is_fp8_marlin_supported():
return False, "FP8 Marlin requires compute capability 7.5 or higher"
if envs.VLLM_BATCH_INVARIANT:
return False, "FP8 Marlin not supported for batch invariant execution."
if (
compute_capability is not None
and compute_capability >= 89
and not envs.VLLM_TEST_FORCE_FP8_MARLIN
):
return (
False,
"To apply FP8 Marlin on high-capability GPUs, please set "
"VLLM_TEST_FORCE_FP8_MARLIN=1",
)
return True, None
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def __init__(
self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
) -> None:
super().__init__(c, layer_param_names)
self.marlin_input_dtype = None
self.block_quant = self.config.weight_quant_key in {kFp8Static128BlockSym}
self.size_k_first = not self.block_quant
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
if self.block_quant:
weight, weight_scale_inv = process_fp8_weight_block_strategy(
layer.weight, layer.weight_scale_inv
)
# Update layer with new values
replace_parameter(layer, "weight", weight.data)
replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
# Non-block: callers must pass weight in (K, N) layout.
layer.input_scale = None
prepare_fp8_layer_for_marlin(
layer, self.size_k_first, input_dtype=self.marlin_input_dtype
)
del layer.input_scale
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
if self.block_quant:
weight_scale = layer.weight_scale_inv
else:
weight_scale = layer.weight_scale
return apply_fp8_marlin_linear(
input=x,
weight=layer.weight,
weight_scale=weight_scale,
workspace=layer.workspace,
size_n=layer.output_size_per_partition,
size_k=layer.input_size_per_partition,
input_dtype=self.marlin_input_dtype,
bias=bias,
)
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
pass
@@ -0,0 +1,242 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
import torch
from vllm.config import CompilationMode, get_current_vllm_config
from vllm.platforms import current_platform
from .ScaledMMLinearKernel import (
FP8ScaledMMLinearKernel,
FP8ScaledMMLinearLayerConfig,
)
def _get_num_tokens(output_shape: list) -> int:
# torch._scaled_mm works with 2D tensors, so input tensors are
# flattened if they are 3D. If output_shape is 3D, num_tokens is
# the product of all dims except the last (hidden dim).
return math.prod(output_shape[:-1])
class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
"""
Base class for FP8 linear kernels using Torch.
Each subclass represents a kernel variant for
specific device capabilities and torch versions.
"""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not (current_platform.is_cuda_alike() or current_platform.is_cpu()):
return False, "requires ROCm, CUDA or CPU."
if compute_capability is not None and compute_capability < 89:
return False, "requires compute capability 89 and above."
return True, None
def get_output_padding(self) -> int | None:
# Note: we pad the input because torch._scaled_mm is more performant
# for matrices with batch dimension > 16.
# This could change in the future.
# We also don't pad when using torch.compile,
# as it breaks with dynamic shapes.
#
# The perf gain is still relevant as of 16/1/2026
# torch version == 2.9.0. More details in the link below:
# https://github.com/vllm-project/vllm/issues/32269
vllm_config = get_current_vllm_config().compilation_config
pad_output = vllm_config.mode < CompilationMode.VLLM_COMPILE
return 17 if pad_output else None
class PerTensorTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
per_tensor_activation_scales = (
c.activation_quant_key.scale.group_shape.is_per_tensor()
)
per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
if not (per_tensor_activation_scales and per_tensor_weight_scales):
return False, "requires per tensor activation and weight scales."
return True, None
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
# torch._scaled_mm under torch.compile does not support 0-D scales
if As.dim() == 0:
As = As.view(1)
if Bs.dim() == 0:
Bs = Bs.view(1)
output = torch._scaled_mm(
A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
)
# A fix for discrepancy in scaled_mm which returns tuple
# for torch < 2.5 and a single value in torch >= 2.5
if type(output) is tuple and len(output) == 2:
output = output[0]
num_tokens = _get_num_tokens(output_shape)
return torch.narrow(output, 0, 0, num_tokens).view(*output_shape)
class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_rocm():
return False, "requires ROCm."
from vllm.platforms.rocm import on_mi3xx
if not on_mi3xx():
return False, "requires MI3xx."
if compute_capability is not None and compute_capability < 94:
return False, "requires compute capability 94 and above."
return True, None
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
per_tensor_activation_scales = (
c.activation_quant_key.scale.group_shape.is_per_tensor()
)
per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
if c.out_dtype == torch.float16:
# hipblaslt rowwise _scaled_mm only supports BFloat16
return False, "supports BFloat16 output data type only."
if per_tensor_activation_scales or per_tensor_weight_scales:
return False, "cannot be used with per tensor activation and weight scales."
return True, None
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
# Note:
# For now it has only been validated on ROCm platform.
# fp8 rowwise scaling in torch._scaled_mm is introduced in
# https://github.com/pytorch/pytorch/pull/144432 using
# hipBLASLt and ROCm 6.3, which only exists in torch 2.7 and above.
#
# For CUDA platform please validate if the torch._scaled_mm supports
# rowwise scaled GEMM before using it
# torch._scaled_mm rowwise requires scale_a = (m, 1), scale_b = (1, n).
# CompressedTensors stores weight_scale as (n, 1), so `.t()` yields (1, n).
# ModelOpt FP8_PER_CHANNEL_PER_TOKEN stores it as 1-D (n,); reshape to
# (1, n) so both paths satisfy the rowwise contract.
scale_b = Bs.view(1, -1) if Bs.dim() == 1 else Bs.t()
if As.dim() == 1:
As = As.view(-1, 1)
# Fused GEMM_DQ Rowwise GEMM
output = torch._scaled_mm(
A,
B,
out_dtype=out_dtype,
scale_a=As,
scale_b=scale_b,
bias=bias,
)
num_tokens = _get_num_tokens(output_shape)
return torch.narrow(output, 0, 0, num_tokens).view(*output_shape)
class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
per_tensor_activation_scales = (
c.activation_quant_key.scale.group_shape.is_per_tensor()
)
per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
if per_tensor_activation_scales and per_tensor_weight_scales:
return False, "cannot be used with per tensor activation and weight scales."
return True, None
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
# Use unfused DQ due to limitations with scaled_mm
# Symmetric quantized GEMM by definition computes the following:
# C = (s_x * X) (s_w * W) + bias
# This is equivalent to dequantizing the weights and activations
# before applying a GEMM.
#
# In order to compute quantized operands, a quantized kernel
# will rewrite the above like so:
# C = s_w * s_x * (X * W) + bias
#
# For the scaled_mm fallback case, we break this down, since it
# does not support s_w being a vector.
# Input scaling factors are no longer optional in _scaled_mm starting
# from pytorch 2.5. Allocating a dummy tensor to pass as scales
dummy_tensor = torch.ones(1, dtype=torch.float32, device=A.device)
# GEMM
# This computes C = (X * W).
# Output in fp32 to allow subsequent ops to happen in-place
output = torch._scaled_mm(
A,
B,
scale_a=dummy_tensor,
scale_b=dummy_tensor,
out_dtype=torch.float32,
)
# A fix for discrepancy in scaled_mm which returns tuple
# for torch < 2.5 and a single value in torch >= 2.5
if type(output) is tuple and len(output) == 2:
output = output[0]
# Unpad (undo num_token_padding)
num_tokens = _get_num_tokens(output_shape)
output = torch.narrow(output, 0, 0, num_tokens)
x_scale = torch.narrow(As, 0, 0, num_tokens)
# DQ
# C = sw * sx * (X * W) + bias
output = output * x_scale * Bs.t()
if bias is not None:
output = output + bias
return output.to(out_dtype).view(*output_shape)
@@ -0,0 +1,118 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import vllm.envs as envs
from vllm import _custom_ops as ops
from vllm.platforms import current_platform
from vllm.utils.platform_utils import num_compute_units
from vllm.utils.torch_utils import direct_register_custom_op
from .ScaledMMLinearKernel import (
FP8ScaledMMLinearKernel,
FP8ScaledMMLinearLayerConfig,
)
def rocm_per_tensor_float_w8a8_scaled_mm_impl(
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor,
) -> torch.Tensor:
if (
A.shape[0] <= 4
and B.shape[0] % 16 == 0 # M TODO: needed?
and B.shape[1] % 16 == 0 # K
and ((bias is None) or (bias.dtype == out_dtype))
):
output = ops.wvSplitKQ(
B.t(),
A,
out_dtype,
As,
Bs,
num_compute_units(),
bias,
)
# Fallback
else:
output = torch._scaled_mm(
A,
B,
out_dtype=out_dtype,
scale_a=As,
scale_b=Bs,
bias=bias,
)
return output
def rocm_per_tensor_float_w8a8_scaled_mm_fake(
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor,
) -> torch.Tensor:
return A.new_empty((*A.shape[:-1], B.shape[1]), dtype=out_dtype)
if current_platform.is_rocm():
direct_register_custom_op(
op_name="rocm_per_tensor_float_w8a8_scaled_mm_impl",
op_func=rocm_per_tensor_float_w8a8_scaled_mm_impl,
fake_impl=rocm_per_tensor_float_w8a8_scaled_mm_fake,
)
class ROCmFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_rocm():
return False, "requires ROCm."
from vllm.platforms.rocm import on_gfx12x, on_mi3xx
if not (on_mi3xx() or on_gfx12x()):
return False, "requires MI3xx or gfx12x"
if not envs.VLLM_ROCM_USE_SKINNY_GEMM:
return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled."
return True, None
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
per_tensor_activation_scales = (
c.activation_quant_key.scale.group_shape.is_per_tensor()
)
per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
if not (per_tensor_activation_scales and per_tensor_weight_scales):
return False, "requires per tensor activation and weight scales."
return True, None
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
output = torch.ops.vllm.rocm_per_tensor_float_w8a8_scaled_mm_impl(
A, B, out_dtype, As, Bs, bias
)
return torch.narrow(output, 0, 0, A.shape[0]).view(*output_shape)
@@ -0,0 +1,220 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm import ( # noqa: E501
triton_scaled_mm,
)
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
convert_to_channelwise,
)
from vllm.platforms import current_platform
from vllm.utils.torch_utils import direct_register_custom_op
from .BlockScaledMMLinearKernel import (
Fp8BlockScaledMMLinearKernel,
)
from .cutlass import CutlassInt8ScaledMMLinearKernel
from .ScaledMMLinearKernel import (
Int8ScaledMMLinearLayerConfig,
)
class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if current_platform.is_cuda_alike():
return True, None
return False, "requires ROCm or CUDA."
@classmethod
def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
w_q, _, i_s, _, _ = self._get_layer_params(layer)
w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names
replace_parameter(
layer,
w_q_name,
torch.nn.Parameter(w_q.t().data, requires_grad=False),
)
# WEIGHT SCALE
# Triton kernel supports only per-tensor and per-channel.
# If we have a fused module (QKV, MLP) with per tensor scales (thus N
# scales being passed to the kernel), convert to the per-channel case.
is_fused_module = len(layer.logical_widths) > 1
weight_scale = getattr(layer, w_s_name)
if is_fused_module and not self.config.is_channelwise:
weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
replace_parameter(
layer,
w_s_name,
torch.nn.Parameter(weight_scale.data, requires_grad=False),
)
# INPUT SCALE
if self.config.is_static_input_scheme:
assert i_s is not None
if self.config.input_symmetric:
replace_parameter(
layer,
i_s_name,
torch.nn.Parameter(i_s.max(), requires_grad=False),
)
setattr(layer, i_zp_name, None)
else:
input_zero_point = getattr(layer, i_zp_name)
# Reconstruct the ranges to find a single scale and azp
int8_traits = torch.iinfo(torch.int8)
azps = input_zero_point.to(dtype=torch.int32)
range_max = (i_s * (int8_traits.max - azps)).max()
range_min = (i_s * (int8_traits.min - azps)).min()
scale = (range_max - range_min) / (int8_traits.max - int8_traits.min)
replace_parameter(
layer,
i_s_name,
torch.nn.Parameter(scale, requires_grad=False),
)
# AZP loaded as int8 but used as int32
azp = (int8_traits.min - range_min / scale).to(dtype=torch.int32)
replace_parameter(
layer,
i_zp_name,
torch.nn.Parameter(azp, requires_grad=False),
)
else:
setattr(layer, i_s_name, None)
setattr(layer, i_zp_name, None)
# azp_adj is the AZP adjustment term, used to account for weights.
# It does not depend on scales or azp, so it is the same for
# static and dynamic quantization.
# See csrc/quantization/w8a8/cutlass/Epilogues.md for the math.
if not self.config.input_symmetric:
weight = getattr(layer, w_q_name)
# weight is already transposed to [K, N], sum over K (dim=0)
azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32)
if self.config.is_static_input_scheme:
# Fold azp into azp_adj for the per-tensor case
azp_adj = getattr(layer, i_zp_name) * azp_adj
setattr(
layer,
azp_adj_name,
torch.nn.Parameter(azp_adj, requires_grad=False),
)
else:
setattr(layer, azp_adj_name, None)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
symmetric = azp_adj is None
x_q, x_s, x_zp = ops.scaled_int8_quant(
x.contiguous(), i_s, i_zp, symmetric=symmetric
)
out = triton_scaled_mm(
x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias
)
if azp_adj is not None:
# Asymmetric quantization: subtract the zero-point correction.
# D = scale_a * scale_b * (A_q @ B_q - azp * azp_adj) + bias
# triton_scaled_mm already computed scale_a * scale_b * (A_q @ B_q) + bias
# so we subtract scale_a * scale_b * azp * azp_adj
#
# x_s: [M, 1] or scalar, w_s: [N, 1] or scalar, azp_adj: [1, N]
# Reshape w_s from [N, 1] to [1, N] for proper broadcasting.
w_s_row = w_s.view(1, -1) if w_s.dim() > 0 else w_s
static = i_zp is not None
if not static and x_zp is not None:
# Dynamic per-token: azp is per-token, azp_adj is per-channel
# x_zp: [M, 1], azp_adj: [1, N]
out -= x_s * w_s_row * (x_zp * azp_adj).to(x.dtype)
else:
# Static per-tensor: azp already folded into azp_adj
out -= (x_s * w_s_row * azp_adj).to(x.dtype)
return out
class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
@classmethod
def is_supported(cls, compute_capability=None):
if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
return False, "only CUDA-alike and XPU devices are supported."
return True, None
def apply_block_scaled_mm(
self,
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
) -> torch.Tensor:
return torch.ops.vllm.w8a8_triton_block_scaled_mm_func(
A,
B,
As,
Bs,
list(self.weight_group_shape),
self.config.out_dtype,
)
# TODO we should be able to change the type of block_size to GroupShape
# after we resolve GroupShape compilation issue
# https://github.com/vllm-project/vllm/issues/25270
def _w8a8_triton_block_scaled_mm_func(
qx: torch.Tensor,
weight: torch.Tensor,
x_scale: torch.Tensor,
weight_scale: torch.Tensor,
block_size: list[int],
output_dtype: torch.dtype,
) -> torch.Tensor:
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
w8a8_triton_block_scaled_mm,
)
return w8a8_triton_block_scaled_mm(
qx, weight, x_scale, weight_scale, block_size, output_dtype
)
def _w8a8_triton_block_scaled_mm_fake(
qx: torch.Tensor,
weight: torch.Tensor,
x_scale: torch.Tensor,
weight_scale: torch.Tensor,
block_size: list[int],
output_dtype: torch.dtype,
) -> torch.Tensor:
return torch.empty(
(qx.size(0), weight.size(0)), dtype=output_dtype, device=qx.device
)
direct_register_custom_op(
"w8a8_triton_block_scaled_mm_func",
_w8a8_triton_block_scaled_mm_func,
fake_impl=_w8a8_triton_block_scaled_mm_fake,
)
@@ -0,0 +1,216 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Sequence
import torch
from vllm.model_executor.layers.quantization.utils.quant_utils import (
kFp8DynamicTensorSym,
kFp8DynamicTokenSym,
kFp8StaticChannelSym,
kFp8StaticTensorSym,
kFp8StaticTokenSym,
)
from vllm.model_executor.utils import replace_parameter
from vllm.platforms import current_platform
from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel, FP8ScaledMMLinearLayerConfig
class XPUW8A8FP8LinearKernel(FP8ScaledMMLinearKernel):
_SUPPORTED_ACT_QUANT_KEYS = {
kFp8DynamicTensorSym,
kFp8DynamicTokenSym,
kFp8StaticTensorSym,
kFp8StaticTokenSym,
}
_SUPPORTED_WEIGHT_QUANT_KEYS = {
kFp8StaticChannelSym,
kFp8StaticTensorSym,
}
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_xpu():
return False, "XPUW8A8FP8Linear only support on XPU"
return True, None
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
if c.weight_quant_key not in cls._SUPPORTED_WEIGHT_QUANT_KEYS:
return (
False,
"XPUW8A8FP8Linear only support per-channel and per-tensor quantization",
)
if c.activation_quant_key not in cls._SUPPORTED_ACT_QUANT_KEYS:
return (
False,
"XPUW8A8FP8Linear only support per-tensor and per-token activation "
"quantization",
)
if c.weight_quant_key.dtype not in {torch.float8_e5m2, torch.float8_e4m3fn}:
return False, "XPUW8A8FP8Linear only support FP8 weight dtype"
if c.activation_quant_key.dtype not in {
torch.float8_e5m2,
torch.float8_e4m3fn,
}:
return False, "XPUW8A8FP8Linear only support FP8 activation dtype"
return True, None
def __init__(
self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
) -> None:
super().__init__(c, layer_param_names)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""Ensure weight is stored as C-contiguous [K, N] (KN layout).
Checkpoints store weight as [N, K]; fp8_gemm requires [K, N],
C-contiguous. Three incoming layouts are possible:
• [N, K] C-contiguous ← direct checkpoint → .t().contiguous()
• [K, N] Fortran-order ← fp8.py's weight.t() → .contiguous()
• [K, N] C-contiguous ← already correct → no-op
For square weights (K == N) the shape is ambiguous; contiguity is used
as a proxy: C-contiguous ≡ checkpoint [N, K] (needs transpose);
Fortran-order ≡ fp8.py already transposed (needs only contiguous).
"""
K = getattr(layer, "input_size_per_partition", self.config.weight_shape[1])
N = getattr(layer, "output_size_per_partition", self.config.weight_shape[0])
w = layer.weight
if w.shape not in {(K, N), (N, K)}:
raise ValueError(
f"XPUFP8ScaledMM expects weight shape ({K},{N}) or ({N},{K}), "
f"but got {tuple(w.shape)}"
)
needs_transpose = w.shape == (N, K) if K != N else w.is_contiguous()
layer_weight = w.t() if needs_transpose else w
replace_parameter(layer, "weight", layer_weight)
ws = layer.weight_scale
if ws.numel() == 1:
replace_parameter(layer, "weight_scale", ws.reshape(1))
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
# B is C-contiguous [K, N] from process_weights_after_loading.
# fp8_gemm routes on scale dtype (float32) and numel:
# As [1] → per-tensor (numel==1 branch)
# As [M,1] → per-token (group={1,K} branch, broadcast across K)
# Bs [1] → per-tensor
# Bs [N] → per-channel (mask=bit1 branch)
# No shape manipulation needed here.
output = torch.ops._xpu_C.fp8_gemm(A, B, out_dtype, As, Bs, bias)
return output.view(*output_shape)
class XPUW8A16FP8LinearKernel(FP8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_xpu():
return False, "XPUW8A16FP8Linear only support on XPU"
return True, None
@classmethod
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
if c.weight_quant_key not in {kFp8StaticChannelSym, kFp8StaticTensorSym}:
return (
False,
"XPUW8A16FP8Linear only support per-channel and per-tensor "
"quantization",
)
if c.weight_quant_key.dtype not in {torch.float8_e5m2, torch.float8_e4m3fn}:
return False, "XPUW8A16FP8Linear only support FP8 weight dtype"
return True, None
def __init__(
self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
) -> None:
assert self.can_implement(c)[0]
assert self.is_supported()[0]
self.config = c
self.layer_param_names = layer_param_names
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# fp8_gemm_w8a16 expects weight in [in, out] layout.
# Transpose if weight is still in [out, in] layout.
# For square matrices, use contiguity as tie-breaker:
# checkpoint weights are contiguous, .t() views are not.
weight = layer.weight
out_features, in_features = self.config.weight_shape
if weight.shape == (out_features, in_features) and (
in_features != out_features or weight.is_contiguous()
):
replace_parameter(layer, "weight", weight.data.t())
# else: already in [in, out] layout — no-op
weight_scale = layer.weight_scale.t().contiguous()
replace_parameter(layer, "weight_scale", weight_scale.data)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
weight = layer.weight
weight_scale = layer.weight_scale
return torch.ops._xpu_C.fp8_gemm_w8a16(x, weight, weight_scale, bias)
def apply_scaled_mm(
self,
*,
A: torch.Tensor,
B: torch.Tensor,
out_dtype: torch.dtype,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None,
output_shape: list,
) -> torch.Tensor:
pass
class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_xpu():
return False, "XPUFp8BlockScaledMM only support on XPU"
return True, None
def apply_block_scaled_mm(
self,
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
) -> torch.Tensor:
# Weight is [N, K]. Use .t() to create a [K, N] view without copying.
# Bs is [N/128, K/128] — transpose to [K/128, N/128] for oneDNN.
return torch.ops._xpu_C.fp8_gemm(
A,
B.t(),
self.config.out_dtype,
As,
Bs.t().contiguous(),
torch.Tensor(),
)
@@ -0,0 +1,98 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Zentorch dynamic-symmetric W8A8 int8 linear kernel for AMD Zen CPUs.
Selected by ``choose_scaled_mm_linear_kernel`` ahead of the generic
oneDNN-backed ``CPUInt8ScaledMMLinearKernel``. When ``is_supported`` or
``can_implement`` rejects a layer, the selector falls through to the next
kernel in ``_POSSIBLE_INT8_KERNELS[PlatformEnum.CPU]``.
"""
import torch
from vllm.logger import init_logger
from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.platforms import current_platform
from .ScaledMMLinearKernel import (
Int8ScaledMMLinearKernel,
Int8ScaledMMLinearLayerConfig,
)
logger = init_logger(__name__)
class ZentorchInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cpu():
return False, "requires CPU."
if not current_platform.is_zen_cpu():
return False, "requires AMD Zen CPU."
if not has_zentorch_op(["zentorch_dynamic_qlinear"]):
return (
False,
"torch.ops.zentorch.zentorch_dynamic_qlinear is not registered.",
)
return True, None
@classmethod
def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
if c.is_static_input_scheme:
return False, "requires dynamic activation quantization."
if not c.input_symmetric:
return False, "requires symmetric activation quantization."
if not c.is_channelwise:
return False, "requires per-channel weight quantization."
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""Prepare weights for ``zentorch_dynamic_qlinear``.
Keeps weight in [N, K] layout (int8, contiguous) and converts the
per-channel weight scale to bf16 with shape ``(N,)``.
"""
w_q_name, w_s_name, _, _, _ = self.layer_param_names
weight = getattr(layer, w_q_name)
n = weight.shape[0]
replace_parameter(
layer,
w_q_name,
torch.nn.Parameter(weight.data.contiguous(), requires_grad=False),
)
weight_scale = getattr(layer, w_s_name)
ws = weight_scale.data
if ws.dim() == 2 and ws.shape[-1] == 1:
ws = ws.squeeze(-1)
ws = ws.to(torch.bfloat16).contiguous()
assert ws.shape == (n,), (
f"[zen_cpu] expected weight scale shape ({n},), got {tuple(ws.shape)}"
)
replace_parameter(
layer,
w_s_name,
torch.nn.Parameter(ws, requires_grad=False),
)
logger.info_once(
"[zen_cpu] Using zentorch_dynamic_qlinear for W8A8 (dynamic-symmetric)"
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
w_q_name, w_s_name, _, _, _ = self.layer_param_names
return torch.ops.zentorch.zentorch_dynamic_qlinear(
x,
getattr(layer, w_q_name),
getattr(layer, w_s_name),
bias,
zentorch_op_name="zentorch::zentorch_dynamic_qlinear",
)
@@ -0,0 +1,23 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Gates zentorch CPU linear dispatch on platform/op availability."""
from __future__ import annotations
import torch
from vllm.platforms import current_platform
__all__ = ["has_zentorch_op"]
def has_zentorch_op(op_names: list[str]) -> bool:
"""Return ``True`` when running on Zen CPU with all named ops registered."""
if not op_names:
raise ValueError("has_zentorch_op requires at least one op name")
if not current_platform.is_zen_cpu():
return False
ns = getattr(torch.ops, "zentorch", None)
if ns is None:
return False
return all(hasattr(ns, op_name) for op_name in op_names)