111 lines
3.8 KiB
Python
111 lines
3.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from collections.abc import Callable
|
|
from functools import partial
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch.nn.parameter import Parameter
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
|
|
dequant_mxfp4,
|
|
quant_dequant_mxfp4,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
|
|
quant_dequant_mxfp6,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
|
QuantKey,
|
|
kMxfp4Dynamic,
|
|
kMxfp6E2M3Dynamic,
|
|
kMxfp6E3M2Dynamic,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.import_utils import has_quark
|
|
|
|
from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_ACTIVATION_QUANT_DEQUANT_FUNCS: dict[
|
|
QuantKey, Callable[[torch.Tensor], torch.Tensor]
|
|
] = {
|
|
kMxfp4Dynamic: quant_dequant_mxfp4,
|
|
kMxfp6E3M2Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e3m2"),
|
|
kMxfp6E2M3Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e2m3"),
|
|
}
|
|
|
|
|
|
class EmulationMxfp4LinearKernel(MxFp4LinearKernel):
|
|
"""Software emulation fallback for OCP MXFP4/MXFP6 (dequant + F.linear)."""
|
|
|
|
def __init__(self, config: MxFp4LinearLayerConfig) -> None:
|
|
super().__init__(config)
|
|
if config.activation_quant_key is None:
|
|
# no input Q/DQ for weight-only
|
|
self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = (
|
|
lambda x: x
|
|
)
|
|
else:
|
|
self.quant_dequant_func = _ACTIVATION_QUANT_DEQUANT_FUNCS[
|
|
config.activation_quant_key
|
|
]
|
|
|
|
@classmethod
|
|
def is_supported(
|
|
cls, compute_capability: int | None = None
|
|
) -> tuple[bool, str | None]:
|
|
return True, None
|
|
|
|
@classmethod
|
|
def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
|
|
if config.activation_quant_key not in (
|
|
None,
|
|
kMxfp4Dynamic,
|
|
kMxfp6E3M2Dynamic,
|
|
kMxfp6E2M3Dynamic,
|
|
):
|
|
return False, "only supports MXFP4 or MXFP6 or unquantized activations"
|
|
|
|
if not has_quark():
|
|
return False, "amd-quark package not available"
|
|
|
|
if (
|
|
current_platform.is_rocm()
|
|
and current_platform.supports_mx()
|
|
and config.activation_quant_key != kMxfp4Dynamic
|
|
):
|
|
logger.warning_once(
|
|
"The current platform supports native MXFP4/MXFP6 computation, "
|
|
f"but kernels for activation_quant_key={config.activation_quant_key} "
|
|
f"are not yet integrated in vLLM. Using EmulationMxfp4LinearKernel, "
|
|
"with simulated weight dequantization and activation "
|
|
"QDQ (quantize and dequantize), with the linear "
|
|
"layers computed in high precision."
|
|
)
|
|
|
|
if not current_platform.supports_mx():
|
|
logger.warning_once(
|
|
"The current platform does not support native MXFP4 "
|
|
"computation. Using EmulationMxfp4LinearKernel, with simulated weight "
|
|
"dequantization and activation QDQ (quantize and dequantize), with "
|
|
"the linear layers computed in high precision."
|
|
)
|
|
|
|
return True, None
|
|
|
|
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
|
layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
|
|
|
|
def apply_weights(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
x: torch.Tensor,
|
|
bias: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype)
|
|
qdq_x = self.quant_dequant_func(x)
|
|
return F.linear(qdq_x, dq_w, bias)
|