137 lines
4.7 KiB
Python
137 lines
4.7 KiB
Python
# 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 torch
|
|
|
|
from vllm.model_executor.layers.quantization.utils.config_utils import (
|
|
find_matching_patterns,
|
|
)
|
|
|
|
QuarkQTensorHint = dict[str, Any] | list[dict[str, Any]] | None
|
|
|
|
|
|
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({}),
|
|
*,
|
|
check_children: bool = False,
|
|
) -> bool:
|
|
if layer_name is None:
|
|
return False
|
|
|
|
# MoE layers are currently all-or-nothing: if any child is ignored,
|
|
# the parent layer must be ignored as well. For example, the
|
|
# amd/GLM-5.2-MXFP4 config ignores children like
|
|
# model.layers.78.mlp.experts.*.down_proj, while the layer checked
|
|
# here is the parent model.layers.N.mlp.experts.
|
|
# See:
|
|
# https://huggingface.co/amd/GLM-5.2-MXFP4/blob/main/config.json#L793-L795
|
|
if check_children and any(
|
|
target == layer_name or target.startswith(layer_name + ".")
|
|
for target in ignore
|
|
if not target.startswith("re:")
|
|
):
|
|
return True
|
|
|
|
# A direct fused-layer pattern takes precedence over expansion. For
|
|
# model.layers.0.self_attn.qkv_proj,
|
|
# ignore=["re:.*qkv_proj.*"] yields [{"re:.*qkv_proj.*"}]. In contrast,
|
|
# ignore=["re:.*[qkv]_proj"] yields one matching set per expanded shard.
|
|
per_shard_matches = find_matching_patterns(layer_name, ignore, fused_mapping)
|
|
shards_ignored = [len(matches) > 0 for matches in per_shard_matches]
|
|
if any(shards_ignored) and not all(shards_ignored):
|
|
raise ValueError(
|
|
f"Found different quantization schemes for the shards of "
|
|
f"{layer_name}. vLLM requires all to use the same scheme."
|
|
)
|
|
return all(shards_ignored)
|
|
|
|
|
|
def parse_w4a16_int4_weight_config(
|
|
weight_config: Mapping[str, Any],
|
|
) -> tuple[int, bool]:
|
|
"""Parse required W4A16 INT4/UINT4 weight fields from Quark config."""
|
|
if "group_size" not in weight_config:
|
|
raise ValueError(
|
|
"Quark W4A16 INT4/UINT4 configs must specify weight.group_size"
|
|
)
|
|
if "symmetric" not in weight_config:
|
|
raise ValueError("Quark W4A16 INT4/UINT4 configs must specify weight.symmetric")
|
|
|
|
group_size = weight_config["group_size"]
|
|
is_symmetric = weight_config["symmetric"]
|
|
if not isinstance(group_size, int) or group_size <= 0:
|
|
raise ValueError(
|
|
f"Quark W4A16 weight.group_size must be a positive int, got {group_size!r}"
|
|
)
|
|
if not isinstance(is_symmetric, bool):
|
|
raise ValueError(
|
|
f"Quark W4A16 weight.symmetric must be a bool, got {is_symmetric!r}"
|
|
)
|
|
return group_size, is_symmetric
|
|
|
|
|
|
def canonicalize_quark_packed_int4(
|
|
packed_weight: torch.Tensor,
|
|
*,
|
|
pack_reorder: bool,
|
|
is_symmetric: bool,
|
|
pack_factor: int = 8,
|
|
) -> torch.Tensor:
|
|
"""Convert Quark export nibble layout to AWQ checkpoint layout."""
|
|
from vllm.model_executor.layers.quantization.auto_awq import (
|
|
_REVERSE_AWQ_PACK_ORDER,
|
|
)
|
|
|
|
if pack_reorder:
|
|
source_order = torch.tensor(
|
|
_REVERSE_AWQ_PACK_ORDER, device=packed_weight.device, dtype=torch.int32
|
|
)
|
|
else:
|
|
source_order = torch.arange(
|
|
pack_factor, device=packed_weight.device, dtype=torch.int32
|
|
)
|
|
target_order = torch.tensor(
|
|
_REVERSE_AWQ_PACK_ORDER, device=packed_weight.device, dtype=torch.int32
|
|
)
|
|
source_shifts = source_order * 4
|
|
target_shifts = target_order * 4
|
|
|
|
values = (packed_weight.to(torch.int32)[..., None] >> source_shifts) & 0xF
|
|
if is_symmetric:
|
|
values = values ^ 0x8
|
|
packed = (values.to(torch.int64) << target_shifts.to(torch.int64)).sum(dim=-1)
|
|
return packed.to(torch.int32)
|
|
|
|
|
|
# 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)
|