chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,781 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import fnmatch
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
import torch
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.attention import Attention
|
||||
from vllm.model_executor.layers.fused_moe import RoutedExperts
|
||||
from vllm.model_executor.layers.linear import (
|
||||
LinearBase,
|
||||
LinearMethodBase,
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
from vllm.model_executor.layers.quantization.base_config import ( # noqa: E501
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
|
||||
from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501
|
||||
QuarkMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.quark.schemes import (
|
||||
QuarkNVFP4,
|
||||
QuarkOCP_MX,
|
||||
QuarkScheme,
|
||||
QuarkW4A8_MXFP4_FP8,
|
||||
QuarkW8A8Fp8,
|
||||
QuarkW8A8Int8,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.quark.utils import (
|
||||
deep_compare,
|
||||
should_ignore_layer,
|
||||
)
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
__all__ = ["QuarkLinearMethod"]
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# model_type values that use dynamic MXFP4 re-quantization for
|
||||
# OCP MX fp4 Quark checkpoints
|
||||
_DEEPSEEK_V3_FAMILY_MODEL_TYPES = frozenset({"deepseek_v3", "deepseek_v32"})
|
||||
|
||||
|
||||
class QuarkConfig(QuantizationConfig):
|
||||
def __init__(
|
||||
self,
|
||||
quant_config: dict[str, Any],
|
||||
kv_cache_group: list[str] | None = None,
|
||||
kv_cache_config: dict[str, Any] | None = None,
|
||||
pack_method: str = "reorder",
|
||||
):
|
||||
super().__init__()
|
||||
if kv_cache_group is None:
|
||||
kv_cache_group = []
|
||||
self.quant_config = quant_config
|
||||
self.kv_cache_group = kv_cache_group
|
||||
self.kv_cache_config = kv_cache_config
|
||||
self.pack_method = pack_method
|
||||
# Note : this flag is kept disabled because the overhead of
|
||||
# dynamic mxfp4 quantization negates the performance gains
|
||||
# that come from shifting to mxfp4. It is left here in case
|
||||
# we want to re-enable it in the future.
|
||||
self.dynamic_mxfp4_quant = False
|
||||
|
||||
def maybe_update_config(
|
||||
self,
|
||||
model_name: str,
|
||||
hf_config: PretrainedConfig | None = None,
|
||||
revision: str | None = None,
|
||||
):
|
||||
"""Enable dynamic MXFP4 only for DeepSeek-V3-family fp4 checkpoints."""
|
||||
|
||||
if hf_config is None:
|
||||
return
|
||||
|
||||
if (
|
||||
getattr(hf_config, "model_type", None)
|
||||
not in _DEEPSEEK_V3_FAMILY_MODEL_TYPES
|
||||
):
|
||||
return
|
||||
|
||||
quant_config = getattr(hf_config, "quantization_config", None)
|
||||
if isinstance(quant_config, dict):
|
||||
quant_dtype = (
|
||||
quant_config.get("global_quant_config", {})
|
||||
.get("weight", {})
|
||||
.get("dtype")
|
||||
)
|
||||
if quant_dtype == "fp4":
|
||||
self.dynamic_mxfp4_quant = True
|
||||
|
||||
def get_linear_method(self) -> "QuarkLinearMethod":
|
||||
return QuarkLinearMethod(self)
|
||||
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.float16, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 70
|
||||
|
||||
def get_name(self) -> QuantizationMethods:
|
||||
return "quark"
|
||||
|
||||
def apply_vllm_mapper( # noqa: B027
|
||||
self, hf_to_vllm_mapper: "WeightsMapper"
|
||||
):
|
||||
"""
|
||||
Interface for models to update module names referenced in
|
||||
quantization configs in order to reflect the vllm model structure
|
||||
|
||||
Args:
|
||||
hf_to_vllm_mapper: maps from hf model structure (the assumed
|
||||
structure of the qconfig) to vllm model structure
|
||||
"""
|
||||
quant_config_with_hf_to_vllm_mapper: dict[str, Any] = {}
|
||||
|
||||
for k, v in self.quant_config.items():
|
||||
if isinstance(v, list):
|
||||
quant_config_with_hf_to_vllm_mapper[k] = hf_to_vllm_mapper.apply_list(v)
|
||||
elif isinstance(v, dict):
|
||||
quant_config_with_hf_to_vllm_mapper[k] = hf_to_vllm_mapper.apply_dict(v)
|
||||
else:
|
||||
if isinstance(v, str):
|
||||
mapped_v_list = hf_to_vllm_mapper.apply_list([v])
|
||||
if mapped_v_list:
|
||||
quant_config_with_hf_to_vllm_mapper[k] = mapped_v_list[0]
|
||||
else:
|
||||
quant_config_with_hf_to_vllm_mapper[k] = v
|
||||
|
||||
self.quant_config = quant_config_with_hf_to_vllm_mapper
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
# Check if the layer is skipped for quantization.
|
||||
exclude_layers = cast(list[str], self.quant_config.get("exclude"))
|
||||
if should_ignore_layer(
|
||||
prefix, ignore=exclude_layers, fused_mapping=self.packed_modules_mapping
|
||||
):
|
||||
if (
|
||||
"self_attn" not in prefix # only quantize attention projections
|
||||
or not getattr(self, "dynamic_mxfp4_quant", False)
|
||||
or not isinstance(layer, LinearBase) # Ignore other methods
|
||||
):
|
||||
return UnquantizedLinearMethod()
|
||||
|
||||
scheme = self.get_scheme(
|
||||
layer=layer,
|
||||
layer_name=prefix,
|
||||
dynamic_mxfp4_quant=True,
|
||||
)
|
||||
layer.scheme = scheme
|
||||
return QuarkLinearMethod(self)
|
||||
if isinstance(layer, LinearBase):
|
||||
scheme = self.get_scheme(layer=layer, layer_name=prefix)
|
||||
layer.scheme = scheme
|
||||
return QuarkLinearMethod(self)
|
||||
if isinstance(layer, Attention):
|
||||
return QuarkKVCacheMethod(self)
|
||||
|
||||
if isinstance(layer, RoutedExperts):
|
||||
return QuarkMoEMethod.get_moe_method(self, module=layer, layer_name=prefix)
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "QuarkConfig":
|
||||
export_config = config.get("export")
|
||||
if export_config is None:
|
||||
raise ValueError(
|
||||
"The export key should be included in "
|
||||
"the configurations of Quark quantized model"
|
||||
)
|
||||
kv_cache_group = cast(list[str], export_config.get("kv_cache_group"))
|
||||
pack_method = cast(str, export_config.get("pack_method"))
|
||||
|
||||
# In the export model of quark, the quantization configuration
|
||||
# of kv_cache is stored in layer_quant_config. First, it is
|
||||
# judged whether kv_cache_group exists, and then it is judged
|
||||
# whether layer_quant_config has a quantization configuration
|
||||
# that matches kv_cache.
|
||||
if len(kv_cache_group) == 0:
|
||||
kv_cache_config = None
|
||||
else:
|
||||
kv_cache_set = set(kv_cache_group)
|
||||
layer_quant_config = cast(dict[str, Any], config.get("layer_quant_config"))
|
||||
layer_quant_names = list(layer_quant_config.keys())
|
||||
layer_quant_set = set(layer_quant_names)
|
||||
|
||||
if not (
|
||||
kv_cache_set.issubset(layer_quant_set)
|
||||
or any(
|
||||
fnmatch.fnmatchcase(layer_quant, pat)
|
||||
for layer_quant in list(layer_quant_set)
|
||||
for pat in list(kv_cache_set)
|
||||
)
|
||||
):
|
||||
raise ValueError(
|
||||
"The Quark quantized model has the "
|
||||
"kv_cache_group parameter setting, "
|
||||
"but no kv_cache quantization settings "
|
||||
"were found in the quantization "
|
||||
"configuration."
|
||||
)
|
||||
|
||||
q_configs = [
|
||||
quant_cfg
|
||||
for name, quant_cfg in layer_quant_config.items()
|
||||
if any(fnmatch.fnmatchcase(name, pattern) for pattern in kv_cache_group)
|
||||
]
|
||||
|
||||
if not all(
|
||||
deep_compare(q_config["output_tensors"], q_configs[0]["output_tensors"])
|
||||
for q_config in q_configs
|
||||
):
|
||||
raise ValueError(
|
||||
"The quantization method used for kv_cache should "
|
||||
"be the same, but the quantization method for the "
|
||||
"kv_cache layer in the config is different."
|
||||
)
|
||||
kv_cache_config = q_configs[0].get("output_tensors")
|
||||
if kv_cache_config is None:
|
||||
raise ValueError("The kv_cache quantization configuration is empty.")
|
||||
|
||||
# Since we have already set kv_cache quantization configurations,
|
||||
# we will remove the quantization configuration for the
|
||||
# output_tensors corresponding to the kv_cache layer.
|
||||
for q_config in q_configs:
|
||||
q_config["output_tensors"] = None
|
||||
|
||||
# In case q_proj output is also quantized, remove the configuration
|
||||
# to keep qkv consistency.
|
||||
q_proj_q_config = cast(dict[str, Any], layer_quant_config.get("*q_proj"))
|
||||
if q_proj_q_config is not None:
|
||||
q_proj_q_config["output_tensors"] = None
|
||||
|
||||
return cls(
|
||||
quant_config=config,
|
||||
kv_cache_group=kv_cache_group,
|
||||
kv_cache_config=kv_cache_config,
|
||||
pack_method=pack_method,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return []
|
||||
|
||||
def _check_scheme_supported(self, min_capability: int, error: bool = True) -> bool:
|
||||
capability_tuple = current_platform.get_device_capability()
|
||||
|
||||
if capability_tuple is not None:
|
||||
capability = capability_tuple.to_int()
|
||||
supported = capability >= min_capability
|
||||
if error and not supported:
|
||||
raise RuntimeError(
|
||||
"Quantization scheme is not supported for ",
|
||||
f"the current GPU. Min capability: {min_capability}. ",
|
||||
f"Current capability: {capability}.",
|
||||
)
|
||||
return supported
|
||||
else:
|
||||
return False
|
||||
|
||||
def _is_fp8_w4a8(
|
||||
self,
|
||||
weight_quant: list[dict[str, Any]] | None,
|
||||
input_quant: dict[str, Any] | None,
|
||||
) -> bool:
|
||||
# Confirm weights and input quantized.
|
||||
if weight_quant is None or input_quant is None:
|
||||
return False
|
||||
|
||||
if not isinstance(weight_quant, list) or len(weight_quant) != 2:
|
||||
return False
|
||||
|
||||
# Confirm weight scheme is supported
|
||||
is_w4a8_dtype = (
|
||||
weight_quant[0].get("dtype") == "fp8_e4m3"
|
||||
and weight_quant[1].get("dtype") == "int4"
|
||||
and input_quant.get("dtype") == "fp8_e4m3"
|
||||
)
|
||||
is_static_weight = not weight_quant[0].get("is_dynamic") and not weight_quant[
|
||||
1
|
||||
].get("is_dynamic")
|
||||
is_per_tensor_fp8_and_per_channel_int4_weight = (
|
||||
weight_quant[0].get("qscheme") == "per_tensor"
|
||||
and weight_quant[1].get("qscheme") == "per_channel"
|
||||
and weight_quant[1].get("symmetric") is True
|
||||
and weight_quant[1].get("ch_axis") == 0
|
||||
)
|
||||
|
||||
if not (
|
||||
is_w4a8_dtype
|
||||
and is_static_weight
|
||||
and is_per_tensor_fp8_and_per_channel_int4_weight
|
||||
):
|
||||
return False
|
||||
|
||||
# Dynamic quantization is always supported if weights supported.
|
||||
if input_quant.get("is_dynamic"):
|
||||
return True
|
||||
|
||||
# Confirm activation scheme is supported.
|
||||
is_per_tensor_activation = input_quant.get("qscheme") == "per_tensor"
|
||||
return is_per_tensor_activation
|
||||
|
||||
def _is_fp8_w8a8(
|
||||
self,
|
||||
weight_quant: dict[str, Any] | None,
|
||||
input_quant: dict[str, Any] | None,
|
||||
) -> bool:
|
||||
# Confirm weights and input quantized.
|
||||
if weight_quant is None or input_quant is None:
|
||||
return False
|
||||
|
||||
# Confirm weight scheme is supported
|
||||
is_fp8_dtype = (
|
||||
weight_quant.get("dtype") == "fp8_e4m3"
|
||||
and input_quant.get("dtype") == "fp8_e4m3"
|
||||
)
|
||||
is_static_weight = not weight_quant.get("is_dynamic")
|
||||
is_per_tensor_or_channel_weight = weight_quant.get("qscheme") in [
|
||||
"per_tensor",
|
||||
"per_channel",
|
||||
]
|
||||
|
||||
if not (is_fp8_dtype and is_static_weight and is_per_tensor_or_channel_weight):
|
||||
return False
|
||||
|
||||
# Dynamic quantization is always supported if weights supported.
|
||||
if input_quant.get("is_dynamic"):
|
||||
return True
|
||||
|
||||
# Confirm activation scheme is supported.
|
||||
is_per_tensor_activation = input_quant.get("qscheme") == "per_tensor"
|
||||
return is_per_tensor_activation
|
||||
|
||||
def _is_static_tensor_w8a8(
|
||||
self,
|
||||
weight_quant: dict[str, Any] | None,
|
||||
input_quant: dict[str, Any] | None,
|
||||
) -> bool:
|
||||
# Confirm weights and input quantized.
|
||||
if weight_quant is None or input_quant is None:
|
||||
return False
|
||||
|
||||
is_int8_dtype = (
|
||||
weight_quant.get("dtype") == "int8" and input_quant.get("dtype") == "int8"
|
||||
)
|
||||
|
||||
is_tensor = (
|
||||
weight_quant.get("qscheme") in ["per_tensor", "per_channel"]
|
||||
and input_quant.get("qscheme") == "per_tensor"
|
||||
)
|
||||
|
||||
is_static = not weight_quant.get("is_dynamic") and not input_quant.get(
|
||||
"is_dynamic"
|
||||
)
|
||||
|
||||
is_weight_symmetric = weight_quant.get("symmetric") is True
|
||||
|
||||
# Both symmetric and asymmetric input quantization supported.
|
||||
# Only symmetric weight quantization supported.
|
||||
return is_int8_dtype and is_tensor and is_weight_symmetric and is_static
|
||||
|
||||
def _is_w4a8_mxfp4_fp8(
|
||||
self,
|
||||
weight_quant: dict[str, Any] | None,
|
||||
input_quant: dict[str, Any] | None,
|
||||
) -> bool:
|
||||
if weight_quant is None or input_quant is None:
|
||||
return False
|
||||
|
||||
is_weight_mxfp4 = (
|
||||
weight_quant.get("dtype") == "fp4"
|
||||
and weight_quant.get("qscheme") == "per_group"
|
||||
and weight_quant.get("group_size") == 32
|
||||
and weight_quant.get("scale_format") == "e8m0"
|
||||
and not weight_quant.get("is_dynamic")
|
||||
)
|
||||
|
||||
is_input_fp8 = (
|
||||
input_quant.get("dtype") == "fp8_e4m3"
|
||||
and input_quant.get("qscheme") == "per_tensor"
|
||||
and not input_quant.get("is_dynamic") # Static per-tensor
|
||||
and input_quant.get("symmetric") is True # Symmetric quantization
|
||||
)
|
||||
|
||||
return is_weight_mxfp4 and is_input_fp8
|
||||
|
||||
def _is_dynamic_per_token_w8a8(
|
||||
self,
|
||||
weight_quant: dict[str, Any] | None,
|
||||
input_quant: dict[str, Any] | None,
|
||||
) -> bool:
|
||||
"""Detect W8A8 INT8 with per-tensor or per-channel
|
||||
weights and dynamic per-token input."""
|
||||
if weight_quant is None or input_quant is None:
|
||||
return False
|
||||
|
||||
is_int8_dtype = (
|
||||
weight_quant.get("dtype") == "int8" and input_quant.get("dtype") == "int8"
|
||||
)
|
||||
|
||||
is_valid_weight_scheme = weight_quant.get("qscheme") in [
|
||||
"per_tensor",
|
||||
"per_channel",
|
||||
]
|
||||
is_per_token_input = input_quant.get("qscheme") == "per_channel"
|
||||
|
||||
is_dynamic_input = input_quant.get("is_dynamic") is True
|
||||
is_weight_symmetric = weight_quant.get("symmetric") is True
|
||||
|
||||
return (
|
||||
is_int8_dtype
|
||||
and is_valid_weight_scheme
|
||||
and is_per_token_input
|
||||
and is_dynamic_input
|
||||
and is_weight_symmetric
|
||||
)
|
||||
|
||||
def _is_nvfp4(
|
||||
self,
|
||||
weight_quant: dict[str, Any] | list[dict[str, Any]] | None,
|
||||
input_quant: dict[str, Any] | list[dict[str, Any]] | None,
|
||||
) -> bool:
|
||||
# Confirm weights and input quantized.
|
||||
if weight_quant is None or input_quant is None:
|
||||
return False
|
||||
|
||||
# Confirm both weight_quant and input_quant are lists with 2 elements
|
||||
if not isinstance(weight_quant, list) or len(weight_quant) != 2:
|
||||
return False
|
||||
if not isinstance(input_quant, list) or len(input_quant) != 2:
|
||||
return False
|
||||
|
||||
# First element should be fp4 with per_group quantization
|
||||
is_fp4_per_group_weight = (
|
||||
weight_quant[0].get("dtype") == "fp4"
|
||||
and weight_quant[0].get("qscheme") == "per_group"
|
||||
and weight_quant[0].get("group_size") == 16
|
||||
and not weight_quant[0].get("is_dynamic")
|
||||
)
|
||||
is_fp4_per_group_input = (
|
||||
input_quant[0].get("dtype") == "fp4"
|
||||
and input_quant[0].get("qscheme") == "per_group"
|
||||
and input_quant[0].get("group_size") == 16
|
||||
and input_quant[0].get("is_dynamic")
|
||||
)
|
||||
|
||||
# Second element should be fp8_e4m3 with per_tensor quantization
|
||||
is_fp8_per_tensor_weight = (
|
||||
weight_quant[1].get("dtype") == "fp8_e4m3"
|
||||
and weight_quant[1].get("qscheme") == "per_tensor"
|
||||
and not weight_quant[1].get("is_dynamic")
|
||||
)
|
||||
is_fp8_per_tensor_input = (
|
||||
input_quant[1].get("dtype") == "fp8_e4m3"
|
||||
and input_quant[1].get("qscheme") == "per_tensor"
|
||||
and not input_quant[1].get("is_dynamic")
|
||||
)
|
||||
|
||||
return (
|
||||
is_fp4_per_group_weight # type: ignore[return-value]
|
||||
and is_fp4_per_group_input
|
||||
and is_fp8_per_tensor_weight
|
||||
and is_fp8_per_tensor_input
|
||||
)
|
||||
|
||||
def _is_w_ocp_mx_a_x(
|
||||
self, weight_quant: dict[str, Any] | None, input_quant: dict[str, Any] | None
|
||||
) -> bool:
|
||||
"""
|
||||
This check returns True only if it is an OCP-MX weight quantization.
|
||||
The activation can be any data type (e.g., FP16/BF16, FP8, or OCP-MX format).
|
||||
The rationale for checking only the weight type is that
|
||||
the model loading concept and process primarily concerns the weights themselves.
|
||||
"""
|
||||
# Confirm weights quantized.
|
||||
if weight_quant is None:
|
||||
logger.debug(
|
||||
"Quark model's weight quantization is incompatible with OCP_MX format: "
|
||||
"weight_quant is not set."
|
||||
)
|
||||
return False
|
||||
|
||||
if isinstance(weight_quant, list):
|
||||
logger.debug(
|
||||
"Quark model's weight quantization is incompatible with OCP_MX format: "
|
||||
"weight_quant is a list (e.g. fp8_w4a8), OCP_MX requires a single dict."
|
||||
)
|
||||
return False
|
||||
|
||||
# Input and weight qscheme needs to be per group.
|
||||
if weight_quant.get("qscheme") != "per_group":
|
||||
logger.debug(
|
||||
"Quark model's weight quantization is incompatible with OCP MX format: "
|
||||
"weight is not per_group."
|
||||
)
|
||||
return False
|
||||
|
||||
# Input and weight group size needs to be 32.
|
||||
if weight_quant.get("group_size") != 32:
|
||||
logger.debug(
|
||||
"Quark model's weight quantization is incompatible with OCP MX format: "
|
||||
"group_size of weight is not 32."
|
||||
)
|
||||
return False
|
||||
|
||||
# Activations and weight scales need to be in e8m0 format.
|
||||
if weight_quant.get("scale_format") != "e8m0":
|
||||
logger.debug(
|
||||
"Quark model's weight quantization is incompatible with OCP MX format: "
|
||||
"scale_format of weight is not e8m0."
|
||||
)
|
||||
return False
|
||||
|
||||
# Input and weight dtypes need to be any of fp4,
|
||||
# fp6_e3m2 or fp6_e3m2, possibly mixed.
|
||||
if weight_quant.get("dtype") not in {
|
||||
"fp4",
|
||||
"fp6_e3m2",
|
||||
"fp6_e2m3",
|
||||
}:
|
||||
logger.debug(
|
||||
"Quark model's weight quantization is incompatible with OCP MX format: "
|
||||
"dtype is not in {fp4, fp6_e3m2, fp6_e2m3}."
|
||||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool:
|
||||
"""
|
||||
For Quark, determine if it's OCP MXFP4 by checking config directly.
|
||||
This allows hidden_size rounding to happen before moe_config creation.
|
||||
"""
|
||||
layer_quant_config = self._find_matched_config(prefix, layer)
|
||||
weight_config = layer_quant_config.get("weight")
|
||||
input_config = layer_quant_config.get("input_tensors")
|
||||
|
||||
return (
|
||||
self._is_w_ocp_mx_a_x(weight_config, input_config)
|
||||
and weight_config is not None
|
||||
and weight_config.get("dtype") == "fp4"
|
||||
and getattr(torch, "float4_e2m1fn_x2", None) is not None
|
||||
)
|
||||
|
||||
def _find_matched_config(
|
||||
self, layer_name: str, module: torch.nn.Module
|
||||
) -> dict[str, Any]:
|
||||
proj_name = layer_name.split(".")[-1]
|
||||
if proj_name in self.packed_modules_mapping:
|
||||
shard_proj_names = self.packed_modules_mapping[proj_name]
|
||||
|
||||
# Convert fused_name --> [shard_names]
|
||||
shard_names = [
|
||||
layer_name.replace(proj_name, shard_proj_name)
|
||||
for shard_proj_name in shard_proj_names
|
||||
]
|
||||
|
||||
shard_configs = []
|
||||
for shard_name in shard_names:
|
||||
if shard_name == layer_name:
|
||||
config = cast(
|
||||
dict[str, Any], self.quant_config.get("global_quant_config")
|
||||
)
|
||||
else:
|
||||
config = self._find_matched_config(shard_name, module)
|
||||
shard_configs.append(config)
|
||||
|
||||
if not all(
|
||||
deep_compare(q_config, shard_configs[0]) for q_config in shard_configs
|
||||
):
|
||||
raise ValueError(
|
||||
f"Found a different quantization configuration for "
|
||||
f"{shard_proj_names} in {layer_name}. vLLM "
|
||||
"requires all to use the same scheme."
|
||||
)
|
||||
return shard_configs[0]
|
||||
else:
|
||||
layer_quant_config = cast(
|
||||
dict[str, Any], self.quant_config.get("layer_quant_config")
|
||||
)
|
||||
|
||||
def _matches_pattern(layer_name, pattern):
|
||||
if "*" not in pattern:
|
||||
return layer_name in pattern
|
||||
return fnmatch.fnmatch(layer_name, pattern)
|
||||
|
||||
for name_pattern, config in layer_quant_config.items():
|
||||
if _matches_pattern(layer_name, name_pattern):
|
||||
return config
|
||||
|
||||
layer_type = cast(str, type(module))
|
||||
layer_type_quant_config = cast(
|
||||
dict[str, Any], self.quant_config.get("layer_type_quant_config")
|
||||
)
|
||||
if layer_type in layer_type_quant_config:
|
||||
return layer_type_quant_config[layer_type]
|
||||
|
||||
global_quant_config = cast(
|
||||
dict[str, Any], self.quant_config.get("global_quant_config")
|
||||
)
|
||||
return global_quant_config
|
||||
|
||||
def _get_scheme_from_config(
|
||||
self, config: dict[str, Any], dynamic_mxfp4_quant: bool = False
|
||||
) -> "QuarkScheme":
|
||||
if config.get("output_tensors") or config.get("bias"):
|
||||
raise NotImplementedError(
|
||||
"Currently, Quark models with output_tensors "
|
||||
"and bias quantized are not supported"
|
||||
)
|
||||
weight_config = cast(dict[str, Any], config.get("weight"))
|
||||
input_config = cast(dict[str, Any], config.get("input_tensors"))
|
||||
|
||||
if self._is_nvfp4(weight_config, input_config):
|
||||
return QuarkNVFP4()
|
||||
elif self._is_fp8_w8a8(weight_config, input_config):
|
||||
is_fp8_w8a8_supported = self._check_scheme_supported(
|
||||
QuarkW8A8Fp8.get_min_capability(), error=False
|
||||
)
|
||||
if is_fp8_w8a8_supported:
|
||||
return QuarkW8A8Fp8(weight_config, input_config)
|
||||
elif self._is_static_tensor_w8a8(weight_config, input_config):
|
||||
weight_qscheme = cast(str, weight_config.get("qscheme"))
|
||||
return QuarkW8A8Int8(
|
||||
qscheme=weight_qscheme,
|
||||
is_static_input_scheme=True,
|
||||
input_symmetric=input_config.get("symmetric"),
|
||||
)
|
||||
elif self._is_w4a8_mxfp4_fp8(weight_config, input_config):
|
||||
is_w4a8_supported = self._check_scheme_supported(
|
||||
QuarkW4A8_MXFP4_FP8.get_min_capability(), error=False
|
||||
)
|
||||
if is_w4a8_supported:
|
||||
return QuarkW4A8_MXFP4_FP8(weight_config, input_config)
|
||||
elif self._is_dynamic_per_token_w8a8(weight_config, input_config):
|
||||
weight_qscheme = cast(str, weight_config.get("qscheme"))
|
||||
return QuarkW8A8Int8(
|
||||
qscheme=weight_qscheme,
|
||||
is_static_input_scheme=False,
|
||||
input_symmetric=input_config.get("symmetric"),
|
||||
)
|
||||
elif self._is_w_ocp_mx_a_x(weight_config, input_config):
|
||||
return QuarkOCP_MX(
|
||||
weight_config, input_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
|
||||
)
|
||||
|
||||
raise NotImplementedError(
|
||||
"No quark compatible scheme was found. "
|
||||
f"Weight config: {weight_config}, "
|
||||
f"Input config: {input_config}"
|
||||
)
|
||||
|
||||
def get_scheme(
|
||||
self, layer: torch.nn.Module, layer_name: str, dynamic_mxfp4_quant: bool = False
|
||||
) -> "QuarkScheme":
|
||||
layer_quant_config = self._find_matched_config(layer_name, layer)
|
||||
|
||||
# Find the quant_scheme
|
||||
scheme = self._get_scheme_from_config(
|
||||
layer_quant_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
|
||||
)
|
||||
# Raise error if device does not support the scheme
|
||||
# (e.g. fp8 needs ada lovelace)
|
||||
self._check_scheme_supported(scheme.get_min_capability())
|
||||
|
||||
return scheme
|
||||
|
||||
@staticmethod
|
||||
def get_cache_scale_mapper() -> "WeightsMapper":
|
||||
"""Map Quark KV-cache scale names to vLLM names."""
|
||||
orig_to_new_suffix = {
|
||||
".k_proj.output_scale": ".attn.k_scale",
|
||||
".v_proj.output_scale": ".attn.v_scale",
|
||||
".q_proj.output_scale": ".attn.q_scale",
|
||||
".self_attn.prob_output_scale": ".self_attn.attn.prob_scale",
|
||||
}
|
||||
cache_scale_mapper = WeightsMapper(orig_to_new_suffix=orig_to_new_suffix)
|
||||
return cache_scale_mapper | QuantizationConfig.get_cache_scale_mapper()
|
||||
|
||||
|
||||
class QuarkLinearMethod(LinearMethodBase):
|
||||
def __init__(self, quantization_config: QuarkConfig):
|
||||
self.quantization_config = quantization_config
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.scheme.process_weights_after_loading(layer)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
"""
|
||||
Use the CompressedTensorsScheme associated with each layer to create
|
||||
the necessary parameters for the layer. See LinearMethodBase for param
|
||||
details
|
||||
"""
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
layer.scheme.create_weights(
|
||||
layer=layer,
|
||||
input_size=input_size,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
output_partition_sizes=output_partition_sizes,
|
||||
output_size=output_size,
|
||||
params_dtype=params_dtype,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
):
|
||||
"""
|
||||
Use the output of create_weights and the CompressedTensorsScheme
|
||||
associated with the layer to apply the forward pass with the
|
||||
layer input. See LinearMethodBase for param details
|
||||
|
||||
"""
|
||||
scheme = layer.scheme
|
||||
if scheme is None:
|
||||
raise ValueError("A scheme must be defined for each layer")
|
||||
|
||||
return scheme.apply_weights(layer, x, bias=bias)
|
||||
|
||||
|
||||
class QuarkKVCacheMethod(BaseKVCacheMethod):
|
||||
"""
|
||||
Supports loading kv-cache scaling factors from quark checkpoints.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: QuarkConfig):
|
||||
self.validate_kv_cache_config(quant_config.kv_cache_config)
|
||||
super().__init__(quant_config)
|
||||
|
||||
@staticmethod
|
||||
def validate_kv_cache_config(kv_cache_config: dict[str, Any] | None):
|
||||
"""
|
||||
Validator for the kv cache configuration. Useful for controlling the
|
||||
kv cache quantization schemes, that are being supported in vLLM
|
||||
|
||||
Args:
|
||||
kv_cache_config: the quark kv cache scheme
|
||||
"""
|
||||
if kv_cache_config is None:
|
||||
return
|
||||
|
||||
dtype = kv_cache_config.get("dtype")
|
||||
if dtype != "fp8_e4m3":
|
||||
raise NotImplementedError(
|
||||
"Currently supported kv cache quantization is "
|
||||
f"dtype=fp8_e4m3, however received {dtype}"
|
||||
)
|
||||
|
||||
qscheme = kv_cache_config.get("qscheme")
|
||||
if qscheme != "per_tensor":
|
||||
raise NotImplementedError(
|
||||
"Only support per-tensor scaling factor "
|
||||
"for quark KV cache. "
|
||||
f"Expected qscheme: per_tensor, found qscheme: {qscheme}"
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,18 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from .quark_nvfp4 import QuarkNVFP4
|
||||
from .quark_ocp_mx import QuarkOCP_MX
|
||||
from .quark_scheme import QuarkScheme
|
||||
from .quark_w4a8_mxfp4_fp8 import QuarkW4A8_MXFP4_FP8
|
||||
from .quark_w8a8_fp8 import QuarkW8A8Fp8
|
||||
from .quark_w8a8_int8 import QuarkW8A8Int8
|
||||
|
||||
__all__ = [
|
||||
"QuarkScheme",
|
||||
"QuarkW8A8Fp8",
|
||||
"QuarkW8A8Int8",
|
||||
"QuarkOCP_MX",
|
||||
"QuarkW4A8_MXFP4_FP8",
|
||||
"QuarkNVFP4",
|
||||
]
|
||||
@@ -0,0 +1,154 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel
|
||||
from vllm.model_executor.kernels.linear.nvfp4.emulation import (
|
||||
EmulationNvFp4LinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.quark.schemes.quark_scheme import (
|
||||
QuarkScheme,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
GroupQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
|
||||
__all__ = ["QuarkNVFP4"]
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class QuarkNVFP4(QuarkScheme):
|
||||
"""
|
||||
Quark NVFP4 quantization scheme.
|
||||
|
||||
Supports loading NVFP4 checkpoints with the following structure:
|
||||
- weight: uint8, shape [out_features, in_features // 2] (packed FP4)
|
||||
- weight_scale: float8_e4m3fn, shape [out_features, in_features // group_size]
|
||||
- weight_scale_2: bfloat16/float32, scalar (global weight scale)
|
||||
- input_scale_2: bfloat16/float32, scalar (global input scale)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
):
|
||||
self.kernel = init_nvfp4_linear_kernel()
|
||||
self.group_size = 16
|
||||
|
||||
if not isinstance(self.kernel, EmulationNvFp4LinearKernel):
|
||||
logger.warning_once(
|
||||
"Only EmulationNvFp4LinearKernel NVFP4 dense implementation is "
|
||||
"tested with QuarkNVFP4, got kernel=%s. Correctness is not validated.",
|
||||
type(self.kernel).__name__,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
# FP4 requires Turing (75) or newer
|
||||
return 75
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
output_partition_sizes: list[int],
|
||||
input_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
weight_loader: Callable,
|
||||
**kwargs,
|
||||
):
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
layer.logical_widths = output_partition_sizes
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
|
||||
if input_size_per_partition % self.group_size != 0:
|
||||
raise ValueError(
|
||||
f"Input size per partition ({input_size_per_partition}) must be "
|
||||
f"divisible by group size ({self.group_size})"
|
||||
)
|
||||
|
||||
# Weight: FP4 packed as uint8 (2 FP4 values per uint8)
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // 2,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# Per-group weight scale (FP8 E4M3)
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // self.group_size,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
# Global weight scale (scalar, per partition)
|
||||
weight_scale_2 = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale_2", weight_scale_2)
|
||||
|
||||
# Global input scale (scalar, per partition)
|
||||
input_scale_2 = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("input_scale_2", input_scale_2)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
input_global_scale = layer.input_scale_2.max().to(torch.float32)
|
||||
layer.input_global_scale = Parameter(input_global_scale, requires_grad=False)
|
||||
del layer.input_scale_2
|
||||
|
||||
weight_global_scale = layer.weight_scale_2.to(torch.float32)
|
||||
|
||||
if torch.unique(weight_global_scale).numel() != 1:
|
||||
logger.warning_once(
|
||||
"In NVFP4 linear, the global scale for weight are different"
|
||||
" for parallel layers (e.g. q_proj, k_proj, v_proj). This"
|
||||
" will likely result in reduced accuracy. Please verify the"
|
||||
" model accuracy. Consider using a checkpoint with a shared"
|
||||
" global NVFP4 scale for fused layers."
|
||||
)
|
||||
|
||||
weight_global_scale = weight_global_scale.max()
|
||||
|
||||
layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False)
|
||||
del layer.weight_scale_2
|
||||
|
||||
layer.alpha = Parameter(
|
||||
layer.input_global_scale * layer.weight_global_scale, requires_grad=False
|
||||
)
|
||||
layer.input_global_scale_inv = Parameter(
|
||||
(1.0 / layer.input_global_scale).to(torch.float32), requires_grad=False
|
||||
)
|
||||
|
||||
# Convert layer to NVFP4 linear kernel format
|
||||
self.kernel.process_weights_after_loading(layer)
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return self.kernel.apply_weights(layer=layer, x=x, bias=bias)
|
||||
@@ -0,0 +1,381 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
from fractions import Fraction
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
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 (
|
||||
dequant_mxfp6,
|
||||
quant_dequant_mxfp6,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
|
||||
OCP_MX_BLOCK_SIZE,
|
||||
OCP_MX_Scheme,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
GroupQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
PackedvLLMParameter,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from .quark_scheme import QuarkScheme
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
try:
|
||||
from aiter.ops.shuffle import shuffle_weight
|
||||
from aiter.ops.triton.gemm_afp4wfp4 import (
|
||||
gemm_afp4wfp4,
|
||||
gemm_afp4wfp4_preshuffled_weight_scales,
|
||||
)
|
||||
from aiter.ops.triton.quant import dynamic_mxfp4_quant
|
||||
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
if rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled():
|
||||
from aiter import gemm_a4w4, per_1x32_f4_quant_hip
|
||||
|
||||
def gemm_with_dynamic_quant(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
rocm_use_aiter_fp4_asm_gemm: bool = False,
|
||||
out_dtype: torch.dtype | None = torch.bfloat16,
|
||||
x_scales: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
M = x.shape[0]
|
||||
N = weight.shape[0]
|
||||
K = weight.shape[1]
|
||||
if rocm_use_aiter_fp4_asm_gemm:
|
||||
if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K):
|
||||
if x_scales is None:
|
||||
# use hip quant kernel for performance
|
||||
if M >= 32:
|
||||
x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
|
||||
else:
|
||||
x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=False)
|
||||
else:
|
||||
x_q = x
|
||||
x_s = x_scales
|
||||
|
||||
if M >= 32:
|
||||
x_s = x_s.view(torch.uint8).view(x_s.shape[0] // 32, -1)
|
||||
else:
|
||||
x_s = x_s[:M, ...].view(torch.uint8)
|
||||
|
||||
y = torch.empty(M, N, device=x_q.device, dtype=out_dtype)
|
||||
gemm_afp4wfp4_preshuffled_weight_scales(
|
||||
x_q.view(torch.uint8),
|
||||
weight.view(torch.uint8).view(weight.shape[0] // 16, -1),
|
||||
x_s,
|
||||
weight_scale.view(torch.uint8).view(
|
||||
weight_scale.shape[0] // 32, -1
|
||||
),
|
||||
out_dtype,
|
||||
y,
|
||||
)
|
||||
else:
|
||||
if x_scales is None:
|
||||
# use hip quant kernel for performance
|
||||
x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
|
||||
else:
|
||||
x_q = x
|
||||
x_s = x_scales
|
||||
|
||||
y = gemm_a4w4(
|
||||
x_q,
|
||||
weight.view(x_q.dtype),
|
||||
x_s,
|
||||
weight_scale.view(x_s.dtype),
|
||||
dtype=out_dtype,
|
||||
bpreshuffle=True,
|
||||
)
|
||||
return y[:M]
|
||||
else:
|
||||
if x_scales is None:
|
||||
x_q, x_s = dynamic_mxfp4_quant(x)
|
||||
else:
|
||||
x_q = x
|
||||
x_s = x_scales
|
||||
y = torch.empty(
|
||||
x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype
|
||||
)
|
||||
|
||||
gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y)
|
||||
return y
|
||||
|
||||
def gemm_with_dynamic_quant_fake(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
x_scales: torch.Tensor = None,
|
||||
rocm_use_aiter_fp4_asm_gemm: bool = False,
|
||||
out_dtype: torch.dtype | None = torch.bfloat16,
|
||||
) -> torch.Tensor:
|
||||
return torch.empty(
|
||||
(*x.shape[:-1], weight.shape[0]), dtype=out_dtype, device=x.device
|
||||
)
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="gemm_with_dynamic_quant",
|
||||
op_func=gemm_with_dynamic_quant,
|
||||
mutates_args=[],
|
||||
fake_impl=gemm_with_dynamic_quant_fake,
|
||||
dispatch_key=current_platform.dispatch_key,
|
||||
)
|
||||
except (ImportError, AttributeError, RuntimeError):
|
||||
if current_platform.is_rocm():
|
||||
logger.warning(
|
||||
"AITER is not found or QuarkOCP_MX is not supported on the current "
|
||||
"platform. QuarkOCP_MX quantization will not be available."
|
||||
)
|
||||
dynamic_mxfp4_quant = gemm_afp4wfp4 = None
|
||||
|
||||
|
||||
class QuarkOCP_MX(QuarkScheme):
|
||||
def __init__(
|
||||
self,
|
||||
weight_quant_spec: dict[str, Any],
|
||||
input_quant_spec: dict[str, Any] | None,
|
||||
dynamic_mxfp4_quant: bool = False,
|
||||
):
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
self.qscheme = "per_group"
|
||||
self.weight_quant_spec = weight_quant_spec
|
||||
self.input_quant_spec = input_quant_spec
|
||||
self.dynamic_mxfp4_quant = dynamic_mxfp4_quant
|
||||
self.weight_dtype = weight_quant_spec["dtype"].replace("fp", "mxfp")
|
||||
self.input_dtype: str | None = None
|
||||
if input_quant_spec is not None:
|
||||
input_quant = input_quant_spec["dtype"]
|
||||
if input_quant == "fp8_e4m3":
|
||||
self.input_dtype = "fp8"
|
||||
else:
|
||||
self.input_dtype = input_quant.replace("fp", "mxfp")
|
||||
|
||||
self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(
|
||||
self.input_dtype, self.weight_dtype
|
||||
)
|
||||
|
||||
if self.weight_dtype == "mxfp4":
|
||||
self.packed_factor: int | Fraction = 2
|
||||
self.dequant_func = dequant_mxfp4
|
||||
else:
|
||||
self.packed_factor = Fraction(numerator=8, denominator=6)
|
||||
self.dequant_func = partial(
|
||||
dequant_mxfp6, quant_dtype=self.weight_dtype.replace("mx", "")
|
||||
)
|
||||
|
||||
if self.input_dtype is None:
|
||||
self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = (
|
||||
lambda x: x
|
||||
) # no input Q/DQ for weight-only
|
||||
elif self.input_dtype == "mxfp4":
|
||||
self.quant_dequant_func = quant_dequant_mxfp4
|
||||
else:
|
||||
self.quant_dequant_func = partial(
|
||||
quant_dequant_mxfp6, quant_dtype=self.input_dtype.replace("mx", "")
|
||||
)
|
||||
|
||||
if input_quant_spec is None:
|
||||
self.static_input_scales = False
|
||||
else:
|
||||
self.static_input_scales = not input_quant_spec.get("is_dynamic")
|
||||
|
||||
if self.static_input_scales:
|
||||
raise NotImplementedError(
|
||||
"QuarkOCP_MX with static input scales is currently not "
|
||||
"implemented. Please open an issue."
|
||||
)
|
||||
|
||||
# TODO: integrate (or test) mixed-precision kernel.
|
||||
self.emulate = not current_platform.supports_mx() or (
|
||||
self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4"
|
||||
)
|
||||
|
||||
self.rocm_use_aiter_fp4_asm_gemm = (
|
||||
rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled()
|
||||
)
|
||||
|
||||
if not self.emulate and (dynamic_mxfp4_quant is None or gemm_afp4wfp4 is None):
|
||||
# Currently need these kernels if not emulating
|
||||
raise NotImplementedError(
|
||||
f"{self.__class__.__name__} requires AITER to be installed "
|
||||
"for non-emulation mode! Please refer to "
|
||||
"https://github.com/ROCm/aiter for installation details."
|
||||
)
|
||||
|
||||
if not current_platform.supports_mx():
|
||||
logger.warning_once(
|
||||
"The current platform does not support native MXFP4/MXFP6 "
|
||||
"computation. Simulated weight dequantization and activation "
|
||||
"QDQ (quantize and dequantize) will be used, with the linear "
|
||||
"layers computed in high precision."
|
||||
)
|
||||
|
||||
if current_platform.supports_mx() and (
|
||||
self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4"
|
||||
):
|
||||
logger.warning_once(
|
||||
"The current platform supports native MXFP4/MXFP6 "
|
||||
f"computation, but kernels for input_dtype={self.input_dtype} "
|
||||
f"and weight_dtype={self.weight_dtype} are not yet integrated "
|
||||
"in vLLM. Simulated weight dequantization and activation "
|
||||
"QDQ (quantize and dequantize) will be used, with the linear "
|
||||
"layers computed in high precision."
|
||||
)
|
||||
|
||||
def get_packed_dim(self, dim: int, quant_dtype: str):
|
||||
if quant_dtype == "mxfp4":
|
||||
assert dim % 2 == 0
|
||||
return dim // 2
|
||||
elif quant_dtype in {"mxfp6_e3m2", "mxfp6_e2m3"}:
|
||||
# FP6 packs 4 * 6 = 24 bits on 3 bytes.
|
||||
assert (dim * 3) % 4 == 0
|
||||
return (dim * 3) // 4
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"Unsupported quant_dtype in QuarkOCP_MX.get_packed_dim, "
|
||||
f"got quant_dtype={quant_dtype}. Something is wrong, please "
|
||||
"open an issue."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 70
|
||||
|
||||
def process_dynamic_mxfp4_weights_after_loading(
|
||||
self, layer: torch.nn.Module
|
||||
) -> None:
|
||||
w_q, w_s = dynamic_mxfp4_quant(layer.weight)
|
||||
layer.weight_scale = torch.nn.Parameter(w_s.T.contiguous(), requires_grad=False)
|
||||
layer.weight = torch.nn.Parameter(w_q, requires_grad=False)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False)
|
||||
|
||||
if self.emulate:
|
||||
if self.dynamic_mxfp4_quant:
|
||||
self.process_dynamic_mxfp4_weights_after_loading(layer)
|
||||
else:
|
||||
layer.weight_scale = torch.nn.Parameter(
|
||||
layer.weight_scale.data, requires_grad=False
|
||||
)
|
||||
else:
|
||||
if self.dynamic_mxfp4_quant:
|
||||
self.process_dynamic_mxfp4_weights_after_loading(layer)
|
||||
elif self.rocm_use_aiter_fp4_asm_gemm:
|
||||
# shuffle weight scale
|
||||
weight_scale_shuffle = layer.weight_scale.data
|
||||
sm, sn = weight_scale_shuffle.shape
|
||||
weight_scale_shuffle = weight_scale_shuffle.view(
|
||||
sm // 32, 2, 16, sn // 8, 2, 4, 1
|
||||
)
|
||||
weight_scale_shuffle = weight_scale_shuffle.permute(
|
||||
0, 3, 5, 2, 4, 1, 6
|
||||
).contiguous()
|
||||
weight_scale_shuffle = weight_scale_shuffle.view(sm, sn)
|
||||
layer.weight_scale = torch.nn.Parameter(
|
||||
weight_scale_shuffle, requires_grad=False
|
||||
)
|
||||
|
||||
# shuffle weight
|
||||
weight_shuffle = layer.weight.data
|
||||
weight_shuffle = shuffle_weight(weight_shuffle, layout=(16, 16))
|
||||
layer.weight = torch.nn.Parameter(weight_shuffle, requires_grad=False)
|
||||
else:
|
||||
layer.weight_scale = torch.nn.Parameter(
|
||||
layer.weight_scale.data.T.contiguous(), requires_grad=False
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
output_partition_sizes: list[int],
|
||||
input_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
weight_loader: Callable,
|
||||
**kwargs,
|
||||
):
|
||||
if self.dynamic_mxfp4_quant:
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
layer.register_parameter("weight", weight)
|
||||
set_weight_attrs(weight, kwargs)
|
||||
else:
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
layer.logical_widths = output_partition_sizes
|
||||
|
||||
# WEIGHT
|
||||
weight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
self.get_packed_dim(input_size_per_partition, self.weight_dtype),
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
packed_dim=1,
|
||||
packed_factor=self.packed_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# WEIGHT SCALE
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // OCP_MX_BLOCK_SIZE,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if self.emulate:
|
||||
dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype)
|
||||
qdq_x = self.quant_dequant_func(x)
|
||||
return F.linear(qdq_x, dq_w, bias)
|
||||
y = torch.ops.vllm.gemm_with_dynamic_quant(
|
||||
x,
|
||||
layer.weight,
|
||||
layer.weight_scale,
|
||||
self.rocm_use_aiter_fp4_asm_gemm,
|
||||
self.out_dtype,
|
||||
)
|
||||
# gemm_with_dynamic_quant has no bias argument; add it here so the
|
||||
# native path matches F.linear (e.g. qkv_proj with qkv_bias=True).
|
||||
if bias is not None:
|
||||
y = y + bias
|
||||
return y
|
||||
@@ -0,0 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
|
||||
__all__ = ["QuarkScheme"]
|
||||
|
||||
|
||||
class QuarkScheme(ABC):
|
||||
"""
|
||||
Abstract class used to describe the weight creation and forward pass
|
||||
of different quantization schemes supported by Quark.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
"""
|
||||
Get minimum device capability.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def create_weights(self, *args, **kwargs):
|
||||
"""
|
||||
Weight creation for the particular scheme. Inputs to this function
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def apply_weights(
|
||||
self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
|
||||
):
|
||||
"""
|
||||
Run the forward pass for the particular scheme. This is where
|
||||
scheme-specific dequant/quant steps/kernels should be applied.
|
||||
|
||||
Args:
|
||||
layer: torch.nn.Module with the registered weights and
|
||||
other parameters relevant to the particular scheme.
|
||||
x: input to the layer
|
||||
bias: bias parameter
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module):
|
||||
"""
|
||||
Called after weight loading is complete for any cleanup that
|
||||
needs to occur.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,218 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
from fractions import Fraction
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
get_fp8_min_max,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
GroupQuantScaleParameter,
|
||||
PackedvLLMParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from .quark_scheme import QuarkScheme
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
__all__ = ["QuarkW4A8_MXFP4_FP8"]
|
||||
|
||||
OCP_MX_BLOCK_SIZE = 32
|
||||
|
||||
|
||||
class QuarkW4A8_MXFP4_FP8(QuarkScheme):
|
||||
"""
|
||||
- Weights: MXFP4 with E8M0 scales per block of 32
|
||||
- Activations: FP8 E4M3 (static per-tensor quantization)
|
||||
|
||||
Uses the AITER Triton kernel and falls back to emulation if AITER not available.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weight_quant_spec: dict[str, Any],
|
||||
input_quant_spec: dict[str, Any],
|
||||
):
|
||||
self.out_dtype = None
|
||||
|
||||
self.weight_dtype = "mxfp4"
|
||||
self.packed_factor: Fraction = Fraction(2, 1) # 2 FP4 values per byte
|
||||
self.weight_block_size = OCP_MX_BLOCK_SIZE
|
||||
|
||||
self.is_static_input_scheme = not input_quant_spec.get("is_dynamic")
|
||||
self.input_qscheme = input_quant_spec.get("qscheme") # "per_tensor"
|
||||
|
||||
self.fp8_min, self.fp8_max = get_fp8_min_max()
|
||||
self.fp8_dtype = current_platform.fp8_dtype()
|
||||
|
||||
if not self.is_static_input_scheme:
|
||||
raise NotImplementedError(
|
||||
"Dynamic FP8 activation quantization is not yet supported "
|
||||
"for W4A8. The current implementation expects static per-tensor "
|
||||
"FP8 scales stored in the checkpoint."
|
||||
)
|
||||
|
||||
kernel_supported_gpu = False
|
||||
if current_platform.is_rocm():
|
||||
from vllm.platforms.rocm import on_gfx950
|
||||
|
||||
kernel_supported_gpu = on_gfx950()
|
||||
|
||||
self.use_aiter_kernel = (
|
||||
is_aiter_found_and_supported()
|
||||
and self.is_static_input_scheme
|
||||
and kernel_supported_gpu
|
||||
)
|
||||
|
||||
if not self.use_aiter_kernel:
|
||||
logger.warning_once(
|
||||
"[W4A8 MXFP4+FP8] Aiter Triton kernel not found. Using emulation mode."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 70
|
||||
|
||||
def get_packed_dim(self, dim: int) -> int:
|
||||
assert dim % 2 == 0, f"Dimension {dim} must be even for MXFP4 packing"
|
||||
return dim // 2
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
output_partition_sizes: list[int],
|
||||
input_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
weight_loader: Callable,
|
||||
**kwargs,
|
||||
):
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
layer.logical_widths = output_partition_sizes
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
|
||||
# MXFP4 WEIGHT (packed, 2 values per byte)
|
||||
weight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
self.get_packed_dim(input_size_per_partition),
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
packed_dim=1,
|
||||
packed_factor=self.packed_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# WEIGHT SCALE (E8M0 format, per block of 32)
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // self.weight_block_size,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
# INPUT SCALE (FP8 per-tensor static scale)
|
||||
if self.is_static_input_scheme:
|
||||
input_scale = PerTensorScaleParameter(
|
||||
data=torch.empty(
|
||||
len(output_partition_sizes),
|
||||
dtype=torch.float32,
|
||||
),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
# Initialize to avoid NaN
|
||||
input_scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("input_scale", input_scale)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
# Ensuring weights & scales are non-trainable
|
||||
layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False)
|
||||
layer.weight_scale = torch.nn.Parameter(
|
||||
layer.weight_scale.data, requires_grad=False
|
||||
)
|
||||
|
||||
if self.is_static_input_scheme:
|
||||
input_scale = layer.input_scale.data
|
||||
# For fused modules (QKV), take the max scale
|
||||
if input_scale.numel() != 1:
|
||||
input_scale = input_scale.max()
|
||||
|
||||
layer.input_scale = torch.nn.Parameter(
|
||||
torch.tensor(input_scale, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if self.use_aiter_kernel:
|
||||
return self._apply_aiter_kernel(layer, x, bias)
|
||||
else:
|
||||
return self._apply_emulation(layer, x, bias)
|
||||
|
||||
def _apply_aiter_kernel(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
M = x.shape[0]
|
||||
out_dtype = x.dtype if self.out_dtype is None else self.out_dtype
|
||||
|
||||
input_scale = layer.input_scale
|
||||
x_fp8 = (x / input_scale).clamp(self.fp8_min, self.fp8_max).to(self.fp8_dtype)
|
||||
|
||||
# Broadcast per-tensor scale to per-row (M, 1) for Aiter kernel
|
||||
x_scales = input_scale.expand(M, 1).to(dtype=torch.float32, device=x.device)
|
||||
|
||||
y = rocm_aiter_ops.gemm_a8wfp4(
|
||||
x_fp8, layer.weight, x_scales, layer.weight_scale, out_dtype
|
||||
)
|
||||
|
||||
if bias is not None:
|
||||
y = y + bias
|
||||
|
||||
return y
|
||||
|
||||
def _apply_emulation(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
|
||||
dequant_mxfp4,
|
||||
)
|
||||
|
||||
weight_dq = dequant_mxfp4(
|
||||
layer.weight,
|
||||
layer.weight_scale,
|
||||
x.dtype,
|
||||
)
|
||||
|
||||
input_scale = layer.input_scale
|
||||
x_fp8 = (x / input_scale).clamp(self.fp8_min, self.fp8_max).to(self.fp8_dtype)
|
||||
x_dq = (x_fp8.to(x.dtype) * input_scale).to(x.dtype)
|
||||
|
||||
return F.linear(x_dq, weight_dq, bias)
|
||||
@@ -0,0 +1,198 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast
|
||||
|
||||
import torch
|
||||
from torch.nn import Parameter
|
||||
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
init_fp8_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8StaticChannelSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
normalize_e4m3fn_to_e4m3fnuz,
|
||||
requantize_with_max_scale,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
ChannelQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
__all__ = ["QuarkW8A8Fp8"]
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class QuarkW8A8Fp8(QuarkScheme):
|
||||
def __init__(
|
||||
self, weight_config: dict[str, Any], input_config: dict[str, Any] | None
|
||||
):
|
||||
self.weight_qscheme = cast(str, weight_config.get("qscheme"))
|
||||
self.is_static_input_scheme: bool = False
|
||||
self.input_qscheme: str | None = None
|
||||
if input_config is not None:
|
||||
self.is_static_input_scheme = not cast(bool, input_config.get("is_dynamic"))
|
||||
self.input_qscheme = cast(str, input_config.get("qscheme"))
|
||||
|
||||
per_token_activation = (
|
||||
not self.is_static_input_scheme and self.input_qscheme == "per_channel"
|
||||
)
|
||||
per_channel_weight = self.weight_qscheme == "per_channel"
|
||||
|
||||
self.activation_quant_key = (
|
||||
kFp8DynamicTokenSym if per_token_activation else kFp8StaticTensorSym
|
||||
)
|
||||
# A per-output-channel weight scale is one fp32 value per weight row
|
||||
# (length N). Tag it as ``GroupShape.PER_CHANNEL`` to match the
|
||||
# canonical compressed-tensors CHANNEL strategy, so kernel selection
|
||||
# (e.g. AITER's pre-shuffled FP8 GEMM) treats it uniformly.
|
||||
self.weight_quant_key = (
|
||||
kFp8StaticChannelSym if per_channel_weight else kFp8StaticTensorSym
|
||||
)
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
self.input_dtype = get_current_vllm_config().model_config.dtype
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
# lovelace and up
|
||||
return 89
|
||||
|
||||
def process_weights_after_loading(self, layer) -> None:
|
||||
# If per tensor, when we have a fused module (e.g. QKV) with per
|
||||
# tensor scales (thus N scales being passed to the kernel),
|
||||
# requantize so we can always run per tensor
|
||||
if self.weight_qscheme == "per_tensor":
|
||||
if current_platform.is_fp8_fnuz():
|
||||
input_scale = getattr(layer, "input_scale", None)
|
||||
weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
|
||||
weight=layer.weight,
|
||||
weight_scale=layer.weight_scale,
|
||||
input_scale=input_scale,
|
||||
)
|
||||
if input_scale is not None:
|
||||
layer.input_scale = Parameter(input_scale, requires_grad=False)
|
||||
else:
|
||||
max_w_scale = layer.weight_scale
|
||||
weight = layer.weight
|
||||
|
||||
max_w_scale, weight = requantize_with_max_scale(
|
||||
weight=weight,
|
||||
weight_scale=max_w_scale,
|
||||
logical_widths=layer.logical_widths,
|
||||
)
|
||||
|
||||
layer.weight = Parameter(weight.t(), requires_grad=False)
|
||||
layer.weight_scale = Parameter(max_w_scale, requires_grad=False)
|
||||
|
||||
# If channelwise, scales are already lined up, so just transpose.
|
||||
elif self.weight_qscheme == "per_channel":
|
||||
weight = layer.weight
|
||||
|
||||
if current_platform.is_fp8_fnuz():
|
||||
input_scale = getattr(layer, "input_scale", None)
|
||||
weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
|
||||
weight=weight,
|
||||
weight_scale=layer.weight_scale,
|
||||
input_scale=input_scale,
|
||||
)
|
||||
if input_scale is not None:
|
||||
layer.input_scale = Parameter(input_scale, requires_grad=False)
|
||||
else:
|
||||
weight_scale = layer.weight_scale.data
|
||||
if self.activation_quant_key.scale.group_shape == GroupShape.PER_TOKEN:
|
||||
weight_scale = weight_scale.view(-1, 1)
|
||||
layer.weight = Parameter(weight.t(), requires_grad=False)
|
||||
# required by torch.compile to be torch.nn.Parameter
|
||||
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown quantization scheme {self.weight_qscheme}")
|
||||
|
||||
# INPUT SCALE
|
||||
if self.is_static_input_scheme:
|
||||
layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
|
||||
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
output_partition_sizes: list[int],
|
||||
input_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
weight_loader: Callable,
|
||||
**kwargs,
|
||||
):
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
layer.logical_widths = output_partition_sizes
|
||||
|
||||
# WEIGHT
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# WEIGHT SCALE
|
||||
# TODO: update create_xxx_parameter functions to return
|
||||
# the newly added parameters
|
||||
if self.weight_qscheme == "per_channel":
|
||||
weight_scale = ChannelQuantScaleParameter(
|
||||
data=torch.empty((sum(output_partition_sizes)), dtype=torch.float32),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
else:
|
||||
assert self.weight_qscheme == "per_tensor"
|
||||
weight_scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
# min requirement for fp8 kernels
|
||||
weight_scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
# INPUT SCALE
|
||||
if self.is_static_input_scheme:
|
||||
input_scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
input_scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("input_scale", input_scale)
|
||||
|
||||
self.fp8_linear = init_fp8_linear_kernel(
|
||||
activation_quant_key=self.activation_quant_key,
|
||||
weight_quant_key=self.weight_quant_key,
|
||||
weight_shape=layer.weight.shape,
|
||||
input_dtype=self.input_dtype,
|
||||
out_dtype=self.out_dtype,
|
||||
module_name=self.__class__.__name__,
|
||||
)
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return self.fp8_linear.apply_weights(layer, x, bias)
|
||||
@@ -0,0 +1,138 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
init_int8_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
|
||||
from vllm.model_executor.parameter import (
|
||||
BasevLLMParameter,
|
||||
ChannelQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class QuarkW8A8Int8(QuarkScheme):
|
||||
def __init__(
|
||||
self,
|
||||
qscheme: str,
|
||||
is_static_input_scheme: bool | None,
|
||||
input_symmetric: bool | None,
|
||||
):
|
||||
self.qscheme = qscheme
|
||||
self.is_static_input_scheme = is_static_input_scheme
|
||||
self.input_symmetric = input_symmetric
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
# turing and up
|
||||
return 75
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
output_partition_sizes: list[int],
|
||||
input_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
weight_loader: Callable,
|
||||
**kwargs,
|
||||
):
|
||||
layer.logical_widths = output_partition_sizes
|
||||
|
||||
# Quark stores per-channel weight_scale as 1D [N]; reshape to [N, 1].
|
||||
def _scale_weight_loader(
|
||||
param: torch.nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if loaded_weight.dim() == 1:
|
||||
loaded_weight = loaded_weight.unsqueeze(-1)
|
||||
return weight_loader(param, loaded_weight, *args, **kwargs)
|
||||
|
||||
self.kernel = init_int8_linear_kernel(
|
||||
is_channelwise=(self.qscheme == "per_channel"),
|
||||
is_static_input_scheme=(self.is_static_input_scheme is True),
|
||||
input_symmetric=(self.input_symmetric is True),
|
||||
module_name=self.__class__.__name__,
|
||||
)
|
||||
|
||||
# WEIGHT
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
sum(output_partition_sizes), input_size_per_partition, dtype=torch.int8
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# WEIGHT SCALE
|
||||
if self.qscheme == "per_channel":
|
||||
weight_scale = ChannelQuantScaleParameter(
|
||||
data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32),
|
||||
output_dim=0,
|
||||
weight_loader=_scale_weight_loader,
|
||||
)
|
||||
ChannelQuantZPParameter = ChannelQuantScaleParameter
|
||||
weight_zero_point = ChannelQuantZPParameter(
|
||||
data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.int8),
|
||||
output_dim=0,
|
||||
weight_loader=_scale_weight_loader,
|
||||
)
|
||||
else:
|
||||
assert self.qscheme == "per_tensor"
|
||||
weight_scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
PerTensorZPParameter = PerTensorScaleParameter
|
||||
weight_zero_point = PerTensorZPParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.int8),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
layer.register_parameter("weight_zero_point", weight_zero_point)
|
||||
|
||||
# INPUT SCALE
|
||||
input_zero_point = None
|
||||
input_scale = None
|
||||
if self.is_static_input_scheme:
|
||||
input_scale = BasevLLMParameter(
|
||||
data=torch.empty(1, dtype=torch.float32), weight_loader=weight_loader
|
||||
)
|
||||
|
||||
input_zero_point = BasevLLMParameter(
|
||||
data=torch.empty(1, dtype=torch.int8), weight_loader=weight_loader
|
||||
)
|
||||
|
||||
layer.register_parameter("input_scale", input_scale)
|
||||
layer.register_parameter("input_zero_point", input_zero_point)
|
||||
if not hasattr(layer, "azp_adj"):
|
||||
layer.register_parameter("azp_adj", None)
|
||||
|
||||
# Checkpoints are serialized in quark format, which is
|
||||
# different from the format the kernel may want. Handle repacking here.
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.register_parameter("weight_zero_point", None)
|
||||
delattr(layer, "weight_zero_point")
|
||||
if self.input_symmetric:
|
||||
layer.register_parameter("input_zero_point", None)
|
||||
delattr(layer, "input_zero_point")
|
||||
|
||||
self.kernel.process_weights_after_loading(layer)
|
||||
|
||||
def apply_weights(
|
||||
self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
|
||||
) -> torch.Tensor:
|
||||
return self.kernel.apply_weights(layer, x, bias)
|
||||
@@ -0,0 +1,120 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Iterable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
|
||||
|
||||
def deep_compare(dict1: Any, dict2: Any) -> bool:
|
||||
if type(dict1) is not type(dict2):
|
||||
return False
|
||||
if isinstance(dict1, dict):
|
||||
if dict1.keys() != dict2.keys():
|
||||
return False
|
||||
return all(deep_compare(dict1[k], dict2[k]) for k in dict1)
|
||||
elif isinstance(dict1, list):
|
||||
# `dict1` may be a list of dict.
|
||||
return all(deep_compare(dict1[i], dict2[i]) for i in range(len(dict1)))
|
||||
else:
|
||||
return dict1 == dict2
|
||||
|
||||
|
||||
def should_ignore_layer(
|
||||
layer_name: str | None,
|
||||
ignore: Iterable[str],
|
||||
fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
|
||||
) -> bool:
|
||||
if layer_name is None:
|
||||
return False
|
||||
|
||||
# layer_name = model.layers.0.self_attn.qkv_proj
|
||||
# proj_name = qkv_proj
|
||||
proj_name = layer_name.split(".")[-1]
|
||||
|
||||
# Fused layers like gate_up_proj or qkv_proj will not be fused
|
||||
# in the safetensors checkpoint. So, we convert the name
|
||||
# from the fused version to unfused + check to make sure that
|
||||
# each shard of the fused layer has the same scheme.
|
||||
if proj_name in fused_mapping:
|
||||
shard_proj_names = fused_mapping[proj_name]
|
||||
|
||||
# Convert fused_name --> [shard_names]
|
||||
shard_names = [
|
||||
layer_name.replace(proj_name, shard_proj_name)
|
||||
for shard_proj_name in shard_proj_names
|
||||
]
|
||||
|
||||
# Layer should be ignored if shards are ignored.
|
||||
should_ignore_layer = None
|
||||
for shard_name in shard_names:
|
||||
should_ignore_shard = check_equal_or_regex_match(
|
||||
layer_name=shard_name, targets=ignore
|
||||
)
|
||||
|
||||
# If shard_idx=0, set layer ignore to match shard.
|
||||
if should_ignore_layer is None:
|
||||
should_ignore_layer = should_ignore_shard
|
||||
|
||||
# If shard_idx=1+ confirm scheme matches prior shards.
|
||||
elif should_ignore_shard != should_ignore_layer:
|
||||
raise ValueError(
|
||||
f"Found a different quantization schemes for "
|
||||
f"{shard_proj_names} in {layer_name}. vLLM "
|
||||
"requires all to use the same scheme."
|
||||
)
|
||||
|
||||
# Unfused layers like down_proj and o_proj will match
|
||||
# the safetensors checkpoint already.
|
||||
else:
|
||||
should_ignore_layer = check_equal_or_regex_match(
|
||||
layer_name=layer_name, targets=ignore
|
||||
)
|
||||
|
||||
assert should_ignore_layer is not None
|
||||
return should_ignore_layer
|
||||
|
||||
|
||||
def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool:
|
||||
"""
|
||||
Checks whether a layer_name is exactly equal or a regex match for
|
||||
if target starts with 're:' to any target in list.
|
||||
"""
|
||||
return any(_is_equal_or_regex_match(layer_name, target) for target in targets)
|
||||
|
||||
|
||||
def _is_equal_or_regex_match(
|
||||
value: str, target: str, check_contains: bool = False
|
||||
) -> bool:
|
||||
"""
|
||||
Checks whether a value is exactly equal or a regex match for target
|
||||
if target starts with 're:'. If check_contains is set to True,
|
||||
additionally checks if the target string is contained within the value.
|
||||
"""
|
||||
|
||||
if target.startswith("re:"):
|
||||
pattern = target[3:]
|
||||
if re.match(pattern, value):
|
||||
return True
|
||||
elif check_contains:
|
||||
if target.lower() in value.lower():
|
||||
return True
|
||||
elif target == value:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# utility for tensor dims > 2 cases
|
||||
def quark_quantize_weight_to_mxfp4(w: torch.Tensor):
|
||||
assert w.dtype == torch.bfloat16, (
|
||||
"Quark dynamic quantization is supported only for fp16 weights and only to MXF4"
|
||||
)
|
||||
|
||||
from aiter.ops.triton.quant import dynamic_mxfp4_quant
|
||||
|
||||
*dims, d = w.shape
|
||||
w, w_scales = dynamic_mxfp4_quant(w.reshape(-1, d))
|
||||
return w.view(*dims, d // 2), w_scales.view(*dims, d // 32)
|
||||
Reference in New Issue
Block a user