chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,191 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Literal, get_args
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
QuantizationMethods = Literal[
|
||||
"awq",
|
||||
"auto_awq",
|
||||
"fp8",
|
||||
"fbgemm_fp8",
|
||||
"fp_quant",
|
||||
"modelopt",
|
||||
"modelopt_fp4",
|
||||
"modelopt_mxfp8",
|
||||
"modelopt_mixed",
|
||||
"auto_gptq",
|
||||
"gptq",
|
||||
"gptq_marlin",
|
||||
"awq_marlin",
|
||||
"humming",
|
||||
"compressed-tensors",
|
||||
"bitsandbytes",
|
||||
"experts_int8",
|
||||
"quark",
|
||||
"moe_wna16",
|
||||
"torchao",
|
||||
"inc",
|
||||
"mxfp4",
|
||||
"gpt_oss_mxfp4",
|
||||
"deepseek_v4_fp8",
|
||||
"online",
|
||||
# Below are online quant shorthand names (see vllm.config.quantization).
|
||||
# Listed here as strings to avoid a circular import; kept in sync with
|
||||
# _ONLINE_SHORTHANDS by the assertion in get_quantization_config().
|
||||
"fp8_per_tensor",
|
||||
"fp8_per_block",
|
||||
"fp8_per_channel",
|
||||
"int8_per_channel_weight_only",
|
||||
"mxfp8",
|
||||
]
|
||||
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
|
||||
|
||||
DEPRECATED_QUANTIZATION_METHODS = [
|
||||
"fbgemm_fp8",
|
||||
"fp_quant",
|
||||
]
|
||||
|
||||
# The customized quantization methods which will be added to this dict.
|
||||
_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {}
|
||||
|
||||
|
||||
def register_quantization_config(quantization: str):
|
||||
"""Register a customized vllm quantization config.
|
||||
|
||||
When a quantization method is not supported by vllm, you can register a customized
|
||||
quantization config to support it.
|
||||
|
||||
Args:
|
||||
quantization (str): The quantization method name.
|
||||
|
||||
Examples:
|
||||
>>> from vllm.model_executor.layers.quantization import (
|
||||
... register_quantization_config,
|
||||
... )
|
||||
>>> from vllm.model_executor.layers.quantization import get_quantization_config
|
||||
>>> from vllm.model_executor.layers.quantization.base_config import (
|
||||
... QuantizationConfig,
|
||||
... )
|
||||
>>>
|
||||
>>> @register_quantization_config("my_quant")
|
||||
... class MyQuantConfig(QuantizationConfig):
|
||||
... pass
|
||||
>>>
|
||||
>>> get_quantization_config("my_quant")
|
||||
<class 'MyQuantConfig'>
|
||||
""" # noqa: E501
|
||||
|
||||
def _wrapper(quant_config_cls):
|
||||
if quantization in QUANTIZATION_METHODS:
|
||||
logger.debug(
|
||||
"The quantization method '%s' already exists and will be "
|
||||
"overwritten by the quantization config %s.",
|
||||
quantization,
|
||||
quant_config_cls,
|
||||
)
|
||||
else:
|
||||
QUANTIZATION_METHODS.append(quantization)
|
||||
# Automatically assume the custom quantization config is supported
|
||||
if sq := current_platform.supported_quantization:
|
||||
sq.append(quantization)
|
||||
|
||||
if not issubclass(quant_config_cls, QuantizationConfig):
|
||||
raise ValueError(
|
||||
"The quantization config must be a subclass of `QuantizationConfig`."
|
||||
)
|
||||
_CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls
|
||||
return quant_config_cls
|
||||
|
||||
return _wrapper
|
||||
|
||||
|
||||
def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
|
||||
if quantization not in QUANTIZATION_METHODS:
|
||||
raise ValueError(f"Invalid quantization method: {quantization}")
|
||||
|
||||
# lazy import to avoid triggering `torch.compile` too early
|
||||
from vllm.config.quantization import _ONLINE_SHORTHANDS
|
||||
from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
|
||||
from vllm.models.deepseek_v4 import DeepseekV4FP8Config
|
||||
|
||||
from .auto_awq import AutoAWQConfig
|
||||
from .auto_gptq import AutoGPTQConfig
|
||||
from .bitsandbytes import BitsAndBytesConfig
|
||||
from .compressed_tensors.compressed_tensors import (
|
||||
CompressedTensorsConfig,
|
||||
)
|
||||
from .experts_int8 import ExpertsInt8Config
|
||||
from .fbgemm_fp8 import FBGEMMFp8Config
|
||||
from .fp8 import Fp8Config
|
||||
from .fp_quant import FPQuantConfig
|
||||
from .humming import HummingConfig
|
||||
from .inc import INCConfig
|
||||
from .modelopt import (
|
||||
ModelOptFp8Config,
|
||||
ModelOptMixedPrecisionConfig,
|
||||
ModelOptMxFp8Config,
|
||||
ModelOptNvFp4Config,
|
||||
)
|
||||
from .moe_wna16 import MoeWNA16Config
|
||||
from .mxfp4 import GptOssMxfp4Config, Mxfp4Config
|
||||
from .online.base import OnlineQuantizationConfig
|
||||
from .torchao import TorchAOConfig
|
||||
|
||||
method_to_config: dict[str, type[QuantizationConfig]] = {
|
||||
"awq": AutoAWQConfig,
|
||||
"awq_marlin": AutoAWQConfig,
|
||||
"auto_awq": AutoAWQConfig,
|
||||
"fp8": Fp8Config,
|
||||
"fbgemm_fp8": FBGEMMFp8Config,
|
||||
"fp_quant": FPQuantConfig,
|
||||
"modelopt": ModelOptFp8Config,
|
||||
"modelopt_fp4": ModelOptNvFp4Config,
|
||||
"modelopt_mxfp8": ModelOptMxFp8Config,
|
||||
"modelopt_mixed": ModelOptMixedPrecisionConfig,
|
||||
"auto_gptq": AutoGPTQConfig,
|
||||
"gptq": AutoGPTQConfig,
|
||||
"gptq_marlin": AutoGPTQConfig,
|
||||
"compressed-tensors": CompressedTensorsConfig,
|
||||
"bitsandbytes": BitsAndBytesConfig,
|
||||
"experts_int8": ExpertsInt8Config,
|
||||
"quark": QuarkConfig,
|
||||
"moe_wna16": MoeWNA16Config,
|
||||
"torchao": TorchAOConfig,
|
||||
"inc": INCConfig,
|
||||
"mxfp4": Mxfp4Config,
|
||||
"gpt_oss_mxfp4": GptOssMxfp4Config,
|
||||
"deepseek_v4_fp8": DeepseekV4FP8Config,
|
||||
"humming": HummingConfig,
|
||||
"online": OnlineQuantizationConfig,
|
||||
# MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the
|
||||
# ModelOpt MXFP8 config (same format). The "mxfp8" online shorthand
|
||||
# below only applies to the `--quantization mxfp8` CLI path.
|
||||
"mxfp8": ModelOptMxFp8Config,
|
||||
}
|
||||
|
||||
# Register online shorthands (e.g. "fp8_per_tensor") as quant methods.
|
||||
# setdefault so a shorthand that is also a checkpoint method (e.g. "mxfp8")
|
||||
# keeps its checkpoint config; the shorthand still works via the
|
||||
# `--quantization` CLI path in `resolve_quantization_config`.
|
||||
for shorthand in _ONLINE_SHORTHANDS:
|
||||
method_to_config.setdefault(shorthand, OnlineQuantizationConfig)
|
||||
|
||||
# Update the `method_to_config` with customized quantization methods.
|
||||
method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
|
||||
|
||||
return method_to_config[quantization]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"QuantizationConfig",
|
||||
"QuantizationMethods",
|
||||
"get_quantization_config",
|
||||
"register_quantization_config",
|
||||
"QUANTIZATION_METHODS",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,852 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from copy import deepcopy
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
import vllm.model_executor.layers.fused_moe # noqa
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
MPLinearLayerConfig,
|
||||
choose_mp_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEExpertsModular,
|
||||
FusedMoEMethodBase,
|
||||
FusedMoEQuantConfig,
|
||||
FusedMoeWeightScaleSupported,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
UnquantizedFusedMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
|
||||
WNA16MoEBackend,
|
||||
convert_to_wna16_moe_kernel_format,
|
||||
make_wna16_moe_kernel,
|
||||
select_wna16_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import LinearMethodBase, set_weight_attrs
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils import replace_parameter
|
||||
from vllm.model_executor.layers.quantization.utils.gptq_utils import (
|
||||
get_dynamic_override,
|
||||
get_linear_quant_method,
|
||||
override_config,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
check_moe_marlin_supports_layer,
|
||||
get_marlin_input_dtype,
|
||||
marlin_make_workspace_new,
|
||||
marlin_repeat_scales_on_all_ranks,
|
||||
verify_marlin_supported,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
QuantKey,
|
||||
kInt4StaticGroupScale,
|
||||
kInt8StaticGroupScale,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
ChannelQuantScaleParameter,
|
||||
GroupQuantScaleParameter,
|
||||
PackedColumnParameter,
|
||||
PackedvLLMParameter,
|
||||
RowvLLMParameter,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
from vllm.transformers_utils.config import get_safetensors_params_metadata
|
||||
from vllm.utils.collection_utils import is_list_of
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def get_moe_quant_method(
|
||||
config: "AutoGPTQConfig",
|
||||
layer: RoutedExperts,
|
||||
prefix: str,
|
||||
moe_method_cls: type,
|
||||
):
|
||||
cloned_config = deepcopy(config)
|
||||
|
||||
assert isinstance(layer, RoutedExperts)
|
||||
# False = skip module, None = no override, else = Positive match
|
||||
if (
|
||||
get_dynamic_override( # noqa: E712
|
||||
cloned_config, # noqa: E712
|
||||
layer_name=prefix,
|
||||
)
|
||||
== False
|
||||
): # noqa: E712
|
||||
return UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
|
||||
if prefix:
|
||||
# Dynamic per module/layer rules may override base config
|
||||
override_config(cloned_config, prefix=prefix)
|
||||
|
||||
return moe_method_cls(cloned_config, layer.moe_config)
|
||||
|
||||
|
||||
class AutoGPTQConfig(QuantizationConfig):
|
||||
"""Config class for AutoGPTQ quantization using Marlin kernels."""
|
||||
|
||||
# (num_bits, is_sym) -> quant_type
|
||||
TYPE_MAP = {
|
||||
(4, True): scalar_types.uint4b8,
|
||||
(8, True): scalar_types.uint8b128,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weight_bits: int,
|
||||
group_size: int,
|
||||
desc_act: bool,
|
||||
is_sym: bool,
|
||||
lm_head_quantized: bool,
|
||||
dynamic: dict[str, dict[str, int | bool]],
|
||||
full_config: dict[str, Any],
|
||||
modules_in_block_to_quantize: list[str] | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if desc_act and group_size == -1:
|
||||
# In this case, act_order == True is the same as act_order == False
|
||||
# (since we have only one group per output channel)
|
||||
desc_act = False
|
||||
|
||||
# GPTQModel use `dynamic` config property to allow per module
|
||||
# quantization config so each module can be individually optimized.
|
||||
# Format is dict[str, dict] where key is a regex string that can
|
||||
# perform both positive ("+:" prefixed) or negative ("-:" prefixed)
|
||||
# matching of a module.
|
||||
# Default to positive match, override base quant config mode, if no
|
||||
# prefix is used. Value is in dict format of field key and override
|
||||
# value.
|
||||
# Negative matching will skip quantization init for this module
|
||||
# entirely:
|
||||
# non-quantized inference. More details and quantization examples can be
|
||||
# found at: https://github.com/ModelCloud/GPTQModel
|
||||
# Example:
|
||||
# # last 1/2 of the layers 10-21 has 8bit vs 4bit for 0-9
|
||||
# # last 1/4 of the layers 16-21 has 8bit and group_size 64
|
||||
# dynamic = {
|
||||
# #`.*\.` matches the layers_node prefix
|
||||
# # positive match layer 10-15
|
||||
# r"+:.*\.(?:1[0-5])\..*": {"bits": 8,},
|
||||
# # positive match layer 16-21
|
||||
# r"+:.*\.(?:1[6-9]|20|21)\..*": {"bits": 8, "group_size": 64,},
|
||||
# r"-:.*\.moe\..*": {}, # negative match (skip) all `moe` layers
|
||||
# }
|
||||
self.dynamic = dynamic
|
||||
|
||||
self.weight_bits = weight_bits
|
||||
self.is_sym = is_sym
|
||||
|
||||
self.pack_factor = 32 // weight_bits # packed into int32
|
||||
self.group_size = group_size
|
||||
self.desc_act = desc_act
|
||||
self.lm_head_quantized = lm_head_quantized
|
||||
self.full_config = full_config
|
||||
|
||||
if (weight_bits, is_sym) not in self.TYPE_MAP:
|
||||
raise ValueError(
|
||||
f"Unsupported quantization config: bits={weight_bits}, sym={is_sym}"
|
||||
)
|
||||
|
||||
self.quant_type = self.TYPE_MAP[(weight_bits, is_sym)]
|
||||
|
||||
self.modules_in_block_to_quantize = modules_in_block_to_quantize or []
|
||||
# used to identify GPTQ model quantized by autoround
|
||||
self.autoround_version = full_config.get("autoround_version", "")
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"AutoGPTQConfig(quant_type={self.quant_type}, "
|
||||
f"group_size={self.group_size}, "
|
||||
f"desc_act={self.desc_act}, "
|
||||
f"lm_head_quantized={self.lm_head_quantized}, "
|
||||
f"dynamic={self.dynamic}, "
|
||||
f"modules_in_block_to_quantize={self.modules_in_block_to_quantize})"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "auto_gptq"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.half, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 60
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return ["quantize_config.json"]
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "AutoGPTQConfig":
|
||||
dynamic = cls.get_from_keys_or(config, ["dynamic"], default={})
|
||||
dynamic = {} if dynamic is None else dynamic
|
||||
|
||||
weight_bits = cls.get_from_keys(config, ["bits"])
|
||||
group_size = cls.get_from_keys(config, ["group_size"])
|
||||
desc_act = cls.get_from_keys(config, ["desc_act"])
|
||||
is_sym = cls.get_from_keys(config, ["sym"])
|
||||
lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
|
||||
modules_in_block_to_quantize = cls.get_from_keys_or(
|
||||
config, ["modules_in_block_to_quantize"], default=None
|
||||
)
|
||||
return cls(
|
||||
weight_bits,
|
||||
group_size,
|
||||
desc_act,
|
||||
is_sym,
|
||||
lm_head_quantized,
|
||||
dynamic,
|
||||
config,
|
||||
modules_in_block_to_quantize,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(
|
||||
cls, hf_quant_cfg, user_quant, hf_config=None
|
||||
) -> QuantizationMethods | None:
|
||||
"""Override to use AutoGPTQ for compatible GPTQ models."""
|
||||
quant_method = hf_quant_cfg.get("quant_method", "").lower()
|
||||
|
||||
if quant_method != "gptq":
|
||||
return None
|
||||
|
||||
is_valid_user_quant = user_quant is None or user_quant in (
|
||||
"gptq",
|
||||
"gptq_marlin",
|
||||
"auto_gptq",
|
||||
"marlin",
|
||||
)
|
||||
|
||||
if is_valid_user_quant:
|
||||
return cls.get_name()
|
||||
|
||||
return None
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
if isinstance(layer, RoutedExperts):
|
||||
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
|
||||
|
||||
if not check_moe_marlin_supports_layer(
|
||||
layer, self.group_size, allow_tile_padding=not self.desc_act
|
||||
):
|
||||
logger.warning_once(
|
||||
f"Layer '{prefix}' is not supported by GPTQMoeMarlin. "
|
||||
"Falling back to Moe WNA16 kernels."
|
||||
)
|
||||
return MoeWNA16Config.from_config(self.full_config).get_quant_method(
|
||||
layer, prefix
|
||||
)
|
||||
moe_quant_method = get_moe_quant_method(
|
||||
self, layer, prefix, AutoGPTQMoEMethod
|
||||
)
|
||||
if moe_quant_method is None:
|
||||
return None
|
||||
moe_quant_method.input_dtype = get_marlin_input_dtype(prefix)
|
||||
return moe_quant_method
|
||||
|
||||
quant_method = get_linear_quant_method(
|
||||
self, layer, prefix, AutoGPTQLinearMethod
|
||||
)
|
||||
if quant_method is None:
|
||||
return None
|
||||
quant_method.input_dtype = get_marlin_input_dtype(prefix)
|
||||
return quant_method
|
||||
|
||||
def apply_vllm_mapper(self, hf_to_vllm_mapper):
|
||||
if self.modules_in_block_to_quantize is not None:
|
||||
self.modules_in_block_to_quantize = hf_to_vllm_mapper.apply_list(
|
||||
self.modules_in_block_to_quantize
|
||||
)
|
||||
|
||||
def maybe_update_config(
|
||||
self,
|
||||
model_name: str,
|
||||
hf_config: PretrainedConfig | None = None,
|
||||
revision: str | None = None,
|
||||
):
|
||||
if self.modules_in_block_to_quantize:
|
||||
if is_list_of(self.modules_in_block_to_quantize, list):
|
||||
# original modules_in_block_to_quantize: list[list[str]]
|
||||
# flatten original modules_in_block_to_quantize
|
||||
self.modules_in_block_to_quantize = [
|
||||
item
|
||||
for sublist in self.modules_in_block_to_quantize
|
||||
for item in sublist
|
||||
]
|
||||
return
|
||||
|
||||
unquant_dtypes = [torch.float16, torch.bfloat16, torch.float32]
|
||||
metadata = get_safetensors_params_metadata(model_name, revision=revision)
|
||||
quant_layers: set[str] = {
|
||||
param_name.rsplit(".", 1)[0]
|
||||
for param_name, info in metadata.items()
|
||||
if (dtype := info.get("dtype", None))
|
||||
and _SAFETENSORS_TO_TORCH_DTYPE[dtype] not in unquant_dtypes
|
||||
}
|
||||
self.modules_in_block_to_quantize = list(quant_layers)
|
||||
|
||||
|
||||
class AutoGPTQLinearMethod(LinearMethodBase):
|
||||
"""Linear method for AutoGPTQ using Marlin kernels.
|
||||
|
||||
Args:
|
||||
quant_config: The AutoGPTQ quantization config.
|
||||
"""
|
||||
|
||||
_kernel_backends_being_used: set[str] = set()
|
||||
|
||||
def __init__(self, quant_config: AutoGPTQConfig) -> None:
|
||||
self.quant_config = quant_config
|
||||
self.input_dtype = None
|
||||
self.quant_type = self.quant_config.quant_type
|
||||
|
||||
# Verify supported on platform.
|
||||
verify_marlin_supported(
|
||||
quant_type=self.quant_config.quant_type,
|
||||
group_size=self.quant_config.group_size,
|
||||
)
|
||||
|
||||
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,
|
||||
) -> None:
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
is_row_parallel = input_size != input_size_per_partition
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
input_dtype = self.input_dtype
|
||||
|
||||
mp_linear_kernel_config = MPLinearLayerConfig(
|
||||
full_weight_shape=(input_size, output_size),
|
||||
partition_weight_shape=(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition,
|
||||
),
|
||||
weight_type=self.quant_config.quant_type,
|
||||
act_type=params_dtype if input_dtype is None else input_dtype,
|
||||
group_size=self.quant_config.group_size,
|
||||
zero_points=False,
|
||||
has_g_idx=self.quant_config.desc_act,
|
||||
)
|
||||
|
||||
kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
|
||||
|
||||
if kernel_type.__name__ not in self._kernel_backends_being_used:
|
||||
logger.info("Using %s for AutoGPTQLinearMethod", kernel_type.__name__)
|
||||
self._kernel_backends_being_used.add(kernel_type.__name__)
|
||||
|
||||
# Normalize group_size
|
||||
if self.quant_config.group_size != -1:
|
||||
group_size = self.quant_config.group_size
|
||||
else:
|
||||
group_size = input_size
|
||||
|
||||
# Determine sharding
|
||||
if marlin_repeat_scales_on_all_ranks(
|
||||
self.quant_config.desc_act, self.quant_config.group_size, is_row_parallel
|
||||
):
|
||||
# By setting scale_dim == None, weight_loader will
|
||||
# repeat the scales on each GPU in TP>1 case.
|
||||
scales_and_zp_input_dim = None
|
||||
scales_and_zp_size = input_size // group_size
|
||||
else:
|
||||
# By setting scale_dim == 0, weight_loader will
|
||||
# shard the scales in TP>1 case.
|
||||
scales_and_zp_input_dim = 0
|
||||
scales_and_zp_size = input_size_per_partition // group_size
|
||||
|
||||
# Quantized weights
|
||||
qweight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition // self.quant_config.pack_factor,
|
||||
output_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=0,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
# Activation order
|
||||
g_idx = RowvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
qzeros_args = {
|
||||
"data": torch.empty(
|
||||
scales_and_zp_size,
|
||||
output_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
"weight_loader": weight_loader,
|
||||
}
|
||||
weight_scale_args = {
|
||||
"data": torch.empty(
|
||||
scales_and_zp_size,
|
||||
output_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
"weight_loader": weight_loader,
|
||||
}
|
||||
|
||||
if scales_and_zp_input_dim is None:
|
||||
scales = ChannelQuantScaleParameter(output_dim=1, **weight_scale_args)
|
||||
qzeros = PackedColumnParameter(
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
**qzeros_args,
|
||||
)
|
||||
|
||||
else:
|
||||
scales = GroupQuantScaleParameter(
|
||||
output_dim=1, input_dim=0, **weight_scale_args
|
||||
)
|
||||
qzeros = PackedvLLMParameter(
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
**qzeros_args,
|
||||
)
|
||||
|
||||
layer.register_parameter("qweight", qweight)
|
||||
layer.register_parameter("g_idx", g_idx)
|
||||
layer.register_parameter("scales", scales)
|
||||
layer.register_parameter("qzeros", qzeros)
|
||||
|
||||
self.kernel = kernel_type(
|
||||
mp_linear_kernel_config,
|
||||
w_q_param_name="qweight",
|
||||
w_s_param_name="scales",
|
||||
w_zp_param_name="qzeros",
|
||||
w_gidx_param_name="g_idx",
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
self.kernel.process_weights_after_loading(layer)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return self.kernel.apply_weights(layer, x, bias)
|
||||
|
||||
|
||||
class AutoGPTQMoEMethod(FusedMoEMethodBase):
|
||||
"""MoE Marlin method with quantization."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
quant_config: AutoGPTQConfig,
|
||||
moe: FusedMoEConfig,
|
||||
) -> None:
|
||||
super().__init__(moe)
|
||||
self.quant_config = quant_config
|
||||
if self.quant_config.quant_type.size_bits == 4:
|
||||
quant_type = scalar_types.uint4b8
|
||||
scale = kInt4StaticGroupScale
|
||||
elif self.quant_config.quant_type.size_bits == 8:
|
||||
quant_type = scalar_types.uint8b128
|
||||
scale = kInt8StaticGroupScale
|
||||
else:
|
||||
raise ValueError("AutoGPTQMoEMethod only supports int4 and int8 now.")
|
||||
self.input_dtype = None
|
||||
self.use_marlin = True
|
||||
weight_key = QuantKey(quant_type, scale)
|
||||
|
||||
self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
|
||||
moe,
|
||||
weight_key,
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.input_dtype = self.input_dtype
|
||||
is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
|
||||
|
||||
if is_a_8bit:
|
||||
assert self.quant_config.quant_type.size_bits == 8, (
|
||||
"W8A8-INT8 is not supported by marlin kernel."
|
||||
)
|
||||
|
||||
intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
|
||||
|
||||
self.is_k_full = (not self.quant_config.desc_act) or (
|
||||
intermediate_size_per_partition == intermediate_size_full
|
||||
)
|
||||
|
||||
if self.quant_config.group_size != -1:
|
||||
scales_size13 = hidden_size // self.quant_config.group_size
|
||||
w2_scales_size = (
|
||||
intermediate_size_full
|
||||
if self.quant_config.desc_act
|
||||
else intermediate_size_per_partition
|
||||
)
|
||||
scales_size2 = w2_scales_size // self.quant_config.group_size
|
||||
strategy = FusedMoeWeightScaleSupported.GROUP.value
|
||||
else:
|
||||
scales_size13 = 1
|
||||
scales_size2 = 1
|
||||
strategy = FusedMoeWeightScaleSupported.CHANNEL.value
|
||||
|
||||
layer.num_groups_w13 = scales_size13
|
||||
layer.num_groups_w2 = scales_size2
|
||||
|
||||
extra_weight_attrs.update({"quant_method": strategy, "is_transposed": True})
|
||||
# Fused gate_up_proj (column parallel)
|
||||
w13_qweight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size // self.quant_config.pack_factor,
|
||||
2 * intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_qweight", w13_qweight)
|
||||
set_weight_attrs(w13_qweight, extra_weight_attrs)
|
||||
# down_proj (row parallel)
|
||||
w2_qweight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition // self.quant_config.pack_factor,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_qweight", w2_qweight)
|
||||
set_weight_attrs(w2_qweight, extra_weight_attrs)
|
||||
# up_proj scales
|
||||
w13_scales = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
scales_size13,
|
||||
2 * intermediate_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_scales", w13_scales)
|
||||
set_weight_attrs(w13_scales, extra_weight_attrs)
|
||||
# down_proj scales
|
||||
w2_scales = torch.nn.Parameter(
|
||||
torch.empty(num_experts, scales_size2, hidden_size, dtype=params_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_scales", w2_scales)
|
||||
set_weight_attrs(w2_scales, extra_weight_attrs)
|
||||
# don't shard the w2 scales when running act order
|
||||
set_weight_attrs(w2_scales, {"load_full_w2": self.quant_config.desc_act})
|
||||
# up_proj scales
|
||||
w13_qzeros = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
scales_size13,
|
||||
2 * intermediate_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_qzeros", w13_qzeros)
|
||||
set_weight_attrs(w13_qzeros, extra_weight_attrs)
|
||||
# down_proj scales
|
||||
w2_qzeros = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
scales_size2,
|
||||
hidden_size // self.quant_config.pack_factor,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_qzeros", w2_qzeros)
|
||||
set_weight_attrs(w2_qzeros, extra_weight_attrs)
|
||||
# don't shard the w2 scales when running act order
|
||||
set_weight_attrs(w2_qzeros, {"load_full_w2": self.quant_config.desc_act})
|
||||
w13_g_idx = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_g_idx", w13_g_idx)
|
||||
set_weight_attrs(w13_g_idx, extra_weight_attrs)
|
||||
w2_g_idx = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_g_idx", w2_g_idx)
|
||||
set_weight_attrs(w2_g_idx, extra_weight_attrs)
|
||||
w13_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
|
||||
set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
|
||||
w2_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
|
||||
set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
|
||||
|
||||
if self.experts_cls is not None and issubclass(
|
||||
self.experts_cls, FusedMoEExpertsModular
|
||||
):
|
||||
device = layer.w13_qweight.device
|
||||
layer.workspace = marlin_make_workspace_new(device, 4)
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
|
||||
|
||||
if is_a_8bit:
|
||||
assert self.quant_config.quant_type.size_bits == 8, (
|
||||
"W8A8-INT8 is not supported by marlin kernel."
|
||||
)
|
||||
|
||||
converted = convert_to_wna16_moe_kernel_format(
|
||||
backend=self.wna16_moe_backend,
|
||||
layer=layer,
|
||||
quant_config=self.quant_config,
|
||||
input_dtype=self.input_dtype,
|
||||
w13=layer.w13_qweight,
|
||||
w2=layer.w2_qweight,
|
||||
w13_scale=layer.w13_scales,
|
||||
w2_scale=layer.w2_scales,
|
||||
w13_g_idx=layer.w13_g_idx,
|
||||
w2_g_idx=layer.w2_g_idx,
|
||||
w13_bias=getattr(layer, "w13_bias", None),
|
||||
w2_bias=getattr(layer, "w2_bias", None),
|
||||
)
|
||||
|
||||
if converted is None:
|
||||
# Backend rewrote the layer's params in place (e.g. Humming).
|
||||
self._setup_kernel(layer)
|
||||
return
|
||||
|
||||
(
|
||||
w13,
|
||||
w2,
|
||||
w13_scale,
|
||||
w2_scale,
|
||||
w13_g_idx,
|
||||
w2_g_idx,
|
||||
w13_g_idx_sort_indices,
|
||||
w2_g_idx_sort_indices,
|
||||
w13_qzeros,
|
||||
w2_qzeros,
|
||||
w13_input_global_scale,
|
||||
w2_input_global_scale,
|
||||
w13_bias,
|
||||
w2_bias,
|
||||
) = converted
|
||||
|
||||
replace_parameter(layer, "w13_qweight", w13)
|
||||
replace_parameter(layer, "w2_qweight", w2)
|
||||
replace_parameter(layer, "w13_scales", w13_scale)
|
||||
replace_parameter(layer, "w2_scales", w2_scale)
|
||||
replace_parameter(layer, "w13_g_idx", w13_g_idx)
|
||||
replace_parameter(layer, "w2_g_idx", w2_g_idx)
|
||||
replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
|
||||
replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
|
||||
if w13_qzeros is not None:
|
||||
replace_parameter(layer, "w13_qzeros", w13_qzeros)
|
||||
if w2_qzeros is not None:
|
||||
replace_parameter(layer, "w2_qzeros", w2_qzeros)
|
||||
if w13_input_global_scale is not None:
|
||||
if hasattr(layer, "w13_input_global_scale"):
|
||||
replace_parameter(
|
||||
layer, "w13_input_global_scale", w13_input_global_scale
|
||||
)
|
||||
else:
|
||||
layer.register_parameter(
|
||||
"w13_input_global_scale",
|
||||
torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
|
||||
)
|
||||
if w2_input_global_scale is not None:
|
||||
if hasattr(layer, "w2_input_global_scale"):
|
||||
replace_parameter(layer, "w2_input_global_scale", w2_input_global_scale)
|
||||
else:
|
||||
layer.register_parameter(
|
||||
"w2_input_global_scale",
|
||||
torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
|
||||
)
|
||||
if w13_bias is not None:
|
||||
if hasattr(layer, "w13_bias"):
|
||||
replace_parameter(layer, "w13_bias", w13_bias)
|
||||
else:
|
||||
layer.register_parameter(
|
||||
"w13_bias", torch.nn.Parameter(w13_bias, requires_grad=False)
|
||||
)
|
||||
if w2_bias is not None:
|
||||
if hasattr(layer, "w2_bias"):
|
||||
replace_parameter(layer, "w2_bias", w2_bias)
|
||||
else:
|
||||
layer.register_parameter(
|
||||
"w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False)
|
||||
)
|
||||
|
||||
# The modular kernel reads w13_weight/w2_weight; marlin keeps *_qweight.
|
||||
layer.w13_weight = layer.w13_qweight
|
||||
layer.w2_weight = layer.w2_qweight
|
||||
|
||||
self._setup_kernel(layer)
|
||||
|
||||
def _setup_kernel(self, layer: RoutedExperts) -> None:
|
||||
"""Build the FusedMoEKernel for this layer."""
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
self.moe_kernel = make_wna16_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
experts_cls=self.experts_cls,
|
||||
backend=self.wna16_moe_backend,
|
||||
layer=layer,
|
||||
is_k_full=self.is_k_full,
|
||||
w13_g_idx=getattr(layer, "w13_g_idx", None),
|
||||
w2_g_idx=getattr(layer, "w2_g_idx", None),
|
||||
w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None),
|
||||
w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None),
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
|
||||
if self.wna16_moe_backend == WNA16MoEBackend.HUMMING:
|
||||
from vllm.model_executor.layers.quantization.utils.humming_utils import (
|
||||
get_humming_moe_quant_config,
|
||||
)
|
||||
|
||||
return get_humming_moe_quant_config(layer)
|
||||
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
gptq_marlin_moe_quant_config,
|
||||
)
|
||||
|
||||
# CPU fused_experts_cpu requires zero points even for symmetric quant
|
||||
use_zp = (
|
||||
not self.quant_config.is_sym
|
||||
or self.wna16_moe_backend == WNA16MoEBackend.CPU
|
||||
)
|
||||
return gptq_marlin_moe_quant_config(
|
||||
w1_scale=layer.w13_scales,
|
||||
w2_scale=layer.w2_scales,
|
||||
weight_bits=self.quant_config.weight_bits,
|
||||
group_size=self.quant_config.group_size,
|
||||
w1_zp=getattr(layer, "w13_qzeros", None) if use_zp else None,
|
||||
w2_zp=getattr(layer, "w2_qzeros", None) if use_zp else None,
|
||||
w1_bias=getattr(layer, "w13_bias", None),
|
||||
w2_bias=getattr(layer, "w2_bias", None),
|
||||
)
|
||||
|
||||
def select_gemm_impl(
|
||||
self,
|
||||
prepare_finalize,
|
||||
layer: RoutedExperts,
|
||||
):
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} uses the new modular kernel "
|
||||
"initialization logic. This function should not be called."
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
hidden_states=x,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
expert_map=layer.expert_map,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
hidden_states=x,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
router_logits=router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
num_expert_group=layer.num_expert_group,
|
||||
topk_group=layer.topk_group,
|
||||
e_score_correction_bias=layer.e_score_correction_bias,
|
||||
routed_scaling_factor=layer.routed_scaling_factor,
|
||||
)
|
||||
@@ -0,0 +1,337 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
AWQ_TRITON_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
|
||||
|
||||
|
||||
@triton.jit
|
||||
def awq_dequantize_kernel(
|
||||
qweight_ptr, # quantized matrix
|
||||
scales_ptr, # scales, per group
|
||||
zeros_ptr, # zeros, per group
|
||||
group_size, # Should always be one of the supported group sizes
|
||||
result_ptr, # Output matrix
|
||||
num_cols, # input num cols in qweight
|
||||
num_rows, # input num rows in qweight
|
||||
BLOCK_SIZE_X: tl.constexpr,
|
||||
BLOCK_SIZE_Y: tl.constexpr,
|
||||
):
|
||||
# Set up the pids.
|
||||
pid_x = tl.program_id(axis=0)
|
||||
pid_y = tl.program_id(axis=1)
|
||||
|
||||
# Compute offsets and masks for qweight_ptr.
|
||||
offsets_y = pid_y * BLOCK_SIZE_Y + tl.arange(0, BLOCK_SIZE_Y)
|
||||
offsets_x = pid_x * BLOCK_SIZE_X + tl.arange(0, BLOCK_SIZE_X)
|
||||
offsets = num_cols * offsets_y[:, None] + offsets_x[None, :]
|
||||
|
||||
masks_y = offsets_y < num_rows
|
||||
masks_x = offsets_x < num_cols
|
||||
|
||||
masks = masks_y[:, None] & masks_x[None, :]
|
||||
|
||||
# Compute offsets and masks for result output ptr.
|
||||
result_offsets_y = pid_y * BLOCK_SIZE_Y + tl.arange(0, BLOCK_SIZE_Y)
|
||||
result_offsets_x = pid_x * BLOCK_SIZE_X * 8 + tl.arange(0, BLOCK_SIZE_X * 8)
|
||||
result_offsets = (
|
||||
8 * num_cols * result_offsets_y[:, None] + result_offsets_x[None, :]
|
||||
)
|
||||
|
||||
result_masks_y = result_offsets_y < num_rows
|
||||
result_masks_x = result_offsets_x < num_cols * 8
|
||||
result_masks = result_masks_y[:, None] & result_masks_x[None, :]
|
||||
|
||||
# Load the weights.
|
||||
iweights = tl.load(qweight_ptr + offsets, masks, 0.0)
|
||||
iweights = tl.interleave(iweights, iweights)
|
||||
iweights = tl.interleave(iweights, iweights)
|
||||
iweights = tl.interleave(iweights, iweights)
|
||||
|
||||
# Create reverse AWQ order as tensor: [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
# that will map given indices to the correct order.
|
||||
reverse_awq_order_tensor = (
|
||||
(tl.arange(0, 2) * 4)[None, :] + tl.arange(0, 4)[:, None]
|
||||
).reshape(8)
|
||||
|
||||
# Use this to compute a set of shifts that can be used to unpack and
|
||||
# reorder the values in iweights and zeros.
|
||||
shifts = reverse_awq_order_tensor * 4
|
||||
shifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_Y * BLOCK_SIZE_X, 8))
|
||||
shifts = tl.reshape(shifts, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8))
|
||||
|
||||
# Unpack and reorder: shift out the correct 4-bit value and mask.
|
||||
iweights = (iweights >> shifts) & 0xF
|
||||
|
||||
# Compute zero offsets and masks.
|
||||
zero_offsets_y = pid_y * BLOCK_SIZE_Y // group_size + tl.arange(0, 1)
|
||||
zero_offsets_x = pid_x * BLOCK_SIZE_X + tl.arange(0, BLOCK_SIZE_X)
|
||||
zero_offsets = num_cols * zero_offsets_y[:, None] + zero_offsets_x[None, :]
|
||||
|
||||
zero_masks_y = zero_offsets_y < num_rows // group_size
|
||||
zero_masks_x = zero_offsets_x < num_cols
|
||||
zero_masks = zero_masks_y[:, None] & zero_masks_x[None, :]
|
||||
|
||||
# Load the zeros.
|
||||
zeros = tl.load(zeros_ptr + zero_offsets, zero_masks, 0.0)
|
||||
zeros = tl.interleave(zeros, zeros)
|
||||
zeros = tl.interleave(zeros, zeros)
|
||||
zeros = tl.interleave(zeros, zeros)
|
||||
zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8))
|
||||
|
||||
# Unpack and reorder: shift out the correct 4-bit value and mask.
|
||||
zeros = (zeros >> shifts) & 0xF
|
||||
|
||||
# Compute scale offsets and masks.
|
||||
scale_offsets_y = pid_y * BLOCK_SIZE_Y // group_size + tl.arange(0, 1)
|
||||
scale_offsets_x = pid_x * BLOCK_SIZE_X * 8 + tl.arange(0, BLOCK_SIZE_X * 8)
|
||||
scale_offsets = num_cols * 8 * scale_offsets_y[:, None] + scale_offsets_x[None, :]
|
||||
scale_masks_y = scale_offsets_y < num_rows // group_size
|
||||
scale_masks_x = scale_offsets_x < num_cols * 8
|
||||
scale_masks = scale_masks_y[:, None] & scale_masks_x[None, :]
|
||||
|
||||
# Load the scales.
|
||||
scales = tl.load(scales_ptr + scale_offsets, scale_masks, 0.0)
|
||||
scales = tl.broadcast_to(scales, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8))
|
||||
|
||||
# Dequantize.
|
||||
iweights = (iweights - zeros) * scales
|
||||
iweights = iweights.to(result_ptr.type.element_ty)
|
||||
|
||||
# Finally, store.
|
||||
tl.store(result_ptr + result_offsets, iweights, result_masks)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def awq_gemm_kernel(
|
||||
a_ptr,
|
||||
b_ptr,
|
||||
c_ptr,
|
||||
zeros_ptr,
|
||||
scales_ptr,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
group_size,
|
||||
BLOCK_SIZE_M: tl.constexpr,
|
||||
BLOCK_SIZE_N: tl.constexpr,
|
||||
BLOCK_SIZE_K: tl.constexpr,
|
||||
SPLIT_K: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(axis=0)
|
||||
pid_z = tl.program_id(1)
|
||||
|
||||
# NOTE: This doesn't work in TRITON_INTERPRET=1 mode. Use below instead.
|
||||
# num_pid_n = (N + BLOCK_SIZE_N - 1) // BLOCK_SIZE_N
|
||||
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
|
||||
|
||||
pid_m = pid // num_pid_n
|
||||
pid_n = pid % num_pid_n
|
||||
|
||||
accumulator_dtype = c_ptr.type.element_ty
|
||||
|
||||
# NOTE: This doesn't work in TRITON_INTERPRET=1 mode. Use below instead.
|
||||
# accumulator = tl.arange(0, BLOCK_SIZE_N)
|
||||
# accumulator = tl.broadcast_to(accumulator[None, :],
|
||||
# (BLOCK_SIZE_M, BLOCK_SIZE_N))
|
||||
# accumulator = accumulator & 0x0
|
||||
# accumulator = accumulator.to(accumulator_dtype)
|
||||
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype)
|
||||
|
||||
# Create reverse AWQ order as tensor: [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
# that will map given indices to the correct order.
|
||||
reverse_awq_order_tensor = (
|
||||
(tl.arange(0, 2) * 4)[None, :] + tl.arange(0, 4)[:, None]
|
||||
).reshape(8)
|
||||
|
||||
# Create the necessary shifts to use to unpack.
|
||||
shifts = reverse_awq_order_tensor * 4
|
||||
shifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_K * (BLOCK_SIZE_N // 8), 8))
|
||||
shifts = tl.reshape(shifts, (BLOCK_SIZE_K, BLOCK_SIZE_N))
|
||||
|
||||
# Offsets and masks.
|
||||
offsets_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
|
||||
masks_am = offsets_am < M
|
||||
|
||||
offsets_bn = pid_n * (BLOCK_SIZE_N // 8) + tl.arange(0, BLOCK_SIZE_N // 8)
|
||||
masks_bn = offsets_bn < N // 8
|
||||
|
||||
offsets_zn = pid_n * (BLOCK_SIZE_N // 8) + tl.arange(0, BLOCK_SIZE_N // 8)
|
||||
masks_zn = offsets_zn < N // 8
|
||||
|
||||
offsets_sn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
|
||||
masks_sn = offsets_sn < N
|
||||
|
||||
offsets_k = pid_z * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
|
||||
offsets_a = K * offsets_am[:, None] + offsets_k[None, :]
|
||||
offsets_b = (N // 8) * offsets_k[:, None] + offsets_bn[None, :]
|
||||
|
||||
a_ptrs = a_ptr + offsets_a
|
||||
b_ptrs = b_ptr + offsets_b
|
||||
|
||||
# NOTE: Use this in TRITON_INTERPRET=1 mode instead of tl.cdiv
|
||||
# block_offset = BLOCK_SIZE_K * SPLIT_K
|
||||
# for k in range(0, (K + block_offset - 1) // (block_offset)):
|
||||
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)):
|
||||
masks_k = offsets_k < K
|
||||
masks_a = masks_am[:, None] & masks_k[None, :]
|
||||
a = tl.load(a_ptrs, mask=masks_a, other=0.0)
|
||||
|
||||
masks_b = masks_k[:, None] & masks_bn[None, :]
|
||||
b = tl.load(b_ptrs, mask=masks_b, other=0.0)
|
||||
b = tl.interleave(b, b)
|
||||
b = tl.interleave(b, b)
|
||||
b = tl.interleave(b, b)
|
||||
|
||||
# Dequantize b.
|
||||
offsets_szk = (
|
||||
BLOCK_SIZE_K * SPLIT_K * k + pid_z * BLOCK_SIZE_K
|
||||
) // group_size + tl.arange(0, 1)
|
||||
offsets_z = (N // 8) * offsets_szk[:, None] + offsets_zn[None, :]
|
||||
masks_zk = offsets_szk < K // group_size
|
||||
masks_z = masks_zk[:, None] & masks_zn[None, :]
|
||||
zeros_ptrs = zeros_ptr + offsets_z
|
||||
zeros = tl.load(zeros_ptrs, mask=masks_z, other=0.0)
|
||||
zeros = tl.interleave(zeros, zeros)
|
||||
zeros = tl.interleave(zeros, zeros)
|
||||
zeros = tl.interleave(zeros, zeros)
|
||||
zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N))
|
||||
|
||||
offsets_s = N * offsets_szk[:, None] + offsets_sn[None, :]
|
||||
masks_sk = offsets_szk < K // group_size
|
||||
masks_s = masks_sk[:, None] & masks_sn[None, :]
|
||||
scales_ptrs = scales_ptr + offsets_s
|
||||
scales = tl.load(scales_ptrs, mask=masks_s, other=0.0)
|
||||
scales = tl.broadcast_to(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N))
|
||||
|
||||
b = (b >> shifts) & 0xF
|
||||
zeros = (zeros >> shifts) & 0xF
|
||||
b = (b - zeros) * scales
|
||||
b = b.to(c_ptr.type.element_ty)
|
||||
|
||||
# Accumulate results.
|
||||
accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
|
||||
|
||||
offsets_k += BLOCK_SIZE_K * SPLIT_K
|
||||
a_ptrs += BLOCK_SIZE_K * SPLIT_K
|
||||
b_ptrs += BLOCK_SIZE_K * SPLIT_K * (N // 8)
|
||||
|
||||
c = accumulator.to(c_ptr.type.element_ty)
|
||||
offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
|
||||
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
|
||||
c_ptrs = c_ptr + pid_z * N * M + N * offs_cm[:, None] + offs_cn[None, :]
|
||||
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
|
||||
tl.store(c_ptrs, c, mask=c_mask)
|
||||
|
||||
|
||||
# qweights - [K , M // 8], int32
|
||||
# scales - [K // G, M ], float16
|
||||
# zeros - [K // G, M // 8], int32
|
||||
def awq_dequantize_triton(
|
||||
qweight: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
zeros: torch.Tensor,
|
||||
block_size_x: int = 32,
|
||||
block_size_y: int = 32,
|
||||
) -> torch.Tensor:
|
||||
K = qweight.shape[0]
|
||||
M = scales.shape[1]
|
||||
group_size = qweight.shape[0] // scales.shape[0]
|
||||
|
||||
assert K > 0 and M > 0
|
||||
assert scales.shape[0] == K // group_size and scales.shape[1] == M
|
||||
assert zeros.shape[0] == K // group_size and zeros.shape[1] == M // 8
|
||||
assert group_size <= K
|
||||
assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K
|
||||
|
||||
# Result tensor:
|
||||
# number of rows = same as input tensor
|
||||
# number of cols = 8 x input tensor num cols
|
||||
result = torch.empty(
|
||||
qweight.shape[0],
|
||||
qweight.shape[1] * 8,
|
||||
device=qweight.device,
|
||||
dtype=scales.dtype,
|
||||
)
|
||||
|
||||
Y = qweight.shape[0] # num rows
|
||||
X = qweight.shape[1] # num cols
|
||||
|
||||
grid = lambda META: (
|
||||
triton.cdiv(X, META["BLOCK_SIZE_X"]),
|
||||
triton.cdiv(Y, META["BLOCK_SIZE_Y"]),
|
||||
)
|
||||
awq_dequantize_kernel[grid](
|
||||
qweight,
|
||||
scales,
|
||||
zeros,
|
||||
group_size,
|
||||
result,
|
||||
X,
|
||||
Y,
|
||||
BLOCK_SIZE_X=block_size_x,
|
||||
BLOCK_SIZE_Y=block_size_y,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# input - [M, K]
|
||||
# qweight - [K, N // 8]
|
||||
# qzeros - [K // G, N // 8]
|
||||
# scales - [K // G, N]
|
||||
# split_k_iters - parallelism along K-dimension, int, power of 2.
|
||||
def awq_gemm_triton(
|
||||
input: torch.Tensor,
|
||||
qweight: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
qzeros: torch.Tensor,
|
||||
split_k_iters: int,
|
||||
block_size_m: int = 32,
|
||||
block_size_n: int = 32,
|
||||
block_size_k: int = 32,
|
||||
) -> torch.Tensor:
|
||||
M, K = input.shape
|
||||
N = qweight.shape[1] * 8
|
||||
group_size = qweight.shape[0] // qzeros.shape[0]
|
||||
|
||||
assert N > 0 and K > 0 and M > 0
|
||||
assert qweight.shape[0] == K and qweight.shape[1] == N // 8
|
||||
assert qzeros.shape[0] == K // group_size and qzeros.shape[1] == N // 8
|
||||
assert scales.shape[0] == K // group_size and scales.shape[1] == N
|
||||
assert split_k_iters & (split_k_iters - 1) == 0 and split_k_iters != 0
|
||||
assert split_k_iters <= 32
|
||||
assert group_size <= K
|
||||
assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K
|
||||
|
||||
grid = lambda META: (
|
||||
triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
|
||||
split_k_iters,
|
||||
)
|
||||
|
||||
result = torch.zeros((split_k_iters, M, N), dtype=scales.dtype, device=input.device)
|
||||
|
||||
# A = input, B = qweight, C = result
|
||||
# A = M x K, B = K x N, C = M x N
|
||||
awq_gemm_kernel[grid](
|
||||
input,
|
||||
qweight,
|
||||
result,
|
||||
qzeros,
|
||||
scales,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
group_size,
|
||||
BLOCK_SIZE_M=block_size_m,
|
||||
BLOCK_SIZE_N=block_size_n,
|
||||
BLOCK_SIZE_K=block_size_k,
|
||||
SPLIT_K=split_k_iters,
|
||||
)
|
||||
|
||||
result = result.sum(0)
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,276 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import inspect
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
else:
|
||||
QuantizationMethods = str
|
||||
|
||||
|
||||
class QuantizeMethodBase(ABC):
|
||||
"""Base class for different quantized methods."""
|
||||
|
||||
uses_meta_device: bool = False
|
||||
"""
|
||||
Whether this method creates weights on meta device for online quantization.
|
||||
When True, weights are created on meta device and quantized layer-wise
|
||||
in process_weights_after_loading, reducing peak memory during loading.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def create_weights(
|
||||
self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
|
||||
):
|
||||
"""Create weights for a layer.
|
||||
|
||||
The weights will be set as attributes of the layer."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
|
||||
"""Apply the weights in layer to the input tensor.
|
||||
|
||||
Expects create_weights to have been called before on the layer."""
|
||||
raise NotImplementedError
|
||||
|
||||
# Not required functions
|
||||
def embedding(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
|
||||
"""Gather embeddings in the layer based on indices in the input tensor.
|
||||
|
||||
Expects create_weights to have been called before on the layer."""
|
||||
raise NotImplementedError
|
||||
|
||||
# Not required functions
|
||||
def tie_weights(self, layer: torch.nn.Module, embed_tokens: torch.nn.Module):
|
||||
"""Tie ``layer``'s weight to ``embed_tokens``' weight.
|
||||
|
||||
The default shares the weight tensor, which is the standard behavior for
|
||||
tied word embeddings and matches what ``ParallelLMHead.tie_weights`` did
|
||||
directly before quantization methods became responsible for it.
|
||||
Quantization methods that need special weight handling (e.g. repacked
|
||||
weights) override this.
|
||||
|
||||
Expects create_weights to have been called before on the layer."""
|
||||
layer.weight = embed_tokens.weight
|
||||
return layer
|
||||
|
||||
def process_weights_after_loading(self, layer: nn.Module) -> None:
|
||||
"""Process the weight after loading.
|
||||
|
||||
This can be used for example, to transpose weights for computation.
|
||||
"""
|
||||
return
|
||||
|
||||
|
||||
def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool:
|
||||
"""
|
||||
Not all quant methods have embedding implemented, so we need to check that
|
||||
it exists for our given method. We check this by making sure the function
|
||||
has been changed from the base implementation.
|
||||
"""
|
||||
base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding", None)
|
||||
class_embedding = inspect.getattr_static(method_class, "embedding", None)
|
||||
|
||||
return class_embedding is not None and class_embedding is not base_embedding
|
||||
|
||||
|
||||
class QuantizationConfig(ABC):
|
||||
"""Base class for quantization configs."""
|
||||
|
||||
_ignore_unexpected_suffixes = (
|
||||
".q_scale",
|
||||
".k_scale",
|
||||
".v_scale",
|
||||
".q_zero_point",
|
||||
".k_zero_point",
|
||||
".v_zero_point",
|
||||
)
|
||||
"""Suffixes of quantization parameters that may be present in the checkpoint but
|
||||
not in the model, and should be ignored if unexpected during loading. These are used
|
||||
after remapping, so should be in vLLM format (e.g. .q_scale, not .q.scale)."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# mapping is updated by models as they initialize
|
||||
self.packed_modules_mapping: dict[str, list[str]] = dict()
|
||||
|
||||
@abstractmethod
|
||||
def get_name(self) -> QuantizationMethods:
|
||||
"""Name of the quantization method."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def get_supported_act_dtypes(self) -> list[torch.dtype]:
|
||||
"""List of supported activation dtypes."""
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
"""Minimum GPU capability to support the quantization method.
|
||||
|
||||
E.g., 70 for Volta, 75 for Turing, 80 for Ampere.
|
||||
This requirement is due to the custom CUDA kernels used by the
|
||||
quantization method.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def get_config_filenames() -> list[str]:
|
||||
"""List of filenames to search for in the model directory."""
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig":
|
||||
"""Create a config class from the model's quantization config."""
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(
|
||||
cls,
|
||||
hf_quant_cfg: dict[str, Any],
|
||||
user_quant: str | None,
|
||||
hf_config: Any = None,
|
||||
) -> QuantizationMethods | None:
|
||||
"""
|
||||
Detects if this quantization method can support a given checkpoint
|
||||
format by overriding the user specified quantization method --
|
||||
this method should only be overwritten by subclasses in exceptional
|
||||
circumstances.
|
||||
|
||||
Args:
|
||||
hf_quant_cfg: The checkpoint's quantization config dict.
|
||||
user_quant: The user-specified quantization method string.
|
||||
hf_config: The HuggingFace model config object (e.g. for
|
||||
model_type checks). May be None if not available.
|
||||
"""
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any:
|
||||
"""Get a value from the model's quantization config."""
|
||||
for key in keys:
|
||||
if key in config:
|
||||
return config[key]
|
||||
raise ValueError(
|
||||
f"Cannot find any of {keys} in the model's quantization config."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any:
|
||||
"""Get an optional value from the model's quantization config."""
|
||||
try:
|
||||
return QuantizationConfig.get_from_keys(config, keys)
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
@abstractmethod
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> QuantizeMethodBase | None:
|
||||
"""Get the quantize method to use for the quantized layer.
|
||||
|
||||
Args:
|
||||
layer: The layer for the quant method.
|
||||
prefix: The full name of the layer in the state dict
|
||||
Returns:
|
||||
The quantize method. None if the given layer doesn't support quant
|
||||
method.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def get_cache_scale_mapper() -> "WeightsMapper":
|
||||
"""Mapping from checkpoint KV-cache scale names to vLLM scale names.
|
||||
|
||||
Returning a mapper here causes `AutoWeightsLoader` to apply it to the
|
||||
weight stream automatically; individual model `load_weights` methods
|
||||
do not need to know about KV-cache scales.
|
||||
"""
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
orig_to_new_regex = {
|
||||
# Deprecated fused kv_scale -> attn.k_scale
|
||||
re.compile(r"\.kv_scale$"): r".attn.k_scale",
|
||||
# ModelOpt: .self_attn.{k,v}_proj.{k,v}_scale -> .self_attn.attn.*
|
||||
re.compile(r"\.self_attn\.[kv]_proj\.([kv])_scale$"): (
|
||||
r".self_attn.attn.\1_scale"
|
||||
),
|
||||
# Fused QKV / qkqkv proj: .self_attn.qk(qk)v_proj.{k,v}_scale -> attn
|
||||
re.compile(r"\.self_attn\.qk(?:qk)?v_proj\.([kv])_scale$"): (
|
||||
r".self_attn.attn.\1_scale"
|
||||
),
|
||||
# NemotronH: .mixer.{k,v}_proj.{k,v}_scale -> .mixer.attn.*
|
||||
re.compile(r"\.mixer\.[kv]_proj\.([kv])_scale$"): r".mixer.attn.\1_scale",
|
||||
# HYV3: .self_attn.q.scale -> .self_attn.attn.q_scale
|
||||
re.compile(r"\.self_attn\.q\.scale$"): r".self_attn.attn.q_scale",
|
||||
# HYV3: .self_attn.{k,v}_cache.scale -> .self_attn.attn.{k,v}_scale
|
||||
re.compile(r"\.self_attn\.([kv])_cache\.scale$"): (
|
||||
r".self_attn.attn.\1_scale"
|
||||
),
|
||||
# Default: .{q,k,v}_scale -> .attn.{q,k,v}_scale (unless already .attn)
|
||||
re.compile(r"(?<!\.attn)\.([qkv])_scale$"): r".attn.\1_scale",
|
||||
re.compile(r"(?<!\.attn)\.([qkv])_zero_point$"): r".attn.\1_zero_point",
|
||||
}
|
||||
return WeightsMapper(orig_to_new_regex=orig_to_new_regex)
|
||||
|
||||
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
|
||||
"""
|
||||
# TODO (@kylesayrs): add implementations for all subclasses
|
||||
pass
|
||||
|
||||
def maybe_update_config( # noqa: B027
|
||||
self,
|
||||
model_name: str,
|
||||
hf_config: PretrainedConfig | None = None,
|
||||
revision: str | None = None,
|
||||
):
|
||||
"""
|
||||
Interface to update values after config initialization.
|
||||
|
||||
Args:
|
||||
model_name: The name of the model
|
||||
hf_config: The Hugging Face config of the model
|
||||
revision: The revision of the model
|
||||
Returns:
|
||||
"""
|
||||
# TODO: revision is never passed currently in vllm.py,
|
||||
# but is used in subclasses, should we remove this parameter?
|
||||
pass
|
||||
|
||||
def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool:
|
||||
"""
|
||||
Determine if mxfp4 quantization will be used for this config.
|
||||
|
||||
This allows hidden_size rounding to happen before moe_config creation
|
||||
without needing to instantiate quant_method first.
|
||||
|
||||
Args:
|
||||
prefix: The layer prefix/name in the model
|
||||
layer: The layer module
|
||||
|
||||
Returns:
|
||||
True if this config uses MXFP4 quantization, False otherwise
|
||||
"""
|
||||
return False
|
||||
@@ -0,0 +1,614 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from functools import cached_property
|
||||
from typing import Any, Union
|
||||
|
||||
import torch
|
||||
from packaging import version
|
||||
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEMethodBase,
|
||||
FusedMoEQuantConfig,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import (
|
||||
LinearBase,
|
||||
LinearMethodBase,
|
||||
UnquantizedLinearMethod,
|
||||
set_weight_attrs,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import (
|
||||
QuantizationConfig,
|
||||
QuantizationMethods,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
|
||||
def _check_bitsandbytes_version():
|
||||
min_version = "0.49.2" if current_platform.is_rocm() else "0.48.1"
|
||||
try:
|
||||
import bitsandbytes
|
||||
|
||||
if version.parse(bitsandbytes.__version__) < version.parse(min_version):
|
||||
raise ImportError(
|
||||
"bitsandbytes version is wrong. Please "
|
||||
f"install bitsandbytes>={min_version}."
|
||||
)
|
||||
except ImportError as err:
|
||||
raise ImportError(
|
||||
f"Please install bitsandbytes>={min_version} via "
|
||||
f"`pip install bitsandbytes>={min_version}` to use "
|
||||
"bitsandbytes quantizer."
|
||||
) from err
|
||||
|
||||
|
||||
class BitsAndBytesConfig(QuantizationConfig):
|
||||
"""Config class for BitsAndBytes Quantization.
|
||||
|
||||
Reference: https://arxiv.org/abs/2305.14314
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
load_in_8bit: bool = False,
|
||||
load_in_4bit: bool = True,
|
||||
bnb_4bit_compute_dtype: str = "float32",
|
||||
bnb_4bit_quant_storage: str = "uint8",
|
||||
bnb_4bit_quant_type: str = "fp4",
|
||||
bnb_4bit_use_double_quant: bool = False,
|
||||
llm_int8_enable_fp32_cpu_offload: bool = False,
|
||||
llm_int8_has_fp16_weight: bool = False,
|
||||
llm_int8_skip_modules: list[str] | None = None,
|
||||
llm_int8_threshold: float = 6.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.load_in_8bit = load_in_8bit
|
||||
self.load_in_4bit = load_in_4bit
|
||||
self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype
|
||||
self.bnb_4bit_quant_storage = bnb_4bit_quant_storage
|
||||
self.bnb_4bit_quant_type = bnb_4bit_quant_type
|
||||
self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant
|
||||
self.llm_int8_enable_fp32_cpu_offload = llm_int8_enable_fp32_cpu_offload
|
||||
self.llm_int8_has_fp16_weight = llm_int8_has_fp16_weight
|
||||
self.llm_int8_skip_modules = llm_int8_skip_modules or []
|
||||
self.llm_int8_threshold = llm_int8_threshold
|
||||
|
||||
if self.bnb_4bit_quant_storage not in ["uint8"]:
|
||||
raise ValueError(
|
||||
f"Unsupported bnb_4bit_quant_storage: {self.bnb_4bit_quant_storage}"
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"BitsAndBytesConfig(load_in_8bit={self.load_in_8bit}, "
|
||||
f"load_in_4bit={self.load_in_4bit}, "
|
||||
f"bnb_4bit_compute_dtype={self.bnb_4bit_compute_dtype}, "
|
||||
f"bnb_4bit_quant_storage={self.bnb_4bit_quant_storage}, "
|
||||
f"bnb_4bit_quant_type={self.bnb_4bit_quant_type}, "
|
||||
f"llm_int8_skip_modules={self.llm_int8_skip_modules})"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_name(self) -> QuantizationMethods:
|
||||
return "bitsandbytes"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(self) -> list[torch.dtype]:
|
||||
return [torch.float32, torch.float16, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 70
|
||||
|
||||
@staticmethod
|
||||
def get_config_filenames() -> list[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "BitsAndBytesConfig":
|
||||
def get_safe_value(config, keys, default_value=None):
|
||||
try:
|
||||
value = cls.get_from_keys(config, keys)
|
||||
return value if value is not None else default_value
|
||||
except ValueError:
|
||||
return default_value
|
||||
|
||||
load_in_8bit = get_safe_value(config, ["load_in_8bit"], default_value=False)
|
||||
load_in_4bit = get_safe_value(config, ["load_in_4bit"], default_value=True)
|
||||
bnb_4bit_compute_dtype = get_safe_value(
|
||||
config, ["bnb_4bit_compute_dtype"], default_value="float32"
|
||||
)
|
||||
bnb_4bit_quant_storage = get_safe_value(
|
||||
config, ["bnb_4bit_quant_storage"], default_value="uint8"
|
||||
)
|
||||
bnb_4bit_quant_type = get_safe_value(
|
||||
config, ["bnb_4bit_quant_type"], default_value="fp4"
|
||||
)
|
||||
bnb_4bit_use_double_quant = get_safe_value(
|
||||
config, ["bnb_4bit_use_double_quant"], default_value=False
|
||||
)
|
||||
llm_int8_enable_fp32_cpu_offload = get_safe_value(
|
||||
config, ["llm_int8_enable_fp32_cpu_offload"], default_value=False
|
||||
)
|
||||
llm_int8_has_fp16_weight = get_safe_value(
|
||||
config, ["llm_int8_has_fp16_weight"], default_value=False
|
||||
)
|
||||
llm_int8_skip_modules = get_safe_value(
|
||||
config, ["llm_int8_skip_modules"], default_value=[]
|
||||
)
|
||||
llm_int8_threshold = get_safe_value(
|
||||
config, ["llm_int8_threshold"], default_value=6.0
|
||||
)
|
||||
|
||||
return cls(
|
||||
load_in_8bit=load_in_8bit,
|
||||
load_in_4bit=load_in_4bit,
|
||||
bnb_4bit_compute_dtype=bnb_4bit_compute_dtype,
|
||||
bnb_4bit_quant_storage=bnb_4bit_quant_storage,
|
||||
bnb_4bit_quant_type=bnb_4bit_quant_type,
|
||||
bnb_4bit_use_double_quant=bnb_4bit_use_double_quant,
|
||||
llm_int8_enable_fp32_cpu_offload=llm_int8_enable_fp32_cpu_offload,
|
||||
llm_int8_has_fp16_weight=llm_int8_has_fp16_weight,
|
||||
llm_int8_skip_modules=llm_int8_skip_modules,
|
||||
llm_int8_threshold=llm_int8_threshold,
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> Union["LinearMethodBase", "BitsAndBytesMoEMethod"] | None:
|
||||
if isinstance(layer, LinearBase):
|
||||
if is_layer_skipped_bnb(prefix, self.llm_int8_skip_modules):
|
||||
return UnquantizedLinearMethod()
|
||||
return BitsAndBytesLinearMethod(self)
|
||||
elif isinstance(layer, RoutedExperts):
|
||||
return BitsAndBytesMoEMethod(self, layer.moe_config)
|
||||
return None
|
||||
|
||||
|
||||
class BitsAndBytesWeightParameter(torch.nn.Parameter):
|
||||
@cached_property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return torch.get_default_dtype()
|
||||
|
||||
|
||||
def is_layer_skipped_bnb(prefix: str, llm_int8_skip_modules: list[str]):
|
||||
# Split the prefix into its dot-separated components
|
||||
components = prefix.split(".")
|
||||
|
||||
# Check if any of the skip modules exactly matches any component
|
||||
substr_check = any(
|
||||
module_name in components for module_name in llm_int8_skip_modules
|
||||
)
|
||||
|
||||
# Allow certain layers to not be quantized
|
||||
set_components = set(".".join(components[: i + 1]) for i in range(len(components)))
|
||||
set_llm_int8_skip_modules = set(llm_int8_skip_modules)
|
||||
prefix_check = len(set_llm_int8_skip_modules & set_components) != 0
|
||||
|
||||
return substr_check or prefix_check
|
||||
|
||||
|
||||
def calculate_quant_ratio(dtype):
|
||||
if dtype.is_floating_point:
|
||||
return torch.finfo(dtype).bits // torch.iinfo(torch.uint8).bits
|
||||
else:
|
||||
return torch.iinfo(dtype).bits // torch.iinfo(torch.uint8).bits
|
||||
|
||||
|
||||
class BitsAndBytesLinearMethod(LinearMethodBase):
|
||||
"""Linear method for BitsAndBytes.
|
||||
|
||||
Args:
|
||||
quant_config: The BitsAndBytes quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: BitsAndBytesConfig):
|
||||
_check_bitsandbytes_version()
|
||||
self.quant_config = quant_config
|
||||
|
||||
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,
|
||||
):
|
||||
from bitsandbytes.nn import Int8Params
|
||||
|
||||
def create_qweight_for_8bit():
|
||||
qweight = Int8Params(
|
||||
data=torch.empty(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition,
|
||||
dtype=torch.int8,
|
||||
),
|
||||
has_fp16_weights=self.quant_config.llm_int8_has_fp16_weight,
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(
|
||||
qweight,
|
||||
{
|
||||
"input_dim": 0,
|
||||
"output_dim": 0,
|
||||
"pack_factor": 1,
|
||||
"use_bitsandbytes_8bit": True,
|
||||
"generation": 0,
|
||||
},
|
||||
)
|
||||
return qweight
|
||||
|
||||
def create_qweight_for_4bit():
|
||||
quant_ratio = calculate_quant_ratio(params_dtype)
|
||||
|
||||
total_size = input_size_per_partition * sum(output_partition_sizes)
|
||||
if total_size % quant_ratio != 0:
|
||||
raise ValueError(
|
||||
"The input size is not aligned with the quantized weight shape."
|
||||
)
|
||||
|
||||
qweight = BitsAndBytesWeightParameter(
|
||||
torch.empty(total_size // quant_ratio, 1, dtype=torch.uint8),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(
|
||||
qweight,
|
||||
{
|
||||
"input_dim": 0,
|
||||
"output_dim": 0,
|
||||
"pack_factor": quant_ratio,
|
||||
"use_bitsandbytes_4bit": True,
|
||||
},
|
||||
)
|
||||
return qweight
|
||||
|
||||
if self.quant_config.load_in_8bit:
|
||||
qweight = create_qweight_for_8bit()
|
||||
else:
|
||||
qweight = create_qweight_for_4bit()
|
||||
# Enable parameters to have the same name as in the BNB
|
||||
# checkpoint format.
|
||||
layer.register_parameter("weight", qweight)
|
||||
set_weight_attrs(qweight, extra_weight_attrs)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if self.quant_config.load_in_8bit:
|
||||
return self._apply_8bit_weight(layer, x, bias)
|
||||
else:
|
||||
return self._apply_4bit_weight(layer, x, bias)
|
||||
|
||||
def _apply_8bit_weight(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
# only load the bitsandbytes module when needed
|
||||
from bitsandbytes import MatmulLtState, matmul
|
||||
|
||||
original_type = x.dtype
|
||||
original_shape = x.shape
|
||||
reshape_after_matmul = False
|
||||
if x.ndim > 2:
|
||||
x = x.reshape(-1, x.size(-1))
|
||||
reshape_after_matmul = True
|
||||
bf_x = x.to(torch.bfloat16)
|
||||
|
||||
qweight = layer.weight
|
||||
offsets = qweight.bnb_shard_offsets
|
||||
quant_states = qweight.bnb_quant_state
|
||||
matmul_states = qweight.matmul_state
|
||||
generation = qweight.generation
|
||||
|
||||
out_dim_0 = x.shape[0]
|
||||
out_dim_1 = sum(
|
||||
[quant_state[1].shape[0] for quant_state in quant_states.items()]
|
||||
)
|
||||
out = torch.empty(out_dim_0, out_dim_1, dtype=torch.float16, device=x.device)
|
||||
|
||||
current_index = 0
|
||||
for i in range(len(quant_states)):
|
||||
output_size = quant_states[i].shape[0]
|
||||
|
||||
# in profile_run or the first generation of inference,
|
||||
# create new matmul_states
|
||||
if generation == 0 or generation == 1:
|
||||
matmul_states[i] = MatmulLtState()
|
||||
matmul_states[i].CB = qweight[offsets[i] : offsets[i + 1]]
|
||||
matmul_states[i].SCB = quant_states[i].to(x.device)
|
||||
matmul_states[i].threshold = self.quant_config.llm_int8_threshold
|
||||
matmul_states[
|
||||
i
|
||||
].has_fp16_weights = self.quant_config.llm_int8_has_fp16_weight
|
||||
matmul_states[i].is_training = False
|
||||
if (
|
||||
matmul_states[i].threshold > 0.0
|
||||
and not matmul_states[i].has_fp16_weights
|
||||
):
|
||||
matmul_states[i].use_pool = True
|
||||
|
||||
new_x = bf_x.unsqueeze(0)
|
||||
|
||||
out[:, current_index : current_index + output_size] = matmul(
|
||||
new_x, qweight[offsets[i] : offsets[i + 1]], state=matmul_states[i]
|
||||
)
|
||||
|
||||
current_index += output_size
|
||||
|
||||
out = out.to(original_type)
|
||||
|
||||
if reshape_after_matmul:
|
||||
out = out.view(*original_shape[:-1], out.size(-1))
|
||||
|
||||
if bias is not None:
|
||||
out += bias
|
||||
|
||||
qweight.generation += 1
|
||||
|
||||
return out
|
||||
|
||||
def _apply_4bit_weight(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
original_type = x.dtype
|
||||
original_shape = x.shape
|
||||
reshape_after_matmul = False
|
||||
if x.ndim > 2:
|
||||
x = x.reshape(-1, x.size(-1))
|
||||
reshape_after_matmul = True
|
||||
bf_x = x.to(torch.bfloat16)
|
||||
|
||||
qweight = layer.weight
|
||||
quant_states = qweight.bnb_quant_state
|
||||
offsets = qweight.bnb_shard_offsets
|
||||
|
||||
out_dim_0 = x.shape[0]
|
||||
out_dim_1 = sum(
|
||||
[quant_state[1].shape[0] for quant_state in quant_states.items()]
|
||||
)
|
||||
out = torch.empty(out_dim_0, out_dim_1, dtype=torch.bfloat16, device=x.device)
|
||||
apply_bnb_4bit(bf_x, qweight, offsets, out)
|
||||
out = out.to(original_type)
|
||||
|
||||
if reshape_after_matmul:
|
||||
out = out.view(*original_shape[:-1], out.size(-1))
|
||||
|
||||
if bias is not None:
|
||||
out += bias
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def _apply_bnb_4bit(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
offsets: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
) -> None:
|
||||
# only load the bitsandbytes module when needed
|
||||
from bitsandbytes import matmul_4bit
|
||||
|
||||
quant_states = weight.bnb_quant_state
|
||||
current_index = 0
|
||||
for i in range(len(quant_states)):
|
||||
output_size = quant_states[i].shape[0]
|
||||
# It is more efficient to use out kwarg like
|
||||
# matmul_4bit(..., out = ...). Infeasible now due to the bug
|
||||
# https://github.com/TimDettmers/bitsandbytes/issues/1235.
|
||||
# Need to change after the bug is fixed.
|
||||
out[:, current_index : current_index + output_size] = matmul_4bit(
|
||||
x, weight[offsets[i] : offsets[i + 1]].t(), quant_states[i]
|
||||
)
|
||||
current_index += output_size
|
||||
|
||||
|
||||
def _apply_bnb_4bit_fake(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
offsets: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
) -> None:
|
||||
return
|
||||
|
||||
|
||||
try:
|
||||
direct_register_custom_op(
|
||||
op_name="apply_bnb_4bit",
|
||||
op_func=_apply_bnb_4bit,
|
||||
mutates_args=["out"],
|
||||
fake_impl=_apply_bnb_4bit_fake,
|
||||
dispatch_key=current_platform.dispatch_key,
|
||||
)
|
||||
apply_bnb_4bit = torch.ops.vllm.apply_bnb_4bit
|
||||
|
||||
except AttributeError as error:
|
||||
raise error
|
||||
|
||||
|
||||
class BitsAndBytesMoEMethod(FusedMoEMethodBase):
|
||||
"""MoE method for BitsAndBytes.
|
||||
|
||||
Args:
|
||||
quant_config: The BitsAndBytes quantization config.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
quant_config: BitsAndBytesConfig,
|
||||
moe: FusedMoEConfig,
|
||||
):
|
||||
super().__init__(moe)
|
||||
_check_bitsandbytes_version()
|
||||
self.quant_config = quant_config
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
if self.quant_config.load_in_8bit:
|
||||
call_fun = self._create_weights_8bit
|
||||
else:
|
||||
call_fun = self._create_weights_4bit
|
||||
call_fun(
|
||||
layer,
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: RoutedExperts
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
return None
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
from vllm.model_executor.layers.fused_moe import fused_experts
|
||||
|
||||
# TODO(bnell): Do these need to be called on the hot path?
|
||||
if self.quant_config.load_in_8bit:
|
||||
w13, w2 = self._apply_8bit_dequant(layer)
|
||||
else:
|
||||
w13, w2 = self._apply_4bit_dequnt(layer)
|
||||
return fused_experts(
|
||||
hidden_states=x,
|
||||
w1=w13,
|
||||
w2=w2,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=layer.activation,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
quant_config=self.moe_quant_config,
|
||||
)
|
||||
|
||||
def _create_weights_4bit(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
quant_ratio = calculate_quant_ratio(params_dtype)
|
||||
# Fused gate_up_proj (column parallel)
|
||||
w13_total_size = (
|
||||
hidden_size * 2 * intermediate_size_per_partition
|
||||
) // quant_ratio
|
||||
w13_qweight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
w13_total_size,
|
||||
1,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_qweight)
|
||||
set_weight_attrs(w13_qweight, extra_weight_attrs)
|
||||
set_weight_attrs(
|
||||
w13_qweight,
|
||||
{
|
||||
"num_experts": num_experts,
|
||||
"input_dim": hidden_size,
|
||||
"output_dim": 2 * intermediate_size_per_partition,
|
||||
"experts_shape": (
|
||||
num_experts,
|
||||
intermediate_size_per_partition * 2,
|
||||
hidden_size,
|
||||
),
|
||||
"pack_factor": quant_ratio,
|
||||
"use_bitsandbytes_4bit": True,
|
||||
},
|
||||
)
|
||||
# down_proj (row parallel)
|
||||
w2_total_size = (hidden_size * intermediate_size_per_partition) // quant_ratio
|
||||
w2_qweight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
w2_total_size,
|
||||
1,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(
|
||||
w2_qweight,
|
||||
{
|
||||
"num_experts": num_experts,
|
||||
"input_dim": intermediate_size_per_partition,
|
||||
"output_dim": hidden_size,
|
||||
"experts_shape": (
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
),
|
||||
"pack_factor": quant_ratio,
|
||||
"use_bitsandbytes_4bit": True,
|
||||
},
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_qweight)
|
||||
set_weight_attrs(w2_qweight, extra_weight_attrs)
|
||||
|
||||
def _create_weights_8bit(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
raise NotImplementedError
|
||||
|
||||
def _apply_4bit_dequnt(
|
||||
self, layer: torch.nn.Module
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
from bitsandbytes.functional import dequantize_4bit
|
||||
|
||||
w13 = dequantize_4bit(
|
||||
layer.w13_weight.reshape(-1, 1),
|
||||
layer.w13_weight.bnb_quant_state,
|
||||
)
|
||||
w2 = dequantize_4bit(
|
||||
layer.w2_weight.reshape(-1, 1),
|
||||
layer.w2_weight.bnb_quant_state,
|
||||
)
|
||||
w13 = w13.reshape(layer.w13_weight.experts_shape)
|
||||
w2 = w2.reshape(layer.w2_weight.experts_shape)
|
||||
return w13, w2
|
||||
|
||||
def _apply_8bit_dequant(
|
||||
self, layer: torch.nn.Module
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+170
@@ -0,0 +1,170 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Quantized embedding method for compressed-tensors.
|
||||
|
||||
Adds dequant-on-lookup support for a pack-quantized ``VocabParallelEmbedding``
|
||||
(2-8 bit INT, channel- or group-quantized). Only the gathered token rows are
|
||||
unpacked and dequantized, so the packed weight is never densified.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
|
||||
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
|
||||
from vllm.model_executor.parameter import (
|
||||
BasevLLMParameter,
|
||||
ChannelQuantScaleParameter,
|
||||
GroupQuantScaleParameter,
|
||||
PackedvLLMParameter,
|
||||
)
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
__all__ = ["CompressedTensorsEmbeddingWNA16Int"]
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _dequant_gather_kernel(
|
||||
ids_ptr,
|
||||
packed_ptr,
|
||||
scale_ptr,
|
||||
out_ptr,
|
||||
hidden,
|
||||
packed_cols,
|
||||
num_groups,
|
||||
NUM_BITS: tl.constexpr,
|
||||
PACK_FACTOR: tl.constexpr,
|
||||
GROUP_SIZE: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
"""Gather embedding rows by token id, unpack int32-packed INT weights, and
|
||||
dequantize to ``out`` dtype in one pass (no int8 intermediate)."""
|
||||
row = tl.program_id(0)
|
||||
col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
|
||||
col_mask = col < hidden
|
||||
tid = tl.load(ids_ptr + row).to(tl.int64)
|
||||
|
||||
packed_idx = col // PACK_FACTOR
|
||||
shift = (col % PACK_FACTOR) * NUM_BITS
|
||||
packed = tl.load(
|
||||
packed_ptr + tid * packed_cols + packed_idx, mask=col_mask, other=0
|
||||
)
|
||||
q = ((packed >> shift) & ((1 << NUM_BITS) - 1)) - (1 << (NUM_BITS - 1))
|
||||
|
||||
if GROUP_SIZE == 0: # channel: one scale per row
|
||||
scale = tl.load(scale_ptr + tid)
|
||||
else: # group: one scale per (row, group)
|
||||
grp = col // GROUP_SIZE
|
||||
scale = tl.load(scale_ptr + tid * num_groups + grp, mask=col_mask, other=0.0)
|
||||
|
||||
out = q.to(tl.float32) * scale.to(tl.float32)
|
||||
tl.store(
|
||||
out_ptr + row * hidden + col, out.to(out_ptr.dtype.element_ty), mask=col_mask
|
||||
)
|
||||
|
||||
|
||||
def _dequant_gather_triton(
|
||||
ids: torch.Tensor,
|
||||
weight_packed: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
hidden: int,
|
||||
num_bits: int,
|
||||
) -> torch.Tensor:
|
||||
n = ids.numel()
|
||||
out = torch.empty(n, hidden, dtype=weight_scale.dtype, device=weight_packed.device)
|
||||
num_groups = weight_scale.shape[1]
|
||||
group_size = 0 if num_groups == 1 else hidden // num_groups
|
||||
block = min(triton.next_power_of_2(hidden), 1024)
|
||||
grid = (n, triton.cdiv(hidden, block))
|
||||
_dequant_gather_kernel[grid](
|
||||
ids,
|
||||
weight_packed,
|
||||
weight_scale,
|
||||
out,
|
||||
hidden,
|
||||
weight_packed.shape[1],
|
||||
num_groups,
|
||||
NUM_BITS=num_bits,
|
||||
PACK_FACTOR=32 // num_bits,
|
||||
GROUP_SIZE=group_size,
|
||||
BLOCK=block,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
class CompressedTensorsEmbeddingWNA16Int(QuantizeMethodBase):
|
||||
def __init__(self, weight_quant: QuantizationArgs):
|
||||
self.num_bits = weight_quant.num_bits
|
||||
self.pack_factor = 32 // self.num_bits
|
||||
self.strategy = weight_quant.strategy
|
||||
self.group_size = weight_quant.group_size
|
||||
self.is_group = (
|
||||
self.strategy == QuantizationStrategy.GROUP.value
|
||||
and self.group_size is not None
|
||||
)
|
||||
|
||||
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,
|
||||
):
|
||||
weight_loader = extra_weight_attrs["weight_loader"]
|
||||
# Embedding weight is [num_embeddings(vocab), embedding_dim(hidden)];
|
||||
# vocab is the output (partitioned) dim, hidden is the input dim.
|
||||
vocab_pp = sum(output_partition_sizes)
|
||||
hidden = input_size_per_partition
|
||||
layer.hidden_size = hidden
|
||||
|
||||
weight_packed = PackedvLLMParameter(
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
packed_dim=1,
|
||||
packed_factor=self.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
data=torch.empty(vocab_pp, hidden // self.pack_factor, dtype=torch.int32),
|
||||
)
|
||||
|
||||
if self.is_group:
|
||||
assert hidden % self.group_size == 0
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
output_dim=0,
|
||||
input_dim=1,
|
||||
weight_loader=weight_loader,
|
||||
data=torch.empty(
|
||||
vocab_pp, hidden // self.group_size, dtype=params_dtype
|
||||
),
|
||||
)
|
||||
else:
|
||||
weight_scale = ChannelQuantScaleParameter(
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
data=torch.empty(vocab_pp, 1, dtype=params_dtype),
|
||||
)
|
||||
|
||||
weight_shape = BasevLLMParameter(
|
||||
data=torch.empty(2, dtype=torch.int64), weight_loader=weight_loader
|
||||
)
|
||||
|
||||
layer.register_parameter("weight_packed", weight_packed)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
layer.register_parameter("weight_shape", weight_shape)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
pass
|
||||
|
||||
def embedding(self, layer: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor:
|
||||
ids = input_.reshape(-1).contiguous()
|
||||
hidden = layer.hidden_size
|
||||
deq = _dequant_gather_triton(
|
||||
ids, layer.weight_packed, layer.weight_scale, hidden, self.num_bits
|
||||
)
|
||||
return deq.reshape(*input_.shape, hidden)
|
||||
|
||||
def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
|
||||
raise NotImplementedError(
|
||||
"CompressedTensorsEmbeddingWNA16Int supports embedding lookup only"
|
||||
)
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe import ( # noqa: E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CompressedTensorsMoEMethod",
|
||||
]
|
||||
+218
@@ -0,0 +1,218 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import torch
|
||||
from compressed_tensors import CompressionFormat
|
||||
from compressed_tensors.quantization import (
|
||||
ActivationOrdering,
|
||||
QuantizationStrategy,
|
||||
QuantizationType,
|
||||
)
|
||||
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEMethodBase,
|
||||
UnquantizedFusedMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
|
||||
WNA16_SUPPORTED_BITS,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
check_moe_marlin_supports_layer,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CompressedTensorsMoEMethod(FusedMoEMethodBase):
|
||||
@staticmethod
|
||||
def get_moe_method(
|
||||
quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
|
||||
layer: torch.nn.Module,
|
||||
layer_name: str,
|
||||
) -> FusedMoEMethodBase:
|
||||
# RoutedExperts was made by combining multiple Linears so need to
|
||||
# make sure quantization config for Linear can target it
|
||||
quant_config._add_fused_moe_to_target_scheme_map()
|
||||
unfused_names = [
|
||||
layer_name + proj_name
|
||||
for proj_name in [".0.gate_proj", ".0.up_proj", ".0.down_proj"]
|
||||
]
|
||||
# TODO: refactor this to use expert_mapping and check all layer numbers
|
||||
all_scheme_dicts = [
|
||||
quant_config.get_scheme_dict(layer, name) for name in unfused_names
|
||||
]
|
||||
scheme_dict = all_scheme_dicts.pop()
|
||||
|
||||
# multiple schemes found
|
||||
if not all([cur_dict == scheme_dict for cur_dict in all_scheme_dicts]):
|
||||
raise ValueError(
|
||||
"All MoE projections need to have same "
|
||||
"quantization scheme but found multiple"
|
||||
)
|
||||
|
||||
if scheme_dict is None: # ignored layer
|
||||
return UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
|
||||
# TODO: @dsikka: refactor this to use schemes as other kernels
|
||||
# are supported + check if the layer is being ignored.
|
||||
weight_quant = scheme_dict.get("weights")
|
||||
input_quant = scheme_dict.get("input_activations")
|
||||
format = scheme_dict.get("format")
|
||||
|
||||
if quant_config._is_mxfp4(weight_quant):
|
||||
from .compressed_tensors_moe_w4a4_mxfp4 import (
|
||||
CompressedTensorsW4A4Mxfp4MoEMethod,
|
||||
)
|
||||
|
||||
return CompressedTensorsW4A4Mxfp4MoEMethod(layer.moe_config)
|
||||
|
||||
if quant_config._is_mxfp8(weight_quant):
|
||||
from .compressed_tensors_moe_w8a8_mxfp8 import (
|
||||
CompressedTensorsW8A8Mxfp8MoEMethod,
|
||||
)
|
||||
|
||||
return CompressedTensorsW8A8Mxfp8MoEMethod(layer.moe_config)
|
||||
|
||||
if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
|
||||
# group_size=None means channelwise
|
||||
group_size = weight_quant.group_size or -1
|
||||
|
||||
valid_format_and_bits = (
|
||||
weight_quant.num_bits in WNA16_SUPPORTED_BITS
|
||||
and format == CompressionFormat.pack_quantized.value
|
||||
)
|
||||
|
||||
if not valid_format_and_bits:
|
||||
raise ValueError(
|
||||
"For Fused MoE layers, only format: ",
|
||||
f"{CompressionFormat.pack_quantized.value} ",
|
||||
f" and bits: {WNA16_SUPPORTED_BITS} is supported ",
|
||||
f"but got format: {CompressionFormat.pack_quantized.value} "
|
||||
f" and bits: {weight_quant.num_bits}",
|
||||
)
|
||||
|
||||
# Prefer to use the MarlinMoE kernel when it is supported.
|
||||
is_actorder = (
|
||||
weight_quant.strategy == QuantizationStrategy.GROUP
|
||||
and weight_quant.actorder
|
||||
in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC)
|
||||
)
|
||||
if (
|
||||
not check_moe_marlin_supports_layer(
|
||||
layer, group_size, allow_tile_padding=not is_actorder
|
||||
)
|
||||
or current_platform.is_rocm()
|
||||
):
|
||||
if is_actorder:
|
||||
raise ValueError(
|
||||
"WNA16MoE is not supported with actorder=group/dynamic."
|
||||
)
|
||||
|
||||
# Native ROCm HIP kernels (RDNA3, etc.)
|
||||
if current_platform.is_rocm():
|
||||
from . import rocm_moe_rdna
|
||||
|
||||
if rocm_moe_rdna.is_supported(weight_quant):
|
||||
return rocm_moe_rdna.make_method(
|
||||
weight_quant, input_quant, layer.moe_config
|
||||
)
|
||||
from vllm.platforms.rocm import on_gfx950
|
||||
|
||||
vllm_config = get_current_vllm_config()
|
||||
is_lora_disabled = vllm_config.lora_config is None
|
||||
moe_backend = vllm_config.kernel_config.moe_backend
|
||||
if (
|
||||
weight_quant.strategy == QuantizationStrategy.GROUP
|
||||
and weight_quant.type == QuantizationType.INT
|
||||
and group_size == 32
|
||||
and weight_quant.num_bits == 4
|
||||
and is_lora_disabled
|
||||
and on_gfx950()
|
||||
and moe_backend == "flydsl"
|
||||
):
|
||||
from .compressed_tensors_moe_w4a16_flydsl import (
|
||||
CompressedTensorsW4A16FlydslMoEMethod,
|
||||
)
|
||||
|
||||
logger.info_once("Using CompressedTensorsW4A16FlydslMoEMethod")
|
||||
return CompressedTensorsW4A16FlydslMoEMethod(
|
||||
weight_quant, input_quant, layer.moe_config
|
||||
)
|
||||
from .compressed_tensors_moe_wna16 import (
|
||||
CompressedTensorsWNA16MoEMethod,
|
||||
)
|
||||
|
||||
logger.info_once("Using CompressedTensorsWNA16MoEMethod")
|
||||
return CompressedTensorsWNA16MoEMethod(
|
||||
weight_quant, input_quant, layer.moe_config
|
||||
)
|
||||
else:
|
||||
from .compressed_tensors_moe_wna16_marlin import (
|
||||
CompressedTensorsWNA16MarlinMoEMethod,
|
||||
)
|
||||
|
||||
logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
|
||||
return CompressedTensorsWNA16MarlinMoEMethod(
|
||||
weight_quant, input_quant, layer.moe_config
|
||||
)
|
||||
elif quant_config._is_nvfp4_format(weight_quant):
|
||||
from .compressed_tensors_moe_w4a4_nvfp4 import (
|
||||
CompressedTensorsW4A4Nvfp4MoEMethod,
|
||||
)
|
||||
|
||||
_is_valid_nvfp4_activations = (
|
||||
quant_config._is_nvfp4_format(input_quant) or input_quant is None
|
||||
)
|
||||
if not _is_valid_nvfp4_activations:
|
||||
raise ValueError(
|
||||
"For NVFP4 weights, input quantization must also be NVFP4 format ",
|
||||
f"or None for NVFP4A16, found {input_quant}",
|
||||
)
|
||||
return CompressedTensorsW4A4Nvfp4MoEMethod(
|
||||
layer.moe_config, layer_name, use_a16=(input_quant is None)
|
||||
)
|
||||
elif (
|
||||
quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
|
||||
or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)
|
||||
or quant_config._is_fp8_w8a8(weight_quant, input_quant)
|
||||
):
|
||||
from .compressed_tensors_moe_w8a8_fp8 import (
|
||||
CompressedTensorsW8A8Fp8MoEMethod,
|
||||
)
|
||||
|
||||
return CompressedTensorsW8A8Fp8MoEMethod(
|
||||
weight_quant, input_quant, layer.moe_config
|
||||
)
|
||||
elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant):
|
||||
from .compressed_tensors_moe_w8a8_int8 import (
|
||||
CompressedTensorsW8A8Int8MoEMethod,
|
||||
)
|
||||
|
||||
return CompressedTensorsW8A8Int8MoEMethod(
|
||||
weight_quant, input_quant, layer.moe_config
|
||||
)
|
||||
elif quant_config._is_fp8_w4a8_sm90(weight_quant, input_quant):
|
||||
from .compressed_tensors_moe_w4a8_fp8 import (
|
||||
CompressedTensorsW4A8Fp8MoEMethod,
|
||||
)
|
||||
|
||||
logger.info_once("Using CompressedTensorsW4A8Fp8MoEMethod")
|
||||
return CompressedTensorsW4A8Fp8MoEMethod(
|
||||
weight_quant, input_quant, layer.moe_config
|
||||
)
|
||||
elif quant_config._is_dynamic_token_w4a8_int(weight_quant, input_quant):
|
||||
from .compressed_tensors_moe_w4a8_int8 import (
|
||||
CompressedTensorsW4A8Int8MoEMethod,
|
||||
)
|
||||
|
||||
return CompressedTensorsW4A8Int8MoEMethod(
|
||||
weight_quant, input_quant, layer.moe_config
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}"
|
||||
)
|
||||
+348
@@ -0,0 +1,348 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import torch
|
||||
from aiter.ops.shuffle import shuffle_weight
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
)
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEQuantConfig,
|
||||
int4_w4a16_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _pack_shuffled_int8_to_packed_int4_no_perm(x_shuf_i8: torch.Tensor) -> torch.Tensor:
|
||||
"""Pack a preshuffled int8 tensor (values in [-8, 7]) into packed int4 bytes.
|
||||
Each contiguous 8-value block [v0..v7] -> 4 bytes:
|
||||
b0=(v4<<4)|v0, b1=(v5<<4)|v1, b2=(v6<<4)|v2, b3=(v7<<4)|v3.
|
||||
This matches the 7-op in-kernel unpack sequence and avoids any v_perm.
|
||||
"""
|
||||
flat = x_shuf_i8.contiguous().view(-1).to(torch.int16)
|
||||
assert flat.numel() % 8 == 0
|
||||
u = (flat & 0xF).to(torch.uint8).view(-1, 8)
|
||||
out = torch.empty((u.shape[0], 4), device=u.device, dtype=torch.uint8)
|
||||
out[:, 0] = u[:, 0] | (u[:, 4] << 4)
|
||||
out[:, 1] = u[:, 1] | (u[:, 5] << 4)
|
||||
out[:, 2] = u[:, 2] | (u[:, 6] << 4)
|
||||
out[:, 3] = u[:, 3] | (u[:, 7] << 4)
|
||||
return out.view(-1).to(torch.int8)
|
||||
|
||||
|
||||
def _unpack_gptq_int32_to_signed_int4(w_int32):
|
||||
"""Unpack GPTQ int32 [E, K//8, N] to signed int4 values [E, N, K] (as int8).
|
||||
Shared by both the packed-int4 and bf16-dequant paths.
|
||||
"""
|
||||
E = w_int32.shape[0]
|
||||
# [E, K//8, N] -> transpose -> [E, N, K//8]
|
||||
w = w_int32.transpose(1, 2).contiguous()
|
||||
N = w.shape[1]
|
||||
K_div8 = w.shape[2]
|
||||
K = K_div8 * 8
|
||||
|
||||
# Unpack int32 -> 8 x uint4 values along K
|
||||
w_expanded = w.unsqueeze(-1).expand(E, N, K_div8, 8) # [E, N, K//8, 8]
|
||||
shifts = torch.arange(8, device=w.device) * 4 # [0, 4, 8, ..., 28]
|
||||
nibbles = ((w_expanded >> shifts) & 0xF).to(torch.int8) # [E, N, K//8, 8]
|
||||
nibbles = nibbles.reshape(E, N, K) # [E, N, K] unsigned int4 as int8
|
||||
|
||||
# Convert unsigned [0,15] to signed [-8,7]
|
||||
signed = nibbles.to(torch.int16) - 8
|
||||
signed = signed.to(torch.int8) # [E, N, K] signed int4 as int8
|
||||
return signed
|
||||
|
||||
|
||||
def _gptq_int32_to_flydsl_packed(w_int32):
|
||||
"""Convert GPTQ int32 [E, K//8, N] to FlyDSL shuffled packed int4 [E, N, K//2].
|
||||
Steps:
|
||||
1. Unpack int32 to individual signed int4 values (as int8)
|
||||
2. Apply FlyDSL preshuffle (on individual int8 values)
|
||||
3. Pack with FlyDSL's interleaved int4 packing
|
||||
"""
|
||||
signed = _unpack_gptq_int32_to_signed_int4(w_int32)
|
||||
E, N, K = signed.shape
|
||||
|
||||
# FlyDSL preshuffle (operates on individual values)
|
||||
shuffled = shuffle_weight(signed, layout=(16, 16))
|
||||
|
||||
# FlyDSL interleaved int4 packing
|
||||
packed = _pack_shuffled_int8_to_packed_int4_no_perm(shuffled).contiguous()
|
||||
return packed.view(E, N, K // 2)
|
||||
|
||||
|
||||
class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod):
|
||||
def __init__(
|
||||
self,
|
||||
weight_quant: QuantizationArgs,
|
||||
input_quant: QuantizationArgs | None,
|
||||
moe: FusedMoEConfig,
|
||||
layer_name: str | None = None,
|
||||
):
|
||||
super().__init__(moe)
|
||||
self.weight_quant = weight_quant
|
||||
self.input_quant = input_quant
|
||||
# Extract properties from weight_quant
|
||||
assert weight_quant.num_bits == 4
|
||||
self.num_bits = weight_quant.num_bits
|
||||
self.packed_factor = 32 // weight_quant.num_bits
|
||||
self.strategy = weight_quant.strategy
|
||||
# channelwise is not supported by this kernel
|
||||
assert weight_quant.strategy == "group"
|
||||
assert weight_quant.group_size == 32
|
||||
self.group_size = weight_quant.group_size
|
||||
# grouped actorder isn't supported by this kernel
|
||||
assert weight_quant.actorder != "group"
|
||||
assert weight_quant.symmetric, (
|
||||
"Only symmetric quantization is supported for MoE"
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
self.num_experts = num_experts
|
||||
self.inter_dim = intermediate_size_per_partition
|
||||
# Will transpose the loaded weight along the
|
||||
# intermediate and hidden dim sizes. Will
|
||||
# shard for TP along the transposed dims
|
||||
extra_weight_attrs.update(
|
||||
{"is_transposed": True, "quant_method": self.strategy}
|
||||
)
|
||||
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size // self.packed_factor,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_packed", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition // self.packed_factor,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_packed", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
w2_scales_size = intermediate_size_per_partition
|
||||
|
||||
if self.strategy == "channel":
|
||||
num_groups_w2 = num_groups_w13 = 1
|
||||
self.group_size = -1
|
||||
else:
|
||||
num_groups_w2 = w2_scales_size // self.group_size
|
||||
num_groups_w13 = hidden_size // self.group_size
|
||||
|
||||
w13_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
num_experts,
|
||||
num_groups_w13,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_scale)
|
||||
set_weight_attrs(w13_scale, extra_weight_attrs)
|
||||
|
||||
w2_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_scale)
|
||||
set_weight_attrs(w2_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_scale, {"load_full_w2": False})
|
||||
|
||||
w2_weight_shape = torch.nn.Parameter(
|
||||
torch.empty(num_experts, 2), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w2_weight_shape", w2_weight_shape)
|
||||
set_weight_attrs(w2_weight_shape, extra_weight_attrs)
|
||||
w13_weight_shape = torch.nn.Parameter(
|
||||
torch.empty(num_experts, 2), requires_grad=False
|
||||
)
|
||||
|
||||
layer.register_parameter("w13_weight_shape", w13_weight_shape)
|
||||
set_weight_attrs(w13_weight_shape, extra_weight_attrs)
|
||||
|
||||
w13_g_idx = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_g_idx", w13_g_idx)
|
||||
set_weight_attrs(w13_g_idx, extra_weight_attrs)
|
||||
|
||||
w2_g_idx = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_g_idx", w2_g_idx)
|
||||
set_weight_attrs(w2_g_idx, extra_weight_attrs)
|
||||
|
||||
w13_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
|
||||
set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
|
||||
|
||||
w2_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
|
||||
set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
|
||||
|
||||
layer.a13_scale = None
|
||||
layer.a2_scale = None
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
# Reconfigure packed weights and scales to match flydsl_w4a16 format
|
||||
|
||||
# Convert w13 weights
|
||||
w13 = layer.w13_weight_packed.data
|
||||
w13 = _gptq_int32_to_flydsl_packed(w13)
|
||||
w13 = w13.view(-1).contiguous()
|
||||
layer.w13_weight_packed = torch.nn.Parameter(w13, requires_grad=False)
|
||||
|
||||
# Convert w2 weights
|
||||
w2 = layer.w2_weight_packed.data
|
||||
w2 = _gptq_int32_to_flydsl_packed(w2)
|
||||
w2 = w2.view(-1).contiguous()
|
||||
layer.w2_weight_packed = torch.nn.Parameter(w2, requires_grad=False)
|
||||
|
||||
# Convert scales for FlyDSL:
|
||||
# per-row: [E, 1, N] -> squeeze -> [E, N]
|
||||
# groupwise: [E, K//gs, N] -> keep as-is (Opt 0: cache-friendly layout)
|
||||
w13_scale = layer.w13_weight_scale.data
|
||||
if self.group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1:
|
||||
E, G, N = w13_scale.shape
|
||||
w13_scale = (
|
||||
w13_scale.view(E, G // 2, 2, N)
|
||||
.permute(0, 1, 3, 2)
|
||||
.contiguous()
|
||||
.view(-1)
|
||||
.contiguous()
|
||||
)
|
||||
elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1:
|
||||
# Per-row: squeeze [E, 1, N] -> [E, N]
|
||||
w13_scale = w13_scale.squeeze(1)
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
w13_scale.contiguous(), requires_grad=False
|
||||
)
|
||||
|
||||
w2_scale = layer.w2_weight_scale.data
|
||||
if self.group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1:
|
||||
E, G, N = w2_scale.shape
|
||||
w2_scale = (
|
||||
w2_scale.view(E, G // 2, 2, N)
|
||||
.permute(0, 1, 3, 2)
|
||||
.contiguous()
|
||||
.view(-1)
|
||||
.contiguous()
|
||||
)
|
||||
elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1:
|
||||
# Per-row: squeeze [E, 1, N] -> [E, N]
|
||||
w2_scale = w2_scale.squeeze(1)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
w2_scale.contiguous(), requires_grad=False
|
||||
)
|
||||
|
||||
layer.w13_weight_packed.is_shuffled = True
|
||||
layer.w2_weight_packed.is_shuffled = True
|
||||
layer.is_aiter_converted = True
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: torch.nn.Module
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
assert self.num_bits == 4
|
||||
return int4_w4a16_moe_quant_config(
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
w1_zp=None,
|
||||
w2_zp=None,
|
||||
block_shape=[0, self.group_size],
|
||||
)
|
||||
|
||||
def select_gemm_impl(
|
||||
self,
|
||||
prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
|
||||
layer: torch.nn.Module,
|
||||
) -> mk.FusedMoEExpertsModular:
|
||||
raise NotImplementedError
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import (
|
||||
fused_flydsl_moe,
|
||||
)
|
||||
|
||||
assert self.moe_quant_config is not None
|
||||
|
||||
return fused_flydsl_moe(
|
||||
x,
|
||||
layer.w13_weight_packed,
|
||||
layer.w2_weight_packed,
|
||||
self.num_experts,
|
||||
self.inter_dim,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
w1_scale=self.moe_quant_config.w1_scale,
|
||||
w2_scale=self.moe_quant_config.w2_scale,
|
||||
topk=topk_weights.shape[-1],
|
||||
group_size=self.group_size,
|
||||
doweight_stage1=layer.apply_router_weight_on_input,
|
||||
scale_is_bf16=True,
|
||||
)
|
||||
+234
@@ -0,0 +1,234 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoeWeightScaleSupported,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
mxfp4_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
|
||||
CutlassExpertsMxfp4,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
|
||||
MarlinExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
|
||||
XPUExpertsMxFp4,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
|
||||
Mxfp4MoeBackend,
|
||||
make_mxfp4_moe_kernel,
|
||||
make_mxfp4_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
|
||||
prepare_moe_fp4_layer_for_marlin,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
|
||||
def __init__(self, moe):
|
||||
super().__init__(moe)
|
||||
self.group_size = 32
|
||||
self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
|
||||
# use cutlass if supported, otherwise fallback to marlin for weight-only FP4
|
||||
self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
|
||||
self.experts_cls: type[mk.FusedMoEExperts]
|
||||
if self.use_cutlass_mxfp4:
|
||||
logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE")
|
||||
self.experts_cls = CutlassExpertsMxfp4
|
||||
elif current_platform.is_xpu():
|
||||
self.mxfp4_backend = Mxfp4MoeBackend.XPU
|
||||
self.experts_cls = XPUExpertsMxFp4
|
||||
logger.info_once("Using XPUExpertsMxFp4 for MXFP4 MoE on XPU platform")
|
||||
else:
|
||||
logger.info_once("Using MarlinExperts for MXFP4 MoE")
|
||||
self.experts_cls = MarlinExperts
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.num_experts = num_experts
|
||||
layer.params_dtype = params_dtype
|
||||
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
# 2 fp4 items are packed in the input dimension
|
||||
hidden_size // 2,
|
||||
requires_grad=False,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_packed", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
# 2 fp4 items are packed in the input dimension
|
||||
intermediate_size_per_partition // 2,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_packed", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
# 2 fp4 items are packed in the input dimension
|
||||
hidden_size // self.group_size,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
# 2 fp4 items are packed in the input dimension
|
||||
intermediate_size_per_partition // self.group_size,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: torch.nn.Module
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
if self.use_cutlass_mxfp4:
|
||||
# W4A4: both weights and activations quantized to MXFP4
|
||||
return mxfp4_moe_quant_config(
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
)
|
||||
else:
|
||||
# W4A16: weight-only via Marlin
|
||||
return make_mxfp4_moe_quant_config(
|
||||
mxfp4_backend=self.mxfp4_backend,
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
layer.w13_weight = torch.nn.Parameter(
|
||||
layer.w13_weight_packed.data, requires_grad=False
|
||||
)
|
||||
delattr(layer, "w13_weight_packed")
|
||||
|
||||
layer.w2_weight = torch.nn.Parameter(
|
||||
layer.w2_weight_packed.data, requires_grad=False
|
||||
)
|
||||
delattr(layer, "w2_weight_packed")
|
||||
|
||||
if self.use_cutlass_mxfp4:
|
||||
# Swizzle weight scales from flat checkpoint layout [E, N, K//32]
|
||||
# to CUTLASS tiled layout [E, numMTiles*numKTiles*512].
|
||||
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
|
||||
swizzle_mxfp4_scales,
|
||||
)
|
||||
|
||||
E = layer.w13_weight_scale.shape[0]
|
||||
w13_N = layer.w13_weight_scale.shape[1]
|
||||
w13_scale_K = layer.w13_weight_scale.shape[2]
|
||||
w13_K = w13_scale_K * 32
|
||||
|
||||
w2_M = layer.w2_weight_scale.shape[1]
|
||||
w2_scale_N = layer.w2_weight_scale.shape[2]
|
||||
w2_N = w2_scale_N * 32
|
||||
|
||||
swizzled_w13 = []
|
||||
swizzled_w2 = []
|
||||
for e_idx in range(E):
|
||||
s13 = layer.w13_weight_scale[e_idx]
|
||||
sw13 = swizzle_mxfp4_scales(s13, w13_N, w13_K)
|
||||
swizzled_w13.append(sw13.reshape(w13_N, w13_scale_K))
|
||||
s2 = layer.w2_weight_scale[e_idx]
|
||||
sw2 = swizzle_mxfp4_scales(s2, w2_M, w2_N)
|
||||
swizzled_w2.append(sw2.reshape(w2_M, w2_scale_N))
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
torch.stack(swizzled_w13), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
torch.stack(swizzled_w2), requires_grad=False
|
||||
)
|
||||
elif current_platform.is_xpu():
|
||||
pass
|
||||
else:
|
||||
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_moe_fp4_layer_for_marlin(layer)
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
if self.moe_quant_config is not None:
|
||||
self.moe_kernel = make_mxfp4_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
experts_cls=self.experts_cls,
|
||||
mxfp4_backend=self.mxfp4_backend,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
+313
@@ -0,0 +1,313 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoeWeightScaleSupported,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
|
||||
convert_to_nvfp4_moe_kernel_format,
|
||||
is_global_sf_supported_for_nvfp4_backend,
|
||||
make_nvfp4_moe_kernel,
|
||||
make_nvfp4_moe_quant_config,
|
||||
select_nvfp4_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kNvfp4Dynamic,
|
||||
kNvfp4Static,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
|
||||
def __init__(
|
||||
self,
|
||||
moe: FusedMoEConfig,
|
||||
layer_name: str | None = None,
|
||||
use_a16: bool = False,
|
||||
):
|
||||
super().__init__(moe)
|
||||
self.group_size = 16
|
||||
|
||||
# Select experts implementation.
|
||||
self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
|
||||
config=self.moe,
|
||||
weight_key=kNvfp4Static,
|
||||
activation_key=None if use_a16 else kNvfp4Dynamic,
|
||||
)
|
||||
|
||||
self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
|
||||
self.nvfp4_backend
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.num_experts = num_experts
|
||||
layer.params_dtype = params_dtype
|
||||
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
|
||||
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
# 2 fp4 items are packed in the input dimension
|
||||
hidden_size // 2,
|
||||
requires_grad=False,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_packed", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
# 2 fp4 items are packed in the input dimension
|
||||
intermediate_size_per_partition // 2,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_packed", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
# Weight Scales
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
# 2 fp4 items are packed in the input dimension
|
||||
hidden_size // self.group_size,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
# 2 fp4 items are packed in the input dimension
|
||||
intermediate_size_per_partition // self.group_size,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
|
||||
)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
# Weight Global Scales
|
||||
w13_weight_scale_2 = torch.nn.Parameter(
|
||||
torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_global_scale", w13_weight_scale_2)
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale_2, extra_weight_attrs)
|
||||
|
||||
w2_weight_scale_2 = torch.nn.Parameter(
|
||||
torch.empty(num_experts, dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w2_weight_global_scale", w2_weight_scale_2)
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
|
||||
)
|
||||
set_weight_attrs(w2_weight_scale_2, extra_weight_attrs)
|
||||
|
||||
# Input Global Scales
|
||||
w13_input_scale = torch.nn.Parameter(
|
||||
torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_input_global_scale", w13_input_scale)
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
|
||||
)
|
||||
set_weight_attrs(w13_input_scale, extra_weight_attrs)
|
||||
|
||||
w2_input_scale = torch.nn.Parameter(
|
||||
torch.empty(num_experts, dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w2_input_global_scale", w2_input_scale)
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
|
||||
)
|
||||
set_weight_attrs(w2_input_scale, extra_weight_attrs)
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
"""
|
||||
Convert NVFP4 MoE weights into kernel format and setup the kernel.
|
||||
"""
|
||||
# NOTE(rob): wN_weight_packed -> wN_weight is because ModularKernelMethod
|
||||
# requires this naming convention. However, the name change breaks
|
||||
# reloading because the state dict no longer matches disk. Once we
|
||||
# remove MKM, we should revert this change to ensure compatibility.
|
||||
layer.w13_weight = torch.nn.Parameter(
|
||||
layer.w13_weight_packed.data, requires_grad=False
|
||||
)
|
||||
delattr(layer, "w13_weight_packed")
|
||||
|
||||
layer.w2_weight = torch.nn.Parameter(
|
||||
layer.w2_weight_packed.data, requires_grad=False
|
||||
)
|
||||
delattr(layer, "w2_weight_packed")
|
||||
|
||||
# Use a single gscale for w13.
|
||||
if self.moe.is_act_and_mul and not torch.allclose(
|
||||
layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1]
|
||||
):
|
||||
logger.warning_once(
|
||||
"w1_weight_global_scale must match w3_weight_global_scale. "
|
||||
"Accuracy may be affected.",
|
||||
)
|
||||
w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous()
|
||||
|
||||
# Shuffle weights into the NvFp4 kernel format.
|
||||
(
|
||||
w13,
|
||||
w13_scale,
|
||||
w13_scale_2,
|
||||
a13_scale,
|
||||
w2,
|
||||
w2_scale,
|
||||
w2_scale_2,
|
||||
a2_scale,
|
||||
) = convert_to_nvfp4_moe_kernel_format(
|
||||
nvfp4_backend=self.nvfp4_backend,
|
||||
layer=layer,
|
||||
w13=layer.w13_weight,
|
||||
w13_scale=layer.w13_weight_scale,
|
||||
w13_scale_2=(1.0 / w13_weight_global_scale),
|
||||
a13_scale=(1.0 / layer.w13_input_global_scale),
|
||||
w2=layer.w2_weight,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
w2_scale_2=(1.0 / layer.w2_weight_global_scale),
|
||||
a2_scale=(1.0 / layer.w2_input_global_scale),
|
||||
is_act_and_mul=self.moe.is_act_and_mul,
|
||||
)
|
||||
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w13_weight_scale", w13_scale)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
replace_parameter(layer, "w2_weight_scale", w2_scale)
|
||||
layer.w13_weight_scale_2 = w13_scale_2
|
||||
layer.w2_weight_scale_2 = w2_scale_2
|
||||
layer.w13_input_scale = a13_scale
|
||||
layer.w2_input_scale = a2_scale
|
||||
|
||||
# Setup modular kernel.
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_nvfp4_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
experts_cls=self.experts_cls,
|
||||
backend=self.nvfp4_backend,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
self.moe_kernel.fused_experts.process_weights_after_loading(layer)
|
||||
|
||||
def maybe_make_prepare_finalize(
|
||||
self,
|
||||
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> mk.FusedMoEPrepareAndFinalizeModular | None:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} uses the new modular kernel initialization "
|
||||
"logic. This function should not be called."
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
|
||||
return make_nvfp4_moe_quant_config(
|
||||
backend=self.nvfp4_backend,
|
||||
w13_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
w13_scale_2=layer.w13_weight_scale_2,
|
||||
w2_scale_2=layer.w2_weight_scale_2,
|
||||
a13_scale=layer.w13_input_scale,
|
||||
a2_scale=layer.w2_input_scale,
|
||||
swiglu_limit=getattr(layer, "swiglu_limit", None),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
num_expert_group=layer.num_expert_group,
|
||||
topk_group=layer.topk_group,
|
||||
e_score_correction_bias=layer.e_score_correction_bias,
|
||||
routed_scaling_factor=layer.routed_scaling_factor,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
+242
@@ -0,0 +1,242 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
)
|
||||
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoeWeightScaleSupported,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.w4a8 import (
|
||||
convert_to_w4a8_moe_kernel_format,
|
||||
make_w4a8_moe_kernel,
|
||||
make_w4a8_moe_quant_config,
|
||||
select_w4a8_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
|
||||
|
||||
|
||||
class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
|
||||
def __init__(
|
||||
self,
|
||||
weight_quant: QuantizationArgs,
|
||||
input_quant: QuantizationArgs,
|
||||
moe: FusedMoEConfig,
|
||||
layer_name: str | None = None,
|
||||
):
|
||||
super().__init__(moe)
|
||||
self.weight_quant = weight_quant
|
||||
self.input_quant = input_quant
|
||||
|
||||
self.group_size = self.weight_quant.group_size
|
||||
self.num_bits = self.weight_quant.num_bits
|
||||
self.packed_factor = 32 // self.num_bits
|
||||
|
||||
assert self.weight_quant.symmetric, (
|
||||
"Only symmetric quantization is supported for W4A8 MoE"
|
||||
)
|
||||
assert self.weight_quant.actorder != "group"
|
||||
assert self.group_size == 128, "Only group size 128 supported for W4A8 MoE"
|
||||
|
||||
self.w4a8_backend, self.experts_cls = select_w4a8_moe_backend(
|
||||
config=self.moe,
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.num_experts = num_experts
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
# requirement for CUTLASS reorder_tensor
|
||||
assert hidden_size % 256 == 0, f"{hidden_size=} must be divisible by 256"
|
||||
assert intermediate_size_per_partition % 256 == 0, (
|
||||
f"{intermediate_size_per_partition=} must be divisible by 256"
|
||||
)
|
||||
# storage type, pack 8xint4 into int32
|
||||
params_dtype = torch.int32
|
||||
|
||||
# WEIGHTS
|
||||
w13_weight_packed = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // self.packed_factor,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_packed", w13_weight_packed)
|
||||
set_weight_attrs(w13_weight_packed, extra_weight_attrs)
|
||||
|
||||
w2_weight_packed = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // self.packed_factor,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_packed", w2_weight_packed)
|
||||
set_weight_attrs(w2_weight_packed, extra_weight_attrs)
|
||||
|
||||
# SCALES
|
||||
# weight_scale refers to the group-wise scales
|
||||
# they are initially loaded as bf16, we will convert to fp8
|
||||
# after loading
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // self.group_size,
|
||||
dtype=layer.orig_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // self.group_size,
|
||||
dtype=layer.orig_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
# Add PER-GROUP quantization for RoutedExperts.weight_loader.
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
# weight shapes
|
||||
w2_weight_shape = torch.nn.Parameter(
|
||||
torch.empty(num_experts, 2), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w2_weight_shape", w2_weight_shape)
|
||||
set_weight_attrs(w2_weight_shape, extra_weight_attrs)
|
||||
w13_weight_shape = torch.nn.Parameter(
|
||||
torch.empty(num_experts, 2), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w13_weight_shape", w13_weight_shape)
|
||||
set_weight_attrs(w13_weight_shape, extra_weight_attrs)
|
||||
|
||||
w13_weight_chan_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_chan_scale", w13_weight_chan_scale)
|
||||
|
||||
w2_weight_chan_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, hidden_size, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_chan_scale", w2_weight_chan_scale)
|
||||
|
||||
# don't use input scales
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
(
|
||||
w13_weight_packed,
|
||||
w2_weight_packed,
|
||||
w13_weight_scale,
|
||||
w2_weight_scale,
|
||||
w13_weight_chan_scale,
|
||||
w2_weight_chan_scale,
|
||||
b_strides1,
|
||||
b_strides2,
|
||||
) = convert_to_w4a8_moe_kernel_format(
|
||||
w13_weight_packed=layer.w13_weight_packed,
|
||||
w2_weight_packed=layer.w2_weight_packed,
|
||||
w13_weight_scale=layer.w13_weight_scale,
|
||||
w2_weight_scale=layer.w2_weight_scale,
|
||||
)
|
||||
|
||||
replace_parameter(layer, "w13_weight_packed", w13_weight_packed)
|
||||
replace_parameter(layer, "w2_weight_packed", w2_weight_packed)
|
||||
replace_parameter(layer, "w13_weight_scale", w13_weight_scale)
|
||||
replace_parameter(layer, "w2_weight_scale", w2_weight_scale)
|
||||
replace_parameter(layer, "w13_weight_chan_scale", w13_weight_chan_scale)
|
||||
replace_parameter(layer, "w2_weight_chan_scale", w2_weight_chan_scale)
|
||||
|
||||
self.b_strides1 = b_strides1
|
||||
self.b_strides2 = b_strides2
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
if self.moe_quant_config is not None:
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_w4a8_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
experts_cls=self.experts_cls,
|
||||
b_strides1=self.b_strides1,
|
||||
b_strides2=self.b_strides2,
|
||||
group_size=self.group_size,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
|
||||
return make_w4a8_moe_quant_config(
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
g1_alphas=layer.w13_weight_chan_scale,
|
||||
g2_alphas=layer.w2_weight_chan_scale,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
hidden_states=x,
|
||||
w1=layer.w13_weight_packed,
|
||||
w2=layer.w2_weight_packed,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
|
||||
@property
|
||||
def supports_eplb(self) -> bool:
|
||||
return False
|
||||
+315
@@ -0,0 +1,315 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
)
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.w4a8_int8 import (
|
||||
convert_to_w4a8_int8_moe_format,
|
||||
make_w4a8_int8_moe_kernel,
|
||||
make_w4a8_int8_moe_quant_config,
|
||||
select_w4a8_int8_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
QuantKey,
|
||||
ScaleDesc,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
|
||||
"""
|
||||
CPU-only MoE method using dynamic 4-bit matmul kernels on Arm Platform
|
||||
- Weights: int4 (stored as int8 values in [-8,7], packed to uint8 nibbles)
|
||||
- Scales: Fp32 for Channelwise , bf16 for groupwise quantization
|
||||
- Bias: Same data type as original weights
|
||||
- Activations: FP32/Bf16 dynamic per-token (A8 Int),
|
||||
quantized inside the kernel
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weight_quant: QuantizationArgs,
|
||||
input_quant: QuantizationArgs,
|
||||
moe: FusedMoEConfig,
|
||||
layer_name: str | None = None,
|
||||
):
|
||||
super().__init__(moe)
|
||||
self.has_bias = self.moe.has_bias
|
||||
self.weight_quant = weight_quant
|
||||
self.input_quant = input_quant
|
||||
self.static_input_scales = False # always dynamic per token
|
||||
# Weight can be channel-wise (group_size=None) or group-wise
|
||||
self.group_size = (
|
||||
weight_quant.group_size if (weight_quant.group_size is not None) else -1
|
||||
)
|
||||
|
||||
# make sure group size is valid
|
||||
if self.group_size != -1 and (
|
||||
moe.hidden_dim % self.group_size != 0
|
||||
or moe.intermediate_size_per_partition % self.group_size != 0
|
||||
):
|
||||
raise ValueError(
|
||||
f"Group size ({self.group_size}) must evenly divide both "
|
||||
f"hidden size ({moe.hidden_dim}) and intermediate size per "
|
||||
f"partition ({moe.intermediate_size_per_partition})."
|
||||
)
|
||||
|
||||
# Construct QuantKey for weights from QuantizationArgs
|
||||
# W4A8 INT4: 4-bit weights (stored as int8), static quantization
|
||||
if self.group_size == -1:
|
||||
# Channel-wise quantization
|
||||
group_shape = GroupShape(-1, 1)
|
||||
scale_dtype = torch.float32
|
||||
else:
|
||||
# Group-wise quantization
|
||||
group_shape = GroupShape(1, self.group_size)
|
||||
scale_dtype = torch.bfloat16
|
||||
|
||||
weight_scale_desc = ScaleDesc(scale_dtype, static=True, group_shape=group_shape)
|
||||
weight_key = QuantKey(torch.int8, weight_scale_desc, symmetric=True)
|
||||
|
||||
self.backend, self.experts_cls = select_w4a8_int8_moe_backend(
|
||||
moe,
|
||||
weight_key,
|
||||
activation_key=None, # unquantized inputs
|
||||
)
|
||||
|
||||
# ---- parameter creation ----
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
# Shapes per local rank (TP/EP):
|
||||
# w13: [E, 2*I_local, H] int8 (int4 values in [-8,7])
|
||||
# w2 : [E, H, I_local] int8
|
||||
# Scales:
|
||||
# channel-wise: group_size=-1 -> per-output-row, single scale per row
|
||||
# group-wise : group_size=g ->
|
||||
# per-output-row, (in_features/g) scales
|
||||
|
||||
E = num_experts
|
||||
H = hidden_size
|
||||
IN = intermediate_size_per_partition
|
||||
g = self.group_size
|
||||
|
||||
# Per-row scale columns
|
||||
def _n_scale_cols(in_features: int) -> int:
|
||||
return 1 if g == -1 else (in_features // g)
|
||||
|
||||
# Register unpacked int4-as-int8 weights the loader will fill.
|
||||
w13 = torch.nn.Parameter(
|
||||
torch.empty(E, 2 * IN, H, dtype=torch.int8), requires_grad=False
|
||||
)
|
||||
set_weight_attrs(w13, extra_weight_attrs)
|
||||
layer.register_parameter("w13_weight", w13)
|
||||
|
||||
w2 = torch.nn.Parameter(
|
||||
torch.empty(E, H, IN, dtype=torch.int8), requires_grad=False
|
||||
)
|
||||
set_weight_attrs(w2, extra_weight_attrs)
|
||||
layer.register_parameter("w2_weight", w2)
|
||||
|
||||
# Register scales
|
||||
# KleidiAI groupwise kernels accepts float32 scales
|
||||
# KleidiAI groupwise kernels accepts bfloat16 scales
|
||||
scale_dtype = torch.float32 if g == -1 else torch.bfloat16
|
||||
|
||||
w13_s = torch.nn.Parameter(
|
||||
torch.ones(E, 2 * IN, _n_scale_cols(H), dtype=scale_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(
|
||||
w13_s,
|
||||
{"quant_method": "channel" if g == -1 else "group", **extra_weight_attrs},
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_s)
|
||||
|
||||
w2_s = torch.nn.Parameter(
|
||||
torch.ones(E, H, _n_scale_cols(IN), dtype=scale_dtype), requires_grad=False
|
||||
)
|
||||
set_weight_attrs(
|
||||
w2_s,
|
||||
{"quant_method": "channel" if g == -1 else "group", **extra_weight_attrs},
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_s)
|
||||
|
||||
if self.has_bias:
|
||||
w13_bias = torch.nn.Parameter(
|
||||
torch.zeros(E, 2 * IN, dtype=params_dtype), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w13_bias", w13_bias)
|
||||
set_weight_attrs(w13_bias, extra_weight_attrs)
|
||||
|
||||
w2_bias = torch.nn.Parameter(
|
||||
torch.zeros(num_experts, hidden_size, dtype=params_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_bias", w2_bias)
|
||||
set_weight_attrs(w2_bias, extra_weight_attrs)
|
||||
|
||||
# Placeholders for packed weights (will be replaced after packing)
|
||||
layer.register_parameter(
|
||||
"w13_weight_packed", torch.nn.Parameter(torch.empty(0), requires_grad=False)
|
||||
)
|
||||
set_weight_attrs(layer.w13_weight_packed, extra_weight_attrs)
|
||||
|
||||
layer.register_parameter(
|
||||
"w2_weight_packed", torch.nn.Parameter(torch.empty(0), requires_grad=False)
|
||||
)
|
||||
set_weight_attrs(layer.w2_weight_packed, extra_weight_attrs)
|
||||
|
||||
# dims for 4 bit fused matmuls
|
||||
layer.w13_in_features = H
|
||||
layer.w13_out_features = 2 * IN
|
||||
layer.w2_in_features = IN
|
||||
layer.w2_out_features = H
|
||||
layer.group_size = g
|
||||
|
||||
# post-load packing to dyn-4bit KleidiAI kernel's format
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
# Use oracle to pack weights.
|
||||
w13_packed, w2_packed, w13_weight_scale, w2_weight_scale, w13_bias, w2_bias = (
|
||||
convert_to_w4a8_int8_moe_format(
|
||||
w13_weight=layer.w13_weight,
|
||||
w2_weight=layer.w2_weight,
|
||||
w13_weight_scale=layer.w13_weight_scale,
|
||||
w2_weight_scale=layer.w2_weight_scale,
|
||||
group_size=self.group_size,
|
||||
w13_bias=layer.w13_bias if self.has_bias else None,
|
||||
w2_bias=layer.w2_bias if self.has_bias else None,
|
||||
)
|
||||
)
|
||||
|
||||
# Register packed weights as parameters
|
||||
replace_parameter(
|
||||
layer,
|
||||
"w13_weight_packed",
|
||||
torch.nn.Parameter(w13_packed, requires_grad=False),
|
||||
)
|
||||
replace_parameter(
|
||||
layer,
|
||||
"w2_weight_packed",
|
||||
torch.nn.Parameter(w2_packed, requires_grad=False),
|
||||
)
|
||||
|
||||
replace_parameter(
|
||||
layer, "w13_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False)
|
||||
)
|
||||
replace_parameter(
|
||||
layer, "w2_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False)
|
||||
)
|
||||
replace_parameter(
|
||||
layer,
|
||||
"w13_weight_scale",
|
||||
torch.nn.Parameter(w13_weight_scale, requires_grad=False),
|
||||
)
|
||||
replace_parameter(
|
||||
layer,
|
||||
"w2_weight_scale",
|
||||
torch.nn.Parameter(w2_weight_scale, requires_grad=False),
|
||||
)
|
||||
|
||||
if self.has_bias:
|
||||
replace_parameter(
|
||||
layer,
|
||||
"w13_bias",
|
||||
torch.nn.Parameter(w13_bias, requires_grad=False),
|
||||
)
|
||||
if self.has_bias:
|
||||
replace_parameter(
|
||||
layer,
|
||||
"w2_bias",
|
||||
torch.nn.Parameter(w2_bias, requires_grad=False),
|
||||
)
|
||||
|
||||
quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert quant_config is not None
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_w4a8_int8_moe_kernel(
|
||||
moe_quant_config=quant_config,
|
||||
moe_config=self.moe,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: torch.nn.Module
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
# Determine block shape from group_size
|
||||
# group_size=-1 means channel-wise: (-1, 1)
|
||||
# group_size=N means group-wise: (1, N)
|
||||
block_shape = (-1, 1) if self.group_size == -1 else (1, self.group_size)
|
||||
|
||||
return make_w4a8_int8_moe_quant_config(block_shape=block_shape)
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
x,
|
||||
layer.w13_weight_packed,
|
||||
layer.w2_weight_packed,
|
||||
router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
num_expert_group=layer.num_expert_group,
|
||||
topk_group=layer.topk_group,
|
||||
e_score_correction_bias=layer.e_score_correction_bias,
|
||||
routed_scaling_factor=layer.routed_scaling_factor,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
x,
|
||||
layer.w13_weight_packed,
|
||||
layer.w2_weight_packed,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
+418
@@ -0,0 +1,418 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
QuantizationStrategy,
|
||||
)
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.distributed import get_tensor_model_parallel_world_size
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoeWeightScaleSupported,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
convert_to_fp8_moe_kernel_format,
|
||||
make_fp8_moe_kernel,
|
||||
make_fp8_moe_quant_config,
|
||||
select_fp8_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
process_fp8_input_tensor_strategy_moe,
|
||||
process_fp8_weight_tensor_strategy_moe,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kFp8Dynamic128Sym,
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8Static128BlockSym,
|
||||
kFp8StaticChannelSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
normalize_e4m3fn_to_e4m3fnuz,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
|
||||
"""W8A8 FP8 MoE quantization using compressed tensors."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weight_quant: QuantizationArgs,
|
||||
input_quant: QuantizationArgs,
|
||||
moe: FusedMoEConfig,
|
||||
layer_name: str | None = None,
|
||||
):
|
||||
super().__init__(moe)
|
||||
self.weight_quant = weight_quant
|
||||
self.input_quant = input_quant
|
||||
|
||||
per_tensor = (
|
||||
self.weight_quant.strategy == QuantizationStrategy.TENSOR
|
||||
and self.input_quant.strategy == QuantizationStrategy.TENSOR
|
||||
)
|
||||
per_channel = (
|
||||
self.weight_quant.strategy == QuantizationStrategy.CHANNEL
|
||||
and self.input_quant.strategy == QuantizationStrategy.TOKEN
|
||||
)
|
||||
if not (per_tensor or per_channel):
|
||||
assert self.weight_quant.strategy == QuantizationStrategy.BLOCK
|
||||
self.weight_block_size = self.weight_quant.block_structure
|
||||
assert self.weight_quant.dynamic is not None
|
||||
else:
|
||||
self.weight_block_size = None
|
||||
self.block_quant = self.weight_block_size is not None
|
||||
|
||||
self.static_input_scales = not self.input_quant.dynamic
|
||||
if self.static_input_scales and per_channel:
|
||||
raise ValueError(
|
||||
"For FP8 Fused MoE layer, we require either per tensor or "
|
||||
"channelwise, dynamic per token quantization."
|
||||
)
|
||||
|
||||
ct2vllm_weight = {
|
||||
QuantizationStrategy.CHANNEL: kFp8StaticChannelSym,
|
||||
QuantizationStrategy.TENSOR: kFp8StaticTensorSym,
|
||||
QuantizationStrategy.BLOCK: kFp8Static128BlockSym,
|
||||
}
|
||||
ct2vllm_act = {
|
||||
QuantizationStrategy.TOKEN: kFp8DynamicTokenSym,
|
||||
QuantizationStrategy.TENSOR: (
|
||||
kFp8StaticTensorSym if self.static_input_scales else kFp8Dynamic128Sym
|
||||
),
|
||||
}
|
||||
weight_key = ct2vllm_weight[self.weight_quant.strategy]
|
||||
if weight_key == kFp8Static128BlockSym:
|
||||
activation_key = kFp8Dynamic128Sym
|
||||
else:
|
||||
activation_key = ct2vllm_act[self.input_quant.strategy]
|
||||
|
||||
# Select Fp8 MoE backend
|
||||
self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
|
||||
config=self.moe,
|
||||
weight_key=weight_key,
|
||||
activation_key=activation_key,
|
||||
allow_vllm_cutlass=True,
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.num_experts = num_experts
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
params_dtype = torch.float8_e4m3fn
|
||||
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
|
||||
|
||||
if self.block_quant:
|
||||
assert self.weight_block_size is not None
|
||||
layer.weight_block_size = self.weight_block_size
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
block_n, block_k = (
|
||||
self.weight_block_size[0],
|
||||
self.weight_block_size[1],
|
||||
)
|
||||
# NOTE: To ensure proper alignment of the block-wise quantization
|
||||
# scales, the output_size of the weights for both the gate and up
|
||||
# layers must be divisible by block_n.
|
||||
# Required by column parallel or enabling merged weights
|
||||
if intermediate_size_per_partition % block_n != 0:
|
||||
raise ValueError(
|
||||
f"The output_size of gate's and up's weight = "
|
||||
f"{intermediate_size_per_partition} is not divisible by "
|
||||
f"weight quantization block_n = {block_n}."
|
||||
)
|
||||
if tp_size > 1 and intermediate_size_per_partition % block_k != 0:
|
||||
# Required by row parallel
|
||||
raise ValueError(
|
||||
f"The input_size of down's weight = "
|
||||
f"{intermediate_size_per_partition} is not divisible by "
|
||||
f"weight quantization block_k = {block_k}."
|
||||
)
|
||||
|
||||
# WEIGHTS
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
hidden_size,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
# WEIGHT_SCALES
|
||||
if self.weight_quant.strategy == QuantizationStrategy.TENSOR:
|
||||
# For gated MoE, allocate 2 scales for w1 and w3 respectively.
|
||||
# They will be combined to a single scale after weight loading.
|
||||
# For non-gated MoE, allocate 1 scale for w13.
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, w13_num_shards, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
# Add PER-TENSOR quantization for RoutedExperts.weight_loader.
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
elif self.weight_quant.strategy == QuantizationStrategy.CHANNEL:
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
num_experts,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
1,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, hidden_size, 1, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
# Add PER-CHANNEL quantization for RoutedExperts.weight_loader.
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
elif self.weight_quant.strategy == QuantizationStrategy.BLOCK:
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
num_experts,
|
||||
w13_num_shards
|
||||
* ((intermediate_size_per_partition + block_n - 1) // block_n),
|
||||
(hidden_size + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
num_experts,
|
||||
(hidden_size + block_n - 1) // block_n,
|
||||
(intermediate_size_per_partition + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
# Add PER-CHANNEL quantization for RoutedExperts.weight_loader.
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
# INPUT_SCALES
|
||||
if self.static_input_scales:
|
||||
w13_input_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w13_input_scale", w13_input_scale)
|
||||
set_weight_attrs(w13_input_scale, extra_weight_attrs)
|
||||
|
||||
w2_input_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w2_input_scale", w2_input_scale)
|
||||
set_weight_attrs(w2_input_scale, extra_weight_attrs)
|
||||
else:
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
# Allow for accessing weights and scales in standard way.
|
||||
w13 = layer.w13_weight
|
||||
w2 = layer.w2_weight
|
||||
w13_scale = layer.w13_weight_scale
|
||||
w2_scale = layer.w2_weight_scale
|
||||
w13_input_scale = layer.w13_input_scale
|
||||
w2_input_scale = layer.w2_input_scale
|
||||
|
||||
# MI300x and MI325x use FNUZ format for FP8. Convert if needed.
|
||||
if current_platform.is_fp8_fnuz():
|
||||
w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz(
|
||||
w13, w13_scale, w13_input_scale
|
||||
)
|
||||
w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
|
||||
w2, w2_scale, w2_input_scale
|
||||
)
|
||||
|
||||
# Per tensor kernels require single activation scale. Use the max.
|
||||
if self.static_input_scales:
|
||||
assert self.input_quant.strategy == QuantizationStrategy.TENSOR
|
||||
assert w13_input_scale is not None and w2_input_scale is not None
|
||||
w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
|
||||
w13_input_scale, w2_input_scale
|
||||
)
|
||||
replace_parameter(layer, "w13_input_scale", w13_input_scale)
|
||||
replace_parameter(layer, "w2_input_scale", w2_input_scale)
|
||||
|
||||
# Per-tensor kernels use a single scale, for W13, but on disk there
|
||||
# is a separate scale for W1 and W3. Requantize with the max scale.
|
||||
if self.weight_quant.strategy == QuantizationStrategy.TENSOR:
|
||||
w13, w13_scale = process_fp8_weight_tensor_strategy_moe(
|
||||
w13,
|
||||
w13_scale,
|
||||
shard_size=layer.intermediate_size_per_partition,
|
||||
num_experts=layer.local_num_experts,
|
||||
is_act_and_mul=self.moe.is_act_and_mul,
|
||||
)
|
||||
|
||||
w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
|
||||
fp8_backend=self.fp8_backend,
|
||||
layer=layer,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
w13_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
w13_input_scale=w13_input_scale,
|
||||
w2_input_scale=w2_input_scale,
|
||||
)
|
||||
|
||||
# Replace parameters with updated versions. Note that this helper
|
||||
# function ensures the replacement is compatible with RL weight reloads.
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
replace_parameter(layer, "w13_weight_scale", w13_scale)
|
||||
replace_parameter(layer, "w2_weight_scale", w2_scale)
|
||||
|
||||
# Setup modular kernel for TP case and naive DP/EP case.
|
||||
# In non-naive DP/EP case, we will create a ModularKernelMethod.
|
||||
# TODO(rob): unify these so FP8MoEMethod owns the ModularKernel
|
||||
# in both cases.
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
if self.moe_quant_config:
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_fp8_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
fp8_backend=self.fp8_backend,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def maybe_make_prepare_finalize(
|
||||
self,
|
||||
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> mk.FusedMoEPrepareAndFinalizeModular | None:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} uses the new modular kernel initialization "
|
||||
"logic. This function should not be called."
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
|
||||
is_per_token = self.input_quant.strategy == QuantizationStrategy.TOKEN
|
||||
return make_fp8_moe_quant_config(
|
||||
fp8_backend=self.fp8_backend,
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
a1_scale=getattr(layer, "w13_input_scale", None),
|
||||
a2_scale=getattr(layer, "w2_input_scale", None),
|
||||
per_act_token_quant=is_per_token,
|
||||
per_out_ch_quant=is_per_token,
|
||||
block_shape=self.weight_block_size,
|
||||
swiglu_limit=getattr(layer, "swiglu_limit", None),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
num_expert_group=layer.num_expert_group,
|
||||
topk_group=layer.topk_group,
|
||||
e_score_correction_bias=layer.e_score_correction_bias,
|
||||
routed_scaling_factor=layer.routed_scaling_factor,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
|
||||
@property
|
||||
def supports_eplb(self) -> bool:
|
||||
return True
|
||||
+233
@@ -0,0 +1,233 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
QuantizationStrategy,
|
||||
)
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoeWeightScaleSupported,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.int8 import (
|
||||
convert_to_int8_moe_kernel_format,
|
||||
make_int8_moe_kernel,
|
||||
make_int8_moe_quant_config,
|
||||
select_int8_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kInt8DynamicTokenSym,
|
||||
kInt8StaticChannelSym,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
|
||||
"""W8A8 Int8 MoE quantization using compressed tensors."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weight_quant: QuantizationArgs,
|
||||
input_quant: QuantizationArgs,
|
||||
moe: FusedMoEConfig,
|
||||
layer_name: str | None = None,
|
||||
):
|
||||
super().__init__(moe)
|
||||
self.weight_quant = weight_quant
|
||||
self.input_quant = input_quant
|
||||
|
||||
per_channel = (
|
||||
self.weight_quant.strategy == QuantizationStrategy.CHANNEL
|
||||
and self.input_quant.strategy == QuantizationStrategy.TOKEN
|
||||
)
|
||||
if not per_channel:
|
||||
raise ValueError(
|
||||
"For INT8 Fused MoE layers, we require channelwise, "
|
||||
"dynamic per token quantization. Found "
|
||||
f"{self.weight_quant}, {self.input_quant}"
|
||||
)
|
||||
|
||||
self.static_input_scales = not self.input_quant.dynamic
|
||||
if self.static_input_scales:
|
||||
raise ValueError(
|
||||
"For INT8 Fused MoE layers, we require channelwise, "
|
||||
"dynamic per token quantization. Found static input scales."
|
||||
)
|
||||
|
||||
# Select Int8 MoE backend.
|
||||
self.int8_backend, self.experts_cls = select_int8_moe_backend(
|
||||
config=self.moe,
|
||||
weight_key=kInt8StaticChannelSym,
|
||||
activation_key=kInt8DynamicTokenSym,
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
params_dtype = torch.int8
|
||||
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
|
||||
|
||||
# WEIGHTS
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
hidden_size,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
# WEIGHT_SCALES
|
||||
assert self.weight_quant.strategy == QuantizationStrategy.CHANNEL
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
num_experts,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
1,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, hidden_size, 1, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
# Add PER-CHANNEL quantization for RoutedExperts.weight_loader.
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
# INPUT_SCALES
|
||||
assert not self.static_input_scales
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
w13, w2 = convert_to_int8_moe_kernel_format(
|
||||
int8_backend=self.int8_backend,
|
||||
w13=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
layer=layer,
|
||||
w13_scale=layer.w13_weight_scale,
|
||||
)
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_int8_moe_kernel(
|
||||
int8_backend=self.int8_backend,
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def maybe_make_prepare_finalize(
|
||||
self,
|
||||
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> mk.FusedMoEPrepareAndFinalizeModular | None:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} uses the new modular kernel initialization "
|
||||
"logic. This function should not be called."
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
|
||||
return make_int8_moe_quant_config(
|
||||
int8_backend=self.int8_backend,
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
a1_scale=layer.w13_input_scale,
|
||||
a2_scale=layer.w2_input_scale,
|
||||
per_act_token_quant=True,
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
num_expert_group=layer.num_expert_group,
|
||||
topk_group=layer.topk_group,
|
||||
e_score_correction_bias=layer.e_score_correction_bias,
|
||||
routed_scaling_factor=layer.routed_scaling_factor,
|
||||
)
|
||||
+217
@@ -0,0 +1,217 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoeWeightScaleSupported,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
convert_to_fp8_moe_kernel_format,
|
||||
make_fp8_moe_kernel,
|
||||
make_fp8_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
|
||||
select_mxfp8_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe import ( # noqa: E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
|
||||
MXFP8_BLOCK_SIZE,
|
||||
MXFP8_SCALE_DTYPE,
|
||||
MXFP8_VALUE_DTYPE,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
|
||||
|
||||
|
||||
class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
|
||||
"""Compressed-tensors MoE method for pre-quantized MXFP8 (W8A8) checkpoints.
|
||||
|
||||
Loads FP8 (E4M3) weights with E8M0 uint8 per-group scales (group_size=32)
|
||||
from checkpoint. Activations are dynamically quantized to MXFP8 at runtime.
|
||||
Supports FlashInfer TRT-LLM and Marlin backends (auto-selected).
|
||||
"""
|
||||
|
||||
def __init__(self, moe: FusedMoEConfig):
|
||||
super().__init__(moe)
|
||||
self.weight_block_size = [1, MXFP8_BLOCK_SIZE]
|
||||
self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.num_experts = num_experts
|
||||
layer.params_dtype = params_dtype
|
||||
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
|
||||
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
hidden_size,
|
||||
dtype=MXFP8_VALUE_DTYPE,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
dtype=MXFP8_VALUE_DTYPE,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
hidden_size // MXFP8_BLOCK_SIZE,
|
||||
dtype=MXFP8_SCALE_DTYPE,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // MXFP8_BLOCK_SIZE,
|
||||
dtype=MXFP8_SCALE_DTYPE,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
layer.weight_block_size = self.weight_block_size
|
||||
|
||||
w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
|
||||
fp8_backend=self.fp8_backend,
|
||||
layer=layer,
|
||||
w13=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
w13_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
w13_input_scale=layer.w13_input_scale,
|
||||
w2_input_scale=layer.w2_input_scale,
|
||||
)
|
||||
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
replace_parameter(layer, "w13_weight_scale", w13_scale)
|
||||
replace_parameter(layer, "w2_weight_scale", w2_scale)
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
if self.moe_quant_config is not None:
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_fp8_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
fp8_backend=self.fp8_backend,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: torch.nn.Module
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
return make_fp8_moe_quant_config(
|
||||
fp8_backend=self.fp8_backend,
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
a1_scale=layer.w13_input_scale,
|
||||
a2_scale=layer.w2_input_scale,
|
||||
block_shape=self.weight_block_size,
|
||||
swiglu_limit=getattr(layer, "swiglu_limit", None),
|
||||
gemm1_alpha=getattr(layer, "swiglu_alpha", None),
|
||||
gemm1_beta=getattr(layer, "swiglu_beta", None),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def maybe_make_prepare_finalize(
|
||||
self,
|
||||
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> mk.FusedMoEPrepareAndFinalizeModular | None:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} uses the new modular kernel "
|
||||
"initialization logic. This function should not be called."
|
||||
)
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
num_expert_group=layer.num_expert_group,
|
||||
topk_group=layer.topk_group,
|
||||
e_score_correction_bias=layer.e_score_correction_bias,
|
||||
routed_scaling_factor=layer.routed_scaling_factor,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
+283
@@ -0,0 +1,283 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
)
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEQuantConfig,
|
||||
int4_w4a16_moe_quant_config,
|
||||
int8_w8a16_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
|
||||
def __init__(
|
||||
self,
|
||||
weight_quant: QuantizationArgs,
|
||||
input_quant: QuantizationArgs | None,
|
||||
moe: FusedMoEConfig,
|
||||
layer_name: str | None = None,
|
||||
):
|
||||
super().__init__(moe)
|
||||
self.weight_quant = weight_quant
|
||||
self.input_quant = input_quant
|
||||
# Extract properties from weight_quant
|
||||
self.num_bits = weight_quant.num_bits
|
||||
self.packed_factor = 32 // weight_quant.num_bits
|
||||
self.strategy = weight_quant.strategy
|
||||
# channelwise is not supported by this kernel
|
||||
assert weight_quant.strategy == "group"
|
||||
self.group_size = weight_quant.group_size
|
||||
# grouped actorder isn't supported by this kernel
|
||||
assert weight_quant.actorder != "group"
|
||||
assert weight_quant.symmetric, (
|
||||
"Only symmetric quantization is supported for MoE"
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
# Will transpose the loaded weight along the
|
||||
# intermediate and hidden dim sizes. Will
|
||||
# shard for TP along the transposed dims
|
||||
extra_weight_attrs.update(
|
||||
{"is_transposed": True, "quant_method": self.strategy}
|
||||
)
|
||||
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size // self.packed_factor,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_packed", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition // self.packed_factor,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_packed", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
w2_scales_size = intermediate_size_per_partition
|
||||
|
||||
if self.strategy == "channel":
|
||||
num_groups_w2 = num_groups_w13 = 1
|
||||
self.group_size = -1
|
||||
else:
|
||||
if hidden_size % self.group_size != 0:
|
||||
raise ValueError(
|
||||
"CompressedTensors WNA16 MoE requires hidden_size "
|
||||
f"({hidden_size}) to be divisible by group_size "
|
||||
f"({self.group_size})."
|
||||
)
|
||||
if intermediate_size_per_partition % self.group_size != 0:
|
||||
raise ValueError(
|
||||
"CompressedTensors WNA16 MoE with static group scales "
|
||||
"requires the MoE intermediate size per tensor-parallel "
|
||||
f"partition ({intermediate_size_per_partition}) to be "
|
||||
f"divisible by group_size ({self.group_size}). Scale "
|
||||
"groups would otherwise cross TP shard boundaries; use a "
|
||||
"compatible TP size or enable expert parallelism."
|
||||
)
|
||||
num_groups_w2 = w2_scales_size // self.group_size
|
||||
num_groups_w13 = hidden_size // self.group_size
|
||||
|
||||
w13_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
num_experts,
|
||||
num_groups_w13,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_scale)
|
||||
set_weight_attrs(w13_scale, extra_weight_attrs)
|
||||
|
||||
w2_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_scale)
|
||||
set_weight_attrs(w2_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_scale, {"load_full_w2": False})
|
||||
|
||||
w2_weight_shape = torch.nn.Parameter(
|
||||
torch.empty(num_experts, 2), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w2_weight_shape", w2_weight_shape)
|
||||
set_weight_attrs(w2_weight_shape, extra_weight_attrs)
|
||||
w13_weight_shape = torch.nn.Parameter(
|
||||
torch.empty(num_experts, 2), requires_grad=False
|
||||
)
|
||||
|
||||
layer.register_parameter("w13_weight_shape", w13_weight_shape)
|
||||
set_weight_attrs(w13_weight_shape, extra_weight_attrs)
|
||||
|
||||
w13_g_idx = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_g_idx", w13_g_idx)
|
||||
set_weight_attrs(w13_g_idx, extra_weight_attrs)
|
||||
|
||||
w2_g_idx = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_g_idx", w2_g_idx)
|
||||
set_weight_attrs(w2_g_idx, extra_weight_attrs)
|
||||
|
||||
w13_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
|
||||
set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
|
||||
|
||||
w2_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
|
||||
set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
|
||||
|
||||
layer.a13_scale = None
|
||||
layer.a2_scale = None
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
# Reconfigure packed weights and scales to match moe_wna16 format
|
||||
layer.w13_weight_packed = torch.nn.Parameter(
|
||||
layer.w13_weight_packed.transpose(1, 2).contiguous().view(torch.uint8),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight_packed = torch.nn.Parameter(
|
||||
layer.w2_weight_packed.transpose(1, 2).contiguous().view(torch.uint8),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
layer.w13_weight_scale.transpose(1, 2).contiguous(), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
layer.w2_weight_scale.transpose(1, 2).contiguous(), requires_grad=False
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: torch.nn.Module
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
assert self.num_bits == 4 or self.num_bits == 8
|
||||
config_builder = (
|
||||
int4_w4a16_moe_quant_config
|
||||
if self.num_bits == 4
|
||||
else int8_w8a16_moe_quant_config
|
||||
)
|
||||
|
||||
return config_builder(
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
w1_zp=None,
|
||||
w2_zp=None,
|
||||
block_shape=[0, self.group_size],
|
||||
)
|
||||
|
||||
def select_gemm_impl(
|
||||
self,
|
||||
prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
|
||||
layer: torch.nn.Module,
|
||||
) -> mk.FusedMoEExpertsModular:
|
||||
if self.moe.is_lora_enabled:
|
||||
assert self.moe_quant_config is not None
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
|
||||
if HAS_TRITON:
|
||||
from vllm.model_executor.layers.fused_moe import TritonWNA16Experts
|
||||
|
||||
layer.w13_weight = layer.w13_weight_packed
|
||||
layer.w2_weight = layer.w2_weight_packed
|
||||
return TritonWNA16Experts(
|
||||
moe_config=self.moe, quant_config=self.moe_quant_config
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"TritonExperts requires Triton. "
|
||||
"Install triton or disable LoRA for MoE."
|
||||
)
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
from vllm.model_executor.layers.fused_moe import fused_experts
|
||||
|
||||
return fused_experts(
|
||||
x,
|
||||
layer.w13_weight_packed,
|
||||
layer.w2_weight_packed,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=layer.activation,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
quant_config=self.moe_quant_config,
|
||||
)
|
||||
|
||||
@property
|
||||
def supports_eplb(self) -> bool:
|
||||
return True
|
||||
+580
@@ -0,0 +1,580 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
)
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEExpertsModular,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
|
||||
WNA16MoEBackend,
|
||||
convert_to_wna16_moe_kernel_format,
|
||||
make_wna16_moe_kernel,
|
||||
make_wna16_moe_quant_config,
|
||||
select_wna16_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
|
||||
CompressedTensorsMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
|
||||
WNA16_SUPPORTED_TYPES_MAP,
|
||||
WNA16_ZP_SUPPORTED_TYPES_MAP,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
get_marlin_input_dtype,
|
||||
marlin_make_workspace_new,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
QuantKey,
|
||||
kInt4Static32GroupScale,
|
||||
kInt4StaticGroupScale,
|
||||
kInt8StaticGroupScale,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
|
||||
def __init__(
|
||||
self,
|
||||
weight_quant: QuantizationArgs,
|
||||
input_quant: QuantizationArgs | None,
|
||||
moe: FusedMoEConfig,
|
||||
layer_name: str | None = None,
|
||||
):
|
||||
super().__init__(moe)
|
||||
self.weight_quant = weight_quant
|
||||
self.input_quant = input_quant
|
||||
self.symmetric = weight_quant.symmetric
|
||||
# Extract properties from weight_quant
|
||||
self.num_bits = weight_quant.num_bits
|
||||
self.packed_factor = 32 // weight_quant.num_bits
|
||||
self.strategy = weight_quant.strategy
|
||||
self.group_size = weight_quant.group_size
|
||||
self.actorder = weight_quant.actorder
|
||||
|
||||
self.quant_type = (
|
||||
WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
|
||||
if self.symmetric
|
||||
else WNA16_ZP_SUPPORTED_TYPES_MAP[self.num_bits]
|
||||
)
|
||||
|
||||
self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
|
||||
|
||||
if self.num_bits == 4:
|
||||
if self.group_size == 32:
|
||||
scale = kInt4Static32GroupScale
|
||||
else:
|
||||
scale = kInt4StaticGroupScale
|
||||
elif self.num_bits == 8:
|
||||
scale = kInt8StaticGroupScale
|
||||
else:
|
||||
raise ValueError(
|
||||
"CompressedTensorsWNA16MarlinMoEMethod only supports int4 and int8 now."
|
||||
)
|
||||
|
||||
weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric)
|
||||
|
||||
# Select WNA16 MoE backend via oracle.
|
||||
self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
|
||||
config=self.moe,
|
||||
weight_key=weight_key,
|
||||
)
|
||||
|
||||
def get_weight_shape(
|
||||
self,
|
||||
weight_name: str,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
num_groups_w2: int | None = None,
|
||||
num_groups_w13: int | None = None,
|
||||
) -> tuple[int, int, int]:
|
||||
"""
|
||||
Get the shape of the weight based on the weight name, number of experts
|
||||
hidden size, intermediate size per partition, number of groups for w2,
|
||||
and number of groups for w13. Pass in num_groups_w2 and num_groups_w13
|
||||
for weight scales/zero_points.
|
||||
"""
|
||||
if weight_name in ("w13_scale", "w13_zp"):
|
||||
assert num_groups_w13 is not None, (
|
||||
"num_groups_w13 must be provided for weight scales/zero_points"
|
||||
)
|
||||
if weight_name in ("w2_scale", "w2_zp"):
|
||||
assert num_groups_w2 is not None, (
|
||||
"num_groups_w2 must be provided for weight scales/zero_points"
|
||||
)
|
||||
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
|
||||
is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
|
||||
shape_map = {
|
||||
"w13_weight": {
|
||||
"Flashinfer": (
|
||||
num_experts,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
hidden_size // self.packed_factor,
|
||||
),
|
||||
"Marlin": (
|
||||
num_experts,
|
||||
hidden_size // self.packed_factor,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
),
|
||||
},
|
||||
"w13_scale": {
|
||||
"Flashinfer": (
|
||||
num_experts,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
num_groups_w13,
|
||||
),
|
||||
"Marlin": (
|
||||
num_experts,
|
||||
num_groups_w13,
|
||||
w13_num_shards * intermediate_size_per_partition,
|
||||
),
|
||||
},
|
||||
"w13_zp": {
|
||||
"Marlin": (
|
||||
num_experts,
|
||||
num_groups_w13,
|
||||
w13_num_shards
|
||||
* intermediate_size_per_partition
|
||||
// self.packed_factor,
|
||||
),
|
||||
},
|
||||
"w2_weight": {
|
||||
"Flashinfer": (
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // self.packed_factor,
|
||||
),
|
||||
"Marlin": (
|
||||
num_experts,
|
||||
intermediate_size_per_partition // self.packed_factor,
|
||||
hidden_size,
|
||||
),
|
||||
},
|
||||
"w2_scale": {
|
||||
"Flashinfer": (num_experts, hidden_size, num_groups_w2),
|
||||
"Marlin": (num_experts, num_groups_w2, hidden_size),
|
||||
},
|
||||
"w2_zp": {
|
||||
"Marlin": (
|
||||
num_experts,
|
||||
num_groups_w2,
|
||||
hidden_size // self.packed_factor,
|
||||
),
|
||||
},
|
||||
}
|
||||
backend_key = "Flashinfer" if is_flashinfer else "Marlin"
|
||||
return shape_map[weight_name][backend_key]
|
||||
|
||||
@staticmethod
|
||||
def _w2_scale_sharding(
|
||||
actorder,
|
||||
group_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
intermediate_size_full: int,
|
||||
) -> tuple[bool, int, bool]:
|
||||
"""Decide how to shard w2 group scales across TP for WNA16 Marlin MoE.
|
||||
|
||||
Only ``actorder="group"`` permutes activations by ``g_idx`` at runtime
|
||||
and therefore needs the full-K (unsharded) w2 scales plus ``is_k_full``.
|
||||
``actorder="weight"``/``"static"`` (and ``None``) reorder weights at
|
||||
quantization time, so scales shard normally per TP rank.
|
||||
"""
|
||||
load_full_w2 = (actorder == "group") and group_size != -1
|
||||
w2_scales_size = (
|
||||
intermediate_size_full if load_full_w2 else intermediate_size_per_partition
|
||||
)
|
||||
is_k_full = (actorder != "group") or (
|
||||
intermediate_size_per_partition == intermediate_size_full
|
||||
)
|
||||
return load_full_w2, w2_scales_size, is_k_full
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
|
||||
|
||||
# Will transpose the loaded weight along the
|
||||
# intermediate and hidden dim sizes. Will
|
||||
# shard for TP along the transposed dims
|
||||
is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
|
||||
extra_weight_attrs.update(
|
||||
{"is_transposed": is_transposed, "quant_method": self.strategy}
|
||||
)
|
||||
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
*self.get_weight_shape(
|
||||
"w13_weight",
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
),
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_packed", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
*self.get_weight_shape(
|
||||
"w2_weight",
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
),
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_packed", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
load_full_w2, w2_scales_size, self.is_k_full = self._w2_scale_sharding(
|
||||
self.actorder,
|
||||
self.group_size,
|
||||
intermediate_size_per_partition,
|
||||
intermediate_size_full,
|
||||
)
|
||||
|
||||
if self.strategy == "channel":
|
||||
num_groups_w2 = num_groups_w13 = 1
|
||||
self.group_size = -1
|
||||
else:
|
||||
if hidden_size % self.group_size != 0:
|
||||
raise ValueError(
|
||||
"CompressedTensors WNA16 Marlin MoE requires hidden_size "
|
||||
f"({hidden_size}) to be divisible by group_size "
|
||||
f"({self.group_size})."
|
||||
)
|
||||
if (
|
||||
not load_full_w2
|
||||
and intermediate_size_per_partition % self.group_size != 0
|
||||
):
|
||||
raise ValueError(
|
||||
"CompressedTensors WNA16 Marlin MoE with static group "
|
||||
"scales requires the MoE intermediate size per "
|
||||
"tensor-parallel partition "
|
||||
f"({intermediate_size_per_partition}) to be divisible by "
|
||||
f"group_size ({self.group_size}). Scale groups would "
|
||||
"otherwise cross TP shard boundaries; use a compatible TP "
|
||||
"size or enable expert parallelism."
|
||||
)
|
||||
num_groups_w2 = w2_scales_size // self.group_size
|
||||
num_groups_w13 = hidden_size // self.group_size
|
||||
|
||||
layer.num_groups_w13 = num_groups_w13
|
||||
layer.num_groups_w2 = num_groups_w2
|
||||
|
||||
w13_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
*self.get_weight_shape(
|
||||
"w13_scale",
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
num_groups_w13=num_groups_w13,
|
||||
),
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_scale)
|
||||
set_weight_attrs(w13_scale, extra_weight_attrs)
|
||||
|
||||
w2_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
*self.get_weight_shape(
|
||||
"w2_scale",
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
num_groups_w2=num_groups_w2,
|
||||
),
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_scale)
|
||||
set_weight_attrs(w2_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2})
|
||||
|
||||
if not self.symmetric:
|
||||
w13_zp = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
*self.get_weight_shape(
|
||||
"w13_zp",
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
num_groups_w13=num_groups_w13,
|
||||
),
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_zero_point", w13_zp)
|
||||
set_weight_attrs(w13_zp, extra_weight_attrs)
|
||||
|
||||
w2_zp = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
*self.get_weight_shape(
|
||||
"w2_zp",
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
num_groups_w2=num_groups_w2,
|
||||
),
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_zero_point", w2_zp)
|
||||
set_weight_attrs(w2_zp, extra_weight_attrs)
|
||||
|
||||
w2_weight_shape = torch.nn.Parameter(
|
||||
torch.empty(num_experts, 2), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w2_weight_shape", w2_weight_shape)
|
||||
set_weight_attrs(w2_weight_shape, extra_weight_attrs)
|
||||
w13_weight_shape = torch.nn.Parameter(
|
||||
torch.empty(num_experts, 2), requires_grad=False
|
||||
)
|
||||
|
||||
layer.register_parameter("w13_weight_shape", w13_weight_shape)
|
||||
set_weight_attrs(w13_weight_shape, extra_weight_attrs)
|
||||
|
||||
w13_g_idx = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_g_idx", w13_g_idx)
|
||||
set_weight_attrs(w13_g_idx, extra_weight_attrs)
|
||||
|
||||
w2_g_idx = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_g_idx", w2_g_idx)
|
||||
set_weight_attrs(w2_g_idx, extra_weight_attrs)
|
||||
|
||||
w13_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
|
||||
set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
|
||||
|
||||
w2_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
|
||||
set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
|
||||
|
||||
layer.a13_scale = None
|
||||
layer.a2_scale = None
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
# Process weights using the shared oracle infrastructure
|
||||
is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
|
||||
converted = convert_to_wna16_moe_kernel_format(
|
||||
backend=self.wna16_backend,
|
||||
layer=layer,
|
||||
quant_config=self.weight_quant,
|
||||
input_dtype=self.marlin_input_dtype,
|
||||
w13=layer.w13_weight_packed,
|
||||
w2=layer.w2_weight_packed,
|
||||
w13_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
w13_g_idx=layer.w13_weight_g_idx,
|
||||
w2_g_idx=layer.w2_weight_g_idx,
|
||||
w13_qzeros=getattr(layer, "w13_weight_zero_point", None),
|
||||
w2_qzeros=getattr(layer, "w2_weight_zero_point", None),
|
||||
)
|
||||
if converted is None:
|
||||
# In-place backends (e.g. Humming) are not wired through this
|
||||
# marlin-only method; fail clearly rather than unpacking None.
|
||||
raise NotImplementedError(
|
||||
f"{type(self).__name__} does not support the "
|
||||
f"{self.wna16_backend.value} MoE backend."
|
||||
)
|
||||
(
|
||||
w13_qweight,
|
||||
w2_qweight,
|
||||
w13_scales,
|
||||
w2_scales,
|
||||
w13_g_idx_processed,
|
||||
w2_g_idx_processed,
|
||||
w13_g_idx_sort_indices,
|
||||
w2_g_idx_sort_indices,
|
||||
w13_qzeros,
|
||||
w2_qzeros,
|
||||
w13_input_global_scale,
|
||||
w2_input_global_scale,
|
||||
_, # w13_bias
|
||||
_, # w2_bias
|
||||
) = converted
|
||||
|
||||
# Replace common parameters
|
||||
replace_parameter(layer, "w13_weight_packed", w13_qweight)
|
||||
replace_parameter(layer, "w2_weight_packed", w2_qweight)
|
||||
replace_parameter(layer, "w13_weight_scale", w13_scales)
|
||||
replace_parameter(layer, "w2_weight_scale", w2_scales)
|
||||
|
||||
# CPU fused_experts_cpu requires zero points even for symmetric quant
|
||||
if not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU:
|
||||
replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
|
||||
replace_parameter(layer, "w2_weight_zero_point", w2_qzeros)
|
||||
|
||||
# Marlin-specific parameters (not needed for Flashinfer)
|
||||
if not is_flashinfer:
|
||||
replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
|
||||
replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)
|
||||
replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
|
||||
replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
|
||||
|
||||
# Register input global scales if present
|
||||
if w13_input_global_scale is not None:
|
||||
layer.register_parameter(
|
||||
"w13_input_global_scale",
|
||||
torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
|
||||
)
|
||||
if w2_input_global_scale is not None:
|
||||
layer.register_parameter(
|
||||
"w2_input_global_scale",
|
||||
torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
|
||||
)
|
||||
|
||||
if self.experts_cls is not None and issubclass(
|
||||
self.experts_cls, FusedMoEExpertsModular
|
||||
):
|
||||
layer.workspace = marlin_make_workspace_new(
|
||||
layer.w13_weight_g_idx.device, 4
|
||||
)
|
||||
|
||||
# Alias packed weights to w13_weight/w2_weight for the modular kernel interface
|
||||
layer.w13_weight = layer.w13_weight_packed
|
||||
layer.w2_weight = layer.w2_weight_packed
|
||||
|
||||
assert self.experts_cls is not None
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert self.moe_quant_config is not None
|
||||
|
||||
# Add Marlin-specific arguments
|
||||
marlin_args: dict[str, Any] = {}
|
||||
if not is_flashinfer:
|
||||
marlin_args = {
|
||||
"w13_g_idx": layer.w13_weight_g_idx,
|
||||
"w2_g_idx": layer.w2_weight_g_idx,
|
||||
"w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices,
|
||||
"w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices,
|
||||
"is_k_full": self.is_k_full,
|
||||
}
|
||||
|
||||
self.moe_kernel = make_wna16_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
**marlin_args,
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: torch.nn.Module
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
return make_wna16_moe_quant_config(
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
group_size=self.group_size,
|
||||
num_bits=self.num_bits,
|
||||
w1_zp=getattr(layer, "w13_weight_zero_point", None),
|
||||
w2_zp=getattr(layer, "w2_weight_zero_point", None),
|
||||
)
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
num_expert_group=layer.num_expert_group,
|
||||
topk_group=layer.topk_group,
|
||||
e_score_correction_bias=layer.e_score_correction_bias,
|
||||
routed_scaling_factor=layer.routed_scaling_factor,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
+261
@@ -0,0 +1,261 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""CompressedTensors MoE W4A16 using the fused RDNA3 (gfx1100) HIP kernel.
|
||||
|
||||
Uses ``moe_gptq_gemm_rdna3`` — a single HIP kernel launch per GEMM that
|
||||
handles expert routing + W4A16 dequant + dot product with atomic output.
|
||||
|
||||
Weight format (per expert, same as dense RDNA3 W4A16):
|
||||
- Packed int32 ``[E, K/8, N]`` with exllama shuffle
|
||||
- Scales ``[E, groups, N]`` in activation dtype
|
||||
- Zero points ``[E, groups, N/8]`` packed int32 (synthesized)
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.activation import (
|
||||
MoEActivation,
|
||||
apply_moe_activation,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
|
||||
moe_align_block_size,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501
|
||||
CompressedTensorsWNA16MoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
pack_quantized_values_into_int32,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _synthesize_qzeros(
|
||||
groups: int, out_features: int, device: torch.device
|
||||
) -> torch.Tensor:
|
||||
"""Create packed zero-point tensor for symmetric quant.
|
||||
|
||||
GPTQv1 +1 quirk: kernel adds 1 to stored zeros, so encode
|
||||
(bias - 1) = 7 for uint4b8 (bias=8).
|
||||
"""
|
||||
zeros = torch.full(
|
||||
(groups, out_features),
|
||||
scalar_types.uint4b8.bias - 1,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1)
|
||||
|
||||
|
||||
class CompressedTensorsWNA16RDNA3MoEMethod(CompressedTensorsWNA16MoEMethod):
|
||||
"""W4A16 MoE using the fused RDNA3 HIP kernel (moe_gptq_gemm_rdna3).
|
||||
|
||||
Weights are in RDNA3 format (shuffled int32 [E, K/8, N]),
|
||||
NOT Triton format (transposed uint8). apply() dispatches through
|
||||
the fused HIP kernel directly.
|
||||
"""
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
device = layer.w13_weight_packed.device
|
||||
num_experts = layer.w13_weight_packed.shape[0]
|
||||
empty_g_idx = torch.empty(0, dtype=torch.int32, device=device)
|
||||
|
||||
# Shuffle weights in-place per expert (exllama nibble interleave)
|
||||
for e in range(num_experts):
|
||||
w13_e = layer.w13_weight_packed.data[e].contiguous()
|
||||
ops.gptq_shuffle(w13_e, empty_g_idx, 4)
|
||||
layer.w13_weight_packed.data[e] = w13_e
|
||||
w2_e = layer.w2_weight_packed.data[e].contiguous()
|
||||
ops.gptq_shuffle(w2_e, empty_g_idx, 4)
|
||||
layer.w2_weight_packed.data[e] = w2_e
|
||||
|
||||
# Keep scales as [E, groups, N] in activation dtype
|
||||
act_dtype = layer.w13_weight_scale.dtype
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
layer.w13_weight_scale.to(dtype=act_dtype).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
layer.w2_weight_scale.to(dtype=act_dtype).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
# Synthesize packed zero points: [E, groups, N/8] int32
|
||||
w13_groups = (layer.w13_weight_packed.shape[1] * 8) // self.group_size
|
||||
w13_N = layer.w13_weight_packed.shape[2]
|
||||
w2_groups = (layer.w2_weight_packed.shape[1] * 8) // self.group_size
|
||||
w2_N = layer.w2_weight_packed.shape[2]
|
||||
|
||||
w13_qz = _synthesize_qzeros(w13_groups, w13_N, device)
|
||||
w2_qz = _synthesize_qzeros(w2_groups, w2_N, device)
|
||||
layer.w13_qzeros = torch.nn.Parameter(
|
||||
w13_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_qzeros = torch.nn.Parameter(
|
||||
w2_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
# Pre-allocate reusable buffers for decode (sizes based on top_k=8)
|
||||
N_gate_up = w13_N
|
||||
hidden_size = w2_N
|
||||
intermediate = N_gate_up // 2 # gated activation
|
||||
# Max tokens we expect in decode; prefill will re-allocate if needed
|
||||
max_decode_tokens = 16
|
||||
top_k = 8 # conservative default
|
||||
buf_size = max_decode_tokens * top_k
|
||||
layer.rdna3_w1_buf = torch.zeros(
|
||||
buf_size, N_gate_up, dtype=act_dtype, device=device
|
||||
)
|
||||
layer.rdna3_act_buf = torch.empty(
|
||||
buf_size, intermediate, dtype=act_dtype, device=device
|
||||
)
|
||||
layer.rdna3_out_buf = torch.zeros(
|
||||
max_decode_tokens, hidden_size, dtype=act_dtype, device=device
|
||||
)
|
||||
layer.rdna3_empty_tw = torch.empty(0, device=device)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
activation = (
|
||||
layer.activation
|
||||
if isinstance(layer.activation, MoEActivation)
|
||||
else MoEActivation.from_str(layer.activation)
|
||||
)
|
||||
return _rdna3_fused_moe(
|
||||
x,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
layer=layer,
|
||||
activation=activation,
|
||||
apply_router_weight_on_input=(layer.apply_router_weight_on_input),
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
)
|
||||
|
||||
|
||||
def _rdna3_fused_moe(
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
layer: RoutedExperts,
|
||||
activation: MoEActivation,
|
||||
apply_router_weight_on_input: bool,
|
||||
global_num_experts: int,
|
||||
expert_map: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
"""Fused MoE forward using the RDNA3 W4A16 HIP kernel.
|
||||
|
||||
Optimizations vs naive dispatch:
|
||||
- BLOCK_SIZE_M=1 for decode (no padding waste, bf16 fast path)
|
||||
- Pre-allocated buffers (no torch.zeros per call)
|
||||
- Inline token sorting for small M (skip moe_align_block_size)
|
||||
- moe_sum fused into output accumulation
|
||||
"""
|
||||
num_tokens = hidden_states.shape[0]
|
||||
top_k = topk_ids.shape[1]
|
||||
total_tokens = num_tokens * top_k
|
||||
N_gate_up = layer.w13_weight_packed.shape[2]
|
||||
hidden_size = layer.w2_weight_packed.shape[2]
|
||||
dtype = hidden_states.dtype
|
||||
device = hidden_states.device
|
||||
|
||||
intermediate_size = N_gate_up // 2 if activation.is_gated else N_gate_up
|
||||
|
||||
if global_num_experts <= 0:
|
||||
global_num_experts = layer.w13_weight_packed.shape[0]
|
||||
|
||||
# BLOCK_SIZE_M=1 for decode (small M), 4 for prefill
|
||||
block_size_m = 1 if num_tokens <= 4 else 4
|
||||
|
||||
# --- Token routing ---
|
||||
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
|
||||
topk_ids,
|
||||
block_size_m,
|
||||
global_num_experts,
|
||||
expert_map,
|
||||
)
|
||||
|
||||
# --- Reuse pre-allocated buffers when possible ---
|
||||
if total_tokens <= layer.rdna3_w1_buf.shape[0]:
|
||||
w1_out = layer.rdna3_w1_buf[:total_tokens]
|
||||
w1_out.zero_()
|
||||
act_out = layer.rdna3_act_buf[:total_tokens]
|
||||
else:
|
||||
w1_out = torch.zeros(
|
||||
total_tokens,
|
||||
N_gate_up,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
act_out = torch.empty(
|
||||
total_tokens,
|
||||
intermediate_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# --- topk weights (pre-cast to float32 for kernel) ---
|
||||
topk_w_float = topk_weights.view(-1).float()
|
||||
empty_tw = layer.rdna3_empty_tw
|
||||
|
||||
# --- w1 GEMM: [M, K] -> [M*top_k, N_gate_up] ---
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
hidden_states,
|
||||
w1_out,
|
||||
layer.w13_weight_packed,
|
||||
layer.w13_weight_scale,
|
||||
layer.w13_qzeros,
|
||||
topk_w_float if apply_router_weight_on_input else empty_tw,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
top_k,
|
||||
block_size_m,
|
||||
apply_router_weight_on_input,
|
||||
)
|
||||
|
||||
# --- Activation (silu_and_mul etc.) ---
|
||||
apply_moe_activation(activation, act_out, w1_out)
|
||||
|
||||
# --- w2 GEMM: [M*top_k, intermediate] -> [M, hidden] (fused reduce) ---
|
||||
# output_topk=top_k: kernel writes to out[token_id / top_k] directly,
|
||||
# fusing moe_sum into the atomic accumulation — saves one kernel launch
|
||||
# and the w2_out intermediate buffer.
|
||||
out = torch.zeros(
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
act_out,
|
||||
out,
|
||||
layer.w2_weight_packed,
|
||||
layer.w2_weight_scale,
|
||||
layer.w2_qzeros,
|
||||
topk_w_float if not apply_router_weight_on_input else empty_tw,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
1,
|
||||
block_size_m,
|
||||
not apply_router_weight_on_input,
|
||||
output_topk=top_k,
|
||||
)
|
||||
return out
|
||||
+48
@@ -0,0 +1,48 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""ROCm MoE kernel dispatcher.
|
||||
|
||||
Selects architecture-specific native HIP MoE kernels in priority order.
|
||||
Falls back to the Triton WNA16 path when no native kernel is available.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def is_supported(weight_quant) -> bool:
|
||||
"""Check if a native ROCm MoE kernel is available for this config."""
|
||||
if weight_quant.num_bits != 4:
|
||||
return False
|
||||
|
||||
from vllm.platforms.rocm import on_gfx1100
|
||||
|
||||
# RDNA3 (gfx1100). Future: add RDNA4 (gfx12x), CDNA (gfx94x), etc.
|
||||
return (
|
||||
on_gfx1100()
|
||||
and hasattr(torch.ops, "_rocm_C")
|
||||
and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
|
||||
)
|
||||
|
||||
|
||||
def make_method(weight_quant, input_quant, moe_config):
|
||||
"""Create the native ROCm MoE method. Call only after is_supported()."""
|
||||
from vllm.platforms.rocm import on_gfx1100
|
||||
|
||||
if on_gfx1100():
|
||||
from .compressed_tensors_moe_wna16_rdna3 import (
|
||||
CompressedTensorsWNA16RDNA3MoEMethod,
|
||||
)
|
||||
|
||||
logger.info_once(
|
||||
"Using CompressedTensorsWNA16RDNA3MoEMethod (native RDNA3 HIP kernel)"
|
||||
)
|
||||
return CompressedTensorsWNA16RDNA3MoEMethod(
|
||||
weight_quant, input_quant, moe_config
|
||||
)
|
||||
|
||||
# Future: RDNA4, CDNA, etc.
|
||||
raise RuntimeError("is_supported() returned True but no kernel matched")
|
||||
@@ -0,0 +1,28 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from .compressed_tensors_scheme import CompressedTensorsScheme
|
||||
from .compressed_tensors_w4a4_mxfp4 import CompressedTensorsW4A4Mxfp4
|
||||
from .compressed_tensors_w4a4_nvfp4 import CompressedTensorsW4A4Fp4
|
||||
from .compressed_tensors_w4a8_fp8 import CompressedTensorsW4A8Fp8
|
||||
from .compressed_tensors_w4a8_int import CompressedTensorsW4A8Int
|
||||
from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
|
||||
from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8
|
||||
from .compressed_tensors_w8a8_mxfp8 import CompressedTensorsW8A8Mxfp8
|
||||
from .compressed_tensors_w8a16_fp8 import CompressedTensorsW8A16Fp8
|
||||
from .compressed_tensors_wNa8o8 import CompressedTensorsWNA8O8Int
|
||||
from .compressed_tensors_wNa16 import CompressedTensorsWNA16
|
||||
|
||||
__all__ = [
|
||||
"CompressedTensorsScheme",
|
||||
"CompressedTensorsWNA16",
|
||||
"CompressedTensorsWNA8O8Int",
|
||||
"CompressedTensorsW8A16Fp8",
|
||||
"CompressedTensorsW8A8Int8",
|
||||
"CompressedTensorsW8A8Fp8",
|
||||
"CompressedTensorsW4A4Mxfp4",
|
||||
"CompressedTensorsW4A4Fp4",
|
||||
"CompressedTensorsW4A8Int",
|
||||
"CompressedTensorsW4A8Fp8",
|
||||
"CompressedTensorsW8A8Mxfp8",
|
||||
]
|
||||
+55
@@ -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__ = ["CompressedTensorsScheme"]
|
||||
|
||||
|
||||
class CompressedTensorsScheme(ABC):
|
||||
"""
|
||||
Abstract class used to describe the weight creation and forward pass
|
||||
of different quantization schemes supported by CompressedTensors.
|
||||
"""
|
||||
|
||||
@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()
|
||||
+96
@@ -0,0 +1,96 @@
|
||||
# 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.model_executor.kernels.linear import init_mxfp4_linear_kernel
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
GroupQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
)
|
||||
|
||||
__all__ = ["CompressedTensorsW4A4Mxfp4"]
|
||||
|
||||
|
||||
class CompressedTensorsW4A4Mxfp4(CompressedTensorsScheme):
|
||||
"""
|
||||
Compressed tensors scheme for MXFP4.
|
||||
|
||||
Supports models quantized with the compressed-tensors mxfp4-pack-quantized
|
||||
format.
|
||||
|
||||
MXFP4 format:
|
||||
- 4-bit float weights (E2M1) packed into uint8
|
||||
- Per-group E8M0 scales with group_size=32
|
||||
- No global scale (unlike NVFP4)
|
||||
|
||||
On SM100+ with FlashInfer: true W4A4 (activations dynamically quantized).
|
||||
Otherwise: W4A16 weight-only via Marlin.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.group_size = 32
|
||||
self.kernel = init_mxfp4_linear_kernel()
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 80
|
||||
|
||||
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
|
||||
layer.params_dtype = params_dtype
|
||||
|
||||
# Packed FP4 weights (2 values per byte)
|
||||
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_packed", weight)
|
||||
|
||||
# Per-group E8M0 scales
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // self.group_size,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
|
||||
del layer.weight_packed
|
||||
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, x, bias)
|
||||
+149
@@ -0,0 +1,149 @@
|
||||
# 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.layers.fusion.quant_activation import (
|
||||
expose_input_quant_key,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
GroupQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
__all__ = ["CompressedTensorsW4A4Fp4"]
|
||||
|
||||
|
||||
class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
|
||||
def __init__(self, use_a16: bool = False):
|
||||
self.use_a16 = use_a16
|
||||
self.kernel = init_nvfp4_linear_kernel(use_a16=use_a16)
|
||||
self.group_size = 16
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
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
|
||||
|
||||
# Weight
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition // 2,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_packed", weight)
|
||||
|
||||
# Global Weight Scale
|
||||
weight_global_scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_global_scale", weight_global_scale)
|
||||
|
||||
# Per Group Weight Scale
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
sum(output_partition_sizes),
|
||||
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)
|
||||
|
||||
if not self.use_a16:
|
||||
input_global_scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("input_global_scale", input_global_scale)
|
||||
|
||||
expose_input_quant_key(layer, self.kernel)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
# Rename CT checkpoint names to standardized names
|
||||
layer.weight = layer.weight_packed
|
||||
del layer.weight_packed
|
||||
|
||||
# Check for mismatched weight global scales
|
||||
if torch.unique(layer.weight_global_scale).numel() != 1:
|
||||
logger.warning_once(
|
||||
"In NVFP4 linear, the weight global scale is 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."
|
||||
)
|
||||
|
||||
# Process weight global scale (CT stores as divisors, i.e. 1/scale)
|
||||
weight_global_scale = layer.weight_global_scale.max().to(torch.float32)
|
||||
layer.weight_global_scale = Parameter(
|
||||
1.0 / weight_global_scale, requires_grad=False
|
||||
)
|
||||
|
||||
if not self.use_a16:
|
||||
if torch.unique(layer.input_global_scale).numel() != 1:
|
||||
logger.warning_once(
|
||||
"In NVFP4 linear, the input global scale is 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."
|
||||
)
|
||||
# Process input global scale and pre-compute alpha for W4A4 mode
|
||||
input_global_scale_inv = layer.input_global_scale.max().to(torch.float32)
|
||||
layer.input_global_scale = Parameter(
|
||||
(1.0 / input_global_scale_inv).to(torch.float32), requires_grad=False
|
||||
)
|
||||
|
||||
# Pre-compute alpha and inverse for runtime quantization
|
||||
layer.input_global_scale_inv = Parameter(
|
||||
input_global_scale_inv, requires_grad=False
|
||||
)
|
||||
layer.alpha = Parameter(
|
||||
layer.input_global_scale * layer.weight_global_scale,
|
||||
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)
|
||||
+177
@@ -0,0 +1,177 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import ActivationOrdering
|
||||
|
||||
from vllm.distributed.utils import verify_group_size_divides_partition
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
MPLinearLayerConfig,
|
||||
choose_mp_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
marlin_repeat_scales_on_all_ranks,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
BasevLLMParameter,
|
||||
ChannelQuantScaleParameter,
|
||||
GroupQuantScaleParameter,
|
||||
PackedvLLMParameter,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
__all__ = ["CompressedTensorsW4A8Fp8"]
|
||||
W4A8_SUPPORTED_TYPES_MAP = {
|
||||
4: scalar_types.int4,
|
||||
}
|
||||
W4A8_SUPPORTED_BITS = list(W4A8_SUPPORTED_TYPES_MAP.keys())
|
||||
|
||||
|
||||
class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
|
||||
_kernel_backends_being_used: set[str] = set()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
strategy: str,
|
||||
num_bits: int,
|
||||
group_size: int | None = None,
|
||||
symmetric: bool | None = True,
|
||||
actorder: ActivationOrdering | None = None,
|
||||
):
|
||||
self.pack_factor = 32 // num_bits
|
||||
self.strategy = strategy
|
||||
self.symmetric = symmetric
|
||||
self.group_size = -1 if group_size is None else group_size
|
||||
self.has_g_idx = actorder == ActivationOrdering.GROUP
|
||||
|
||||
if self.group_size != 128 or self.strategy != "group":
|
||||
raise ValueError(
|
||||
"W4A8 kernels require group quantization with group size 128"
|
||||
)
|
||||
|
||||
if num_bits not in W4A8_SUPPORTED_TYPES_MAP:
|
||||
raise ValueError(
|
||||
f"Unsupported num_bits = {num_bits}. "
|
||||
f"Supported num_bits = {W4A8_SUPPORTED_TYPES_MAP.keys()}"
|
||||
)
|
||||
|
||||
self.quant_type = W4A8_SUPPORTED_TYPES_MAP[num_bits]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
# hopper
|
||||
return 90
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
output_size: int,
|
||||
input_size: int,
|
||||
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)
|
||||
|
||||
mp_linear_kernel_config = MPLinearLayerConfig(
|
||||
full_weight_shape=(input_size, output_size),
|
||||
partition_weight_shape=(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition,
|
||||
),
|
||||
weight_type=self.quant_type,
|
||||
act_type=torch.float8_e4m3fn, # always use fp8(e4m3)
|
||||
group_size=self.group_size,
|
||||
zero_points=not self.symmetric,
|
||||
has_g_idx=self.has_g_idx,
|
||||
out_type=params_dtype,
|
||||
)
|
||||
|
||||
kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
|
||||
|
||||
if kernel_type.__name__ not in self._kernel_backends_being_used:
|
||||
logger.info("Using %s for CompressedTensorsW4A8Fp8", kernel_type.__name__)
|
||||
self._kernel_backends_being_used.add(kernel_type.__name__)
|
||||
|
||||
# If group_size is -1, we are in channelwise case.
|
||||
group_size = self.group_size if self.group_size != -1 else input_size
|
||||
row_parallel = input_size != input_size_per_partition
|
||||
partition_scales = not marlin_repeat_scales_on_all_ranks(
|
||||
self.has_g_idx, self.group_size, row_parallel
|
||||
)
|
||||
|
||||
scales_and_zp_size = input_size // group_size
|
||||
|
||||
if partition_scales:
|
||||
verify_group_size_divides_partition(input_size_per_partition, group_size)
|
||||
scales_and_zp_size = input_size_per_partition // group_size
|
||||
|
||||
weight = PackedvLLMParameter(
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
packed_factor=self.pack_factor,
|
||||
packed_dim=1,
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // self.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
)
|
||||
|
||||
# After loading, we will transform bf16 -> fp8 ->
|
||||
# expand by 8x via `cutlass_pack_scale_fp8`
|
||||
# and construct per-channel fp32 scales.
|
||||
weight_scale_args = {
|
||||
"weight_loader": weight_loader,
|
||||
"data": torch.empty(
|
||||
output_size_per_partition,
|
||||
scales_and_zp_size,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
}
|
||||
|
||||
if not partition_scales:
|
||||
weight_scale = ChannelQuantScaleParameter(output_dim=0, **weight_scale_args)
|
||||
else:
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
output_dim=0, input_dim=1, **weight_scale_args
|
||||
)
|
||||
|
||||
# A 2D array defining the original shape of the weights
|
||||
# before packing
|
||||
weight_shape = BasevLLMParameter(
|
||||
data=torch.empty(2, dtype=torch.int64), weight_loader=weight_loader
|
||||
)
|
||||
|
||||
layer.register_parameter("weight_packed", weight)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
layer.register_parameter("weight_shape", weight_shape)
|
||||
|
||||
self.kernel = kernel_type(
|
||||
mp_linear_kernel_config,
|
||||
w_q_param_name="weight_packed",
|
||||
w_s_param_name="weight_scale",
|
||||
w_zp_param_name="weight_zero_point",
|
||||
w_gidx_param_name="weight_g_idx",
|
||||
)
|
||||
|
||||
# Checkpoints are serialized in compressed-tensors 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:
|
||||
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)
|
||||
+153
@@ -0,0 +1,153 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.distributed.utils import verify_group_size_divides_partition
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
MPLinearLayerConfig,
|
||||
choose_mp_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
ChannelQuantScaleParameter,
|
||||
GroupQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
__all__ = ["CompressedTensorsW4A8Int"]
|
||||
W4A8_SUPPORTED_TYPES_MAP = {
|
||||
4: scalar_types.int4,
|
||||
}
|
||||
W4A8_SUPPORTED_BITS = list(W4A8_SUPPORTED_TYPES_MAP.keys())
|
||||
|
||||
|
||||
class CompressedTensorsW4A8Int(CompressedTensorsScheme):
|
||||
_kernel_backends_being_used: set[str] = set()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
strategy: str,
|
||||
num_bits: int,
|
||||
group_size: int | None = None,
|
||||
is_static_input_scheme: bool = False,
|
||||
input_symmetric: bool = True,
|
||||
):
|
||||
self.strategy = strategy
|
||||
self.group_size = -1 if group_size is None else group_size
|
||||
self.is_static_input_scheme = is_static_input_scheme
|
||||
self.input_symmetric = input_symmetric
|
||||
|
||||
if num_bits not in W4A8_SUPPORTED_TYPES_MAP:
|
||||
raise ValueError(
|
||||
f"Unsupported num_bits = {num_bits}."
|
||||
f"Supported num_bits = {W4A8_SUPPORTED_TYPES_MAP.keys()}"
|
||||
)
|
||||
self.quant_type = W4A8_SUPPORTED_TYPES_MAP[num_bits]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 1
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
output_size: int,
|
||||
input_size: int,
|
||||
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)
|
||||
row_parallel = input_size != input_size_per_partition
|
||||
|
||||
# Compute effective group_size
|
||||
if self.group_size == -1:
|
||||
effective_group_size = (
|
||||
input_size_per_partition if row_parallel else input_size
|
||||
)
|
||||
else:
|
||||
effective_group_size = self.group_size
|
||||
|
||||
# Ensure group_size divides input_size_per_partition
|
||||
verify_group_size_divides_partition(
|
||||
input_size_per_partition, effective_group_size
|
||||
)
|
||||
|
||||
# Determine scale partitioning
|
||||
is_channelwise = self.group_size == -1
|
||||
repeat_scales = is_channelwise and row_parallel
|
||||
partition_scales = not repeat_scales
|
||||
|
||||
mp_linear_kernel_config = MPLinearLayerConfig(
|
||||
full_weight_shape=(input_size, output_size),
|
||||
partition_weight_shape=(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition,
|
||||
),
|
||||
weight_type=self.quant_type,
|
||||
act_type=params_dtype,
|
||||
group_size=effective_group_size,
|
||||
zero_points=False,
|
||||
has_g_idx=False,
|
||||
)
|
||||
|
||||
kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
|
||||
if kernel_type.__name__ not in self._kernel_backends_being_used:
|
||||
logger.info("Using %s for CompressedTensorsW4A8Int", kernel_type.__name__)
|
||||
self._kernel_backends_being_used.add(kernel_type.__name__)
|
||||
|
||||
scales_and_zp_size = input_size_per_partition // effective_group_size
|
||||
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition, input_size_per_partition, dtype=torch.int8
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
weight_scale_args = {
|
||||
"weight_loader": weight_loader,
|
||||
"data": torch.empty(
|
||||
output_size_per_partition, scales_and_zp_size, dtype=params_dtype
|
||||
),
|
||||
}
|
||||
|
||||
if partition_scales:
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
output_dim=0, input_dim=1, **weight_scale_args
|
||||
)
|
||||
else:
|
||||
weight_scale = ChannelQuantScaleParameter(output_dim=0, **weight_scale_args)
|
||||
|
||||
layer.register_parameter("weight_packed", weight)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
self.kernel = kernel_type(
|
||||
mp_linear_kernel_config,
|
||||
w_q_param_name="weight_packed",
|
||||
w_s_param_name="weight_scale",
|
||||
w_zp_param_name=None,
|
||||
w_gidx_param_name=None,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
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)
|
||||
+159
@@ -0,0 +1,159 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
|
||||
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
init_wfp8_a16_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
|
||||
STRATEGY_TO_PARAMETER_TYPE,
|
||||
STRATEGY_TO_WEIGHT_QUANT_KEY,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
create_fp8_scale_parameter,
|
||||
create_fp8_weight_parameter,
|
||||
validate_fp8_block_shape,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kFp8DynamicTensorSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
convert_to_channelwise,
|
||||
)
|
||||
from vllm.model_executor.parameter import PerTensorScaleParameter
|
||||
from vllm.model_executor.utils import replace_parameter
|
||||
|
||||
__all__ = ["CompressedTensorsW8A16Fp8"]
|
||||
|
||||
|
||||
class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
|
||||
def __init__(self, weight_quant: QuantizationArgs, is_static_input_scheme: bool):
|
||||
self.weight_quant = weight_quant
|
||||
self.strategy = weight_quant.strategy
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
self.input_dtype = get_current_vllm_config().model_config.dtype
|
||||
self.is_static_input_scheme = is_static_input_scheme
|
||||
self.weight_block_size = self.weight_quant.block_structure
|
||||
|
||||
self.weight_quant_key = STRATEGY_TO_WEIGHT_QUANT_KEY[self.strategy]
|
||||
self.activation_quant_key = (
|
||||
kFp8StaticTensorSym if is_static_input_scheme else kFp8DynamicTensorSym
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
# turing and up
|
||||
return 75
|
||||
|
||||
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,
|
||||
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
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
if self.strategy == QuantizationStrategy.BLOCK:
|
||||
assert self.weight_block_size is not None
|
||||
layer.weight_block_size = self.weight_block_size
|
||||
# Validate block quantization shapes
|
||||
validate_fp8_block_shape(
|
||||
layer,
|
||||
input_size,
|
||||
output_size,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
self.weight_block_size,
|
||||
)
|
||||
|
||||
# WEIGHT
|
||||
weight = create_fp8_weight_parameter(
|
||||
output_size_per_partition, input_size_per_partition, weight_loader
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# WEIGHT SCALE
|
||||
weight_scale = create_fp8_scale_parameter(
|
||||
STRATEGY_TO_PARAMETER_TYPE[self.strategy],
|
||||
output_partition_sizes,
|
||||
input_size_per_partition,
|
||||
layer.weight_block_size,
|
||||
weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
# INPUT SCALE (to deal with converted checkpoints)
|
||||
if self.is_static_input_scheme:
|
||||
input_scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("input_scale", input_scale)
|
||||
|
||||
self.linear_kernel = init_wfp8_a16_linear_kernel(
|
||||
weight_quant_key=self.weight_quant_key,
|
||||
activation_quant_key=self.activation_quant_key,
|
||||
weight_shape=layer.weight.shape,
|
||||
input_dtype=self.input_dtype,
|
||||
out_dtype=self.out_dtype,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
if self.strategy == QuantizationStrategy.BLOCK:
|
||||
assert self.is_static_input_scheme is False
|
||||
# MarlinFP8ScaledMMLinearKernel uses "weight_scale_inv" for block
|
||||
# quant, while CT registers the scale as "weight_scale".
|
||||
# Rename by deleting the old parameter and adding the new one so
|
||||
# that prepare_fp8_layer_for_marlin (which prefers "weight_scale"
|
||||
# over "weight_scale_inv") picks up "weight_scale_inv" correctly.
|
||||
weight_scale_data = layer.weight_scale.data
|
||||
del layer._parameters["weight_scale"]
|
||||
replace_parameter(layer, "weight_scale_inv", weight_scale_data)
|
||||
else:
|
||||
if self.strategy == QuantizationStrategy.TENSOR:
|
||||
# For fused modules with per-tensor scales, expand each scale
|
||||
# to its shard's channels.
|
||||
replace_parameter(
|
||||
layer,
|
||||
"weight_scale",
|
||||
convert_to_channelwise(layer.weight_scale, layer.logical_widths),
|
||||
)
|
||||
self.strategy = QuantizationStrategy.CHANNEL
|
||||
self.weight_quant_key = STRATEGY_TO_WEIGHT_QUANT_KEY[self.strategy]
|
||||
self.linear_kernel.config.weight_quant_key = self.weight_quant_key
|
||||
|
||||
# Canonicalize to (K, N) for the kernel.
|
||||
replace_parameter(layer, "weight", layer.weight.t())
|
||||
# Preserve the dim tags dropped by the transpose so layout-aware
|
||||
# kernels see (K, N).
|
||||
layer.weight.input_dim = 0
|
||||
layer.weight.output_dim = 1
|
||||
|
||||
self.linear_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.linear_kernel.apply_weights(layer, x, bias)
|
||||
+207
@@ -0,0 +1,207 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
|
||||
from torch.nn import Parameter
|
||||
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
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.fusion.quant_activation import (
|
||||
QuantizedActivation,
|
||||
expose_input_quant_key,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
|
||||
STRATEGY_TO_PARAMETER_TYPE,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
create_fp8_input_scale,
|
||||
create_fp8_scale_parameter,
|
||||
create_fp8_weight_parameter,
|
||||
process_fp8_weight_channel_strategy,
|
||||
process_fp8_weight_tensor_strategy,
|
||||
validate_fp8_block_shape,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
create_fp8_quant_key,
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8StaticChannelSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
cutlass_block_fp8_supported,
|
||||
)
|
||||
|
||||
__all__ = ["CompressedTensorsW8A8Fp8"]
|
||||
|
||||
STATIC_QUANT = True
|
||||
DYNAMIC_QUANT = False
|
||||
activation_quant_key_mapping = {
|
||||
STATIC_QUANT: kFp8StaticTensorSym,
|
||||
DYNAMIC_QUANT: kFp8DynamicTokenSym,
|
||||
}
|
||||
weight_quant_key_mapping = {
|
||||
QuantizationStrategy.CHANNEL: kFp8StaticChannelSym,
|
||||
QuantizationStrategy.TENSOR: kFp8StaticTensorSym,
|
||||
}
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
|
||||
def __init__(self, weight_quant: QuantizationArgs, is_static_input_scheme: bool):
|
||||
self.weight_quant = weight_quant
|
||||
self.strategy = weight_quant.strategy
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
self.input_dtype = get_current_vllm_config().model_config.dtype
|
||||
self.is_static_input_scheme = is_static_input_scheme
|
||||
self.weight_block_size = self.weight_quant.block_structure
|
||||
|
||||
if self.weight_block_size is not None:
|
||||
self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
|
||||
self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled()
|
||||
assert not self.is_static_input_scheme
|
||||
self.act_q_group_shape = GroupShape(1, self.weight_block_size[0])
|
||||
self.weight_quant_key = create_fp8_quant_key(
|
||||
static=True, group_shape=GroupShape(*self.weight_block_size)
|
||||
)
|
||||
self.activation_quant_key = create_fp8_quant_key(
|
||||
static=False, group_shape=self.act_q_group_shape
|
||||
)
|
||||
else:
|
||||
self.activation_quant_key = activation_quant_key_mapping[
|
||||
self.is_static_input_scheme
|
||||
]
|
||||
self.weight_quant_key = weight_quant_key_mapping[self.strategy]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
# lovelace and up
|
||||
return 89
|
||||
|
||||
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,
|
||||
weight_loader: Callable,
|
||||
**kwargs,
|
||||
):
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
layer.logical_widths = output_partition_sizes
|
||||
layer.weight_block_size = None
|
||||
layer.orig_dtype = params_dtype
|
||||
|
||||
if self.strategy == QuantizationStrategy.BLOCK:
|
||||
assert self.weight_block_size is not None
|
||||
layer.weight_block_size = self.weight_block_size
|
||||
# Validate block quantization shapes
|
||||
validate_fp8_block_shape(
|
||||
layer,
|
||||
input_size,
|
||||
output_size,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
self.weight_block_size,
|
||||
)
|
||||
|
||||
# WEIGHT
|
||||
weight = create_fp8_weight_parameter(
|
||||
output_size_per_partition, input_size_per_partition, weight_loader
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# WEIGHT SCALE
|
||||
weight_scale = create_fp8_scale_parameter(
|
||||
STRATEGY_TO_PARAMETER_TYPE[self.strategy],
|
||||
output_partition_sizes,
|
||||
input_size_per_partition,
|
||||
layer.weight_block_size,
|
||||
weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
# INPUT SCALE
|
||||
if self.is_static_input_scheme:
|
||||
input_scale = create_fp8_input_scale(output_partition_sizes, weight_loader)
|
||||
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,
|
||||
input_dtype=self.input_dtype,
|
||||
out_dtype=self.out_dtype,
|
||||
weight_shape=(output_size_per_partition, input_size_per_partition),
|
||||
module_name=self.__class__.__name__,
|
||||
)
|
||||
|
||||
expose_input_quant_key(layer, self.fp8_linear)
|
||||
|
||||
def process_weights_after_loading(self, layer) -> None:
|
||||
if self.strategy == QuantizationStrategy.TENSOR:
|
||||
weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy(
|
||||
layer.weight,
|
||||
layer.weight_scale,
|
||||
layer.logical_widths,
|
||||
getattr(layer, "input_scale", None),
|
||||
)
|
||||
weight = weight.t()
|
||||
elif self.strategy == QuantizationStrategy.CHANNEL:
|
||||
weight, weight_scale, input_scale = process_fp8_weight_channel_strategy(
|
||||
layer.weight, layer.weight_scale, getattr(layer, "input_scale", None)
|
||||
)
|
||||
weight = weight.t()
|
||||
|
||||
elif self.strategy == QuantizationStrategy.BLOCK:
|
||||
assert self.is_static_input_scheme is False
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
|
||||
layer.input_scale = None
|
||||
# fp8_linear.process_weights_after_loading applies the post process
|
||||
# and reassigns the weight and weight_scale buffers to layer attributes.
|
||||
return
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown quantization strategy {self.strategy}: "
|
||||
f"should be one of {list(QuantizationStrategy)}"
|
||||
)
|
||||
|
||||
# required by torch.compile to be torch.nn.Parameter
|
||||
layer.weight = Parameter(weight.data, requires_grad=False)
|
||||
# Preserve the dim tags dropped by the transpose so layout-aware
|
||||
# kernels (humming) see (K, N) instead of assuming (N, K).
|
||||
layer.weight.input_dim = 0
|
||||
layer.weight.output_dim = 1
|
||||
layer.weight_scale = Parameter(weight_scale.data, requires_grad=False)
|
||||
if input_scale is not None:
|
||||
layer.input_scale = Parameter(input_scale.data, requires_grad=False)
|
||||
|
||||
# INPUT SCALE
|
||||
if self.is_static_input_scheme and hasattr(layer, "input_scale"):
|
||||
layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
|
||||
else:
|
||||
layer.input_scale = None
|
||||
|
||||
if hasattr(self, "fp8_linear"):
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor | QuantizedActivation,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return self.fp8_linear.apply_weights(layer, x, bias)
|
||||
+112
@@ -0,0 +1,112 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import QuantizationStrategy
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
init_int8_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
BasevLLMParameter,
|
||||
ChannelQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CompressedTensorsW8A8Int8(CompressedTensorsScheme):
|
||||
def __init__(
|
||||
self, strategy: str, is_static_input_scheme: bool, input_symmetric: bool
|
||||
):
|
||||
self.strategy = strategy
|
||||
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
|
||||
|
||||
self.kernel = init_int8_linear_kernel(
|
||||
is_channelwise=(self.strategy == QuantizationStrategy.CHANNEL),
|
||||
is_static_input_scheme=self.is_static_input_scheme,
|
||||
input_symmetric=self.input_symmetric,
|
||||
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.strategy == QuantizationStrategy.CHANNEL:
|
||||
weight_scale = ChannelQuantScaleParameter(
|
||||
data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
else:
|
||||
assert self.strategy == QuantizationStrategy.TENSOR
|
||||
weight_scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
# 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
|
||||
)
|
||||
if not self.input_symmetric:
|
||||
# Note: compressed-tensors stores the zp using the same dtype
|
||||
# as the weights
|
||||
# AZP loaded as int8 but used as int32
|
||||
input_zero_point = BasevLLMParameter(
|
||||
data=torch.empty(1, dtype=torch.int8), weight_loader=weight_loader
|
||||
)
|
||||
|
||||
layer.register_parameter("input_zero_point", input_zero_point)
|
||||
layer.register_parameter("input_scale", input_scale)
|
||||
if not hasattr(layer, "azp_adj"):
|
||||
layer.register_parameter("azp_adj", None)
|
||||
|
||||
# Checkpoints are serialized in compressed-tensors 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:
|
||||
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)
|
||||
+92
@@ -0,0 +1,92 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
|
||||
MXFP8_BLOCK_SIZE,
|
||||
MXFP8_SCALE_DTYPE,
|
||||
MXFP8_VALUE_DTYPE,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
GroupQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
)
|
||||
|
||||
__all__ = ["CompressedTensorsW8A8Mxfp8"]
|
||||
|
||||
|
||||
class CompressedTensorsW8A8Mxfp8(CompressedTensorsScheme):
|
||||
"""
|
||||
Compressed tensors scheme for MXFP8 quantization (W8A8).
|
||||
|
||||
Loads pre-quantized MXFP8 weights from compressed-tensors checkpoints.
|
||||
Activations are dynamically quantized to MXFP8 at runtime.
|
||||
|
||||
MXFP8 format:
|
||||
- 8-bit float weights (E4M3) stored as float8_e4m3fn
|
||||
- Per-group E8M0 scales (uint8) with group_size=32
|
||||
- Activations dynamically quantized to MXFP8 during inference
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.kernel = init_mxfp8_linear_kernel()
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
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
|
||||
layer.params_dtype = params_dtype
|
||||
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition,
|
||||
dtype=MXFP8_VALUE_DTYPE,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // MXFP8_BLOCK_SIZE,
|
||||
dtype=MXFP8_SCALE_DTYPE,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
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, x, bias)
|
||||
+257
@@ -0,0 +1,257 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
from fractions import Fraction
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import ActivationOrdering
|
||||
|
||||
from vllm.distributed.utils import verify_group_size_divides_partition
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
MarlinLinearKernel,
|
||||
MPLinearLayerConfig,
|
||||
choose_mp_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
get_marlin_input_dtype,
|
||||
marlin_repeat_scales_on_all_ranks,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
BasevLLMParameter,
|
||||
ChannelQuantScaleParameter,
|
||||
GroupQuantScaleParameter,
|
||||
PackedColumnParameter,
|
||||
PackedvLLMParameter,
|
||||
RowvLLMParameter,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
__all__ = ["CompressedTensorsWNA16"]
|
||||
WNA16_SUPPORTED_TYPES_MAP = {
|
||||
2: scalar_types.uint2b2,
|
||||
3: scalar_types.uint3b4,
|
||||
4: scalar_types.uint4b8,
|
||||
5: scalar_types.uint5b16,
|
||||
6: scalar_types.uint6b32,
|
||||
7: scalar_types.uint7b64,
|
||||
8: scalar_types.uint8b128,
|
||||
}
|
||||
WNA16_ZP_SUPPORTED_TYPES_MAP = {4: scalar_types.uint4, 8: scalar_types.uint8}
|
||||
WNA16_SUPPORTED_BITS = list(WNA16_SUPPORTED_TYPES_MAP.keys())
|
||||
|
||||
|
||||
class CompressedTensorsWNA16(CompressedTensorsScheme):
|
||||
_kernel_backends_being_used: set[str] = set()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
strategy: str,
|
||||
num_bits: int,
|
||||
group_size: int | None = None,
|
||||
symmetric: bool | None = True,
|
||||
actorder: ActivationOrdering | None = None,
|
||||
layer_name: str | None = None,
|
||||
):
|
||||
self.num_bits = num_bits
|
||||
self.pack_factor = Fraction(32, num_bits)
|
||||
self.strategy = strategy
|
||||
self.symmetric = symmetric
|
||||
self.group_size = -1 if group_size is None else group_size
|
||||
self.has_g_idx = actorder == ActivationOrdering.GROUP
|
||||
self.layer_name = layer_name
|
||||
|
||||
if self.group_size == -1 and self.strategy != "channel":
|
||||
raise ValueError(
|
||||
"Pack-quantized format requires group quantization "
|
||||
"or channelwise quantization, but found no group "
|
||||
"size and strategy is not channelwise."
|
||||
)
|
||||
|
||||
if num_bits not in WNA16_SUPPORTED_TYPES_MAP:
|
||||
raise ValueError(
|
||||
f"Unsupported num_bits = {num_bits}. "
|
||||
f"Supported num_bits = {list(WNA16_SUPPORTED_TYPES_MAP)}"
|
||||
)
|
||||
|
||||
if not self.symmetric and num_bits not in WNA16_ZP_SUPPORTED_TYPES_MAP:
|
||||
raise ValueError(
|
||||
f"Asymmetric quantization not supported for "
|
||||
f"num_bits = {num_bits}. Supported: "
|
||||
f"{list(WNA16_ZP_SUPPORTED_TYPES_MAP)}"
|
||||
)
|
||||
|
||||
self.quant_type = (
|
||||
WNA16_ZP_SUPPORTED_TYPES_MAP[num_bits]
|
||||
if not self.symmetric
|
||||
else WNA16_SUPPORTED_TYPES_MAP[num_bits]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
# Turing and up
|
||||
return 75
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
output_size: int,
|
||||
input_size: int,
|
||||
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.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
layer.output_partition_sizes = output_partition_sizes
|
||||
layer.params_dtype = params_dtype
|
||||
if not hasattr(layer, "has_bias"):
|
||||
layer.has_bias = False
|
||||
|
||||
mp_linear_kernel_config = MPLinearLayerConfig(
|
||||
full_weight_shape=(input_size, output_size),
|
||||
partition_weight_shape=(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition,
|
||||
),
|
||||
weight_type=self.quant_type,
|
||||
act_type=params_dtype,
|
||||
group_size=self.group_size,
|
||||
zero_points=not self.symmetric,
|
||||
has_g_idx=self.has_g_idx,
|
||||
)
|
||||
|
||||
kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
|
||||
|
||||
if kernel_type.__name__ not in self._kernel_backends_being_used:
|
||||
logger.info("Using %s for CompressedTensorsWNA16", kernel_type.__name__)
|
||||
self._kernel_backends_being_used.add(kernel_type.__name__)
|
||||
|
||||
if kernel_type is MarlinLinearKernel:
|
||||
input_dtype = get_marlin_input_dtype(self.layer_name)
|
||||
if input_dtype is not None:
|
||||
mp_linear_kernel_config.act_type = input_dtype
|
||||
|
||||
# If group_size is -1, we are in channelwise case.
|
||||
group_size = self.group_size if self.group_size != -1 else input_size
|
||||
row_parallel = input_size != input_size_per_partition
|
||||
partition_scales = not marlin_repeat_scales_on_all_ranks(
|
||||
self.has_g_idx, self.group_size, row_parallel
|
||||
)
|
||||
|
||||
scales_and_zp_size = input_size // group_size
|
||||
|
||||
if partition_scales:
|
||||
verify_group_size_divides_partition(
|
||||
input_size_per_partition, group_size, self.layer_name
|
||||
)
|
||||
scales_and_zp_size = input_size_per_partition // group_size
|
||||
|
||||
packed_input_dim = math.ceil(input_size_per_partition * self.num_bits / 32)
|
||||
weight = PackedvLLMParameter(
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
packed_factor=self.pack_factor,
|
||||
packed_dim=1,
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
packed_input_dim,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
)
|
||||
|
||||
weight_scale_args = {
|
||||
"weight_loader": weight_loader,
|
||||
"data": torch.empty(
|
||||
output_size_per_partition,
|
||||
scales_and_zp_size,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
}
|
||||
|
||||
packed_output_dim = math.ceil(output_size_per_partition * self.num_bits / 32)
|
||||
zeros_args = {
|
||||
"weight_loader": weight_loader,
|
||||
"data": torch.zeros(
|
||||
packed_output_dim,
|
||||
scales_and_zp_size,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
}
|
||||
|
||||
if not partition_scales:
|
||||
weight_scale = ChannelQuantScaleParameter(output_dim=0, **weight_scale_args)
|
||||
|
||||
if not self.symmetric:
|
||||
qzeros = PackedColumnParameter(
|
||||
output_dim=0,
|
||||
packed_dim=0,
|
||||
packed_factor=self.pack_factor,
|
||||
**zeros_args,
|
||||
)
|
||||
else:
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
output_dim=0, input_dim=1, **weight_scale_args
|
||||
)
|
||||
if not self.symmetric:
|
||||
qzeros = PackedvLLMParameter(
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
packed_dim=0,
|
||||
packed_factor=self.pack_factor,
|
||||
**zeros_args,
|
||||
)
|
||||
|
||||
# A 2D array defining the original shape of the weights
|
||||
# before packing
|
||||
weight_shape = BasevLLMParameter(
|
||||
data=torch.empty(2, dtype=torch.int64), weight_loader=weight_loader
|
||||
)
|
||||
|
||||
layer.register_parameter("weight_packed", weight)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
layer.register_parameter("weight_shape", weight_shape)
|
||||
|
||||
if not self.symmetric:
|
||||
layer.register_parameter("weight_zero_point", qzeros)
|
||||
|
||||
# group index (for activation reordering)
|
||||
if self.has_g_idx:
|
||||
weight_g_idx = RowvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_g_idx", weight_g_idx)
|
||||
|
||||
self.kernel = kernel_type(
|
||||
mp_linear_kernel_config,
|
||||
w_q_param_name="weight_packed",
|
||||
w_s_param_name="weight_scale",
|
||||
w_zp_param_name="weight_zero_point",
|
||||
w_gidx_param_name="weight_g_idx",
|
||||
)
|
||||
|
||||
# Checkpoints are serialized in compressed-tensors 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:
|
||||
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)
|
||||
+260
@@ -0,0 +1,260 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Weight N-bit INT scheme with static INT8 input/output activation quant.
|
||||
|
||||
Handles compressed-tensors INT weight checkpoints that carry static per-tensor
|
||||
INT8 ``input_activations`` and/or ``output_activations``. The activation quant is
|
||||
reproduced as a float fake-quant on the layer input and output, around a
|
||||
weight-only matmul, rather than a fused int8 GEMM.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from compressed_tensors.compressors.pack_quantized.helpers import pack_to_int32
|
||||
|
||||
from vllm.distributed.utils import verify_group_size_divides_partition
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
MPLinearLayerConfig,
|
||||
choose_mp_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
marlin_repeat_scales_on_all_ranks,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
BasevLLMParameter,
|
||||
ChannelQuantScaleParameter,
|
||||
GroupQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
PackedvLLMParameter,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
__all__ = ["CompressedTensorsWNA8O8Int", "fake_quant_static_int8"]
|
||||
|
||||
WNA8O8_SUPPORTED_TYPES_MAP = {
|
||||
2: scalar_types.uint2b2,
|
||||
4: scalar_types.uint4b8,
|
||||
8: scalar_types.uint8b128,
|
||||
}
|
||||
|
||||
|
||||
def fake_quant_static_int8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
|
||||
"""Static per-tensor symmetric INT8 quantize-dequantize, in x's dtype."""
|
||||
scale = scale.to(x.dtype)
|
||||
q = torch.clamp(torch.round(x / scale), -128.0, 127.0)
|
||||
return q * scale
|
||||
|
||||
|
||||
class CompressedTensorsWNA8O8Int(CompressedTensorsScheme):
|
||||
def __init__(
|
||||
self,
|
||||
num_bits: int,
|
||||
strategy: str,
|
||||
group_size: int | None = None,
|
||||
has_input_act: bool = False,
|
||||
has_output_act: bool = False,
|
||||
layer_name: str | None = None,
|
||||
quant_format: str = "pack-quantized",
|
||||
):
|
||||
self.num_bits = num_bits
|
||||
self.pack_factor = 32 // num_bits
|
||||
self.strategy = strategy
|
||||
self.group_size = -1 if group_size is None else group_size
|
||||
self.has_input_act = has_input_act
|
||||
self.has_output_act = has_output_act
|
||||
self.layer_name = layer_name
|
||||
# "pack-quantized" (sub-byte, int32-packed) or "int-quantized" (8-bit int8).
|
||||
self.quant_format = quant_format
|
||||
self.is_int_quantized = quant_format == "int-quantized"
|
||||
if num_bits not in WNA8O8_SUPPORTED_TYPES_MAP:
|
||||
raise ValueError(
|
||||
f"Unsupported num_bits = {num_bits} for WNA8O8Int; "
|
||||
f"supported = {sorted(WNA8O8_SUPPORTED_TYPES_MAP)}"
|
||||
)
|
||||
self.quant_type = WNA8O8_SUPPORTED_TYPES_MAP[num_bits]
|
||||
self._input_scale: torch.Tensor | None = None
|
||||
self._output_scale: torch.Tensor | None = None
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 70
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
output_size: int,
|
||||
input_size: int,
|
||||
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.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
# Set for kernels' weight prep; also covers ParallelLMHead, which does
|
||||
# not set these in __init__.
|
||||
layer.output_partition_sizes = output_partition_sizes
|
||||
layer.params_dtype = params_dtype
|
||||
if not hasattr(layer, "has_bias"):
|
||||
layer.has_bias = False
|
||||
|
||||
mp_config = MPLinearLayerConfig(
|
||||
full_weight_shape=(input_size, output_size),
|
||||
partition_weight_shape=(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition,
|
||||
),
|
||||
weight_type=self.quant_type,
|
||||
act_type=params_dtype, # activation quant applied externally (SRQ)
|
||||
group_size=self.group_size,
|
||||
zero_points=False,
|
||||
has_g_idx=False,
|
||||
)
|
||||
self.kernel = choose_mp_linear_kernel(mp_config)(
|
||||
mp_config,
|
||||
w_q_param_name="weight_packed",
|
||||
w_s_param_name="weight_scale",
|
||||
)
|
||||
|
||||
self._register_weight(
|
||||
layer, input_size, input_size_per_partition, params_dtype, weight_loader
|
||||
)
|
||||
|
||||
def _register_weight(
|
||||
self, layer, input_size, input_size_per_partition, params_dtype, weight_loader
|
||||
):
|
||||
out = layer.output_size_per_partition
|
||||
if self.is_int_quantized:
|
||||
# Plain int8 weight; packed to the canonical int32 layout after load.
|
||||
layer.register_parameter(
|
||||
"weight",
|
||||
ModelWeightParameter(
|
||||
data=torch.empty(out, input_size_per_partition, dtype=torch.int8),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
),
|
||||
)
|
||||
else:
|
||||
layer.register_parameter(
|
||||
"weight_packed",
|
||||
PackedvLLMParameter(
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
packed_dim=1,
|
||||
packed_factor=self.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
data=torch.empty(
|
||||
out,
|
||||
input_size_per_partition // self.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
),
|
||||
)
|
||||
layer.register_parameter(
|
||||
"weight_shape",
|
||||
BasevLLMParameter(
|
||||
data=torch.empty(2, dtype=torch.int64), weight_loader=weight_loader
|
||||
),
|
||||
)
|
||||
|
||||
# Scale: per-output-channel, or per group along the input dim under TP.
|
||||
group_size = self.group_size if self.group_size != -1 else input_size
|
||||
partitioned = not marlin_repeat_scales_on_all_ranks(
|
||||
False, self.group_size, input_size != input_size_per_partition
|
||||
)
|
||||
scales = (input_size_per_partition if partitioned else input_size) // group_size
|
||||
scale_data = torch.empty(out, scales, dtype=params_dtype)
|
||||
if partitioned:
|
||||
verify_group_size_divides_partition(
|
||||
input_size_per_partition, group_size, self.layer_name
|
||||
)
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
data=scale_data, output_dim=0, input_dim=1, weight_loader=weight_loader
|
||||
)
|
||||
else:
|
||||
weight_scale = ChannelQuantScaleParameter(
|
||||
data=scale_data, output_dim=0, weight_loader=weight_loader
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
for name, present in (
|
||||
("input_scale", self.has_input_act),
|
||||
("output_scale", self.has_output_act),
|
||||
):
|
||||
if present:
|
||||
layer.register_parameter(
|
||||
name,
|
||||
BasevLLMParameter(
|
||||
data=torch.empty(1, dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
),
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
# Lift the static activation scales off the layer (applied externally) so
|
||||
# the kernel only sees weight tensors. Drop uncalibrated (zero) scales.
|
||||
self._input_scale = self._take_act_scale(layer, "input_scale")
|
||||
self._output_scale = self._take_act_scale(layer, "output_scale")
|
||||
self.has_input_act = self._input_scale is not None
|
||||
self.has_output_act = self._output_scale is not None
|
||||
|
||||
if self.is_int_quantized:
|
||||
self._pack_int_quantized_weight(layer)
|
||||
|
||||
self.kernel.process_weights_after_loading(layer)
|
||||
|
||||
def _pack_int_quantized_weight(self, layer: torch.nn.Module) -> None:
|
||||
"""Normalize an int-quantized (plain int8) weight to the canonical
|
||||
``weight_packed`` int32 + ``weight_shape`` layout the MP kernels expect."""
|
||||
weight = layer.weight
|
||||
out_features, in_features = weight.shape
|
||||
packed = pack_to_int32(weight.data.contiguous(), self.num_bits)
|
||||
delattr(layer, "weight")
|
||||
|
||||
def _noop_loader(*_, **__):
|
||||
return None
|
||||
|
||||
layer.register_parameter(
|
||||
"weight_packed",
|
||||
PackedvLLMParameter(
|
||||
data=packed.contiguous(),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
packed_dim=1,
|
||||
packed_factor=self.pack_factor,
|
||||
weight_loader=_noop_loader,
|
||||
),
|
||||
)
|
||||
layer.register_parameter(
|
||||
"weight_shape",
|
||||
BasevLLMParameter(
|
||||
data=torch.tensor([out_features, in_features], dtype=torch.int64),
|
||||
weight_loader=_noop_loader,
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _take_act_scale(layer, name: str) -> torch.Tensor | None:
|
||||
param = getattr(layer, name, None)
|
||||
if param is None:
|
||||
return None
|
||||
scale = param.data.clone()
|
||||
delattr(layer, name)
|
||||
return None if float(scale.reshape(-1)[0]) == 0.0 else scale
|
||||
|
||||
def apply_weights(
|
||||
self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
|
||||
) -> torch.Tensor:
|
||||
if self.has_input_act:
|
||||
x = fake_quant_static_int8(x, self._input_scale)
|
||||
out = self.kernel.apply_weights(layer, x, bias)
|
||||
if self.has_output_act:
|
||||
out = fake_quant_static_int8(out, self._output_scale)
|
||||
return out
|
||||
@@ -0,0 +1,260 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from collections.abc import Callable, Generator
|
||||
from itertools import accumulate
|
||||
|
||||
import torch
|
||||
from compressed_tensors.transform import (
|
||||
TransformArgs,
|
||||
TransformConfig,
|
||||
TransformLocation,
|
||||
TransformScheme,
|
||||
)
|
||||
from compressed_tensors.utils import is_match
|
||||
|
||||
from vllm.model_executor.layers.linear import (
|
||||
WEIGHT_LOADER_V2_SUPPORTED,
|
||||
LinearMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
|
||||
CompressedTensorsScheme,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.transform.module import ( # noqa: E501
|
||||
HadamardTransform,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.transform.utils import ( # noqa: E501
|
||||
TransformTuple,
|
||||
)
|
||||
|
||||
|
||||
class CompressedTensorsLinearTransformMethod(LinearMethodBase):
|
||||
"""
|
||||
Wraps `CompressedTensorsLinearMethod` or `UnquantizedLinearMethod` and adds
|
||||
input and output transforms to either side of the original apply method
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def from_schemes(
|
||||
cls,
|
||||
quant_method: LinearMethodBase,
|
||||
quant_scheme: CompressedTensorsScheme | None,
|
||||
input_tfms: dict[int, TransformTuple],
|
||||
output_tfms: dict[int, TransformTuple],
|
||||
) -> "CompressedTensorsLinearTransformMethod":
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.transform.schemes.linear_qutlass_nvfp4 import ( # noqa: E501
|
||||
QutlassNvFP4LinearMethod,
|
||||
is_qutlass_fp4_scheme,
|
||||
)
|
||||
|
||||
assert input_tfms or output_tfms
|
||||
|
||||
if is_qutlass_fp4_scheme(quant_scheme, input_tfms):
|
||||
return QutlassNvFP4LinearMethod(quant_method, input_tfms, output_tfms)
|
||||
|
||||
# hadacore or dense gemm is selected by Transform module
|
||||
|
||||
return cls(quant_method, input_tfms, output_tfms)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
quant_method: LinearMethodBase,
|
||||
input_tfms: dict[int, TransformTuple],
|
||||
output_tfms: dict[int, TransformTuple],
|
||||
):
|
||||
self.quant_method = quant_method
|
||||
self.input_tfms = input_tfms
|
||||
self.output_tfms = output_tfms
|
||||
|
||||
self.input_transform: HadamardTransform | None = None
|
||||
self.output_transform: HadamardTransform | None = None
|
||||
|
||||
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,
|
||||
):
|
||||
# get weight loader for transforms
|
||||
weight_loader: Callable = extra_weight_attrs.get("weight_loader") # type: ignore[assignment]
|
||||
|
||||
# HACK: UnquantizedLinearMethod does not support weight loader v2, but
|
||||
# transforms (specifically SharedWeightParameter) requires
|
||||
# weight loader v2. Until UnquantizedLinearMethod supports v2, we must
|
||||
# hack around this by getting weight loader v1 so ULM can load correctly
|
||||
quant_method_name = self.quant_method.__class__.__name__
|
||||
if quant_method_name not in WEIGHT_LOADER_V2_SUPPORTED:
|
||||
weight_loader_v1 = layer.weight_loader
|
||||
extra_weight_attrs["weight_loader"] = weight_loader_v1
|
||||
|
||||
self.quant_method.create_weights(
|
||||
layer=layer,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
output_partition_sizes=output_partition_sizes,
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
params_dtype=params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
|
||||
# validate schemes
|
||||
num_partitions = len(output_partition_sizes)
|
||||
self._validate_tfm_schemes(num_partitions)
|
||||
|
||||
# create submodules for weight loading
|
||||
if len(self.input_tfms) > 0:
|
||||
scheme_name = list(self.input_tfms.values())[0].scheme_name
|
||||
location = list(self.input_tfms.values())[0].args.location
|
||||
transform_name = f"{scheme_name}_{location}"
|
||||
|
||||
transform = HadamardTransform(
|
||||
self.input_tfms,
|
||||
layer,
|
||||
weight_loader,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
)
|
||||
layer.register_module(transform_name, transform)
|
||||
self.input_transform = transform
|
||||
|
||||
if len(self.output_tfms) > 0:
|
||||
scheme_name = list(self.output_tfms.values())[0].scheme_name
|
||||
location = list(self.output_tfms.values())[0].args.location
|
||||
transform_name = f"{scheme_name}_{location}"
|
||||
|
||||
transform = HadamardTransform(
|
||||
self.output_tfms,
|
||||
layer,
|
||||
weight_loader,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
)
|
||||
layer.register_module(transform_name, transform)
|
||||
self.output_transform = transform
|
||||
|
||||
# compute partition ranges for slicing activations
|
||||
starts = [0] + list(accumulate(output_partition_sizes))[:-1]
|
||||
self.partition_ranges = list(zip(starts, output_partition_sizes))
|
||||
|
||||
def process_weights_after_loading(self, layer):
|
||||
self.quant_method.process_weights_after_loading(layer)
|
||||
|
||||
for submodule in layer.children():
|
||||
if isinstance(submodule, HadamardTransform):
|
||||
submodule.process_weights_after_loading()
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if self.input_transform is not None:
|
||||
x = self.input_transform(x)
|
||||
|
||||
assert bias is None
|
||||
x = self.quant_method.apply(layer, x, bias)
|
||||
|
||||
# In most cases, input transforms are preferred over output transforms
|
||||
# (@ksayers): confirm that this is done concurrently
|
||||
if self.output_transform is not None:
|
||||
for part_id, (start, length) in enumerate(self.partition_ranges):
|
||||
x[:, start : start + length] = self.output_transform(
|
||||
x[:, start : start + length].clone(), part_id=part_id
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
def _validate_tfm_schemes(self, num_partitions: int):
|
||||
if len(self.input_tfms) > 0:
|
||||
if 0 not in self.input_tfms:
|
||||
raise ValueError("Must have same input")
|
||||
|
||||
for part_index in range(num_partitions):
|
||||
if self.input_tfms[part_index] != self.input_tfms[0]:
|
||||
raise ValueError("Must have same input")
|
||||
|
||||
if len(self.output_tfms) > 0:
|
||||
scheme_name = list(self.output_tfms.values())[0].scheme_name
|
||||
location = list(self.output_tfms.values())[0].args.location
|
||||
|
||||
for tfm in self.output_tfms.values():
|
||||
if tfm.scheme_name != scheme_name:
|
||||
raise ValueError("Must have same scheme name")
|
||||
if tfm.args.location != location:
|
||||
raise ValueError("Must have same location")
|
||||
|
||||
return self.input_tfms, self.output_tfms
|
||||
|
||||
|
||||
def get_linear_transform_schemes(
|
||||
layer: torch.nn.Module,
|
||||
layer_name: str,
|
||||
transform_config: TransformConfig | None,
|
||||
packed_modules_mapping: dict[str, list[str]],
|
||||
) -> tuple[
|
||||
dict[int, TransformTuple], dict[int, TransformTuple]
|
||||
]: # [input_transform, [output_transform, ...]]
|
||||
# there can only be one transform input scheme per (fused) module
|
||||
input_tfms = {}
|
||||
output_tfms = {}
|
||||
|
||||
partition_names = get_layer_partition_names(layer_name, packed_modules_mapping)
|
||||
|
||||
for scheme_name, scheme, args in get_schemes_args(transform_config):
|
||||
for part_index, part_name in enumerate(partition_names):
|
||||
if (
|
||||
is_match(part_name, layer, args.targets, args.ignore)
|
||||
and args.is_online()
|
||||
):
|
||||
if args.location == TransformLocation.INPUT:
|
||||
input_tfms[part_index] = TransformTuple(scheme_name, scheme, args)
|
||||
|
||||
elif args.location == TransformLocation.OUTPUT:
|
||||
output_tfms[part_index] = TransformTuple(scheme_name, scheme, args)
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Cannot apply `{args.location}` transform to `{layer_name}`"
|
||||
)
|
||||
|
||||
return (input_tfms, output_tfms)
|
||||
|
||||
|
||||
def get_schemes_args(
|
||||
transform_config: TransformConfig | None,
|
||||
) -> Generator[tuple[str, TransformScheme, TransformArgs]]:
|
||||
if transform_config is None:
|
||||
return
|
||||
|
||||
for scheme_name, scheme in transform_config.config_groups.items():
|
||||
for args in scheme.apply:
|
||||
yield (scheme_name, scheme, args)
|
||||
|
||||
|
||||
def get_layer_partition_names(
|
||||
layer_name: str, packed_modules_mapping: dict[str, list[str]]
|
||||
) -> list[str]:
|
||||
"""
|
||||
Get all partition names associated with this layer.
|
||||
Names are returned in order of their partition indices.
|
||||
|
||||
```python
|
||||
mapping = {"gate_up_proj", "gate_proj", "up_proj"}
|
||||
|
||||
assert get_layer_partition_names("mlp.gate_up_proj", mapping) == [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
]
|
||||
assert get_layer_partition_names("mlp.down_proj", mapping) == ["down_proj"]"""
|
||||
for fused_suffix, part_suffixes in packed_modules_mapping.items():
|
||||
if layer_name.endswith(fused_suffix):
|
||||
return [
|
||||
layer_name.removesuffix(fused_suffix) + part_suffix
|
||||
for part_suffix in part_suffixes
|
||||
]
|
||||
|
||||
return [layer_name]
|
||||
@@ -0,0 +1,173 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import math
|
||||
from collections.abc import Callable, Hashable
|
||||
|
||||
import torch
|
||||
from compressed_tensors.transform import (
|
||||
TransformArgs,
|
||||
TransformLocation,
|
||||
TransformScheme,
|
||||
)
|
||||
from torch import Tensor
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
|
||||
from vllm.model_executor.layers.linear import LinearBase
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.transform.utils import ( # noqa: E501
|
||||
TransformTuple,
|
||||
)
|
||||
from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
|
||||
from vllm.model_executor.parameter import SharedWeightParameter
|
||||
|
||||
|
||||
class HadamardTransform(torch.nn.Module):
|
||||
"""
|
||||
Class which handles weight loading, postprocessing, and application of
|
||||
transforms. Meant to be used with `CompressedTensorsLinearTransformMethod`
|
||||
and attention transforms method (not implemented yet)
|
||||
"""
|
||||
|
||||
transforms: dict[int, TransformTuple] # info parsed from transforms config
|
||||
weight: SharedWeightParameter # container for shared tensors
|
||||
|
||||
scales: dict[int, float] # hadamard scale, usually sqrt(matrix.size(0))
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
transforms: dict[int, TransformTuple],
|
||||
layer: torch.nn.Module,
|
||||
weight_loader: Callable,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
):
|
||||
super().__init__()
|
||||
self.transforms = transforms
|
||||
self.scales = {}
|
||||
|
||||
if get_tensor_model_parallel_world_size() > 1:
|
||||
raise NotImplementedError(
|
||||
"Online transforms with tensor parallelism is not supported"
|
||||
)
|
||||
|
||||
# Similar to row/col parallel params, but tensors are separate
|
||||
# to allow for loading with shared memory
|
||||
self.weight = SharedWeightParameter(weight_loader=weight_loader)
|
||||
|
||||
# create shared partition data for each partition of the original weight
|
||||
input_size = input_size_per_partition
|
||||
for part_index, (_scheme_name, scheme, args) in self.transforms.items():
|
||||
output_size = output_partition_sizes[part_index]
|
||||
weight_size = self._get_weight_size(
|
||||
layer, scheme, args, input_size, output_size
|
||||
)
|
||||
|
||||
data_key = self._get_data_key(scheme, weight_size)
|
||||
self.weight.add_partition(
|
||||
part_index,
|
||||
data_key,
|
||||
size=(weight_size, weight_size),
|
||||
dtype=scheme.precision,
|
||||
)
|
||||
|
||||
# validate that shared tensors and schemes are correct
|
||||
self._validate_input_transforms()
|
||||
|
||||
def process_weights_after_loading(self):
|
||||
for part_id in self.weight.partitions:
|
||||
data = self.weight.partitions[part_id].data
|
||||
|
||||
# required by torch.compile
|
||||
self.weight.process_weights_after_loading()
|
||||
|
||||
# precompute scale as a runtime multiply, not division
|
||||
# do not fold into weight in order to utilize FWHT
|
||||
self.scales[part_id] = 1 / math.sqrt(data.size(0))
|
||||
|
||||
# FUTURE: avoid runtime transpose by processing weights
|
||||
# prior to apply
|
||||
|
||||
def forward(self, value: Tensor, part_id: int = 0) -> Tensor:
|
||||
if part_id not in self.weight.partitions:
|
||||
return value
|
||||
|
||||
# use hadacore if possible
|
||||
if self.transforms[part_id].scheme.type == "hadamard":
|
||||
if self.transforms[part_id].scheme.head_dim is not None:
|
||||
weight_size = self.transforms[part_id].scheme.head_dim
|
||||
value = value.unflatten(-1, (-1, weight_size))
|
||||
value = ops.hadacore_transform(value)
|
||||
value = value.flatten(-2, -1)
|
||||
|
||||
return value
|
||||
|
||||
# sylvester transforms are symmetric, inv => transpose => original
|
||||
return ops.hadacore_transform(value)
|
||||
|
||||
# fall back to dense
|
||||
else:
|
||||
weight = self.weight.partitions[part_id]
|
||||
weight = (
|
||||
weight if self.transforms[part_id].args.inverse else weight.T
|
||||
) # linear := x(W.T)
|
||||
scale = self.scales[part_id]
|
||||
|
||||
if self.transforms[part_id].scheme.head_dim is not None:
|
||||
value = value.unflatten(-1, (-1, weight.size(0)))
|
||||
value = (
|
||||
dispatch_unquantized_gemm()(
|
||||
self, value.to(weight.dtype), weight, None
|
||||
).to(value.dtype)
|
||||
* scale
|
||||
)
|
||||
value = value.flatten(-2, -1)
|
||||
|
||||
return value
|
||||
|
||||
return (
|
||||
dispatch_unquantized_gemm()(
|
||||
self, value.to(weight.dtype), weight, None
|
||||
).to(value.dtype)
|
||||
* scale
|
||||
)
|
||||
|
||||
def _get_data_key(self, scheme: TransformScheme, weight_size: int) -> Hashable:
|
||||
return (id(scheme), weight_size)
|
||||
|
||||
def _get_weight_size(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
scheme: TransformScheme,
|
||||
args: TransformArgs,
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
) -> int:
|
||||
if scheme.head_dim is not None:
|
||||
return scheme.head_dim
|
||||
|
||||
if isinstance(layer, LinearBase):
|
||||
if args.location == TransformLocation.INPUT:
|
||||
return input_size
|
||||
|
||||
elif args.location == TransformLocation.OUTPUT:
|
||||
return output_size
|
||||
|
||||
elif isinstance(layer, VocabParallelEmbedding):
|
||||
if args.location == TransformLocation.INPUT:
|
||||
return output_size
|
||||
|
||||
elif args.location == TransformLocation.OUTPUT:
|
||||
return input_size
|
||||
|
||||
raise ValueError()
|
||||
|
||||
def _validate_input_transforms(self):
|
||||
assert len(self.transforms) > 0
|
||||
location = list(self.transforms.values())[0].args.location
|
||||
|
||||
if location == TransformLocation.INPUT:
|
||||
first_data = self.weight.partitions[0].data
|
||||
for partition in self.weight.partitions.values():
|
||||
if partition.data.data_ptr() != first_data.data_ptr():
|
||||
raise ValueError("")
|
||||
+64
@@ -0,0 +1,64 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
|
||||
CompressedTensorsScheme,
|
||||
CompressedTensorsW4A4Fp4,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.transform.linear import ( # noqa: E501
|
||||
CompressedTensorsLinearTransformMethod,
|
||||
TransformTuple,
|
||||
)
|
||||
|
||||
__all__ = ["is_qutlass_fp4_scheme", "QutlassNvFP4LinearMethod"]
|
||||
|
||||
|
||||
def is_qutlass_fp4_scheme(
|
||||
quant_scheme: CompressedTensorsScheme | None,
|
||||
input_tfms: dict[int, TransformTuple],
|
||||
) -> bool:
|
||||
return (
|
||||
isinstance(quant_scheme, (CompressedTensorsW4A4Fp4,))
|
||||
and len(input_tfms) == 1
|
||||
and input_tfms[0].scheme.head_dim == quant_scheme.group_size
|
||||
)
|
||||
|
||||
|
||||
class QutlassNvFP4LinearMethod(CompressedTensorsLinearTransformMethod):
|
||||
def create_weights(
|
||||
self,
|
||||
layer,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
input_size,
|
||||
output_size,
|
||||
params_dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
# initializes fp4 qparams
|
||||
assert isinstance(layer.scheme, (CompressedTensorsW4A4Fp4,))
|
||||
ret = super().create_weights(
|
||||
layer,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
input_size,
|
||||
output_size,
|
||||
params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
|
||||
assert self.input_transform is not None
|
||||
assert len(self.input_transform.weight) == 1
|
||||
assert self.input_transform.weight[0].size(0) == layer.scheme.group_size
|
||||
|
||||
return ret
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import NamedTuple
|
||||
|
||||
from compressed_tensors.transform import TransformArgs, TransformScheme
|
||||
|
||||
__all__ = ["TransformTuple"]
|
||||
|
||||
|
||||
class TransformTuple(NamedTuple):
|
||||
scheme_name: str
|
||||
scheme: TransformScheme
|
||||
args: TransformArgs
|
||||
@@ -0,0 +1,224 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
|
||||
def is_weak_contiguous(x: torch.Tensor):
|
||||
strides = x.stride()
|
||||
sizes = x.shape
|
||||
is_not_transpose = strides[0] == 1 and (strides[1] >= max(1, sizes[0]))
|
||||
is_transpose = strides[1] == 1 and (strides[0] >= max(1, sizes[1]))
|
||||
return is_transpose or is_not_transpose
|
||||
|
||||
|
||||
@triton.jit
|
||||
def scaled_mm_kernel(
|
||||
a_ptr,
|
||||
b_ptr,
|
||||
scale_a_ptr,
|
||||
scale_b_ptr,
|
||||
c_ptr,
|
||||
bias_ptr,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
stride_am,
|
||||
stride_ak,
|
||||
stride_bk,
|
||||
stride_bn,
|
||||
stride_cm,
|
||||
stride_cn,
|
||||
ACCUMULATOR_DTYPE: tl.constexpr,
|
||||
BLOCK_SIZE_M: tl.constexpr,
|
||||
BLOCK_SIZE_N: tl.constexpr,
|
||||
BLOCK_SIZE_K: tl.constexpr,
|
||||
BLOCK_SIZE_SCALE_A: tl.constexpr,
|
||||
BLOCK_SIZE_SCALE_B: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(axis=0)
|
||||
|
||||
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
|
||||
|
||||
pid_m = pid // num_pid_n
|
||||
pid_n = pid % num_pid_n
|
||||
|
||||
accumulator_dtype = ACCUMULATOR_DTYPE
|
||||
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype)
|
||||
|
||||
# NOTE: Some tensor inputs are so large, they will cause int32 overflow
|
||||
# so it is necessary to use tl.int64 for all the offsets, else SEGV will
|
||||
# eventually occur.
|
||||
|
||||
# Offsets and masks.
|
||||
offsets_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
|
||||
masks_am = offsets_am < M
|
||||
|
||||
offsets_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
|
||||
masks_bn = offsets_bn < N
|
||||
|
||||
offsets_k = tl.arange(0, BLOCK_SIZE_K).to(tl.int64)
|
||||
offsets_a = stride_am * offsets_am[:, None] + stride_ak * offsets_k[None, :]
|
||||
offsets_b = stride_bk * offsets_k[:, None] + stride_bn * offsets_bn[None, :]
|
||||
|
||||
# NOTE: BLOCK_SIZE_SCALE_A could be 1 or BLOCK_SIZE_M, so need to create
|
||||
# appropriate offsets and masks for each case. Same goes for
|
||||
# BLOCK_SIZE_SCALE_B.
|
||||
offsets_scale_am = (
|
||||
tl.arange(0, BLOCK_SIZE_SCALE_A)
|
||||
+ (BLOCK_SIZE_SCALE_A > 1) * pid_m * BLOCK_SIZE_M
|
||||
)
|
||||
masks_scale_am = offsets_scale_am < M
|
||||
|
||||
offsets_scale_bn = (
|
||||
tl.arange(0, BLOCK_SIZE_SCALE_B)
|
||||
+ (BLOCK_SIZE_SCALE_B > 1) * pid_n * BLOCK_SIZE_N
|
||||
)
|
||||
masks_scale_bn = offsets_scale_bn < N
|
||||
|
||||
a_ptrs = a_ptr + offsets_a
|
||||
b_ptrs = b_ptr + offsets_b
|
||||
|
||||
scale_a_ptrs = scale_a_ptr + offsets_scale_am
|
||||
scale_b_ptrs = scale_b_ptr + offsets_scale_bn
|
||||
|
||||
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
|
||||
masks_k = offsets_k < K
|
||||
masks_a = masks_am[:, None] & masks_k[None, :]
|
||||
a = tl.load(a_ptrs, mask=masks_a)
|
||||
|
||||
masks_b = masks_k[:, None] & masks_bn[None, :]
|
||||
b = tl.load(b_ptrs, mask=masks_b)
|
||||
|
||||
# Accumulate results.
|
||||
accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
|
||||
|
||||
offsets_k += BLOCK_SIZE_K
|
||||
a_ptrs += BLOCK_SIZE_K * stride_ak
|
||||
b_ptrs += BLOCK_SIZE_K * stride_bk
|
||||
|
||||
# Apply scale at end.
|
||||
masks_scale_a = masks_scale_am[:, None] & (tl.arange(0, 1) < 1)[:, None]
|
||||
scale_a = tl.load(scale_a_ptrs[:, None], masks_scale_a)
|
||||
# Need to broadcast to the appropriate size, if scale_a is already
|
||||
# (BLOCK_SIZE_M, 1) then it will broadcast to its own shape. Same goes
|
||||
# for scale_b below.
|
||||
scale_a = scale_a.broadcast_to((BLOCK_SIZE_M, 1))
|
||||
accumulator = scale_a * accumulator.to(tl.float32)
|
||||
|
||||
masks_scale_b = masks_scale_bn[:, None] & (tl.arange(0, 1) < 1)[None, :]
|
||||
scale_b = tl.load(scale_b_ptrs[:, None], masks_scale_b)
|
||||
scale_b = scale_b.broadcast_to((BLOCK_SIZE_N, 1))
|
||||
accumulator = scale_b.T * accumulator.to(tl.float32)
|
||||
|
||||
# Convert to output format.
|
||||
c = accumulator.to(c_ptr.type.element_ty)
|
||||
|
||||
# Add bias, it's already in output format, so add it after conversion.
|
||||
if bias_ptr:
|
||||
offsets_bias = offsets_bn
|
||||
bias_ptrs = bias_ptr + offsets_bias
|
||||
bias_mask = offsets_bias < N
|
||||
bias = tl.load(bias_ptrs, bias_mask)
|
||||
c += bias
|
||||
|
||||
# Save output
|
||||
offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
|
||||
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
|
||||
offs_cm = offs_cm.to(tl.int64)
|
||||
offs_cn = offs_cn.to(tl.int64)
|
||||
c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
|
||||
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
|
||||
|
||||
tl.store(c_ptrs, c, mask=c_mask)
|
||||
|
||||
|
||||
# input - [M, K]
|
||||
# weight - [K, N]
|
||||
def triton_scaled_mm(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
scale_a: torch.Tensor,
|
||||
scale_b: torch.Tensor,
|
||||
out_dtype: type[torch.dtype],
|
||||
bias: torch.Tensor | None = None,
|
||||
block_size_m: int = 32,
|
||||
block_size_n: int = 32,
|
||||
block_size_k: int = 32,
|
||||
use_heuristic=True,
|
||||
) -> torch.Tensor:
|
||||
M, K = input.shape
|
||||
N = weight.shape[1]
|
||||
|
||||
assert N > 0 and K > 0 and M > 0
|
||||
assert weight.shape[0] == K
|
||||
assert input.dtype == weight.dtype
|
||||
|
||||
scale_a = scale_a.reshape(-1, 1) if scale_a.dim() <= 1 else scale_a
|
||||
scale_b = scale_b.reshape(-1, 1) if scale_b.dim() <= 1 else scale_b
|
||||
|
||||
assert scale_a.dtype == scale_b.dtype and scale_a.is_floating_point()
|
||||
assert scale_a.shape[1] == 1 and (scale_a.shape[0] == 1 or scale_a.shape[0] == M)
|
||||
assert scale_b.shape[1] == 1 and (scale_b.shape[0] == 1 or scale_b.shape[0] == N)
|
||||
assert out_dtype.is_floating_point
|
||||
assert bias is None or bias.is_floating_point()
|
||||
assert is_weak_contiguous(input)
|
||||
assert is_weak_contiguous(weight)
|
||||
|
||||
grid = lambda META: (
|
||||
triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
|
||||
)
|
||||
|
||||
result = torch.empty((M, N), dtype=out_dtype, device=input.device)
|
||||
|
||||
has_scalar = lambda x: x.shape[0] == 1 and x.shape[1] == 1
|
||||
|
||||
if use_heuristic:
|
||||
is_small_N = N < 8192
|
||||
next_power_of_2_M = max(32, triton.next_power_of_2(M))
|
||||
if next_power_of_2_M <= 32:
|
||||
tile_shape = (64, 64, 256) if is_small_N else (64, 128, 256)
|
||||
elif next_power_of_2_M <= 64:
|
||||
tile_shape = (64, 64, 256)
|
||||
elif next_power_of_2_M <= 128:
|
||||
tile_shape = (64, 128, 128)
|
||||
else:
|
||||
tile_shape = (128, 128, 128)
|
||||
|
||||
block_size_m, block_size_n, block_size_k = tile_shape
|
||||
|
||||
block_size_sa = 1 if has_scalar(scale_a) else block_size_m
|
||||
block_size_sb = 1 if has_scalar(scale_b) else block_size_n
|
||||
|
||||
accumulator_dtype = tl.float32 if input.is_floating_point() else tl.int32
|
||||
|
||||
# A = input, B = weight, C = result
|
||||
# A = M x K, B = K x N, C = M x N
|
||||
scaled_mm_kernel[grid](
|
||||
input,
|
||||
weight,
|
||||
scale_a,
|
||||
scale_b,
|
||||
result,
|
||||
bias,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
input.stride(0),
|
||||
input.stride(1),
|
||||
weight.stride(0),
|
||||
weight.stride(1),
|
||||
result.stride(0),
|
||||
result.stride(1),
|
||||
accumulator_dtype,
|
||||
BLOCK_SIZE_M=block_size_m,
|
||||
BLOCK_SIZE_N=block_size_n,
|
||||
BLOCK_SIZE_K=block_size_k,
|
||||
BLOCK_SIZE_SCALE_A=block_size_sa,
|
||||
BLOCK_SIZE_SCALE_B=block_size_sb,
|
||||
)
|
||||
|
||||
return result.to(out_dtype)
|
||||
@@ -0,0 +1,239 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Iterable, Mapping
|
||||
from types import MappingProxyType
|
||||
|
||||
import regex as re
|
||||
from compressed_tensors import CompressionFormat
|
||||
from compressed_tensors.quantization import QuantizationStrategy
|
||||
from torch.nn import Module
|
||||
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kFp8Static128BlockSym,
|
||||
kFp8StaticChannelSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
BlockQuantScaleParameter,
|
||||
ChannelQuantScaleParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
|
||||
# Maps quantization strategy to the corresponding scale parameter type.
|
||||
# Shared across compressed-tensor scheme classes (w8a16_fp8, w8a8_fp8, …).
|
||||
STRATEGY_TO_PARAMETER_TYPE = {
|
||||
QuantizationStrategy.BLOCK: BlockQuantScaleParameter,
|
||||
QuantizationStrategy.CHANNEL: ChannelQuantScaleParameter,
|
||||
QuantizationStrategy.TENSOR: PerTensorScaleParameter,
|
||||
}
|
||||
|
||||
# Maps quantization strategy to the vLLM weight-quant key used for
|
||||
# kernel selection. Shared across compressed-tensor scheme classes.
|
||||
STRATEGY_TO_WEIGHT_QUANT_KEY = {
|
||||
QuantizationStrategy.BLOCK: kFp8Static128BlockSym,
|
||||
QuantizationStrategy.CHANNEL: kFp8StaticChannelSym,
|
||||
QuantizationStrategy.TENSOR: kFp8StaticTensorSym,
|
||||
}
|
||||
|
||||
|
||||
def is_activation_quantization_format(format: str) -> bool:
|
||||
_ACTIVATION_QUANTIZATION_FORMATS = [
|
||||
CompressionFormat.naive_quantized.value,
|
||||
CompressionFormat.int_quantized.value,
|
||||
CompressionFormat.float_quantized.value,
|
||||
CompressionFormat.nvfp4_pack_quantized.value,
|
||||
]
|
||||
return format in _ACTIVATION_QUANTIZATION_FORMATS
|
||||
|
||||
|
||||
def should_ignore_layer(
|
||||
layer_name: str | None,
|
||||
ignore: Iterable[str] = tuple(),
|
||||
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 and layer_name not in ignore:
|
||||
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 find_matched_target(
|
||||
layer_name: str | None,
|
||||
module: Module,
|
||||
targets: Iterable[str],
|
||||
fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
|
||||
) -> str | None:
|
||||
"""
|
||||
Helper function to look up which "target" in the compressed-tensors
|
||||
config that a layer corresponds to.
|
||||
|
||||
Recall that a compressed-tensors configs has a concept of
|
||||
config_groups, where each layer can be quantized with a different
|
||||
scheme.
|
||||
|
||||
targets in each config_group will be a list of either layer names
|
||||
(or regexes corresponding to layer names) or names of torch Modules.
|
||||
|
||||
First, we try to match the layer_name with a target
|
||||
Second, we try to match the module's name with a target
|
||||
Third, we try to map the layer_name to a list of fused module names.
|
||||
*All* component module names must match in order for a match to be
|
||||
successful. A successful match returns the first component target
|
||||
|
||||
Args:
|
||||
layer_name: layer name
|
||||
module: torch.nn.Module
|
||||
targets: list of targets to match the layer against
|
||||
fused_mapping: map from fused layer names to its components
|
||||
"""
|
||||
|
||||
if layer_name is None:
|
||||
layer_name = ""
|
||||
|
||||
matched_target = (
|
||||
_find_first_match(layer_name, targets)
|
||||
or _find_first_match(module.__class__.__name__, targets, True)
|
||||
or _match_fused_layer(layer_name, targets, fused_mapping)
|
||||
)
|
||||
|
||||
return matched_target
|
||||
|
||||
|
||||
def _find_first_match(
|
||||
value: str, targets: Iterable[str], check_contains: bool = False
|
||||
) -> str | None:
|
||||
"""
|
||||
Returns first element of target that matches value either
|
||||
exactly or as a regex after 're:'. If check_contains is set to True,
|
||||
additionally checks if the target string is contained within the value.
|
||||
|
||||
Args:
|
||||
value: string to compare the list of targets against
|
||||
targets: list of targets to match the layer against
|
||||
check_contains: whether or not to do a substring match
|
||||
"""
|
||||
|
||||
for target in targets:
|
||||
if _is_equal_or_regex_match(value, target, check_contains=check_contains):
|
||||
return target
|
||||
return None
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def _match_fused_layer(
|
||||
layer_name: str,
|
||||
target_layers: Iterable[str],
|
||||
fused_mapping: Mapping[str, list[str]],
|
||||
) -> str | None:
|
||||
"""
|
||||
Match a fused layer name to its corresponding individual layer in
|
||||
target_layers. Returns first value in fused_mapping which matches targets
|
||||
|
||||
Implements an "all" matching strategy where a fused layer matches iff
|
||||
"all" of its components match
|
||||
|
||||
Args:
|
||||
layer_name: layer name
|
||||
target_layers: list of targets to match the layer against
|
||||
fused_mapping: map from fused layer names to its components
|
||||
|
||||
Examples:
|
||||
layer_name = "model.layers.0.self_attn.qkv_proj"
|
||||
target_layers = ["model.layers.0.self_attn.q_proj",
|
||||
"model.layers.0.self_attn.k_proj",
|
||||
"model.layers.0.self_attn.v_proj"]
|
||||
"""
|
||||
# find layer_name in mapping
|
||||
fused = next((key for key in fused_mapping if layer_name.endswith(key)), None)
|
||||
if fused is None:
|
||||
return None
|
||||
|
||||
# expand path of unfused components
|
||||
unfused_paths = [
|
||||
layer_name.replace(fused, unfused) for unfused in fused_mapping[fused]
|
||||
]
|
||||
|
||||
# for each unfused component, find a match in targets
|
||||
unfused_matches: list[str | None] = []
|
||||
for unfused in unfused_paths:
|
||||
for target in target_layers:
|
||||
if _is_equal_or_regex_match(unfused, target):
|
||||
unfused_matches.append(target)
|
||||
break
|
||||
else:
|
||||
unfused_matches.append(None)
|
||||
|
||||
return unfused_matches[0] if all(unfused_matches) else None
|
||||
@@ -0,0 +1,60 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
RoutedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.online.int8 import (
|
||||
Int8OnlineMoEMethod,
|
||||
)
|
||||
|
||||
|
||||
class ExpertsInt8Config(QuantizationConfig):
|
||||
"""Online int8 quantization for MoE expert weights.
|
||||
Linear layers are left unquantized.
|
||||
|
||||
Backward-compatible config for ``--quantization experts_int8``.
|
||||
Prefer ``--quantization int8_per_channel``
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "experts_int8"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.bfloat16, torch.half]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 80
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "ExpertsInt8Config":
|
||||
return cls()
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
if isinstance(layer, LinearBase):
|
||||
return UnquantizedLinearMethod()
|
||||
elif isinstance(layer, RoutedExperts):
|
||||
return Int8OnlineMoEMethod(layer=layer)
|
||||
return None
|
||||
@@ -0,0 +1,186 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.nn import Module
|
||||
from torch.nn.parameter 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.linear import (
|
||||
LinearBase,
|
||||
LinearMethodBase,
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
|
||||
prepare_fp8_layer_for_marlin,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
is_layer_skipped,
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8StaticTokenSym,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
normalize_e4m3fn_to_e4m3fnuz,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
ChannelQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class FBGEMMFp8Config(QuantizationConfig):
|
||||
"""Config class for FBGEMM Fp8."""
|
||||
|
||||
def __init__(self, ignore_list: list[str], input_scale_ub: float):
|
||||
super().__init__()
|
||||
self.ignore_list = ignore_list if ignore_list else []
|
||||
self.input_scale_ub = input_scale_ub
|
||||
|
||||
# For GPUs that lack FP8 hardware support, we can leverage the Marlin
|
||||
# kernel for fast weight-only FP8 quantization
|
||||
self.use_marlin = not current_platform.has_device_capability(89)
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "fbgemm_fp8"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.bfloat16, torch.float16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 80
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "FBGEMMFp8Config":
|
||||
ignore_list = cls.get_from_keys(config, ["modules_to_not_convert"])
|
||||
input_scale_ub = cls.get_from_keys(config, ["activation_scale_ub"])
|
||||
return cls(ignore_list=ignore_list, input_scale_ub=input_scale_ub)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
if isinstance(layer, LinearBase):
|
||||
if is_layer_skipped(
|
||||
prefix=prefix,
|
||||
ignored_layers=self.ignore_list,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedLinearMethod()
|
||||
return FBGEMMFp8LinearMethod(self)
|
||||
return None
|
||||
|
||||
|
||||
class FBGEMMFp8LinearMethod(LinearMethodBase):
|
||||
def __init__(self, quant_config: FBGEMMFp8Config):
|
||||
self.quant_config = quant_config
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
self.input_dtype = get_current_vllm_config().model_config.dtype
|
||||
|
||||
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,
|
||||
):
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
del input_size, output_size
|
||||
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
|
||||
layer.orig_dtype = params_dtype
|
||||
|
||||
# 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
|
||||
weight_scale = ChannelQuantScaleParameter(
|
||||
data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
weight_scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
# INPUT SCALE UPPER BOUND
|
||||
input_scale_ub = torch.nn.Parameter(
|
||||
torch.tensor((self.quant_config.input_scale_ub), dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.input_scale_ub = input_scale_ub
|
||||
|
||||
self.fp8_linear = init_fp8_linear_kernel(
|
||||
activation_quant_key=kFp8DynamicTokenSym,
|
||||
weight_quant_key=kFp8StaticTokenSym,
|
||||
weight_shape=layer.weight.shape,
|
||||
input_dtype=self.input_dtype,
|
||||
out_dtype=self.out_dtype,
|
||||
module_name=self.__class__.__name__,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
# required by torch.compile
|
||||
layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
|
||||
layer.weight = Parameter(layer.weight.data, requires_grad=False)
|
||||
|
||||
weight = layer.weight
|
||||
|
||||
if current_platform.is_fp8_fnuz():
|
||||
weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
|
||||
weight=weight, weight_scale=layer.weight_scale, input_scale=None
|
||||
)
|
||||
if input_scale is not None:
|
||||
layer.input_scale = Parameter(input_scale, requires_grad=False)
|
||||
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
|
||||
|
||||
layer.weight = Parameter(weight.t(), requires_grad=False)
|
||||
if self.quant_config.use_marlin:
|
||||
prepare_fp8_layer_for_marlin(layer)
|
||||
# Activations not quantized for marlin.
|
||||
del layer.input_scale_ub
|
||||
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
|
||||
def apply(
|
||||
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,865 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
|
||||
import vllm.envs as envs
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.distributed import get_tensor_model_parallel_world_size
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
init_fp8_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.scaled_mm import (
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
MarlinFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.attention import Attention
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEMethodBase,
|
||||
FusedMoeWeightScaleSupported,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
UnquantizedFusedMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
Fp8MoeBackend,
|
||||
convert_to_fp8_moe_kernel_format,
|
||||
make_fp8_moe_kernel,
|
||||
make_fp8_moe_quant_config,
|
||||
select_fp8_moe_backend,
|
||||
)
|
||||
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 (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
create_fp8_input_scale,
|
||||
create_fp8_scale_parameter,
|
||||
create_fp8_weight_parameter,
|
||||
process_fp8_input_tensor_strategy_moe,
|
||||
process_fp8_weight_tensor_strategy,
|
||||
process_fp8_weight_tensor_strategy_moe,
|
||||
validate_fp8_block_shape,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
get_marlin_input_dtype,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
create_fp8_quant_key,
|
||||
is_layer_skipped,
|
||||
kFp8Dynamic128Sym,
|
||||
kFp8DynamicTensorSym,
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8Static128BlockSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
cutlass_block_fp8_supported,
|
||||
cutlass_fp8_supported,
|
||||
normalize_e4m3fn_to_e4m3fnuz,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
BlockQuantScaleParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.deep_gemm import (
|
||||
is_deep_gemm_supported,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
ACTIVATION_SCHEMES = ["static", "dynamic"]
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class Fp8Config(QuantizationConfig):
|
||||
"""Config class for FP8."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
is_checkpoint_fp8_serialized: bool = False,
|
||||
activation_scheme: str = "dynamic",
|
||||
ignored_layers: list[str] | None = None,
|
||||
weight_block_size: list[int] | None = None,
|
||||
store_dtype: str | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
|
||||
|
||||
if activation_scheme not in ACTIVATION_SCHEMES:
|
||||
raise ValueError(f"Unsupported activation scheme {activation_scheme}")
|
||||
self.activation_scheme = activation_scheme
|
||||
self.ignored_layers = ignored_layers or []
|
||||
self.store_dtype = store_dtype
|
||||
if weight_block_size is not None:
|
||||
if not is_checkpoint_fp8_serialized:
|
||||
raise ValueError(
|
||||
"The block-wise quantization only supports fp8-serialized "
|
||||
"checkpoint for now."
|
||||
)
|
||||
if len(weight_block_size) != 2:
|
||||
raise ValueError(
|
||||
"The quantization block size of weight must have 2 "
|
||||
f"dimensions, but got {len(weight_block_size)} dimensions"
|
||||
)
|
||||
if activation_scheme != "dynamic":
|
||||
raise ValueError(
|
||||
"The block-wise quantization only supports "
|
||||
"dynamic activation scheme for now, but got "
|
||||
f"{activation_scheme} activation scheme."
|
||||
)
|
||||
self.weight_block_size = weight_block_size
|
||||
self.use_deep_gemm: bool | None = None
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "fp8"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.bfloat16, torch.half]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 75
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return []
|
||||
|
||||
def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
|
||||
if self.ignored_layers is not None:
|
||||
self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers)
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "Fp8Config":
|
||||
quant_method = cls.get_from_keys(config, ["quant_method"])
|
||||
is_checkpoint_fp8_serialized = "fp8" in quant_method
|
||||
activation_scheme = cls.get_from_keys(config, ["activation_scheme"])
|
||||
ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None)
|
||||
weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None)
|
||||
store_dtype = cls.get_from_keys_or(config, ["store_dtype"], None)
|
||||
if not ignored_layers:
|
||||
ignored_layers = cls.get_from_keys_or(
|
||||
config, ["modules_to_not_convert"], None
|
||||
)
|
||||
return cls(
|
||||
is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
|
||||
activation_scheme=activation_scheme,
|
||||
ignored_layers=ignored_layers,
|
||||
weight_block_size=weight_block_size,
|
||||
store_dtype=store_dtype,
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
if isinstance(layer, LinearBase):
|
||||
if is_layer_skipped(
|
||||
prefix=prefix,
|
||||
ignored_layers=self.ignored_layers,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedLinearMethod()
|
||||
if not self.is_checkpoint_fp8_serialized:
|
||||
from vllm.model_executor.layers.quantization.online.fp8 import (
|
||||
Fp8PerTensorOnlineLinearMethod,
|
||||
)
|
||||
|
||||
online_method = Fp8PerTensorOnlineLinearMethod()
|
||||
online_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
|
||||
return online_method
|
||||
else:
|
||||
offline_method = Fp8LinearMethod(self)
|
||||
offline_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
|
||||
return offline_method
|
||||
elif isinstance(layer, RoutedExperts):
|
||||
if is_layer_skipped(
|
||||
prefix=prefix,
|
||||
ignored_layers=self.ignored_layers,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
if self.store_dtype == "mxfp4":
|
||||
from vllm.model_executor.layers.quantization.mxfp4 import (
|
||||
Mxfp4MoEMethod,
|
||||
)
|
||||
|
||||
return Mxfp4MoEMethod(layer.moe_config)
|
||||
if self.is_checkpoint_fp8_serialized:
|
||||
return Fp8MoEMethod(self, layer)
|
||||
else:
|
||||
from vllm.model_executor.layers.quantization.online.fp8 import (
|
||||
Fp8PerTensorOnlineMoEMethod,
|
||||
)
|
||||
|
||||
return Fp8PerTensorOnlineMoEMethod(layer=layer)
|
||||
elif isinstance(layer, Attention):
|
||||
return Fp8KVCacheMethod(self)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_cache_scale_mapper() -> "WeightsMapper":
|
||||
"""Map compressed-tensors KV-cache scale names to vLLM names."""
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
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 CopyNumelCounter(TorchDispatchMode):
|
||||
"""
|
||||
Tracks total number of elements modified with `copy_`. Useful for keeping
|
||||
track of weight loading where underlying weights can be arbitrarily
|
||||
transformed (such as with `narrow`) before calling copy.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.copied_numel = 0
|
||||
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
out = func(*args, **kwargs)
|
||||
if func == torch.ops.aten.copy_.default:
|
||||
self.copied_numel += args[0].numel()
|
||||
return out
|
||||
|
||||
|
||||
def _copy_missing_attrs(old: torch.Tensor, new: torch.Tensor) -> None:
|
||||
"""Copies any attrs present in `old` but not in `new` to `new`"""
|
||||
new_attrs = set(dir(new))
|
||||
attrs_to_set = {}
|
||||
for attr in dir(old):
|
||||
if attr not in new_attrs:
|
||||
attrs_to_set[attr] = getattr(old, attr)
|
||||
set_weight_attrs(new, attrs_to_set)
|
||||
|
||||
|
||||
class Fp8LinearMethod(LinearMethodBase):
|
||||
"""Linear method for FP8.
|
||||
Supports loading FP8 checkpoints with static weight scale and
|
||||
dynamic/static activation scale.
|
||||
|
||||
Limitations:
|
||||
1. Only support float8_e4m3fn data type due to the limitation of
|
||||
torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856)
|
||||
|
||||
Args:
|
||||
quant_config: The quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: Fp8Config):
|
||||
self.quant_config = quant_config
|
||||
self.is_scale_e8m0 = getattr(quant_config, "is_scale_e8m0", False)
|
||||
self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
self.input_dtype = get_current_vllm_config().model_config.dtype
|
||||
|
||||
# For GPUs that lack FP8 hardware support, we can leverage the Marlin
|
||||
# kernel for fast weight-only FP8 quantization
|
||||
self.marlin_input_dtype = None
|
||||
self.use_marlin = False
|
||||
|
||||
if self.quant_config.use_deep_gemm is not None:
|
||||
self.use_deep_gemm = self.quant_config.use_deep_gemm
|
||||
else:
|
||||
self.use_deep_gemm = is_deep_gemm_supported()
|
||||
|
||||
self.weight_block_size = self.quant_config.weight_block_size
|
||||
self.block_quant = self.weight_block_size is not None
|
||||
self.act_q_static = self.quant_config.activation_scheme == "static"
|
||||
|
||||
if self.block_quant:
|
||||
assert not self.act_q_static
|
||||
assert self.weight_block_size is not None
|
||||
|
||||
self.activation_quant_key = create_fp8_quant_key(
|
||||
static=self.act_q_static,
|
||||
group_shape=GroupShape(1, self.weight_block_size[0]),
|
||||
)
|
||||
self.weight_quant_key = create_fp8_quant_key(
|
||||
static=True, group_shape=GroupShape(*self.weight_block_size)
|
||||
)
|
||||
else:
|
||||
self.weight_quant_key = kFp8StaticTensorSym
|
||||
# Use per-token quantization for better perf if dynamic and cutlass
|
||||
if self.act_q_static:
|
||||
self.activation_quant_key = kFp8StaticTensorSym
|
||||
elif cutlass_fp8_supported():
|
||||
self.activation_quant_key = kFp8DynamicTokenSym
|
||||
else:
|
||||
self.activation_quant_key = kFp8DynamicTensorSym
|
||||
|
||||
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,
|
||||
):
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
layer.logical_widths = output_partition_sizes
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
if self.block_quant:
|
||||
assert self.weight_block_size is not None
|
||||
layer.weight_block_size = self.weight_block_size
|
||||
validate_fp8_block_shape(
|
||||
layer,
|
||||
input_size,
|
||||
output_size,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
self.weight_block_size,
|
||||
)
|
||||
|
||||
weight = create_fp8_weight_parameter(
|
||||
output_size_per_partition, input_size_per_partition, weight_loader
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# WEIGHT SCALE
|
||||
if not self.block_quant:
|
||||
scale = create_fp8_scale_parameter(
|
||||
PerTensorScaleParameter,
|
||||
output_partition_sizes,
|
||||
input_size_per_partition,
|
||||
None,
|
||||
weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", scale)
|
||||
else:
|
||||
assert not self.act_q_static
|
||||
assert self.weight_block_size is not None
|
||||
scale = create_fp8_scale_parameter(
|
||||
BlockQuantScaleParameter,
|
||||
output_partition_sizes,
|
||||
input_size_per_partition,
|
||||
self.weight_block_size,
|
||||
weight_loader,
|
||||
scale_dtype=(torch.float8_e8m0fnu if self.is_scale_e8m0 else None),
|
||||
)
|
||||
# The weight_scale_inv name is intentional for deepseekv3
|
||||
layer.register_parameter("weight_scale_inv", scale)
|
||||
|
||||
# INPUT ACTIVATION SCALE
|
||||
if self.act_q_static:
|
||||
scale = create_fp8_input_scale(output_partition_sizes, weight_loader)
|
||||
set_weight_attrs(scale, {"scale_type": "input_scale"})
|
||||
layer.register_parameter("input_scale", 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__,
|
||||
)
|
||||
|
||||
self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
if self.use_marlin:
|
||||
if not self.block_quant:
|
||||
# Canonicalize to (K, N) for the kernel.
|
||||
replace_parameter(layer, "weight", layer.weight.t())
|
||||
# Only Marlin kernels support `marlin_input_dtype`; guard to avoid
|
||||
# AttributeError if backend selection changes.
|
||||
if hasattr(self.fp8_linear, "marlin_input_dtype"):
|
||||
self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
return
|
||||
|
||||
input_scale = None
|
||||
# TODO(rob): refactor block quant into separate class.
|
||||
if self.block_quant:
|
||||
assert not self.act_q_static
|
||||
|
||||
# If checkpoint not serialized fp8, quantize the weights.
|
||||
else:
|
||||
# If checkpoint is fp8 per-tensor, handle that there are N scales for N
|
||||
# shards in a fused module
|
||||
weight = layer.weight
|
||||
weight_scale = layer.weight_scale
|
||||
|
||||
# If using w8a8, torch._scaled_mm needs per tensor, so
|
||||
# requantize the logical shards as a single weight.
|
||||
weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy(
|
||||
weight,
|
||||
weight_scale,
|
||||
layer.logical_widths,
|
||||
getattr(layer, "input_scale", None),
|
||||
)
|
||||
if self.act_q_static:
|
||||
assert input_scale is not None
|
||||
input_scale = input_scale.max()
|
||||
weight = weight.t()
|
||||
|
||||
# Update layer with new values.
|
||||
replace_parameter(layer, "weight", weight.data)
|
||||
replace_parameter(layer, "weight_scale", weight_scale.data)
|
||||
|
||||
if input_scale is not None:
|
||||
replace_parameter(layer, "input_scale", input_scale)
|
||||
else:
|
||||
layer.input_scale = None
|
||||
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
# if batch invariant mode is enabled, prefer direct FP8 path
|
||||
# we will use BF16 dequant when direct FP8 is not supported.
|
||||
if envs.VLLM_BATCH_INVARIANT:
|
||||
if self.block_quant:
|
||||
assert self.weight_block_size is not None
|
||||
return self.fp8_linear.apply_weights(
|
||||
layer,
|
||||
x,
|
||||
bias,
|
||||
)
|
||||
else:
|
||||
if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel):
|
||||
return self.fp8_linear.apply_weights(layer, x, bias)
|
||||
|
||||
# per-tensor/channel: dequant to BF16 and run GEMM
|
||||
weight_fp8 = layer.weight.to(torch.bfloat16)
|
||||
weight_scale = layer.weight_scale.to(torch.bfloat16)
|
||||
if weight_scale.numel() == 1:
|
||||
# Per-tensor: simple scalar multiplication
|
||||
weight_bf16 = weight_fp8 * weight_scale
|
||||
else:
|
||||
# Multiple scales (fused modules like QKV)
|
||||
# Try to infer correct broadcasting
|
||||
# weight is [K, N], scale could be [num_logical_weights]
|
||||
# Need to figure out how to broadcast - for now just try
|
||||
# direct multiplication
|
||||
if (
|
||||
weight_scale.dim() == 1
|
||||
and weight_scale.shape[0] == weight_fp8.shape[0]
|
||||
):
|
||||
# Per-row scaling
|
||||
weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1)
|
||||
else:
|
||||
# Fallback
|
||||
weight_bf16 = weight_fp8 * weight_scale
|
||||
return torch.nn.functional.linear(x, weight_bf16.t(), bias)
|
||||
|
||||
return self.fp8_linear.apply_weights(layer, x, bias)
|
||||
|
||||
|
||||
class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
"""MoE method for FP8.
|
||||
Supports loading FP8 checkpoints with static weight scale and
|
||||
dynamic/static activation scale.
|
||||
|
||||
Also supports loading quantized FP16/BF16 model checkpoints with dynamic
|
||||
activation scaling. The weight scaling factor will be initialized after
|
||||
the model weights are loaded.
|
||||
|
||||
Args:
|
||||
quant_config: The quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: Fp8Config, layer: RoutedExperts):
|
||||
super().__init__(layer.moe_config)
|
||||
self.quant_config = quant_config
|
||||
self.weight_block_size = self.quant_config.weight_block_size
|
||||
self.block_quant: bool = self.weight_block_size is not None
|
||||
self.weight_scale_name = (
|
||||
"weight_scale_inv" if self.block_quant else "weight_scale"
|
||||
)
|
||||
|
||||
# Set weight key and activation key for kernel compatibility
|
||||
if self.block_quant:
|
||||
weight_key = kFp8Static128BlockSym
|
||||
activation_key = kFp8Dynamic128Sym
|
||||
else:
|
||||
weight_key = kFp8StaticTensorSym
|
||||
activation_key = (
|
||||
kFp8StaticTensorSym
|
||||
if self.quant_config.activation_scheme == "static"
|
||||
else kFp8DynamicTensorSym
|
||||
)
|
||||
|
||||
# Select Fp8 MoE backend
|
||||
self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
|
||||
config=self.moe,
|
||||
weight_key=weight_key,
|
||||
activation_key=activation_key,
|
||||
allow_vllm_cutlass=False,
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.num_experts = num_experts
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
assert self.quant_config.is_checkpoint_fp8_serialized
|
||||
params_dtype = torch.float8_e4m3fn
|
||||
|
||||
if self.block_quant:
|
||||
assert self.weight_block_size is not None
|
||||
layer.weight_block_size = self.weight_block_size
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
block_n, block_k = (
|
||||
self.weight_block_size[0],
|
||||
self.weight_block_size[1],
|
||||
)
|
||||
# NOTE: To ensure proper alignment of the block-wise quantization
|
||||
# scales, the output_size of the weights for both the gate and up
|
||||
# layers must be divisible by block_n.
|
||||
# Required by column parallel or enabling merged weights
|
||||
if intermediate_size_per_partition % block_n != 0:
|
||||
raise ValueError(
|
||||
f"The output_size of gate's and up's weight = "
|
||||
f"{intermediate_size_per_partition} is not divisible by "
|
||||
f"weight quantization block_n = {block_n}."
|
||||
)
|
||||
if tp_size > 1 and intermediate_size_per_partition % block_k != 0:
|
||||
# Required by row parallel
|
||||
raise ValueError(
|
||||
f"The input_size of down's weight = "
|
||||
f"{intermediate_size_per_partition} is not divisible by "
|
||||
f"weight quantization block_k = {block_k}."
|
||||
)
|
||||
|
||||
# WEIGHTS
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
# BIASES (for models like GPT-OSS that have biased MoE)
|
||||
if self.moe.has_bias:
|
||||
w13_bias = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
dtype=layer.orig_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_bias", w13_bias)
|
||||
set_weight_attrs(w13_bias, extra_weight_attrs)
|
||||
w2_bias = torch.nn.Parameter(
|
||||
torch.zeros(num_experts, hidden_size, dtype=layer.orig_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_bias", w2_bias)
|
||||
set_weight_attrs(w2_bias, extra_weight_attrs)
|
||||
|
||||
# WEIGHT_SCALES
|
||||
if not self.block_quant:
|
||||
# For per-tensor quant, the scales are per expert and weight.
|
||||
w13_scale_data = torch.ones(num_experts, 2, dtype=torch.float32)
|
||||
w2_scale_data = torch.ones(num_experts, dtype=torch.float32)
|
||||
else:
|
||||
# For block quant, the scales are per block (typically 128x128).
|
||||
w13_scale_data = torch.ones(
|
||||
num_experts,
|
||||
2 * ((intermediate_size_per_partition + block_n - 1) // block_n),
|
||||
(hidden_size + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
w2_scale_data = torch.ones(
|
||||
num_experts,
|
||||
(hidden_size + block_n - 1) // block_n,
|
||||
(intermediate_size_per_partition + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
w13_weight_scale = torch.nn.Parameter(w13_scale_data, requires_grad=False)
|
||||
w2_weight_scale = torch.nn.Parameter(w2_scale_data, requires_grad=False)
|
||||
# Note: name is weight_scale for tensor, weight_scale_inv for block.
|
||||
layer.register_parameter(f"w13_{self.weight_scale_name}", w13_weight_scale)
|
||||
layer.register_parameter(f"w2_{self.weight_scale_name}", w2_weight_scale)
|
||||
|
||||
# Add the quantization method used (per tensor/grouped/channel)
|
||||
# to ensure the weight scales are loaded in properly
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
|
||||
if self.block_quant
|
||||
else {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
# INPUT_SCALES
|
||||
if self.quant_config.activation_scheme == "static":
|
||||
assert not self.block_quant
|
||||
w13_input_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w13_input_scale", w13_input_scale)
|
||||
set_weight_attrs(w13_input_scale, extra_weight_attrs)
|
||||
|
||||
w2_input_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w2_input_scale", w2_input_scale)
|
||||
set_weight_attrs(w2_input_scale, extra_weight_attrs)
|
||||
|
||||
else:
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
def _setup_kernel(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
w13_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
w13_input_scale: torch.Tensor | None,
|
||||
w2_input_scale: torch.Tensor | None,
|
||||
) -> None:
|
||||
# Shuffle weights to runtime format.
|
||||
w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
|
||||
fp8_backend=self.fp8_backend,
|
||||
layer=layer,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
w13_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
w13_input_scale=w13_input_scale,
|
||||
w2_input_scale=w2_input_scale,
|
||||
)
|
||||
|
||||
# Replace parameters with updated versions. Note that this helper
|
||||
# function ensures the replacement is compatible with RL weight reloads.
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
|
||||
replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
|
||||
|
||||
# AITER backend requires weights to be marked as shuffled.
|
||||
if self.fp8_backend == Fp8MoeBackend.AITER:
|
||||
layer.w13_weight.is_shuffled = True
|
||||
layer.w2_weight.is_shuffled = True
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert self.moe_quant_config is not None
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_fp8_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
fp8_backend=self.fp8_backend,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
# Allow for accessing weights and scales in standard way.
|
||||
w13 = layer.w13_weight
|
||||
w2 = layer.w2_weight
|
||||
w13_scale = getattr(layer, f"w13_{self.weight_scale_name}")
|
||||
w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
|
||||
w13_input_scale = layer.w13_input_scale
|
||||
w2_input_scale = layer.w2_input_scale
|
||||
|
||||
# MI300x and MI325x use FNUZ format for FP8. Convert if needed.
|
||||
if current_platform.is_fp8_fnuz():
|
||||
w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz(
|
||||
w13,
|
||||
w13_scale,
|
||||
w13_input_scale,
|
||||
)
|
||||
w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
|
||||
w2,
|
||||
w2_scale,
|
||||
w2_input_scale,
|
||||
)
|
||||
|
||||
# Per tensor kernels require single activation scale. Use the max.
|
||||
if self.quant_config.activation_scheme == "static":
|
||||
assert not self.block_quant
|
||||
assert w13_input_scale is not None and w2_input_scale is not None
|
||||
w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
|
||||
w13_input_scale, w2_input_scale
|
||||
)
|
||||
replace_parameter(layer, "w13_input_scale", w13_input_scale)
|
||||
replace_parameter(layer, "w2_input_scale", w2_input_scale)
|
||||
|
||||
# Per tensor kernels require single weight scale for w13 per expert, but
|
||||
# on disk there is a scale for w1 and w3. Use the max to requantize.
|
||||
if not self.block_quant:
|
||||
shard_size = layer.intermediate_size_per_partition
|
||||
w13, w13_scale = process_fp8_weight_tensor_strategy_moe(
|
||||
w13, w13_scale, shard_size, layer.local_num_experts
|
||||
)
|
||||
|
||||
# Shuffle weights to runtime format and setup kernel.
|
||||
self._setup_kernel(
|
||||
layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale
|
||||
)
|
||||
|
||||
def maybe_make_prepare_finalize(
|
||||
self,
|
||||
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> mk.FusedMoEPrepareAndFinalizeModular | None:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} uses the new modular kernel initialization "
|
||||
"logic. This function should not be called."
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
|
||||
w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
|
||||
w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
|
||||
a1_scale = layer.w13_input_scale
|
||||
a2_scale = layer.w2_input_scale
|
||||
|
||||
quant_config = make_fp8_moe_quant_config(
|
||||
fp8_backend=self.fp8_backend,
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
a1_scale=a1_scale,
|
||||
a2_scale=a2_scale,
|
||||
block_shape=self.weight_block_size,
|
||||
swiglu_limit=getattr(layer, "swiglu_limit", None),
|
||||
gemm1_alpha=getattr(layer, "swiglu_alpha", None),
|
||||
gemm1_beta=getattr(layer, "swiglu_beta", None),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
# Inject biases into the quant config if the model has them
|
||||
# (e.g. GPT-OSS biased MoE)
|
||||
if quant_config is not None and self.moe.has_bias:
|
||||
w13_bias = getattr(layer, "w13_bias", None)
|
||||
w2_bias = getattr(layer, "w2_bias", None)
|
||||
if w13_bias is not None:
|
||||
quant_config._w1.bias = w13_bias
|
||||
if w2_bias is not None:
|
||||
quant_config._w2.bias = w2_bias
|
||||
|
||||
return quant_config
|
||||
|
||||
@property
|
||||
def supports_eplb(self) -> bool:
|
||||
return True
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
num_expert_group=layer.num_expert_group,
|
||||
topk_group=layer.topk_group,
|
||||
e_score_correction_bias=layer.e_score_correction_bias,
|
||||
routed_scaling_factor=layer.routed_scaling_factor,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
|
||||
|
||||
class Fp8KVCacheMethod(BaseKVCacheMethod):
|
||||
"""
|
||||
Supports loading kv-cache scaling factors from FP8 checkpoints.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: Fp8Config):
|
||||
super().__init__(quant_config)
|
||||
@@ -0,0 +1,399 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Supports FP-Quant compression, see https://arxiv.org/abs/2509.23202
|
||||
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from vllm._custom_ops import (
|
||||
cutlass_scaled_fp4_mm,
|
||||
fusedQuantizeMx,
|
||||
fusedQuantizeNv,
|
||||
matmul_mxf4_bf16_tn,
|
||||
)
|
||||
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 QuantizationConfig
|
||||
from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
|
||||
class FPQuantConfig(QuantizationConfig):
|
||||
"""Config class for FPQuant."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hadamard_group_size: int = 32,
|
||||
forward_dtype: str = "mxfp4",
|
||||
forward_method: str = "abs_max",
|
||||
modules_to_not_convert: list[str] | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.hadamard_group_size = hadamard_group_size
|
||||
self.forward_dtype = forward_dtype
|
||||
self.forward_method = forward_method
|
||||
self.modules_to_not_convert = modules_to_not_convert
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"FPQuantConfig(hadamard_group_size={self.hadamard_group_size}, "
|
||||
f"forward_dtype={self.forward_dtype}, "
|
||||
f"forward_method={self.forward_method}, "
|
||||
f"modules_to_not_convert={self.modules_to_not_convert})"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "fp_quant"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 100
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return [] # no extra configs.
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "FPQuantConfig":
|
||||
hadamard_group_size = cls.get_from_keys(config, ["hadamard_group_size"])
|
||||
forward_dtype = cls.get_from_keys(config, ["forward_dtype"])
|
||||
forward_method = cls.get_from_keys(config, ["forward_method"])
|
||||
modules_to_not_convert = cls.get_from_keys(config, ["modules_to_not_convert"])
|
||||
return cls(
|
||||
hadamard_group_size,
|
||||
forward_dtype,
|
||||
forward_method,
|
||||
modules_to_not_convert,
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> LinearMethodBase | None:
|
||||
if self.modules_to_not_convert is not None and any(
|
||||
prefix.endswith(module) for module in self.modules_to_not_convert
|
||||
):
|
||||
return UnquantizedLinearMethod()
|
||||
|
||||
if isinstance(layer, LinearBase):
|
||||
return FPQuantLinearMethod(self)
|
||||
return None
|
||||
|
||||
|
||||
class FPQuantLinearMethod(LinearMethodBase):
|
||||
"""Linear method for FPQuant.
|
||||
|
||||
Args:
|
||||
quant_config: The FPQuant quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: FPQuantConfig):
|
||||
self.quant_config = quant_config
|
||||
|
||||
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,
|
||||
):
|
||||
del output_size # Unused.
|
||||
del input_size # Unused.
|
||||
|
||||
if params_dtype != torch.bfloat16:
|
||||
raise ValueError("Only bfloat16 is currently supported by FPQuant")
|
||||
if input_size_per_partition % self.quant_config.hadamard_group_size != 0: # noqa: E501
|
||||
raise ValueError(
|
||||
"The input size is not aligned with the quantized "
|
||||
"weight shape. This can be caused by too large "
|
||||
"tensor parallel size. Or other skill issues."
|
||||
)
|
||||
|
||||
assert self.quant_config.forward_dtype in ["mxfp4", "nvfp4"], (
|
||||
"Only mxfp4 and nvfp4 are supported for now"
|
||||
)
|
||||
if self.quant_config.forward_dtype == "mxfp4":
|
||||
group_size = 32
|
||||
elif self.quant_config.forward_dtype == "nvfp4":
|
||||
group_size = 16
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported forward_dtype: {self.quant_config.forward_dtype}"
|
||||
)
|
||||
|
||||
qweight = Parameter(
|
||||
torch.empty(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition // 2,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(
|
||||
qweight,
|
||||
{
|
||||
"input_dim": 1,
|
||||
"output_dim": 0,
|
||||
"packed_dim": 1,
|
||||
"pack_factor": 2,
|
||||
}
|
||||
| extra_weight_attrs,
|
||||
)
|
||||
layer.register_parameter("qweight", qweight)
|
||||
|
||||
scales = Parameter(
|
||||
torch.empty(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition // group_size,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(
|
||||
scales,
|
||||
{
|
||||
"input_dim": 1,
|
||||
"output_dim": 0,
|
||||
"packed_dim": 1,
|
||||
"pack_factor": group_size,
|
||||
}
|
||||
| extra_weight_attrs,
|
||||
)
|
||||
layer.register_parameter("scales", scales)
|
||||
|
||||
weight_global_scale = Parameter(
|
||||
torch.empty(1, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(
|
||||
weight_global_scale, {"ignore_warning": True} | extra_weight_attrs
|
||||
)
|
||||
layer.register_parameter("weight_global_scale", weight_global_scale)
|
||||
|
||||
act_global_scale = Parameter(
|
||||
torch.empty(1, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(
|
||||
act_global_scale, {"ignore_warning": True} | extra_weight_attrs
|
||||
)
|
||||
layer.register_parameter("act_global_scale", act_global_scale)
|
||||
|
||||
forward_hadamard_matrix = Parameter(
|
||||
torch.empty(
|
||||
self.quant_config.hadamard_group_size,
|
||||
self.quant_config.hadamard_group_size,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(
|
||||
forward_hadamard_matrix, {"ignore_warning": True} | extra_weight_attrs
|
||||
)
|
||||
layer.register_parameter("forward_hadamard_matrix", forward_hadamard_matrix)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return quantized_forward(
|
||||
x,
|
||||
layer.qweight,
|
||||
layer.scales,
|
||||
layer.weight_global_scale,
|
||||
layer.act_global_scale,
|
||||
bias,
|
||||
layer.forward_hadamard_matrix,
|
||||
self.quant_config.forward_method,
|
||||
self.quant_config.forward_dtype,
|
||||
)
|
||||
|
||||
|
||||
def fused_quantize_mx(
|
||||
x_flat: torch.Tensor, hadamard_matrix: torch.Tensor, forward_method: str
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return fusedQuantizeMx(
|
||||
x_flat,
|
||||
hadamard_matrix,
|
||||
method=cast(Literal["quest", "abs_max"], forward_method),
|
||||
)
|
||||
|
||||
|
||||
def fused_quantize_mx_fake(x_flat, hadamard_matrix, forward_method):
|
||||
rows, cols = x_flat.size(0), x_flat.size(1) // 32
|
||||
padded_rows = ((rows + 128 - 1) // 128) * 128
|
||||
padded_cols = ((cols + 4 - 1) // 4) * 4
|
||||
|
||||
xh_e2m1 = torch.empty(
|
||||
x_flat.size(0), x_flat.size(1) // 2, dtype=torch.uint8, device=x_flat.device
|
||||
)
|
||||
xh_e8m0 = torch.empty(
|
||||
padded_rows, padded_cols, dtype=torch.float8_e8m0fnu, device=x_flat.device
|
||||
)
|
||||
|
||||
return xh_e2m1, xh_e8m0
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="fused_quantize_mx",
|
||||
op_func=fused_quantize_mx,
|
||||
mutates_args=[],
|
||||
fake_impl=fused_quantize_mx_fake,
|
||||
dispatch_key=current_platform.dispatch_key,
|
||||
)
|
||||
|
||||
|
||||
def matmul_mxf4_bf16(
|
||||
x: torch.Tensor,
|
||||
w: torch.Tensor,
|
||||
xs: torch.Tensor,
|
||||
ws: torch.Tensor,
|
||||
alpha: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return matmul_mxf4_bf16_tn(
|
||||
x,
|
||||
w,
|
||||
to_blocked(xs, backend="triton").view(torch.float8_e8m0fnu),
|
||||
to_blocked(ws, backend="triton").view(torch.float8_e8m0fnu),
|
||||
alpha,
|
||||
)
|
||||
|
||||
|
||||
def matmul_mxf4_bf16_fake(x, w, xs, ws, alpha):
|
||||
return torch.empty(*x.shape[:-1], w.shape[0], dtype=torch.bfloat16, device=x.device)
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="matmul_mxf4_bf16",
|
||||
op_func=matmul_mxf4_bf16,
|
||||
mutates_args=[],
|
||||
fake_impl=matmul_mxf4_bf16_fake,
|
||||
dispatch_key=current_platform.dispatch_key,
|
||||
)
|
||||
|
||||
|
||||
def fused_quantize_nv(
|
||||
x_flat: torch.Tensor, hadamard_matrix: torch.Tensor, global_scale: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return fusedQuantizeNv(x_flat, hadamard_matrix, global_scale)
|
||||
|
||||
|
||||
def fused_quantize_nv_fake(x_flat, hadamard_matrix, global_scale):
|
||||
rows, cols = x_flat.size(0), x_flat.size(1) // 16
|
||||
padded_rows = ((rows + 128 - 1) // 128) * 128
|
||||
padded_cols = ((cols + 4 - 1) // 4) * 4
|
||||
|
||||
xh_e2m1 = torch.empty(
|
||||
x_flat.size(0), x_flat.size(1) // 2, dtype=torch.uint8, device=x_flat.device
|
||||
)
|
||||
xh_e8m0 = torch.empty(
|
||||
padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=x_flat.device
|
||||
)
|
||||
|
||||
return xh_e2m1, xh_e8m0
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="fused_quantize_nv",
|
||||
op_func=fused_quantize_nv,
|
||||
mutates_args=[],
|
||||
fake_impl=fused_quantize_nv_fake,
|
||||
dispatch_key=current_platform.dispatch_key,
|
||||
)
|
||||
|
||||
|
||||
def matmul_nvf4_bf16(
|
||||
x: torch.Tensor,
|
||||
w: torch.Tensor,
|
||||
xs: torch.Tensor,
|
||||
ws: torch.Tensor,
|
||||
alpha: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return cutlass_scaled_fp4_mm(
|
||||
x,
|
||||
w,
|
||||
to_blocked(xs, backend="triton")
|
||||
.view(torch.float8_e4m3fn)
|
||||
.view(-1, x.shape[1] // 8), # *2//16
|
||||
to_blocked(ws, backend="triton")
|
||||
.view(torch.float8_e4m3fn)
|
||||
.view(-1, x.shape[1] // 8),
|
||||
alpha,
|
||||
torch.bfloat16,
|
||||
)
|
||||
|
||||
|
||||
def matmul_nvf4_bf16_fake(x, w, xs, ws, alpha):
|
||||
return torch.empty(*x.shape[:-1], w.shape[0], dtype=torch.bfloat16, device=x.device)
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="matmul_nvf4_bf16",
|
||||
op_func=matmul_nvf4_bf16,
|
||||
mutates_args=[],
|
||||
fake_impl=matmul_nvf4_bf16_fake,
|
||||
dispatch_key=current_platform.dispatch_key,
|
||||
)
|
||||
|
||||
|
||||
def quantized_forward(
|
||||
x: torch.Tensor,
|
||||
qweight: torch.Tensor,
|
||||
weight_scales: torch.Tensor,
|
||||
weight_global_scale: torch.Tensor,
|
||||
act_global_scale: torch.Tensor,
|
||||
bias: torch.Tensor | None,
|
||||
forward_hadamard_matrix: torch.Tensor,
|
||||
forward_method: str,
|
||||
forward_dtype: str,
|
||||
) -> torch.Tensor:
|
||||
x_flat = x.contiguous().flatten(end_dim=-2)
|
||||
|
||||
if forward_dtype == "mxfp4":
|
||||
x_flat_q, x_flat_scales = torch.ops.vllm.fused_quantize_mx(
|
||||
x_flat, forward_hadamard_matrix, forward_method
|
||||
)
|
||||
y = torch.ops.vllm.matmul_mxf4_bf16(
|
||||
x_flat_q,
|
||||
qweight,
|
||||
x_flat_scales,
|
||||
weight_scales,
|
||||
1 / (weight_global_scale * act_global_scale),
|
||||
)
|
||||
elif forward_dtype == "nvfp4":
|
||||
x_flat_q, x_flat_scales = torch.ops.vllm.fused_quantize_nv(
|
||||
x_flat, forward_hadamard_matrix, act_global_scale
|
||||
)
|
||||
y = torch.ops.vllm.matmul_nvf4_bf16(
|
||||
x_flat_q,
|
||||
qweight,
|
||||
x_flat_scales,
|
||||
weight_scales,
|
||||
1 / (weight_global_scale * act_global_scale),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported forward_dtype: {forward_dtype}")
|
||||
|
||||
y = y.view(*x.shape[:-1], y.shape[-1])
|
||||
if bias is not None:
|
||||
y += bias
|
||||
|
||||
return y
|
||||
@@ -0,0 +1,829 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
|
||||
import vllm.utils.humming as _hm
|
||||
from vllm import envs
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEMethodBase,
|
||||
FusedMoEQuantConfig,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
|
||||
UnquantizedFusedMoEMethod,
|
||||
)
|
||||
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 (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.humming_utils import (
|
||||
convert_to_humming_moe_kernel_format,
|
||||
get_humming_moe_quant_config,
|
||||
input_schema_to_quant_key,
|
||||
make_humming_moe_kernel,
|
||||
select_humming_moe_experts,
|
||||
weight_schema_to_quant_key,
|
||||
)
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
from vllm.model_executor.parameter import (
|
||||
BasevLLMParameter,
|
||||
BlockQuantScaleParameter,
|
||||
ChannelQuantScaleParameter,
|
||||
GroupQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
PackedvLLMParameter,
|
||||
PerTensorScaleParameter,
|
||||
RowvLLMParameter,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
from vllm.utils.humming import (
|
||||
BaseInputSchema,
|
||||
BaseWeightSchema,
|
||||
HummingInputSchema,
|
||||
HummingWeightSchema,
|
||||
)
|
||||
|
||||
|
||||
def prepare_padded_shape(shape, x):
|
||||
padded_shape = math.ceil(shape / x) * x
|
||||
return padded_shape, padded_shape - shape
|
||||
|
||||
|
||||
def prepare_param(tensor, name, extra_attrs):
|
||||
extra_attrs = extra_attrs.copy()
|
||||
scale_type = extra_attrs.pop("scale_type", None)
|
||||
param_cls_name_map = {
|
||||
"block": BlockQuantScaleParameter,
|
||||
"tensor": PerTensorScaleParameter,
|
||||
"group": GroupQuantScaleParameter,
|
||||
"channel": ChannelQuantScaleParameter,
|
||||
"input_scale": PerTensorScaleParameter,
|
||||
}
|
||||
|
||||
param_cls: type[BasevLLMParameter]
|
||||
if "packed_dim" in extra_attrs:
|
||||
param_cls = PackedvLLMParameter
|
||||
elif scale_type in param_cls_name_map:
|
||||
param_cls = param_cls_name_map[scale_type]
|
||||
elif "output_dim" in extra_attrs and "input_dim" in extra_attrs:
|
||||
param_cls = ModelWeightParameter
|
||||
elif "input_dim" in extra_attrs:
|
||||
param_cls = RowvLLMParameter
|
||||
elif "output_dim" in extra_attrs:
|
||||
param_cls = ChannelQuantScaleParameter
|
||||
else:
|
||||
param_cls = BasevLLMParameter
|
||||
|
||||
kwargs_keys = [
|
||||
"input_dim",
|
||||
"output_dim",
|
||||
"packed_dim",
|
||||
"packed_factor",
|
||||
"weight_loader",
|
||||
]
|
||||
cls_kwargs = {}
|
||||
for key in extra_attrs.copy():
|
||||
if key in kwargs_keys:
|
||||
cls_kwargs[key] = extra_attrs.pop(key)
|
||||
|
||||
param = param_cls(data=tensor, **cls_kwargs)
|
||||
set_weight_attrs(param, extra_attrs)
|
||||
|
||||
param.param_name = name
|
||||
param.ignore_warning = True
|
||||
if scale_type in ["tensor", "input_scale"]:
|
||||
param.needs_scalar_to_array = True
|
||||
|
||||
return param
|
||||
|
||||
|
||||
def prepare_moe_param(tensor: torch.Tensor, name: str, extra_attrs: dict[str, Any]):
|
||||
param = torch.nn.Parameter(tensor, requires_grad=False)
|
||||
if "scale_type" in extra_attrs:
|
||||
extra_attrs["quant_method"] = extra_attrs["scale_type"]
|
||||
|
||||
if "input_dim" in extra_attrs and "output_dim" in extra_attrs:
|
||||
input_dim = extra_attrs["input_dim"]
|
||||
output_dim = extra_attrs["output_dim"]
|
||||
extra_attrs["is_transposed"] = input_dim < output_dim
|
||||
|
||||
set_weight_attrs(param, extra_attrs)
|
||||
param.param_name = name
|
||||
return param
|
||||
|
||||
|
||||
def may_pad_loaded_weight(param, loaded_weight):
|
||||
pad_shape = getattr(param, "pad_shape", None)
|
||||
if pad_shape is None:
|
||||
return loaded_weight
|
||||
value = 1 if loaded_weight.dtype == torch.float8_e8m0fnu else 0
|
||||
padding = []
|
||||
for x in pad_shape[::-1][: loaded_weight.ndim]:
|
||||
padding += [0, x]
|
||||
loaded_weight = torch.nn.functional.pad(
|
||||
input=loaded_weight,
|
||||
pad=padding,
|
||||
value=value,
|
||||
)
|
||||
return loaded_weight
|
||||
|
||||
|
||||
def compressed_tensors_get_config(config: dict[str, Any], key: str):
|
||||
assert key in ["weights", "input_activations"]
|
||||
target_group_config = None
|
||||
for group_config in config["config_groups"].values():
|
||||
if "Linear" in group_config["targets"]:
|
||||
if "weights" not in group_config:
|
||||
return None
|
||||
if key not in group_config or group_config[key] is None:
|
||||
return None
|
||||
target_group_config = group_config[key].copy()
|
||||
break
|
||||
|
||||
if target_group_config is None:
|
||||
return None
|
||||
target_group_config["quant_method"] = config["quant_method"]
|
||||
if config["quant_method"] == "compressed-tensors":
|
||||
target_group_config["format"] = config["format"]
|
||||
elif config["quant_method"] == "modelopt":
|
||||
target_group_config["quant_algo"] = config["quant_algo"]
|
||||
return target_group_config
|
||||
|
||||
|
||||
class HummingConfig(QuantizationConfig):
|
||||
packed_modules_mapping: dict[str, list[str]] = {}
|
||||
|
||||
def __init__(self, full_config: dict[str, Any] | None = None):
|
||||
self.full_config: dict[str, Any] = full_config or {}
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "humming"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.bfloat16, torch.half]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 75
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "HummingConfig":
|
||||
return cls(full_config=config)
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(
|
||||
cls, hf_quant_cfg, user_quant, hf_config=None
|
||||
) -> QuantizationMethods | None:
|
||||
if user_quant == "humming" and hf_config is not None:
|
||||
model_type = hf_config.model_type
|
||||
quant_method = hf_quant_cfg.get("quant_method", None)
|
||||
if model_type == "gpt_oss" and quant_method == "mxfp4":
|
||||
msg = (
|
||||
"For gpt-oss model, use '--moe-backend humming' "
|
||||
"instead of '--quantization humming'."
|
||||
)
|
||||
raise ValueError(msg)
|
||||
return "humming" if user_quant == "humming" else None
|
||||
|
||||
def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
|
||||
self.hf_to_vllm_mapper = hf_to_vllm_mapper
|
||||
|
||||
def is_layer_skipped(self, config: dict[str, Any], prefix: str):
|
||||
keys = ["ignored_layers", "ignore", "modules_to_not_convert"]
|
||||
ignored_layers = self.get_from_keys_or(config, keys, []) or []
|
||||
if hasattr(self, "hf_to_vllm_mapper"):
|
||||
ignored_layers = self.hf_to_vllm_mapper.apply_list(ignored_layers)
|
||||
|
||||
if any(module_name in prefix for module_name in ignored_layers):
|
||||
return True
|
||||
if "lm_head" in prefix:
|
||||
return True
|
||||
|
||||
for regex in config.get("dynamic", {}):
|
||||
if regex[:1] != "-":
|
||||
continue
|
||||
if re.match(regex[2:], prefix):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def get_layer_weight_schema(self, config: dict[str, Any], prefix: str):
|
||||
if self.is_layer_skipped(config, prefix):
|
||||
return None
|
||||
|
||||
if config["quant_method"] in ["compressed-tensors", "modelopt"]:
|
||||
group_config = compressed_tensors_get_config(config, "weights")
|
||||
if group_config is None:
|
||||
return None
|
||||
config = group_config
|
||||
|
||||
layer_config = config
|
||||
layer_dynamic = config.get("dynamic", {})
|
||||
if not isinstance(layer_dynamic, dict):
|
||||
layer_dynamic = {}
|
||||
for regex, override_config in layer_dynamic.items():
|
||||
if regex[:1] != "+":
|
||||
continue
|
||||
if re.match(regex[2:], prefix):
|
||||
layer_config = config.copy()
|
||||
layer_config.update(override_config)
|
||||
break
|
||||
|
||||
if "quant_method" in layer_config:
|
||||
return _hm.BaseWeightSchema.from_config(layer_config)
|
||||
return None
|
||||
|
||||
def get_layer_input_schema(self, config: dict[str, Any], prefix: str):
|
||||
if self.is_layer_skipped(config, prefix):
|
||||
return None
|
||||
if config["quant_method"] in ["compressed-tensors", "modelopt"]:
|
||||
group_config = compressed_tensors_get_config(config, "input_activations")
|
||||
if group_config is None:
|
||||
return None
|
||||
config = group_config
|
||||
|
||||
if config.get("quant_method", None) in _hm.BaseInputSchema.INPUT_SCHEMA_MAP:
|
||||
return _hm.BaseInputSchema.from_config(config)
|
||||
return None
|
||||
|
||||
def get_quant_config_for_layer(
|
||||
self, prefix: str, layer_type: str
|
||||
) -> "HummingLayerQuantizationConfig | None":
|
||||
weight_schema: BaseWeightSchema | None = None
|
||||
force_weight_schema: HummingWeightSchema | None = None
|
||||
|
||||
if self.full_config:
|
||||
weight_schema = self.get_layer_weight_schema(self.full_config, prefix)
|
||||
|
||||
is_online_quant = False
|
||||
online_quant_config = envs.VLLM_HUMMING_ONLINE_QUANT_CONFIG or {}
|
||||
if not self.full_config or online_quant_config.get("force_requant", False):
|
||||
online_quant_config["quant_method"] = "humming"
|
||||
schema = self.get_layer_weight_schema(online_quant_config, prefix)
|
||||
if not self.full_config:
|
||||
weight_schema = schema
|
||||
is_online_quant = True
|
||||
else:
|
||||
force_weight_schema = schema
|
||||
|
||||
if weight_schema is not None:
|
||||
input_schema = None
|
||||
force_input_schema = None
|
||||
|
||||
if self.full_config:
|
||||
input_schema = self.get_layer_input_schema(self.full_config, prefix)
|
||||
|
||||
if envs.VLLM_HUMMING_INPUT_QUANT_CONFIG:
|
||||
quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG.copy()
|
||||
quant_config["quant_method"] = "humming"
|
||||
force_input_schema = self.get_layer_input_schema(quant_config, prefix)
|
||||
if input_schema is None:
|
||||
input_schema = force_input_schema
|
||||
|
||||
if force_weight_schema is not None and force_input_schema is None:
|
||||
force_input_schema = _hm.HummingInputSchema()
|
||||
|
||||
return HummingLayerQuantizationConfig(
|
||||
weight_schema=weight_schema,
|
||||
input_schema=input_schema,
|
||||
force_weight_schema=force_weight_schema,
|
||||
force_input_schema=force_input_schema,
|
||||
is_online_quant=is_online_quant,
|
||||
)
|
||||
return None
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
layer_type = "other"
|
||||
if isinstance(layer, RoutedExperts):
|
||||
layer_type = "moe"
|
||||
elif isinstance(layer, LinearBase):
|
||||
layer_type = "linear"
|
||||
|
||||
quant_config = self.get_quant_config_for_layer(prefix, layer_type)
|
||||
if quant_config is None:
|
||||
if isinstance(layer, RoutedExperts):
|
||||
return UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
elif isinstance(layer, LinearBase):
|
||||
return UnquantizedLinearMethod()
|
||||
elif isinstance(layer, LinearBase):
|
||||
return HummingLinearMethod(quant_config)
|
||||
elif isinstance(layer, RoutedExperts):
|
||||
return HummingMoEMethod(quant_config, layer.moe_config)
|
||||
return None
|
||||
|
||||
|
||||
class HummingLayerQuantizationConfig(HummingConfig):
|
||||
def __init__(
|
||||
self,
|
||||
weight_schema: "BaseWeightSchema",
|
||||
input_schema: "BaseInputSchema | None" = None,
|
||||
force_weight_schema: "HummingWeightSchema | None" = None,
|
||||
force_input_schema: "HummingInputSchema | None" = None,
|
||||
is_online_quant: bool = False,
|
||||
):
|
||||
self.weight_schema = weight_schema
|
||||
if input_schema is None:
|
||||
input_schema = _hm.HummingInputSchema()
|
||||
self.input_schema = input_schema
|
||||
self.force_weight_schema = force_weight_schema
|
||||
self.force_input_schema = force_input_schema
|
||||
self.is_online_quant = is_online_quant
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config):
|
||||
weight_schema = _hm.BaseWeightSchema.from_config(config)
|
||||
return cls(weight_schema)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> QuantizeMethodBase | None:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class HummingLinearMethod(LinearMethodBase):
|
||||
def __init__(self, quant_config: HummingLayerQuantizationConfig):
|
||||
self.quant_config = quant_config
|
||||
self.weight_schema = quant_config.weight_schema
|
||||
self.input_schema = quant_config.input_schema
|
||||
self.force_weight_schema = quant_config.force_weight_schema
|
||||
self.force_input_schema = quant_config.force_input_schema
|
||||
self.is_online_quant = self.quant_config.is_online_quant
|
||||
|
||||
def prepare_weight_loader(self, layer: torch.nn.Module, weight_loader: Callable):
|
||||
def new_weight_loader(
|
||||
param: torch.nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
shard_id: str | int | None = None,
|
||||
):
|
||||
name = param.param_name
|
||||
float_dtypes = [torch.float16, torch.bfloat16, torch.float32]
|
||||
is_unquantized = name == "weight" and loaded_weight.dtype in float_dtypes
|
||||
if is_unquantized and self.is_online_quant:
|
||||
# online quant (fp16/bf16 -> quant_type)
|
||||
assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
|
||||
f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype)
|
||||
has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type)
|
||||
tensor_list = _hm.quantize_weight(
|
||||
weight=loaded_weight,
|
||||
dtype=self.weight_schema.b_dtype,
|
||||
scale_dtype=self.weight_schema.bs_dtype or f16_dtype,
|
||||
group_size=self.weight_schema.weight_scale_group_size,
|
||||
has_zero_point=self.weight_schema.has_zero_point,
|
||||
has_global_scale=has_global_scale,
|
||||
is_fp_zero_point=self.weight_schema.is_fp_zero_point,
|
||||
pack=True,
|
||||
)
|
||||
|
||||
key_list = ["weight", "weight_scale", "zero_point", "global_scale"]
|
||||
for key, tensor in zip(key_list, tensor_list):
|
||||
if tensor is None or tensor.nelement() == 0:
|
||||
continue
|
||||
param = getattr(layer, key)
|
||||
param.weight_loader(param, tensor, shard_id)
|
||||
|
||||
return None
|
||||
elif is_unquantized and not self.is_online_quant:
|
||||
# fallback to unquantized linear
|
||||
# some model skip some layer when quantizing model, but
|
||||
# don't mark the layer as unquantized.
|
||||
if not layer.is_fallback:
|
||||
layer.is_fallback = True
|
||||
for name, _ in list(layer.named_parameters()):
|
||||
if name != "bias":
|
||||
delattr(layer, name)
|
||||
delattr(layer, "locks")
|
||||
self.__class__ = UnquantizedLinearMethod # type: ignore
|
||||
tensor = torch.empty(
|
||||
(
|
||||
layer.output_partition_sizes_sum,
|
||||
layer.input_size_per_partition,
|
||||
),
|
||||
dtype=layer.param_dtype,
|
||||
device=param.device,
|
||||
)
|
||||
extra_weight_attrs = layer.extra_weight_attrs.copy()
|
||||
orig_weight_loader = extra_weight_attrs.pop("weight_loader")
|
||||
layer.weight = ModelWeightParameter(
|
||||
data=tensor,
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=orig_weight_loader,
|
||||
)
|
||||
layer.weight.tp_size = layer.tp_size
|
||||
layer.weight.tp_rank = layer.tp_rank
|
||||
set_weight_attrs(layer.weight, extra_weight_attrs)
|
||||
|
||||
param = layer.weight
|
||||
if shard_id is not None:
|
||||
return layer.weight.weight_loader(param, loaded_weight, shard_id)
|
||||
return layer.weight.weight_loader(param, loaded_weight)
|
||||
|
||||
# weight processing logic for specific quantization schema
|
||||
loaded_weight = self.weight_schema.process_loaded_weight(
|
||||
tensor=loaded_weight,
|
||||
name=name,
|
||||
)
|
||||
if shard_id is not None:
|
||||
return weight_loader(param, loaded_weight, shard_id)
|
||||
return weight_loader(param, loaded_weight)
|
||||
|
||||
return new_weight_loader
|
||||
|
||||
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,
|
||||
):
|
||||
layer.is_fallback = False
|
||||
layer.param_dtype = params_dtype
|
||||
layer.input_size = input_size
|
||||
layer.output_size = output_size
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_partition_sizes_sum = sum(output_partition_sizes)
|
||||
layer.output_partition_sizes = output_partition_sizes
|
||||
layer.extra_weight_attrs = extra_weight_attrs.copy()
|
||||
|
||||
weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader)
|
||||
new_weight_loader = self.prepare_weight_loader(layer, weight_loader)
|
||||
extra_weight_attrs["weight_loader"] = new_weight_loader
|
||||
|
||||
for key in ["weight_block_size", "block_structure"]:
|
||||
block_size = getattr(self.weight_schema, key, None)
|
||||
if block_size is not None:
|
||||
layer.weight_block_size = block_size
|
||||
|
||||
weight_tensor_attrs = self.weight_schema.get_tensors_attrs(
|
||||
shape_n=layer.output_partition_sizes_sum,
|
||||
shape_k=layer.input_size_per_partition,
|
||||
param_dtype=params_dtype,
|
||||
stack_size=len(layer.output_partition_sizes),
|
||||
)
|
||||
|
||||
input_tensor_attrs = self.input_schema.get_tensors_attrs(
|
||||
shape_k=layer.input_size_per_partition,
|
||||
param_dtype=params_dtype,
|
||||
stack_size=len(layer.output_partition_sizes),
|
||||
)
|
||||
|
||||
tensors_attrs = weight_tensor_attrs | input_tensor_attrs
|
||||
|
||||
for name, attrs in tensors_attrs.items():
|
||||
tensor = torch.empty(attrs["shape"], dtype=attrs["dtype"])
|
||||
extra_attrs = attrs.get("extra_attrs", {}).copy()
|
||||
extra_attrs.update(extra_weight_attrs)
|
||||
param = prepare_param(tensor, name, extra_attrs)
|
||||
setattr(layer, name, param)
|
||||
|
||||
locks = torch.zeros(1024, dtype=torch.int32)
|
||||
layer.register_buffer("locks", locks)
|
||||
|
||||
if self.force_input_schema is not None:
|
||||
self.input_schema = self.force_input_schema
|
||||
|
||||
if not hasattr(layer, "weight"):
|
||||
param = prepare_param(torch.tensor(0), "weight", extra_weight_attrs)
|
||||
layer.weight = param
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
if layer.is_fallback:
|
||||
return None
|
||||
|
||||
# convert from checkpoint format to humming format
|
||||
if not isinstance(self.weight_schema, _hm.HummingWeightSchema):
|
||||
self.weight_schema, tensors = self.weight_schema.convert_humming(
|
||||
tensors=layer.state_dict(),
|
||||
shape_n_stacks=layer.output_partition_sizes,
|
||||
shape_k_stacks=[layer.input_size_per_partition],
|
||||
param_dtype=layer.param_dtype,
|
||||
)
|
||||
|
||||
self.input_schema, _ = self.input_schema.convert_humming(
|
||||
tensors=layer.state_dict(),
|
||||
shape_n_stacks=layer.output_partition_sizes,
|
||||
shape_k_stacks=[layer.input_size_per_partition],
|
||||
param_dtype=layer.param_dtype,
|
||||
)
|
||||
|
||||
for name, _ in list(layer.named_parameters()):
|
||||
delattr(layer, name)
|
||||
|
||||
for name, tensor in tensors.items():
|
||||
param = torch.nn.Parameter(tensor, requires_grad=False)
|
||||
setattr(layer, name, param)
|
||||
|
||||
del tensors
|
||||
|
||||
# force requant (origin quant setting -> fp16/bf16 -> new_quant setting)
|
||||
assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
|
||||
force_requant = self.force_weight_schema is not None
|
||||
if force_requant and self.weight_schema != self.force_weight_schema:
|
||||
tensors = self.weight_schema.requant_tensors(
|
||||
tensors=layer.state_dict(),
|
||||
target_weight_schema=self.force_weight_schema,
|
||||
param_dtype=layer.param_dtype,
|
||||
)
|
||||
|
||||
self.weight_schema = self.force_weight_schema
|
||||
|
||||
for name, _ in list(layer.named_parameters()):
|
||||
if name != "bias":
|
||||
delattr(layer, name)
|
||||
|
||||
for name, tensor in tensors.items():
|
||||
param = torch.nn.Parameter(tensor, requires_grad=False)
|
||||
setattr(layer, name, param)
|
||||
|
||||
del tensors
|
||||
|
||||
# prepare layer config from humming kernel
|
||||
_hm.HummingMethod.prepare_layer_meta(
|
||||
layer=layer,
|
||||
shape_n=layer.output_partition_sizes_sum,
|
||||
shape_k=layer.input_size_per_partition,
|
||||
weight_schema=self.weight_schema,
|
||||
input_schema=self.input_schema,
|
||||
pad_n_to_multiple=256,
|
||||
pad_k_to_multiple=128,
|
||||
has_bias=layer.has_bias,
|
||||
torch_dtype=layer.param_dtype,
|
||||
)
|
||||
|
||||
# preprocess weight for inference
|
||||
_hm.HummingMethod.transform_humming_layer(layer)
|
||||
|
||||
# compute_config: kernel configs that do not directly affect weights
|
||||
# but significantly impact kernel behavior or computation precision.
|
||||
# see https://github.com/inclusionAI/humming/blob/main/docs/config.md
|
||||
compute_config = {
|
||||
"use_batch_invariant": envs.VLLM_BATCH_INVARIANT,
|
||||
"use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,
|
||||
"gemm_type": "dense",
|
||||
}
|
||||
self.compute_config = json.dumps(compute_config)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
flatten_inputs = x.view(-1, x.size(-1))
|
||||
output = _hm.HummingMethod.forward_layer(
|
||||
layer=layer,
|
||||
inputs=flatten_inputs,
|
||||
compute_config=self.compute_config,
|
||||
)
|
||||
output = output.view(*x.shape[:-1], output.size(-1))
|
||||
return output
|
||||
|
||||
|
||||
class HummingMoEMethod(FusedMoEMethodBase):
|
||||
def __init__(
|
||||
self, quant_config: HummingLayerQuantizationConfig, moe: "FusedMoEConfig"
|
||||
) -> None:
|
||||
super().__init__(moe)
|
||||
self.quant_config = quant_config
|
||||
self.weight_schema = quant_config.weight_schema
|
||||
self.input_schema = quant_config.input_schema
|
||||
self.force_weight_schema = quant_config.force_weight_schema
|
||||
self.force_input_schema = quant_config.force_input_schema
|
||||
|
||||
# Derive QuantKeys from humming schemas.
|
||||
# Prefer force schemas (the final format after requant) over base.
|
||||
weight_key = weight_schema_to_quant_key(
|
||||
self.force_weight_schema or self.weight_schema
|
||||
)
|
||||
activation_key = input_schema_to_quant_key(
|
||||
self.force_input_schema or self.input_schema
|
||||
)
|
||||
|
||||
# Select Humming MoE experts
|
||||
self.experts_cls = select_humming_moe_experts(
|
||||
config=self.moe,
|
||||
weight_key=weight_key,
|
||||
activation_key=activation_key,
|
||||
)
|
||||
|
||||
def prepare_weight_loader(self, layer, weight_loader):
|
||||
def new_weight_loader(
|
||||
param: torch.nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
weight_name: str,
|
||||
shard_id: str,
|
||||
expert_id: int | None = None,
|
||||
return_success: bool = False,
|
||||
):
|
||||
name = param.param_name
|
||||
float_dtypes = [torch.float16, torch.bfloat16, torch.float32]
|
||||
is_unquantized = name == "weight" and loaded_weight.dtype in float_dtypes
|
||||
# online quant (fp16/bf16 -> quant_type)
|
||||
if is_unquantized:
|
||||
assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
|
||||
f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype)
|
||||
has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type)
|
||||
tensor_list = _hm.quantize_weight(
|
||||
weight=loaded_weight,
|
||||
dtype=self.weight_schema.b_dtype,
|
||||
scale_dtype=self.weight_schema.bs_dtype or f16_dtype,
|
||||
group_size=self.weight_schema.weight_scale_group_size,
|
||||
has_zero_point=self.weight_schema.has_zero_point,
|
||||
has_global_scale=has_global_scale,
|
||||
is_fp_zero_point=self.weight_schema.is_fp_zero_point,
|
||||
pack=True,
|
||||
)
|
||||
|
||||
key_list = ["weight", "weight_scale", "zero_point", "global_scale"]
|
||||
success = True
|
||||
for key, tensor in zip(key_list, tensor_list):
|
||||
if tensor is None or tensor.nelement() == 0:
|
||||
continue
|
||||
sublayer_name = "w2" if shard_id == "w2" else "w13"
|
||||
|
||||
param = getattr(layer, sublayer_name + "_" + key)
|
||||
part_success = param.weight_loader(
|
||||
param=param,
|
||||
loaded_weight=tensor.cpu(),
|
||||
weight_name=shard_id + "_" + key,
|
||||
shard_id=shard_id,
|
||||
expert_id=expert_id,
|
||||
return_success=return_success,
|
||||
)
|
||||
success = success and part_success
|
||||
|
||||
return success if return_success else None
|
||||
|
||||
# weight processing logic for specific quantization schema
|
||||
loaded_weight = self.weight_schema.process_loaded_weight(
|
||||
tensor=loaded_weight,
|
||||
name=name,
|
||||
)
|
||||
return weight_loader(
|
||||
param,
|
||||
loaded_weight,
|
||||
weight_name,
|
||||
shard_id=shard_id,
|
||||
expert_id=expert_id,
|
||||
return_success=return_success,
|
||||
)
|
||||
|
||||
return new_weight_loader
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.num_experts = num_experts
|
||||
layer.param_dtype = params_dtype
|
||||
layer.intermediate_size = intermediate_size_per_partition
|
||||
weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader)
|
||||
weight_loader = self.prepare_weight_loader(layer, weight_loader)
|
||||
extra_weight_attrs["weight_loader"] = weight_loader
|
||||
|
||||
# sublayer: a layer contains multiple sets of weights for quantized GEMM
|
||||
# (e.g., weight, weight_scale, etc.).
|
||||
# The weight names of sublayer start with the prefix "{sublayer_name}_"
|
||||
layer.sublayer_configs = {
|
||||
"w13": {
|
||||
"shape_n": intermediate_size_per_partition * 2,
|
||||
"shape_k": hidden_size,
|
||||
"tensors_attrs": self.weight_schema.get_padded_tensors_attrs(
|
||||
shape_n=intermediate_size_per_partition * 2,
|
||||
shape_k=hidden_size,
|
||||
num_experts=num_experts,
|
||||
param_dtype=params_dtype,
|
||||
has_bias=self.moe.has_bias,
|
||||
),
|
||||
},
|
||||
"w2": {
|
||||
"shape_n": hidden_size,
|
||||
"shape_k": intermediate_size_per_partition,
|
||||
"tensors_attrs": self.weight_schema.get_padded_tensors_attrs(
|
||||
shape_n=hidden_size,
|
||||
shape_k=intermediate_size_per_partition,
|
||||
num_experts=num_experts,
|
||||
param_dtype=params_dtype,
|
||||
has_bias=self.moe.has_bias,
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
for sublayer_name, configs in layer.sublayer_configs.items():
|
||||
for name, attrs in configs["tensors_attrs"].items():
|
||||
tensor = torch.empty(attrs["shape"], dtype=attrs["dtype"])
|
||||
param = torch.nn.Parameter(tensor, requires_grad=False)
|
||||
extra_attrs = attrs.get("extra_attrs", {}).copy()
|
||||
extra_attrs.update(extra_weight_attrs)
|
||||
param = prepare_moe_param(tensor, name, extra_attrs)
|
||||
setattr(layer, f"{sublayer_name}_{name}", param)
|
||||
|
||||
if self.force_input_schema is not None:
|
||||
self.input_schema = self.force_input_schema
|
||||
|
||||
locks = torch.zeros(1024, dtype=torch.int32)
|
||||
layer.register_buffer("locks", locks)
|
||||
|
||||
def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
|
||||
return get_humming_moe_quant_config(layer)
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
if getattr(self, "processed", False):
|
||||
return
|
||||
self.processed = True
|
||||
|
||||
# Convert weights to Humming kernel format
|
||||
convert_to_humming_moe_kernel_format(
|
||||
layer=layer,
|
||||
sublayer_configs=layer.sublayer_configs,
|
||||
weight_schema=self.weight_schema,
|
||||
input_schema=self.input_schema,
|
||||
force_weight_schema=self.force_weight_schema,
|
||||
)
|
||||
|
||||
# Build the MoE kernel
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert self.moe_quant_config is not None
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_humming_moe_kernel(
|
||||
self.moe_quant_config,
|
||||
self.moe,
|
||||
self.experts_cls,
|
||||
layer=layer,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Apply Humming-quantized MoE computation using the standard kernel flow.
|
||||
|
||||
This method uses FusedMoEKernel.apply() which orchestrates:
|
||||
1. Preparation (quantization if needed - skipped for Humming via
|
||||
expects_unquantized_inputs=True to prevent double quantization)
|
||||
2. Expert computation (via experts.apply())
|
||||
3. Finalization (weight application & reduction - no-op for Humming
|
||||
since it's already done internally)
|
||||
|
||||
Humming handles all quantization, weight application, and reduction
|
||||
internally in the experts.apply() method via HummingMethod calls.
|
||||
|
||||
Note: Although w1/w2 weights are passed to the kernel for interface
|
||||
consistency, Humming's experts.apply() reads weights directly from
|
||||
the layer object via HummingMethod.forward_layer() and ignores the
|
||||
w1/w2 parameters.
|
||||
"""
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
hidden_states=x,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
topk_ids=topk_ids,
|
||||
topk_weights=topk_weights,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=False,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
@@ -0,0 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from .inc import INCConfig
|
||||
|
||||
__all__ = ["INCConfig"]
|
||||
@@ -0,0 +1,188 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import regex as re
|
||||
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
from .inc import INCConfig
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class INCLayerConfig:
|
||||
bits: int
|
||||
group_size: int
|
||||
sym: bool
|
||||
packing_format: str
|
||||
backend: str
|
||||
data_type: str
|
||||
quantized: bool
|
||||
|
||||
@property
|
||||
def is_gptq(self) -> bool:
|
||||
return "gptq" in self.packing_format or "gptq" in self.backend
|
||||
|
||||
@property
|
||||
def is_awq(self) -> bool:
|
||||
return "awq" in self.packing_format or "awq" in self.backend
|
||||
|
||||
@property
|
||||
def is_wna16_int(self) -> bool:
|
||||
return self.data_type == "int" and self.quantized
|
||||
|
||||
@property
|
||||
def is_mxfp4(self) -> bool:
|
||||
return self.data_type == "mx_fp" and self.bits == 4
|
||||
|
||||
@property
|
||||
def is_mxfp8(self) -> bool:
|
||||
return self.data_type == "mx_fp" and self.bits == 8
|
||||
|
||||
|
||||
class INCConfigParser:
|
||||
def __init__(self, config: "INCConfig") -> None:
|
||||
self._config = config
|
||||
|
||||
def resolve(self, layer: "torch.nn.Module", layer_name: str) -> INCLayerConfig:
|
||||
bits, group_size, sym = self._resolve_raw(layer, layer_name)
|
||||
return INCLayerConfig(
|
||||
bits=bits,
|
||||
group_size=group_size,
|
||||
sym=sym,
|
||||
packing_format=self._config.packing_format,
|
||||
backend=self._config.backend,
|
||||
data_type=self._config.data_type,
|
||||
quantized=bits < 16,
|
||||
)
|
||||
|
||||
def get_layer_config(
|
||||
self, layer: "torch.nn.Module", layer_name: str
|
||||
) -> tuple[int, int, bool]:
|
||||
layer_config = self.resolve(layer, layer_name)
|
||||
return layer_config.bits, layer_config.group_size, layer_config.sym
|
||||
|
||||
def _resolve_raw(
|
||||
self, layer: "torch.nn.Module", layer_name: str
|
||||
) -> tuple[int, int, bool]:
|
||||
REGEX_SPECIAL_CHARS = set(r"*+?^$()[]{}|\\")
|
||||
|
||||
def is_explicitly_configured(name: str) -> bool:
|
||||
"""Return True if *name* has an explicit entry in extra_config,
|
||||
either via exact key match or via a regex pattern key."""
|
||||
if not self._config.extra_config:
|
||||
return False
|
||||
if name in self._config.extra_config:
|
||||
return True
|
||||
for pattern in self._config.extra_config:
|
||||
if not isinstance(pattern, str) or not any(
|
||||
c in REGEX_SPECIAL_CHARS for c in pattern
|
||||
):
|
||||
continue
|
||||
try:
|
||||
if re.search(re.compile(pattern), name) is not None:
|
||||
return True
|
||||
except re.error:
|
||||
continue
|
||||
return False
|
||||
|
||||
def get_config(name: str, quantized: bool = True) -> tuple[int, int, bool]:
|
||||
if not self._config.extra_config:
|
||||
return (
|
||||
self._config.weight_bits if quantized else 16,
|
||||
self._config.group_size if quantized else -1,
|
||||
self._config.sym if quantized else True,
|
||||
)
|
||||
|
||||
if name in self._config.extra_config:
|
||||
cfg = self._config.extra_config[name]
|
||||
return (
|
||||
cfg.get("bits", self._config.weight_bits if quantized else 16),
|
||||
cfg.get(
|
||||
"group_size",
|
||||
self._config.group_size if quantized else -1,
|
||||
),
|
||||
cfg.get("sym", self._config.sym if quantized else True),
|
||||
)
|
||||
|
||||
regex_special_chars = set(r"*+?^$()[]{}|\\")
|
||||
for pattern, cfg in self._config.extra_config.items():
|
||||
if not isinstance(pattern, str) or not any(
|
||||
c in regex_special_chars for c in pattern
|
||||
):
|
||||
continue
|
||||
|
||||
try:
|
||||
if re.search(re.compile(pattern), name) is not None:
|
||||
return (
|
||||
cfg.get(
|
||||
"bits",
|
||||
self._config.weight_bits if quantized else 16,
|
||||
),
|
||||
cfg.get(
|
||||
"group_size",
|
||||
self._config.group_size if quantized else -1,
|
||||
),
|
||||
cfg.get("sym", self._config.sym if quantized else True),
|
||||
)
|
||||
except re.error:
|
||||
continue
|
||||
|
||||
return (
|
||||
self._config.weight_bits if quantized else 16,
|
||||
self._config.group_size if quantized else -1,
|
||||
self._config.sym if quantized else True,
|
||||
)
|
||||
|
||||
if self._config.extra_config and layer_name in self._config.extra_config:
|
||||
return get_config(layer_name)
|
||||
|
||||
quantized = not isinstance(layer, ParallelLMHead)
|
||||
if self._config.block_name_to_quantize:
|
||||
quantized = any(
|
||||
layer_name.startswith(name)
|
||||
for name in self._config.block_name_to_quantize
|
||||
)
|
||||
|
||||
if self._config.extra_config and "fusedmoe" in layer.__class__.__name__.lower():
|
||||
moe_configs = [
|
||||
get_config(name, quantized)
|
||||
for name in self._config.extra_config
|
||||
if name.startswith(layer_name)
|
||||
]
|
||||
if moe_configs:
|
||||
if len(set(moe_configs)) == 1:
|
||||
return moe_configs[0]
|
||||
raise ValueError(
|
||||
f"Fused MoE layer '{layer_name}' requires "
|
||||
f"consistent quant config for all sub-layers"
|
||||
)
|
||||
|
||||
if self._config.extra_config:
|
||||
for fusion_key, sub_keys in self._config.packed_modules_mapping.items():
|
||||
if fusion_key in layer_name and layer_name.count(fusion_key) == 1:
|
||||
sub_names = [
|
||||
layer_name.replace(fusion_key, sub_key) for sub_key in sub_keys
|
||||
]
|
||||
# Only trigger if at least one sub_name is explicitly
|
||||
# configured in extra_config (via exact match or regex).
|
||||
# This prevents false matches when a short fusion_key
|
||||
# (e.g. "qkv") is merely a substring of a longer layer
|
||||
# name (e.g. "in_proj_qkvz") and none of the generated
|
||||
# sub_names are actually configured.
|
||||
if not any(is_explicitly_configured(n) for n in sub_names):
|
||||
continue
|
||||
sub_configs = [get_config(name, quantized) for name in sub_names]
|
||||
if len(set(sub_configs)) == 1:
|
||||
return sub_configs[0]
|
||||
raise ValueError(
|
||||
f"Fused module '{layer_name}' requires "
|
||||
f"consistent quant config for {sub_names}"
|
||||
)
|
||||
|
||||
return get_config(layer_name, quantized)
|
||||
@@ -0,0 +1,192 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from fractions import Fraction
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
RoutedExperts,
|
||||
UnquantizedFusedMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import (
|
||||
LinearBase,
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import (
|
||||
QuantizationConfig,
|
||||
QuantizationMethods,
|
||||
)
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
|
||||
|
||||
from .config_parser import INCConfigParser
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class INCConfig(QuantizationConfig):
|
||||
"""Config class for Intel Neural Compressor (INC).
|
||||
Repo: https://github.com/intel/neural-compressor
|
||||
"""
|
||||
|
||||
SUPPORTED_BITS = {2, 3, 4, 8}
|
||||
SUPPORTED_DTYPES = {"int"}
|
||||
SUPPORTED_FORMATS = {"auto_round:auto_gptq", "auto_round:auto_awq"}
|
||||
SUPPORTED_BACKENDS = {
|
||||
"auto",
|
||||
"gptq",
|
||||
"gptq:marlin",
|
||||
"awq",
|
||||
"awq:marlin",
|
||||
"marlin",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weight_bits: int,
|
||||
group_size: int,
|
||||
sym: bool = True,
|
||||
packing_format: str = "auto_round:auto_gptq",
|
||||
block_name_to_quantize: str | list[str] | None = None,
|
||||
extra_config: dict[str, Any] | None = None,
|
||||
data_type: str = "int",
|
||||
backend: str = "auto",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if weight_bits not in self.SUPPORTED_BITS:
|
||||
raise ValueError(
|
||||
f"Unsupported weight_bits: {weight_bits}, "
|
||||
f"currently only support {self.SUPPORTED_BITS}."
|
||||
)
|
||||
if data_type not in self.SUPPORTED_DTYPES:
|
||||
raise ValueError(
|
||||
f"Unsupported data_type: {data_type},"
|
||||
f" currently only support {self.SUPPORTED_DTYPES}."
|
||||
)
|
||||
if packing_format not in self.SUPPORTED_FORMATS:
|
||||
raise ValueError(
|
||||
f"Unsupported packing_format: {packing_format}, "
|
||||
f"currently only support {self.SUPPORTED_FORMATS}."
|
||||
)
|
||||
if backend not in self.SUPPORTED_BACKENDS:
|
||||
raise ValueError(
|
||||
f"Unsupported backend: {backend}, "
|
||||
f"currently only support {self.SUPPORTED_BACKENDS}."
|
||||
)
|
||||
|
||||
self.weight_bits = weight_bits
|
||||
self.group_size = group_size
|
||||
self.sym = sym
|
||||
self.packing_format = packing_format
|
||||
self.block_name_to_quantize = (
|
||||
block_name_to_quantize.split(",")
|
||||
if isinstance(block_name_to_quantize, str)
|
||||
else block_name_to_quantize
|
||||
)
|
||||
self.extra_config = extra_config
|
||||
self.data_type = data_type
|
||||
self.backend = backend
|
||||
self.pack_factor = Fraction(32, weight_bits)
|
||||
self.config_parser = INCConfigParser(self)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"INCConfig(weight_bits={self.weight_bits}, "
|
||||
f"group_size={self.group_size}, sym={self.sym})"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "inc"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.half, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 60
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return ["quantization_config.json"]
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "INCConfig":
|
||||
return cls(
|
||||
weight_bits=cls.get_from_keys(config, ["bits"]),
|
||||
group_size=cls.get_from_keys(config, ["group_size"]),
|
||||
sym=cls.get_from_keys(config, ["sym"]),
|
||||
packing_format=cls.get_from_keys_or(
|
||||
config, ["packing_format"], "auto_round:auto_gptq"
|
||||
),
|
||||
block_name_to_quantize=cls.get_from_keys_or(
|
||||
config, ["block_name_to_quantize", "to_quant_block_names"], None
|
||||
),
|
||||
extra_config=cls.get_from_keys_or(config, ["extra_config"], None),
|
||||
data_type=cls.get_from_keys_or(config, ["data_type"], "int"),
|
||||
backend=cls.get_from_keys_or(config, ["backend", "vllm_backend"], "auto"),
|
||||
)
|
||||
|
||||
def get_layer_config(self, layer, layer_name: str):
|
||||
return self.config_parser.get_layer_config(layer, layer_name)
|
||||
|
||||
def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
|
||||
if self.block_name_to_quantize is not None:
|
||||
self.block_name_to_quantize = hf_to_vllm_mapper.apply_list(
|
||||
self.block_name_to_quantize
|
||||
)
|
||||
if self.extra_config is not None:
|
||||
self.extra_config = hf_to_vllm_mapper.apply_dict(self.extra_config)
|
||||
|
||||
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
|
||||
from .schemes.factory import resolve_scheme
|
||||
|
||||
# Match original: check model.-prefixed names for unquantized layers
|
||||
if prefix and self.extra_config:
|
||||
for layer_name in self.extra_config:
|
||||
if (
|
||||
layer_name == prefix or layer_name == f"model.{prefix}"
|
||||
) and self.extra_config[layer_name].get("bits", 16) >= 16:
|
||||
if isinstance(layer, RoutedExperts):
|
||||
return UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
return UnquantizedLinearMethod()
|
||||
|
||||
layer_config = self.config_parser.resolve(layer, prefix)
|
||||
if not layer_config.quantized:
|
||||
if isinstance(layer, (LinearBase, ParallelLMHead)):
|
||||
return UnquantizedLinearMethod()
|
||||
if isinstance(layer, RoutedExperts):
|
||||
return UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
return None
|
||||
|
||||
logger.debug(
|
||||
"[%s] Type: %s, Bits: %s, Group Size: %s, Sym: %s",
|
||||
prefix,
|
||||
layer.__class__.__name__,
|
||||
layer_config.bits,
|
||||
layer_config.group_size,
|
||||
layer_config.sym,
|
||||
)
|
||||
|
||||
scheme = resolve_scheme(layer_config)
|
||||
if isinstance(layer, (LinearBase, ParallelLMHead)):
|
||||
return scheme.get_linear_method(self, layer, prefix, layer_config)
|
||||
if isinstance(layer, RoutedExperts):
|
||||
return scheme.get_moe_method(self, layer, prefix, layer_config)
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(
|
||||
cls, hf_quant_cfg, user_quant, hf_config=None
|
||||
) -> "QuantizationMethods | None":
|
||||
"""Override the `auto-round` method to `inc`."""
|
||||
is_auto_round_format = hf_quant_cfg.get("quant_method", None) == "auto-round"
|
||||
if is_auto_round_format:
|
||||
return cls.get_name()
|
||||
return None
|
||||
@@ -0,0 +1,47 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.linear import LinearMethodBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .schemes.inc_scheme import INCLinearScheme
|
||||
|
||||
|
||||
class INCLinearMethod(LinearMethodBase):
|
||||
def __init__(self, scheme: "INCLinearScheme") -> None:
|
||||
self.scheme = scheme
|
||||
|
||||
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,
|
||||
):
|
||||
return self.scheme.create_weights(
|
||||
layer=layer,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
output_partition_sizes=output_partition_sizes,
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
params_dtype=params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
return self.scheme.process_weights_after_loading(layer)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return self.scheme.apply_weights(layer, x, bias)
|
||||
@@ -0,0 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from .factory import resolve_scheme
|
||||
from .inc_scheme import INCLinearScheme, INCScheme
|
||||
from .inc_wna16_scheme import INCWna16Scheme
|
||||
|
||||
__all__ = [
|
||||
"INCScheme",
|
||||
"INCLinearScheme",
|
||||
"INCWna16Scheme",
|
||||
"resolve_scheme",
|
||||
]
|
||||
@@ -0,0 +1,22 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..config_parser import INCLayerConfig
|
||||
from .inc_scheme import INCScheme
|
||||
|
||||
|
||||
def resolve_scheme(layer_config: "INCLayerConfig") -> "INCScheme":
|
||||
from .inc_wna16_scheme import INCWna16Scheme
|
||||
|
||||
scheme_list: list[type[INCScheme]] = [
|
||||
INCWna16Scheme,
|
||||
]
|
||||
|
||||
for scheme_cls in scheme_list:
|
||||
if scheme_cls.can_handle(layer_config):
|
||||
return scheme_cls()
|
||||
|
||||
raise NotImplementedError(f"No INC scheme found for layer config: {layer_config}")
|
||||
@@ -0,0 +1,123 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_OPS_REGISTERED = False
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_ark_state() -> tuple[bool, str | None, Any | None, Any | None]:
|
||||
"""Return ARK availability, error details, cached module, and QuantLinear."""
|
||||
try:
|
||||
import auto_round_kernel as ark
|
||||
from auto_round_kernel.qlinear import QuantLinear
|
||||
|
||||
logger.info("Successfully imported auto_round_kernel.")
|
||||
except ImportError as error:
|
||||
return False, str(error), None, None
|
||||
|
||||
if getattr(ark, "cpu_lib", None) is None and getattr(ark, "xpu_lib", None) is None:
|
||||
return (
|
||||
False,
|
||||
"No ARK backend library is available.",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
|
||||
logger.info("Successfully loaded auto_round_kernel backend library.")
|
||||
return True, None, ark, QuantLinear
|
||||
|
||||
|
||||
def _inc_ark_woq_linear_impl(
|
||||
x: torch.Tensor,
|
||||
qweight: torch.Tensor,
|
||||
bias: torch.Tensor | None,
|
||||
out_features: int,
|
||||
in_features: int,
|
||||
group_size: int,
|
||||
compute_type: str,
|
||||
weight_type: str,
|
||||
scale_type: str,
|
||||
asym: bool,
|
||||
) -> torch.Tensor:
|
||||
ark = get_ark_state()[2]
|
||||
assert ark is not None
|
||||
|
||||
return ark.woqgemm_linear(
|
||||
x,
|
||||
qweight,
|
||||
bias,
|
||||
out_features,
|
||||
in_features,
|
||||
group_size,
|
||||
compute_type,
|
||||
weight_type,
|
||||
scale_type,
|
||||
asym,
|
||||
)
|
||||
|
||||
|
||||
def _inc_ark_woq_linear_fake(
|
||||
x: torch.Tensor,
|
||||
qweight: torch.Tensor,
|
||||
bias: torch.Tensor | None,
|
||||
out_features: int,
|
||||
in_features: int,
|
||||
group_size: int,
|
||||
compute_type: str,
|
||||
weight_type: str,
|
||||
scale_type: str,
|
||||
asym: bool,
|
||||
) -> torch.Tensor:
|
||||
del qweight
|
||||
del bias
|
||||
del in_features
|
||||
del group_size
|
||||
del compute_type
|
||||
del weight_type
|
||||
del scale_type
|
||||
del asym
|
||||
return torch.empty(
|
||||
(*x.shape[:-1], out_features),
|
||||
dtype=x.dtype,
|
||||
device=x.device,
|
||||
)
|
||||
|
||||
|
||||
class ark_ops:
|
||||
@staticmethod
|
||||
def register_ops_once() -> None:
|
||||
global _OPS_REGISTERED
|
||||
if _OPS_REGISTERED:
|
||||
return
|
||||
|
||||
is_available, error_str, _, _ = get_ark_state()
|
||||
if not is_available:
|
||||
logger.debug(
|
||||
"Skip registering ark op because ARK is unavailable: %s",
|
||||
error_str or "unknown error",
|
||||
)
|
||||
return
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="inc_ark_woq_linear",
|
||||
op_func=_inc_ark_woq_linear_impl,
|
||||
fake_impl=_inc_ark_woq_linear_fake,
|
||||
dispatch_key=current_platform.dispatch_key,
|
||||
)
|
||||
_OPS_REGISTERED = True
|
||||
|
||||
|
||||
ark_ops.register_ops_once()
|
||||
|
||||
__all__ = ["get_ark_state"]
|
||||
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.fused_moe import FusedMoEMethodBase
|
||||
from vllm.model_executor.layers.linear import LinearMethodBase
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
|
||||
from ..config_parser import INCLayerConfig
|
||||
from ..inc import INCConfig
|
||||
|
||||
|
||||
class INCScheme(ABC):
|
||||
"""One class per quant type. Single registration point for the factory.
|
||||
|
||||
Each subclass defines:
|
||||
- can_handle(): when does this scheme apply?
|
||||
- get_linear_method(): required — how to quantize Linear layers
|
||||
- get_moe_method(): optional — how to quantize MoE layers
|
||||
- get_kvcache_method(): optional — how to quantize KV cache
|
||||
|
||||
Schemes that don't support MoE/KVCache inherit the default raise.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def can_handle(layer_config: "INCLayerConfig") -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def get_linear_method(
|
||||
self,
|
||||
config: "INCConfig",
|
||||
layer: "torch.nn.Module",
|
||||
prefix: str,
|
||||
layer_config: "INCLayerConfig",
|
||||
) -> "LinearMethodBase":
|
||||
raise NotImplementedError
|
||||
|
||||
def get_moe_method(
|
||||
self,
|
||||
config: "INCConfig",
|
||||
layer: "torch.nn.Module",
|
||||
prefix: str,
|
||||
layer_config: "INCLayerConfig",
|
||||
) -> "FusedMoEMethodBase | None":
|
||||
"""Optional. Override if this scheme supports MoE.
|
||||
Default raises NotImplementedError."""
|
||||
raise NotImplementedError(
|
||||
f"{type(self).__name__} does not support MoE layers. "
|
||||
f"Layer config: {layer_config}"
|
||||
)
|
||||
|
||||
def get_kvcache_method(
|
||||
self,
|
||||
config: "INCConfig",
|
||||
layer: "torch.nn.Module",
|
||||
prefix: str,
|
||||
layer_config: "INCLayerConfig",
|
||||
) -> "QuantizationMethods":
|
||||
"""Optional. Override if this scheme supports KV cache quantization.
|
||||
Default raises NotImplementedError."""
|
||||
raise NotImplementedError(
|
||||
f"{type(self).__name__} does not support KV cache quantization. "
|
||||
f"Layer config: {layer_config}"
|
||||
)
|
||||
|
||||
|
||||
class INCLinearScheme(ABC):
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
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,
|
||||
) -> 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,463 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
|
||||
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
check_marlin_supported,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
GroupQuantScaleParameter,
|
||||
PackedvLLMParameter,
|
||||
RowvLLMParameter,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
from .inc_scheme import INCLinearScheme
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..config_parser import INCLayerConfig
|
||||
|
||||
|
||||
class INCWNA16LinearScheme(INCLinearScheme):
|
||||
def __init__(self, layer_config: "INCLayerConfig") -> None:
|
||||
self.layer_config = layer_config
|
||||
self.inner_method = self._build_inner_method()
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 60
|
||||
|
||||
def _build_inner_method(self):
|
||||
if self.layer_config.is_gptq:
|
||||
return self._build_gptq_method()
|
||||
if self.layer_config.is_awq:
|
||||
return self._build_awq_method()
|
||||
raise NotImplementedError(
|
||||
f"WNA16 linear scheme does not support {self.layer_config}"
|
||||
)
|
||||
|
||||
def _build_gptq_method(self):
|
||||
gptq_type_map = {
|
||||
(4, True): scalar_types.uint4b8,
|
||||
(8, True): scalar_types.uint8b128,
|
||||
}
|
||||
use_marlin = (
|
||||
self.layer_config.backend == "auto" or "marlin" in self.layer_config.backend
|
||||
) and (self.layer_config.bits, self.layer_config.sym) in gptq_type_map
|
||||
if use_marlin:
|
||||
use_marlin = check_marlin_supported(
|
||||
gptq_type_map[(self.layer_config.bits, self.layer_config.sym)],
|
||||
self.layer_config.group_size,
|
||||
has_zp=not self.layer_config.sym,
|
||||
)
|
||||
|
||||
if use_marlin:
|
||||
from vllm.model_executor.layers.quantization.auto_gptq import (
|
||||
AutoGPTQLinearMethod,
|
||||
)
|
||||
|
||||
return AutoGPTQLinearMethod(
|
||||
AutoGPTQConfig(
|
||||
weight_bits=self.layer_config.bits,
|
||||
group_size=self.layer_config.group_size,
|
||||
desc_act=False,
|
||||
is_sym=self.layer_config.sym,
|
||||
lm_head_quantized=False,
|
||||
dynamic={},
|
||||
full_config={},
|
||||
)
|
||||
)
|
||||
|
||||
raise NotImplementedError(
|
||||
f"INC quantization with bits={self.layer_config.bits}, "
|
||||
f"sym={self.layer_config.sym} is not supported. "
|
||||
"Only 4-bit and 8-bit symmetric quantization is supported "
|
||||
"with Marlin kernels."
|
||||
)
|
||||
|
||||
def _build_awq_method(self):
|
||||
awq_type_map = {
|
||||
4: scalar_types.uint4,
|
||||
8: scalar_types.uint8,
|
||||
}
|
||||
use_marlin = (
|
||||
self.layer_config.backend == "auto" or "marlin" in self.layer_config.backend
|
||||
) and self.layer_config.bits in awq_type_map
|
||||
if use_marlin:
|
||||
use_marlin = check_marlin_supported(
|
||||
awq_type_map[self.layer_config.bits],
|
||||
self.layer_config.group_size,
|
||||
not self.layer_config.sym,
|
||||
)
|
||||
|
||||
if use_marlin:
|
||||
from vllm.model_executor.layers.quantization.auto_awq import (
|
||||
AutoAWQMarlinLinearMethod,
|
||||
)
|
||||
|
||||
return AutoAWQMarlinLinearMethod(
|
||||
AutoAWQConfig(
|
||||
weight_bits=self.layer_config.bits,
|
||||
group_size=self.layer_config.group_size,
|
||||
zero_point=not self.layer_config.sym,
|
||||
lm_head_quantized=False,
|
||||
modules_to_not_convert=[],
|
||||
full_config={},
|
||||
)
|
||||
)
|
||||
|
||||
from vllm.model_executor.layers.quantization.auto_awq import (
|
||||
AutoAWQLinearMethod,
|
||||
)
|
||||
|
||||
return AutoAWQLinearMethod(
|
||||
AutoAWQConfig(
|
||||
weight_bits=self.layer_config.bits,
|
||||
group_size=self.layer_config.group_size,
|
||||
zero_point=not self.layer_config.sym,
|
||||
lm_head_quantized=False,
|
||||
)
|
||||
)
|
||||
|
||||
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,
|
||||
) -> None:
|
||||
return self.inner_method.create_weights(
|
||||
layer=layer,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
output_partition_sizes=output_partition_sizes,
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
params_dtype=params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: "torch.nn.Module") -> None:
|
||||
return self.inner_method.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.inner_method.apply(layer, x, bias)
|
||||
|
||||
|
||||
class INCXPULinearBase(INCLinearScheme):
|
||||
# AWQ packs nibbles within each int32 in the order [0, 2, 4, 6, 1, 3, 5, 7];
|
||||
# this permutation undoes that ordering so values can be repacked in
|
||||
# standard sequential (GPTQ) order.
|
||||
_REVERSE_AWQ_PACK_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
|
||||
def __init__(self, layer_config: "INCLayerConfig") -> None:
|
||||
self.weight_bits = layer_config.bits
|
||||
self.group_size = layer_config.group_size
|
||||
self.sym = layer_config.sym
|
||||
self.pack_factor = 32 // self.weight_bits
|
||||
self.is_awq_packed = layer_config.is_awq
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 0
|
||||
|
||||
def _create_inc_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
params_dtype: torch.dtype,
|
||||
weight_loader: Any,
|
||||
) -> None:
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
scales_and_zp_size = input_size_per_partition // self.group_size
|
||||
|
||||
if self.is_awq_packed:
|
||||
# AWQ: qweight [in, out // pack_factor] packed along output dim
|
||||
qweight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition // self.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
else:
|
||||
# GPTQ: qweight [in // pack_factor, out] packed along input dim
|
||||
qweight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition // self.pack_factor,
|
||||
output_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=0,
|
||||
packed_factor=self.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
scales = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
scales_and_zp_size,
|
||||
output_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
# Both AWQ and GPTQ checkpoints store qzeros with this shape; for
|
||||
# symmetric quantization the values are ignored downstream.
|
||||
qzeros = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
scales_and_zp_size,
|
||||
output_size_per_partition // self.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
layer.register_parameter("qweight", qweight)
|
||||
layer.register_parameter("scales", scales)
|
||||
layer.register_parameter("qzeros", qzeros)
|
||||
|
||||
g_idx = RowvLLMParameter(
|
||||
data=torch.tensor(
|
||||
[i // self.group_size for i in range(input_size_per_partition)],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("g_idx", g_idx)
|
||||
|
||||
def _convert_awq_qweight_to_gptq(self, qw: torch.Tensor) -> torch.Tensor:
|
||||
"""Convert AWQ qweight [K, N // pf] to GPTQ qweight [K // pf, N].
|
||||
|
||||
AWQ packs along the output dim with a non-standard nibble order; GPTQ
|
||||
packs along the input dim with sequential nibble order. The conversion
|
||||
is lossless — it only reshuffles bits.
|
||||
"""
|
||||
size_bits = self.weight_bits
|
||||
pack_factor = self.pack_factor
|
||||
mask = (1 << size_bits) - 1
|
||||
device = qw.device
|
||||
reverse_order = torch.tensor(
|
||||
self._REVERSE_AWQ_PACK_ORDER, dtype=torch.long, device=device
|
||||
)
|
||||
shifts = torch.arange(0, 32, size_bits, dtype=torch.int32, device=device)
|
||||
|
||||
K, N_packed = qw.shape
|
||||
N = N_packed * pack_factor
|
||||
|
||||
# Unpack int32 → individual values, fix AWQ nibble ordering
|
||||
unpacked = (qw.unsqueeze(-1) >> shifts) & mask # (K, N_packed, pf)
|
||||
unpacked = unpacked[:, :, reverse_order]
|
||||
unpacked = unpacked.reshape(K, N) # (K, N)
|
||||
|
||||
# Repack along input dim (dim 0) in sequential nibble order
|
||||
unpacked = unpacked.reshape(K // pack_factor, pack_factor, N)
|
||||
new_qw = (unpacked.to(torch.int32) << shifts[None, :, None]).sum(
|
||||
dim=1, dtype=torch.int32
|
||||
)
|
||||
return new_qw.contiguous()
|
||||
|
||||
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,
|
||||
) -> None:
|
||||
del input_size, output_size
|
||||
self._create_inc_weights(
|
||||
layer=layer,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
output_partition_sizes=output_partition_sizes,
|
||||
params_dtype=params_dtype,
|
||||
weight_loader=extra_weight_attrs.get("weight_loader"),
|
||||
)
|
||||
|
||||
|
||||
class INCXPULinearMethod(INCXPULinearBase):
|
||||
"""XPU linear method for INC w4a16 quantization (symmetric only).
|
||||
|
||||
Supports both GPTQ-packed (``auto_round:auto_gptq``) and AWQ-packed
|
||||
(``auto_round:auto_awq``) AutoRound checkpoints. AWQ-packed qweights are
|
||||
losslessly repacked into the GPTQ-style nibble layout during
|
||||
``process_weights_after_loading``, before the final oneDNN "NT" transpose
|
||||
that ``torch.ops._xpu_C.int4_gemm_w4a16`` expects.
|
||||
"""
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
device = layer.qweight.data.device
|
||||
|
||||
qweight_data = layer.qweight.data
|
||||
if self.is_awq_packed:
|
||||
# Lossless repack: AWQ [K, N // pf] → GPTQ [K // pf, N]
|
||||
qweight_data = self._convert_awq_qweight_to_gptq(qweight_data)
|
||||
|
||||
qweight_ct = qweight_data.t().contiguous()
|
||||
layer.qweight = Parameter(qweight_ct.t(), requires_grad=False)
|
||||
layer.scales = Parameter(layer.scales.data, requires_grad=False)
|
||||
layer.qzeros = Parameter(
|
||||
torch.tensor([8], dtype=torch.int8, device=device),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
out_shape = x.shape[:-1] + (layer.qweight.shape[1],)
|
||||
reshaped_x = x.reshape(-1, x.shape[-1])
|
||||
out = torch.ops._xpu_C.int4_gemm_w4a16(
|
||||
reshaped_x,
|
||||
layer.qweight,
|
||||
bias,
|
||||
layer.scales,
|
||||
layer.qzeros,
|
||||
self.group_size,
|
||||
None,
|
||||
)
|
||||
return out.reshape(out_shape)
|
||||
|
||||
|
||||
class INCARKLinearMethod(INCXPULinearBase):
|
||||
def __init__(self, layer_config: "INCLayerConfig") -> None:
|
||||
super().__init__(layer_config)
|
||||
|
||||
from .inc_ark_ops import get_ark_state
|
||||
|
||||
is_available, error_str, _, quant_linear_cls = get_ark_state()
|
||||
if not is_available or quant_linear_cls is None:
|
||||
reason = error_str or "unknown error"
|
||||
raise ImportError(f"Failed to import auto_round_kernel. {reason}")
|
||||
|
||||
self.quant_linear_cls = quant_linear_cls
|
||||
|
||||
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,
|
||||
) -> None:
|
||||
super().create_weights(
|
||||
layer=layer,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
output_partition_sizes=output_partition_sizes,
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
params_dtype=params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
layer.in_features = input_size_per_partition
|
||||
layer.out_features = sum(output_partition_sizes)
|
||||
layer.params_dtype = params_dtype
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
if hasattr(layer, "input_size_per_partition"):
|
||||
in_features = layer.input_size_per_partition
|
||||
elif hasattr(layer, "input_size"):
|
||||
in_features = layer.input_size
|
||||
else:
|
||||
raise AttributeError("Cannot determine in_features for layer.")
|
||||
|
||||
if hasattr(layer, "output_partition_sizes"):
|
||||
out_features = sum(layer.output_partition_sizes)
|
||||
elif hasattr(layer, "output_size_per_partition"):
|
||||
out_features = layer.output_size_per_partition
|
||||
elif hasattr(layer, "output_size"):
|
||||
out_features = layer.output_size
|
||||
else:
|
||||
out_features = layer.scales.shape[-1]
|
||||
|
||||
ark_linear = self.quant_linear_cls(
|
||||
bits=self.weight_bits,
|
||||
group_size=self.group_size,
|
||||
sym=self.sym,
|
||||
in_features=in_features,
|
||||
out_features=out_features,
|
||||
bias=layer.bias is not None,
|
||||
weight_dtype=layer.params_dtype,
|
||||
)
|
||||
ark_linear.to(layer.qweight.device)
|
||||
|
||||
with torch.no_grad():
|
||||
qweight_src = layer.qweight.detach()
|
||||
if self.is_awq_packed:
|
||||
# ARK consumes GPTQ-style packed nibbles; convert AWQ losslessly.
|
||||
qweight_src = self._convert_awq_qweight_to_gptq(qweight_src)
|
||||
ark_linear.qweight.copy_(qweight_src)
|
||||
if hasattr(layer, "qzeros") and layer.qzeros is not None:
|
||||
ark_linear.qzeros.copy_(layer.qzeros.detach())
|
||||
else:
|
||||
ark_linear.qzeros = None
|
||||
ark_linear.scales.copy_(layer.scales.detach())
|
||||
if hasattr(layer, "bias") and layer.bias is not None:
|
||||
ark_linear.bias.copy_(layer.bias.detach())
|
||||
|
||||
ark_linear.post_init()
|
||||
|
||||
layer.qweight = Parameter(ark_linear.qweight.detach(), requires_grad=False)
|
||||
layer.ark_bias = ark_linear.bias
|
||||
layer.ark_compute_type = ark_linear.cdt
|
||||
layer.ark_weight_type = ark_linear.wdt
|
||||
layer.ark_scale_type = ark_linear.sdt
|
||||
|
||||
if hasattr(layer, "qzeros"):
|
||||
del layer.qzeros
|
||||
del layer.scales
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return torch.ops.vllm.inc_ark_woq_linear.default(
|
||||
x,
|
||||
layer.qweight,
|
||||
layer.ark_bias,
|
||||
layer.out_features,
|
||||
layer.in_features,
|
||||
self.group_size,
|
||||
layer.ark_compute_type,
|
||||
layer.ark_weight_type,
|
||||
layer.ark_scale_type,
|
||||
not self.sym,
|
||||
)
|
||||
|
||||
|
||||
class INCXPUW4A16LinearScheme(INCXPULinearMethod):
|
||||
pass
|
||||
@@ -0,0 +1,201 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
|
||||
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
from ..inc_linear import INCLinearMethod
|
||||
from .inc_scheme import INCScheme
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
from ..config_parser import INCLayerConfig
|
||||
from ..inc import INCConfig
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class INCWna16Scheme(INCScheme):
|
||||
@staticmethod
|
||||
def can_handle(layer_config: "INCLayerConfig") -> bool:
|
||||
return layer_config.is_wna16_int
|
||||
|
||||
def get_linear_method(
|
||||
self,
|
||||
config: "INCConfig",
|
||||
layer: "torch.nn.Module",
|
||||
prefix: str,
|
||||
layer_config: "INCLayerConfig",
|
||||
):
|
||||
del config, layer
|
||||
if current_platform.is_xpu():
|
||||
if layer_config.bits == 4 and layer_config.sym:
|
||||
from .inc_ark_ops import get_ark_state
|
||||
from .inc_wna16_linear import (
|
||||
INCARKLinearMethod,
|
||||
INCXPULinearMethod,
|
||||
)
|
||||
|
||||
is_ark_available, ark_error, _, _ = get_ark_state()
|
||||
if is_ark_available:
|
||||
return INCLinearMethod(INCARKLinearMethod(layer_config))
|
||||
|
||||
logger.debug(
|
||||
"ARK backend is unavailable for layer %s; "
|
||||
"falling back to the default XPU INC path. Error: %s",
|
||||
prefix,
|
||||
ark_error or "unknown error",
|
||||
)
|
||||
return INCLinearMethod(INCXPULinearMethod(layer_config))
|
||||
raise NotImplementedError(f"INC on XPU: unsupported config {layer_config}")
|
||||
|
||||
if current_platform.is_cpu() and layer_config.is_gptq:
|
||||
if layer_config.bits == 4 and layer_config.sym:
|
||||
from .inc_ark_ops import get_ark_state
|
||||
from .inc_wna16_linear import (
|
||||
INCARKLinearMethod,
|
||||
INCWNA16LinearScheme,
|
||||
)
|
||||
|
||||
is_ark_available, ark_error, _, _ = get_ark_state()
|
||||
if is_ark_available:
|
||||
return INCLinearMethod(INCARKLinearMethod(layer_config))
|
||||
|
||||
logger.debug(
|
||||
"ARK backend is unavailable for layer %s; "
|
||||
"falling back to the default CPU INC path. Error: %s",
|
||||
prefix,
|
||||
ark_error or "unknown error",
|
||||
)
|
||||
return INCLinearMethod(INCWNA16LinearScheme(layer_config))
|
||||
raise NotImplementedError(f"INC on CPU: unsupported config {layer_config}")
|
||||
|
||||
from .inc_wna16_linear import INCWNA16LinearScheme
|
||||
|
||||
return INCLinearMethod(INCWNA16LinearScheme(layer_config))
|
||||
|
||||
def get_moe_method(
|
||||
self,
|
||||
config: "INCConfig",
|
||||
layer: "torch.nn.Module",
|
||||
prefix: str,
|
||||
layer_config: "INCLayerConfig",
|
||||
):
|
||||
del config, prefix
|
||||
# XPU and CPU do not support MoE quantization yet
|
||||
if current_platform.is_xpu() or current_platform.is_cpu():
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
UnquantizedFusedMoEMethod,
|
||||
)
|
||||
|
||||
return UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
if layer_config.is_gptq:
|
||||
return _resolve_gptq_moe(layer, layer_config)
|
||||
if layer_config.is_awq:
|
||||
return _resolve_awq_moe(layer, layer_config)
|
||||
raise NotImplementedError(f"WNA16 MoE does not support config {layer_config}")
|
||||
|
||||
|
||||
def _resolve_gptq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"):
|
||||
from vllm.model_executor.layers.quantization.auto_gptq import (
|
||||
AutoGPTQMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.moe_wna16 import (
|
||||
MoeWNA16Config,
|
||||
MoeWNA16Method,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
check_marlin_supported,
|
||||
check_moe_marlin_supports_layer,
|
||||
)
|
||||
|
||||
gptq_type_map = {
|
||||
(4, True): scalar_types.uint4b8,
|
||||
(8, True): scalar_types.uint8b128,
|
||||
}
|
||||
use_marlin = (layer_config.bits, layer_config.sym) in gptq_type_map
|
||||
if use_marlin:
|
||||
use_marlin = check_marlin_supported(
|
||||
gptq_type_map[(layer_config.bits, layer_config.sym)],
|
||||
layer_config.group_size,
|
||||
has_zp=not layer_config.sym,
|
||||
) and check_moe_marlin_supports_layer(layer, layer_config.group_size)
|
||||
|
||||
if use_marlin:
|
||||
return AutoGPTQMoEMethod(
|
||||
AutoGPTQConfig(
|
||||
weight_bits=layer_config.bits,
|
||||
group_size=layer_config.group_size,
|
||||
desc_act=False,
|
||||
is_sym=layer_config.sym,
|
||||
lm_head_quantized=False,
|
||||
dynamic={},
|
||||
full_config={},
|
||||
),
|
||||
layer.moe_config,
|
||||
)
|
||||
|
||||
moe_config = MoeWNA16Config.from_config(
|
||||
{
|
||||
"quant_method": "gptq",
|
||||
"bits": layer_config.bits,
|
||||
"group_size": layer_config.group_size,
|
||||
"sym": layer_config.sym,
|
||||
"lm_head": False,
|
||||
}
|
||||
)
|
||||
return MoeWNA16Method(moe_config, layer.moe_config)
|
||||
|
||||
|
||||
def _resolve_awq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"):
|
||||
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQMoEMethod
|
||||
from vllm.model_executor.layers.quantization.moe_wna16 import (
|
||||
MoeWNA16Config,
|
||||
MoeWNA16Method,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
check_marlin_supported,
|
||||
check_moe_marlin_supports_layer,
|
||||
)
|
||||
|
||||
awq_type_map = {
|
||||
4: scalar_types.uint4,
|
||||
8: scalar_types.uint8,
|
||||
}
|
||||
use_marlin = layer_config.bits in awq_type_map
|
||||
if use_marlin:
|
||||
use_marlin = check_marlin_supported(
|
||||
awq_type_map[layer_config.bits],
|
||||
layer_config.group_size,
|
||||
not layer_config.sym,
|
||||
) and check_moe_marlin_supports_layer(layer, layer_config.group_size)
|
||||
|
||||
if use_marlin:
|
||||
return AutoAWQMoEMethod(
|
||||
AutoAWQConfig(
|
||||
weight_bits=layer_config.bits,
|
||||
group_size=layer_config.group_size,
|
||||
zero_point=not layer_config.sym,
|
||||
lm_head_quantized=False,
|
||||
modules_to_not_convert=[],
|
||||
full_config={},
|
||||
),
|
||||
layer.moe_config,
|
||||
)
|
||||
|
||||
moe_config = MoeWNA16Config.from_config(
|
||||
{
|
||||
"quant_method": "awq",
|
||||
"bits": layer_config.bits,
|
||||
"group_size": layer_config.group_size,
|
||||
"zero_point": not layer_config.sym,
|
||||
"lm_head": False,
|
||||
}
|
||||
)
|
||||
return MoeWNA16Method(moe_config, layer.moe_config)
|
||||
@@ -0,0 +1,266 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.model_executor.custom_op import CustomOp
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
get_fp8_min_max,
|
||||
group_broadcast,
|
||||
prep_scale_for_group_broadcast,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.deep_gemm import (
|
||||
DeepGemmQuantScaleFMT,
|
||||
is_deep_gemm_e8m0_used,
|
||||
is_deep_gemm_supported,
|
||||
)
|
||||
|
||||
_FP8_DTYPE = current_platform.fp8_dtype()
|
||||
_FP8_MIN, _FP8_MAX = get_fp8_min_max()
|
||||
_FP8_MIN_SCALING_FACTOR = 1.0 / (_FP8_MAX * 512.0)
|
||||
|
||||
|
||||
# --8<-- [start:quant_fp8]
|
||||
@CustomOp.register("quant_fp8")
|
||||
class QuantFP8(CustomOp):
|
||||
"""
|
||||
Quantize input tensor to FP8 (per-tensor, per-token, per-channel, or per-group).
|
||||
This CustomOp supports both static and dynamic quantization.
|
||||
"""
|
||||
|
||||
# --8<-- [end:quant_fp8]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
static: bool,
|
||||
group_shape: GroupShape,
|
||||
num_token_padding: int | None = None,
|
||||
column_major_scales: bool = False,
|
||||
tma_aligned_scales: bool = False,
|
||||
use_ue8m0: bool | None = None, # for Torch compile
|
||||
compile_native: bool = True,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
static: static or dynamic quantization
|
||||
group_shape: quantization group shape (PER_TOKEN, PER_TENSOR,
|
||||
PER_CHANNEL, or arbitrary block size)
|
||||
num_token_padding: Pad the token dimension of output to this
|
||||
size
|
||||
tma_aligned_scales: For group quantization, output scales in
|
||||
TMA-aligned layout
|
||||
column_major_scales: For group quantization, output scales in
|
||||
column major format
|
||||
compile_native: Manually compile forward_native if compile mode > None
|
||||
"""
|
||||
super().__init__(compile_native=compile_native)
|
||||
self.static = static
|
||||
self.group_shape = group_shape
|
||||
self.use_per_token_if_dynamic = group_shape == GroupShape.PER_TOKEN
|
||||
self.num_token_padding = num_token_padding
|
||||
self.column_major_scales = column_major_scales
|
||||
self.tma_aligned_scales = tma_aligned_scales
|
||||
self.use_ue8m0 = is_deep_gemm_e8m0_used() if use_ue8m0 is None else use_ue8m0
|
||||
self.use_deep_gemm_supported = is_deep_gemm_supported()
|
||||
|
||||
self.use_aiter = rocm_aiter_ops.is_linear_fp8_enabled()
|
||||
|
||||
self.is_group_quant = group_shape.is_per_group()
|
||||
if self.is_group_quant:
|
||||
self.group_size = group_shape.col
|
||||
else:
|
||||
self.use_per_token_if_dynamic = group_shape == GroupShape.PER_TOKEN
|
||||
if not static:
|
||||
assert group_shape in (GroupShape.PER_TOKEN, GroupShape.PER_TENSOR), (
|
||||
"Only per-token or per-tensor scales are supported for dynamic "
|
||||
"non-group quantization."
|
||||
)
|
||||
|
||||
def forward_cuda(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor | None = None,
|
||||
scale_ub: torch.Tensor | None = None,
|
||||
use_triton: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
from vllm.model_executor.layers.quantization.utils import fp8_utils
|
||||
|
||||
if (
|
||||
self.is_group_quant
|
||||
and self.use_ue8m0
|
||||
and self.use_deep_gemm_supported
|
||||
and (DeepGemmQuantScaleFMT.from_oracle() == DeepGemmQuantScaleFMT.UE8M0)
|
||||
):
|
||||
return fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
|
||||
x,
|
||||
group_size=self.group_size,
|
||||
use_ue8m0=True,
|
||||
)
|
||||
|
||||
if self.is_group_quant and not self.static:
|
||||
assert scale is None, "Dynamic group quantization does not use scale"
|
||||
|
||||
return fp8_utils.per_token_group_quant_fp8(
|
||||
x,
|
||||
group_size=self.group_size,
|
||||
column_major_scales=self.column_major_scales,
|
||||
tma_aligned_scales=self.tma_aligned_scales,
|
||||
dtype=_FP8_DTYPE,
|
||||
use_ue8m0=self.use_ue8m0,
|
||||
)
|
||||
|
||||
assert (scale is not None) == self.static
|
||||
assert scale_ub is None or (
|
||||
not self.static
|
||||
and self.group_shape == GroupShape.PER_TOKEN
|
||||
and scale_ub.numel() == 1
|
||||
)
|
||||
|
||||
return ops.scaled_fp8_quant(
|
||||
x,
|
||||
scale,
|
||||
num_token_padding=self.num_token_padding,
|
||||
scale_ub=scale_ub,
|
||||
use_per_token_if_dynamic=self.use_per_token_if_dynamic,
|
||||
group_shape=(self.group_shape.row, self.group_shape.col)
|
||||
if self.static
|
||||
else None,
|
||||
)
|
||||
|
||||
def forward_hip(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor | None = None,
|
||||
scale_ub: torch.Tensor | None = None,
|
||||
use_triton: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if self.is_group_quant and use_triton:
|
||||
assert scale is None, "Dynamic group quantization does not use scale"
|
||||
|
||||
return torch.ops.vllm.triton_per_token_group_quant_fp8(x, self.group_size)
|
||||
|
||||
use_aiter_quant = self.use_aiter and scale_ub is None and x.is_contiguous()
|
||||
use_aiter_per_tensor_quant = (
|
||||
use_aiter_quant and self.group_shape.is_per_tensor()
|
||||
)
|
||||
use_aiter_per_token_quant = use_aiter_quant and self.group_shape.is_per_token()
|
||||
|
||||
use_aiter_per_group_quant = use_aiter_quant and self.group_shape.is_per_group()
|
||||
|
||||
if use_aiter_per_group_quant:
|
||||
return rocm_aiter_ops.group_fp8_quant(x, self.group_size)
|
||||
if use_aiter_per_tensor_quant:
|
||||
return rocm_aiter_ops.per_tensor_quant(x, _FP8_DTYPE, scale)
|
||||
if use_aiter_per_token_quant:
|
||||
return rocm_aiter_ops.per_token_quant(x, _FP8_DTYPE, scale)
|
||||
|
||||
# Fallback to CUDA implementation
|
||||
return self.forward_cuda(x, scale, scale_ub)
|
||||
|
||||
def forward_xpu(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor | None = None,
|
||||
scale_ub: torch.Tensor | None = None,
|
||||
use_triton: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if self.is_group_quant and not self.static:
|
||||
from vllm.model_executor.layers.quantization.utils import fp8_utils
|
||||
|
||||
return fp8_utils.per_token_group_quant_fp8(
|
||||
x,
|
||||
group_size=self.group_size,
|
||||
column_major_scales=self.column_major_scales,
|
||||
dtype=_FP8_DTYPE,
|
||||
use_ue8m0=self.use_ue8m0,
|
||||
)
|
||||
return self.forward_cuda(x, scale, scale_ub, use_triton)
|
||||
|
||||
def forward_native(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor | None = None,
|
||||
scale_ub: torch.Tensor | None = None,
|
||||
use_triton: bool = False,
|
||||
):
|
||||
if self.is_group_quant and not self.static:
|
||||
assert scale is None, "Dynamic group quantization does not use scale"
|
||||
return self._quantize_group_native(x)
|
||||
|
||||
assert (scale is not None) == self.static
|
||||
assert scale_ub is None or (
|
||||
not self.static
|
||||
and self.group_shape == GroupShape.PER_TOKEN
|
||||
and scale_ub.numel() == 1
|
||||
)
|
||||
|
||||
if scale is None:
|
||||
if self.group_shape == GroupShape.PER_TOKEN:
|
||||
x_max, _ = x.abs().max(dim=-1)
|
||||
x_max = x_max.unsqueeze(-1).to(torch.float32)
|
||||
if scale_ub is not None:
|
||||
x_max = x_max.clamp(max=scale_ub)
|
||||
else:
|
||||
x_max = x.abs().max().unsqueeze(-1).to(torch.float32)
|
||||
|
||||
scale = (x_max / _FP8_MAX).clamp(min=_FP8_MIN_SCALING_FACTOR)
|
||||
else:
|
||||
scale = prep_scale_for_group_broadcast(scale, x, self.group_shape)
|
||||
|
||||
# Even for dynamic per-token scales,
|
||||
# reciprocal performs slightly better than division
|
||||
out = (
|
||||
x.to(torch.float32)
|
||||
* group_broadcast(scale.to(torch.float32), x.shape[-2:]).reciprocal()
|
||||
)
|
||||
out = out.clamp(_FP8_MIN, _FP8_MAX).to(_FP8_DTYPE)
|
||||
|
||||
# This currently generates an extra Triton kernel in compilation.
|
||||
# Fortunately, we don't use padding if compiling.
|
||||
# TODO(luka): benchmark torch._scaled_mm to hopefully remove padding
|
||||
# in general.
|
||||
if self.num_token_padding is not None:
|
||||
padding = max(self.num_token_padding - out.size(0), 0)
|
||||
out = F.pad(out, (0, 0, 0, padding), "constant", 0.0)
|
||||
|
||||
return out, scale
|
||||
|
||||
def _quantize_group_native(
|
||||
self, x: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
orig_shape = x.shape
|
||||
hidden_dim = x.shape[-1]
|
||||
num_groups = (hidden_dim + self.group_size - 1) // self.group_size
|
||||
padded_dim = num_groups * self.group_size
|
||||
|
||||
if padded_dim != hidden_dim:
|
||||
padding = padded_dim - hidden_dim
|
||||
x = F.pad(x, (0, padding), mode="constant", value=0.0)
|
||||
|
||||
x_grouped = x.view(-1, num_groups, self.group_size)
|
||||
absmax = x_grouped.abs().max(dim=-1, keepdim=True)[0].float()
|
||||
scales_raw = absmax / _FP8_MAX
|
||||
if self.use_ue8m0:
|
||||
scales_raw = torch.exp2(torch.ceil(torch.log2(scales_raw)))
|
||||
scales = (scales_raw).clamp(min=_FP8_MIN_SCALING_FACTOR)
|
||||
|
||||
x_scaled = x_grouped / scales
|
||||
x_quant = x_scaled.clamp(_FP8_MIN, _FP8_MAX).to(_FP8_DTYPE)
|
||||
|
||||
x_quant = x_quant.view(-1, padded_dim)
|
||||
if padded_dim != hidden_dim:
|
||||
x_quant = x_quant[..., :hidden_dim]
|
||||
x_quant = x_quant.view(orig_shape)
|
||||
|
||||
scales = scales.squeeze(-1)
|
||||
scales = scales.reshape(orig_shape[:-1] + (num_groups,))
|
||||
|
||||
if self.column_major_scales:
|
||||
scales = scales.transpose(-2, -1).contiguous().transpose(-1, -2)
|
||||
|
||||
return x_quant, scales
|
||||
@@ -0,0 +1,197 @@
|
||||
# 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.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import is_quantized_kv_cache
|
||||
from vllm.v1.kv_cache_interface import kv_cache_uses_per_token_head_scales
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class KVCacheScaleParameter(torch.nn.Parameter):
|
||||
"""Scalar parameter for KV-cache scales.
|
||||
|
||||
Initialized to -1.0 (an invalid sentinel) so call sites just write
|
||||
`KVCacheScaleParameter()`. The `weight_loader` accepts shape `()` or
|
||||
`(1,)` and rejects anything else — per-head scales go through a separate
|
||||
path (compressed-tensors' `_tp_aware_loader`), not this one. Per-instance
|
||||
overrides still work because instance attribute assignment shadows this
|
||||
class-level loader.
|
||||
"""
|
||||
|
||||
def __new__(cls) -> "KVCacheScaleParameter":
|
||||
return super().__new__(cls, torch.tensor(-1.0), requires_grad=False)
|
||||
|
||||
@staticmethod
|
||||
def weight_loader(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> None:
|
||||
if loaded_weight.numel() != 1:
|
||||
raise ValueError(
|
||||
f"KV-cache scale expects a scalar weight, got shape "
|
||||
f"{tuple(loaded_weight.shape)}"
|
||||
)
|
||||
param.data.copy_(loaded_weight.reshape(()))
|
||||
|
||||
|
||||
class BaseKVCacheMethod(QuantizeMethodBase):
|
||||
"""
|
||||
Quant method that adds `_k_scale` and `_v_scale` attributes to the
|
||||
Attention layer to support loading those scaling factors from checkpoints.
|
||||
The k/v_scale will be used to:
|
||||
- quantize k/v_cache entries before saving them to the cache
|
||||
- dequantize k/v_cache entries before fetching them from the cache
|
||||
|
||||
Args:
|
||||
quant_config: the appropriate QuantizationConfig
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: QuantizationConfig):
|
||||
self.quant_config = quant_config
|
||||
|
||||
def create_weights(self, layer: torch.nn.Module):
|
||||
"""
|
||||
Create "weight" (aka q_scale, k_scale and v_scale)
|
||||
for an attention layer.
|
||||
"""
|
||||
# Initialize the Q and KV cache scales to -1.0, an invalid value.
|
||||
# If the q and k/v_scales appear in the checkpoint, it will be
|
||||
# overwritten when loading weights.
|
||||
layer.q_scale = KVCacheScaleParameter()
|
||||
layer.k_scale = KVCacheScaleParameter()
|
||||
layer.v_scale = KVCacheScaleParameter()
|
||||
# Initialize P = softmax(QK^T) scales
|
||||
layer.prob_scale = KVCacheScaleParameter()
|
||||
|
||||
def apply(self, layer: torch.nn.Module) -> torch.Tensor:
|
||||
raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.")
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
# skip if there are no weights to process (for example, weight reloading)
|
||||
if not hasattr(layer, "q_scale"):
|
||||
assert not hasattr(layer, "k_scale")
|
||||
assert not hasattr(layer, "v_scale")
|
||||
assert not hasattr(layer, "prob_scale")
|
||||
return
|
||||
|
||||
# Per-token-head quantized KV cache: scales are computed dynamically
|
||||
# per (token, head) in the kernel at cache-write time. Checkpoint
|
||||
# scales are never used regardless of calculate_kv_scales.
|
||||
if kv_cache_uses_per_token_head_scales(layer.kv_cache_dtype):
|
||||
layer._k_scale.copy_(1.0)
|
||||
layer._v_scale.copy_(1.0)
|
||||
layer._k_scale_float = 1.0
|
||||
layer._v_scale_float = 1.0
|
||||
del layer.k_scale
|
||||
del layer.v_scale
|
||||
del layer.q_scale
|
||||
del layer.prob_scale
|
||||
return
|
||||
|
||||
# If the kv-cache is not quantized, we enforce the k/v_scale to be 1.0
|
||||
# regardless whether the kv-scale is available in the checkpoint.
|
||||
# No need to process kv scales after loading if we are going to
|
||||
# calculate them on the fly.
|
||||
if (
|
||||
is_quantized_kv_cache(layer.kv_cache_dtype)
|
||||
and not layer.calculate_kv_scales
|
||||
):
|
||||
if layer.k_scale > 0.0 and layer.v_scale > 0.0:
|
||||
# We prefer to use separate k_scale and v_scale if present
|
||||
k_scale = layer.k_scale.to("cpu").tolist()
|
||||
v_scale = layer.v_scale.to("cpu").tolist()
|
||||
if current_platform.is_fp8_fnuz():
|
||||
k_scale *= 2
|
||||
v_scale *= 2
|
||||
elif layer.k_scale < 0.0 and layer.v_scale < 0.0:
|
||||
# If no scales were loaded (both scales are invalid negative
|
||||
# values), use the default value of 1.0
|
||||
k_scale = 1.0
|
||||
v_scale = 1.0
|
||||
else:
|
||||
# If we find a single kv_scale in the checkpoint, we remap
|
||||
# kv_scale to k_scale during weight loading, and duplicate
|
||||
# k_scale to v_scale here
|
||||
assert layer.k_scale > 0.0
|
||||
scale_to_duplicate = max(layer.k_scale, layer.v_scale)
|
||||
k_scale = scale_to_duplicate.to("cpu").tolist()
|
||||
v_scale = scale_to_duplicate.to("cpu").tolist()
|
||||
if current_platform.is_fp8_fnuz():
|
||||
k_scale *= 2
|
||||
v_scale *= 2
|
||||
|
||||
if not isinstance(k_scale, float) or not isinstance(v_scale, float):
|
||||
raise ValueError(
|
||||
"Only support per-tensor scaling factor for fp8 KV cache"
|
||||
)
|
||||
|
||||
if layer.q_scale < 0.0:
|
||||
logger.warning_once(
|
||||
"Checkpoint does not provide a q scaling factor. "
|
||||
"Setting it to k_scale. This only matters for "
|
||||
"FP8 Attention backends (flash-attn or flashinfer)."
|
||||
)
|
||||
layer._q_scale.copy_(k_scale)
|
||||
layer._q_scale_float = k_scale
|
||||
|
||||
# These are used in the final Attention.forward()
|
||||
layer._k_scale.copy_(k_scale)
|
||||
layer._v_scale.copy_(v_scale)
|
||||
layer._k_scale_float = k_scale
|
||||
layer._v_scale_float = v_scale
|
||||
if k_scale == 1.0 and v_scale == 1.0 and "e5m2" not in layer.kv_cache_dtype:
|
||||
logger.warning_once(
|
||||
"Using KV cache scaling factor 1.0 for fp8_e4m3. "
|
||||
"If this is unintended, verify that k/v_scale "
|
||||
"scaling factors are properly set in the checkpoint."
|
||||
)
|
||||
|
||||
if layer.q_scale > 0.0:
|
||||
q_scale = layer.q_scale
|
||||
if current_platform.is_fp8_fnuz():
|
||||
q_scale *= 2
|
||||
layer.calculate_kv_scales = False
|
||||
else:
|
||||
q_scale = 1.0
|
||||
if layer.prob_scale > 0.0:
|
||||
prob_scale = layer.prob_scale
|
||||
if current_platform.is_fp8_fnuz():
|
||||
prob_scale *= 2
|
||||
else:
|
||||
prob_scale = 1.0
|
||||
|
||||
is_singleton_float = (
|
||||
lambda x: isinstance(x, float)
|
||||
or isinstance(x, torch.Tensor)
|
||||
and x.numel() == 1
|
||||
and x.is_floating_point()
|
||||
)
|
||||
if not is_singleton_float(q_scale) or not is_singleton_float(prob_scale):
|
||||
raise ValueError(
|
||||
"Only support per-tensor scaling factorfor fp8-quantized Q/prob"
|
||||
)
|
||||
|
||||
# These are used in the final Attention.forward()
|
||||
layer._q_scale.copy_(q_scale)
|
||||
layer._q_scale_float = (
|
||||
q_scale.item() if isinstance(q_scale, torch.Tensor) else q_scale
|
||||
)
|
||||
|
||||
layer._prob_scale.copy_(prob_scale)
|
||||
if layer.kv_cache_dtype == "fp8" and (q_scale == 1.0 or prob_scale == 1.0):
|
||||
logger.warning_once(
|
||||
f"Using uncalibrated q_scale {q_scale} and/or prob_scale "
|
||||
f"{prob_scale} with fp8 attention. This may cause accuracy "
|
||||
"issues. Please make sure q/prob scaling factors are "
|
||||
"available in the fp8 checkpoint."
|
||||
)
|
||||
|
||||
del layer.k_scale
|
||||
del layer.v_scale
|
||||
del layer.q_scale
|
||||
del layer.prob_scale
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,491 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEMethodBase,
|
||||
FusedMoeWeightScaleSupported,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
int4_w4a16_moe_quant_config,
|
||||
int8_w8a16_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
|
||||
UnquantizedFusedMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
class MoeWNA16Config(QuantizationConfig):
|
||||
"""Config class for MOE WNA16 (W8A16/W4A16) quantization."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
linear_quant_method: str,
|
||||
weight_bits: int,
|
||||
group_size: int,
|
||||
has_zp: bool,
|
||||
lm_head_quantized: bool,
|
||||
modules_to_not_convert: list[str] | None,
|
||||
full_config: dict[str, Any],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.weight_bits = weight_bits
|
||||
self.group_size = group_size
|
||||
self.has_zp = has_zp
|
||||
self.bit8_pack_factor = 8 // self.weight_bits
|
||||
self.lm_head_quantized = lm_head_quantized
|
||||
self.linear_quant_method = linear_quant_method
|
||||
self.full_config = full_config
|
||||
# Avoid circular import
|
||||
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
|
||||
|
||||
if self.linear_quant_method == "gptq":
|
||||
pass
|
||||
elif self.linear_quant_method in ("awq", "awq_marlin"):
|
||||
capability_tuple = current_platform.get_device_capability()
|
||||
device_capability = (
|
||||
-1 if capability_tuple is None else capability_tuple.to_int()
|
||||
)
|
||||
awq_min_capability = AutoAWQConfig.get_min_capability()
|
||||
if device_capability < awq_min_capability:
|
||||
raise ValueError(
|
||||
"The quantization method moe_wna16 + awq is not supported "
|
||||
"for the current GPU. "
|
||||
f"Minimum capability: {awq_min_capability}. "
|
||||
f"Current capability: {device_capability}."
|
||||
)
|
||||
else:
|
||||
raise ValueError("moe_wna16 only support gptq and awq.")
|
||||
|
||||
if modules_to_not_convert is None:
|
||||
self.modules_to_not_convert = []
|
||||
else:
|
||||
self.modules_to_not_convert = modules_to_not_convert
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "moe_wna16"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.bfloat16, torch.half]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 70
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return ["quantize_config.json"]
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "MoeWNA16Config":
|
||||
linear_quant_method = cls.get_from_keys(config, ["quant_method"])
|
||||
weight_bits = cls.get_from_keys(config, ["bits"])
|
||||
group_size = cls.get_from_keys(config, ["group_size"])
|
||||
lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
|
||||
if linear_quant_method == "gptq":
|
||||
has_zp = not cls.get_from_keys(config, ["sym"])
|
||||
modules_to_not_convert = []
|
||||
elif linear_quant_method in ("awq", "awq_marlin"):
|
||||
has_zp = cls.get_from_keys(config, ["zero_point"])
|
||||
modules_to_not_convert = cls.get_from_keys_or(
|
||||
config, ["modules_to_not_convert"], None
|
||||
)
|
||||
else:
|
||||
raise ValueError("moe_wna16 only support gptq and awq.")
|
||||
|
||||
return cls(
|
||||
linear_quant_method,
|
||||
weight_bits,
|
||||
group_size,
|
||||
has_zp,
|
||||
lm_head_quantized,
|
||||
modules_to_not_convert,
|
||||
config,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(
|
||||
cls, hf_quant_cfg, user_quant, hf_config=None
|
||||
) -> QuantizationMethods | None:
|
||||
can_convert = cls.is_moe_wna16_compatible(hf_quant_cfg)
|
||||
if can_convert and user_quant == "moe_wna16":
|
||||
return cls.get_name()
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def is_moe_wna16_compatible(cls, quant_config: dict[str, Any]):
|
||||
# Extract data from quant config.
|
||||
quant_method = quant_config.get("quant_method", "").lower()
|
||||
num_bits = quant_config.get("bits")
|
||||
desc_act = quant_config.get("desc_act")
|
||||
|
||||
capability_tuple = current_platform.get_device_capability()
|
||||
device_capability = (
|
||||
-1 if capability_tuple is None else capability_tuple.to_int()
|
||||
)
|
||||
# Avoid circular import
|
||||
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
|
||||
|
||||
awq_min_capability = AutoAWQConfig.get_min_capability()
|
||||
|
||||
gptq_compatible = quant_method == "gptq" and not desc_act and num_bits in [4, 8]
|
||||
awq_compatible = (
|
||||
quant_method == "awq"
|
||||
and num_bits == 4
|
||||
and device_capability >= awq_min_capability
|
||||
)
|
||||
|
||||
return gptq_compatible or awq_compatible
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
if is_layer_skipped_quant(prefix, self.modules_to_not_convert):
|
||||
if isinstance(layer, RoutedExperts):
|
||||
return UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
return UnquantizedLinearMethod()
|
||||
elif isinstance(layer, LinearBase):
|
||||
# Avoid circular import
|
||||
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
|
||||
from vllm.model_executor.layers.quantization.auto_gptq import (
|
||||
AutoGPTQConfig,
|
||||
)
|
||||
|
||||
if self.linear_quant_method == "gptq":
|
||||
return AutoGPTQConfig.from_config(self.full_config).get_quant_method(
|
||||
layer, prefix
|
||||
)
|
||||
elif self.linear_quant_method in ("awq", "awq_marlin"):
|
||||
return AutoAWQConfig.from_config(self.full_config).get_quant_method(
|
||||
layer, prefix
|
||||
)
|
||||
else:
|
||||
raise ValueError("moe_wna16 only support gptq and awq.")
|
||||
elif isinstance(layer, RoutedExperts):
|
||||
return MoeWNA16Method(self, layer.moe_config)
|
||||
return None
|
||||
|
||||
|
||||
def is_layer_skipped_quant(prefix: str, modules_to_not_convert: list[str]):
|
||||
return any(module_name in prefix for module_name in modules_to_not_convert)
|
||||
|
||||
|
||||
class MoeWNA16Method(FusedMoEMethodBase):
|
||||
"""Linear method for MOE WNA16 (W8A16/W4A16) quantization.
|
||||
|
||||
Args:
|
||||
quant_config: The MOE WNA16 (W8A16/W4A16) quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None:
|
||||
super().__init__(moe)
|
||||
self.quant_config = quant_config
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.quant_config = self.quant_config
|
||||
bit8_pack_factor = self.quant_config.bit8_pack_factor
|
||||
group_size = self.quant_config.group_size
|
||||
group_size_div_factor = 1
|
||||
|
||||
# make intermediate_size and hidden_size divisible by group_size
|
||||
# we reduce the group size to ensure that
|
||||
# and we would repeat the loaded_weight later
|
||||
while intermediate_size_per_partition % group_size or hidden_size % group_size:
|
||||
group_size = group_size // 2
|
||||
group_size_div_factor *= 2
|
||||
assert group_size >= 32
|
||||
layer.group_size = group_size
|
||||
layer.group_size_div_factor = group_size_div_factor
|
||||
|
||||
strategy = FusedMoeWeightScaleSupported.GROUP.value
|
||||
extra_weight_attrs.update({"quant_method": strategy, "is_transposed": False})
|
||||
|
||||
assert "weight_loader" in extra_weight_attrs
|
||||
weight_loader = extra_weight_attrs["weight_loader"]
|
||||
wrapped_weight_loader = MoeWNA16Method.get_weight_loader(layer, weight_loader)
|
||||
extra_weight_attrs["weight_loader"] = wrapped_weight_loader
|
||||
|
||||
# Fused gate_up_proj (column parallel)
|
||||
w13_qweight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // bit8_pack_factor,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_qweight", w13_qweight)
|
||||
set_weight_attrs(w13_qweight, extra_weight_attrs)
|
||||
|
||||
# down_proj (row parallel)
|
||||
w2_qweight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // bit8_pack_factor,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_qweight", w2_qweight)
|
||||
set_weight_attrs(w2_qweight, extra_weight_attrs)
|
||||
|
||||
w13_scales = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // group_size,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_scales", w13_scales)
|
||||
set_weight_attrs(w13_scales, extra_weight_attrs)
|
||||
|
||||
w2_scales = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // group_size,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_scales", w2_scales)
|
||||
set_weight_attrs(w2_scales, extra_weight_attrs)
|
||||
|
||||
if self.quant_config.has_zp:
|
||||
w13_qzeros = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition // bit8_pack_factor,
|
||||
hidden_size // group_size,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_qzeros", w13_qzeros)
|
||||
set_weight_attrs(w13_qzeros, extra_weight_attrs)
|
||||
|
||||
w2_qzeros = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
hidden_size // bit8_pack_factor,
|
||||
intermediate_size_per_partition // group_size,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_qzeros", w2_qzeros)
|
||||
set_weight_attrs(w2_qzeros, extra_weight_attrs)
|
||||
|
||||
if self.quant_config.linear_quant_method == "gptq":
|
||||
# some param are unused, but we need to init them in order to
|
||||
# load weights
|
||||
invalid_param_keys = ["w13_g_idx", "w2_g_idx"]
|
||||
if not self.quant_config.has_zp:
|
||||
invalid_param_keys += ["w13_qzeros", "w2_qzeros"]
|
||||
for key in invalid_param_keys:
|
||||
param = torch.nn.Parameter(
|
||||
torch.empty((0,), dtype=torch.int32), requires_grad=False
|
||||
)
|
||||
layer.register_parameter(key, param)
|
||||
set_weight_attrs(param, extra_weight_attrs)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: RoutedExperts
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
weight_bits = self.quant_config.weight_bits
|
||||
has_zp = self.quant_config.has_zp
|
||||
assert weight_bits == 4 or weight_bits == 8
|
||||
config_builder = (
|
||||
int4_w4a16_moe_quant_config
|
||||
if weight_bits == 4
|
||||
else int8_w8a16_moe_quant_config
|
||||
)
|
||||
|
||||
return config_builder(
|
||||
w1_scale=layer.w13_scales,
|
||||
w2_scale=layer.w2_scales,
|
||||
w1_zp=layer.w13_qzeros if has_zp else None,
|
||||
w2_zp=layer.w2_qzeros if has_zp else None,
|
||||
block_shape=[0, layer.group_size],
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
from vllm.model_executor.layers.fused_moe import fused_experts
|
||||
|
||||
return fused_experts(
|
||||
x,
|
||||
layer.w13_qweight,
|
||||
layer.w2_qweight,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=layer.activation,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
quant_config=self.moe_quant_config,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_weight_loader(layer, weight_loader):
|
||||
def convert_awq_tensor(tensor, tensor_type):
|
||||
# convert awq qweight/qzeros to a standard format (assume int4)
|
||||
# qweight: (k, n // pack_factor_bit32) -> (n, k // pack_factor_bit8)
|
||||
# qzeros: (k // group_size, n // pack_factor_bit32) ->
|
||||
# (n // pack_factor_bit8, k // group_size)
|
||||
# pack_factor_bit32 = 32 // weight_bits
|
||||
# pack_factor_bit8 = 8 // weight_bits
|
||||
|
||||
# 0. suppose origin shape (a, b), dtype int32
|
||||
# 1. convert to uint8, shape (a, b) -> (a, 4 * b)
|
||||
size0 = tensor.size(0)
|
||||
tensor = tensor.view(torch.uint8)
|
||||
|
||||
# 2. unpack to uint4 (only when weight_bits == 4)
|
||||
# shape (a, 4 * b) -> (a, 4 * b, 2)
|
||||
shifter = torch.tensor([0, 4], dtype=torch.uint8, device=tensor.device)
|
||||
tensor = (tensor[:, :, None] >> shifter) & 0xF
|
||||
|
||||
# 3. change order, see
|
||||
# https://github.com/casper-hansen/AutoAWQ/blob/v0.2.8/awq/utils/quant_utils.py
|
||||
# shape -> (a, 4 * b * pack_factor_bit8)
|
||||
reverse_awq_pack_order = [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
tensor = tensor.view(-1, 8)[:, reverse_awq_pack_order]
|
||||
tensor = tensor.view(size0, -1)
|
||||
|
||||
# 4. transpose, shape -> (4 * b * pack_factor_bit8, a)
|
||||
tensor = tensor.T.contiguous()
|
||||
|
||||
# 5. repack (only when weight_bits == 4)
|
||||
# qweight shape -> (4 * b * pack_factor_bit8, a // pack_factor_bit8)
|
||||
# qzeros shape -> (4 * b, a)
|
||||
|
||||
if tensor_type == "qweight":
|
||||
tensor = tensor[:, 1::2] * 16 + tensor[:, ::2]
|
||||
elif tensor_type == "qzeros":
|
||||
tensor = tensor[1::2, :] * 16 + tensor[::2, :]
|
||||
return tensor
|
||||
|
||||
def convert_gptq_int4_qzeros(tensor):
|
||||
tensor = tensor.view(torch.uint8)
|
||||
shifter = torch.tensor([0, 4], dtype=torch.uint8, device=tensor.device)
|
||||
tensor = (tensor[:, :, None] >> shifter) & 0xF
|
||||
tensor = tensor + 1
|
||||
tensor = tensor[:, :, 0] + tensor[:, :, 1] * 16
|
||||
return tensor
|
||||
|
||||
def moe_wna16_weight_loader(
|
||||
param: torch.nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
weight_name: str,
|
||||
shard_id: str,
|
||||
expert_id: int,
|
||||
return_success: bool = False,
|
||||
):
|
||||
if "g_idx" in weight_name:
|
||||
return False if return_success else None
|
||||
if not layer.quant_config.has_zp and "qzeros" in weight_name:
|
||||
return False if return_success else None
|
||||
|
||||
device = get_tp_group().device
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
loaded_weight = loaded_weight.to(device)
|
||||
shard_size = layer.intermediate_size_per_partition
|
||||
|
||||
# convert gptq and awq weight to a standard format
|
||||
# awq_marlin uses the same weight format as awq
|
||||
if layer.quant_config.linear_quant_method in ("awq", "awq_marlin"):
|
||||
assert layer.quant_config.weight_bits == 4
|
||||
if "weight" in weight_name:
|
||||
loaded_weight = convert_awq_tensor(loaded_weight, "qweight")
|
||||
elif "zeros" in weight_name:
|
||||
loaded_weight = convert_awq_tensor(loaded_weight, "qzeros")
|
||||
else:
|
||||
loaded_weight = loaded_weight.T
|
||||
elif layer.quant_config.linear_quant_method == "gptq":
|
||||
assert layer.quant_config.weight_bits in [4, 8]
|
||||
if "weight" in weight_name:
|
||||
loaded_weight = loaded_weight.T.contiguous().view(torch.uint8)
|
||||
elif "zeros" in weight_name:
|
||||
# add 1 to gptq qzeros to align with awq
|
||||
loaded_weight = loaded_weight.view(torch.uint8)
|
||||
if layer.quant_config.weight_bits == 4:
|
||||
loaded_weight = convert_gptq_int4_qzeros(loaded_weight).T
|
||||
else:
|
||||
loaded_weight = loaded_weight.T + 1
|
||||
else:
|
||||
loaded_weight = loaded_weight.T
|
||||
|
||||
# repeat the qzeros/scales to fit new group size
|
||||
if (
|
||||
layer.group_size_div_factor > 1
|
||||
and "qzeros" in weight_name
|
||||
or "scales" in weight_name
|
||||
):
|
||||
loaded_weight = loaded_weight.repeat_interleave(
|
||||
layer.group_size_div_factor, 1
|
||||
)
|
||||
|
||||
if "w13_qzeros" in weight_name:
|
||||
tensor = loaded_weight.view(
|
||||
layer.moe_config.tp_size, -1, loaded_weight.size(1)
|
||||
)[tp_rank]
|
||||
if shard_id == "w1":
|
||||
param.data[expert_id, : shard_size // 2] = tensor
|
||||
else:
|
||||
param.data[expert_id, shard_size // 2 :] = tensor
|
||||
return True if return_success else None
|
||||
elif "w2_qzeros" in weight_name:
|
||||
param.data[expert_id] = loaded_weight.view(
|
||||
loaded_weight.size(0), layer.moe_config.tp_size, -1
|
||||
)[:, tp_rank]
|
||||
return True if return_success else None
|
||||
else:
|
||||
# Delegate to the original loader, passing return_success
|
||||
return weight_loader(
|
||||
param,
|
||||
loaded_weight,
|
||||
weight_name,
|
||||
shard_id,
|
||||
expert_id,
|
||||
return_success=return_success,
|
||||
)
|
||||
|
||||
return moe_wna16_weight_loader
|
||||
@@ -0,0 +1,822 @@
|
||||
# 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.attention import Attention
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEMethodBase,
|
||||
FusedMoEParallelConfig,
|
||||
FusedMoEQuantConfig,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe import modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
|
||||
TRITON_BACKENDS,
|
||||
Mxfp4MoeBackend,
|
||||
convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
|
||||
convert_weight_to_mxfp4_moe_kernel_format,
|
||||
make_mxfp4_moe_kernel,
|
||||
make_mxfp4_moe_quant_config,
|
||||
mxfp4_round_up_hidden_size_and_intermediate_size,
|
||||
select_deepseek_v4_mxfp4_moe_backend,
|
||||
select_mxfp4_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
|
||||
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class Mxfp4Config(QuantizationConfig):
|
||||
"""Canonical base config for MXFP4 quantization.
|
||||
|
||||
Subclasses override get_name() and override_quantization_method() to
|
||||
register themselves as the handler for a specific checkpoint format.
|
||||
"""
|
||||
|
||||
def __init__(self, ignored_layers: list[str] | None = None):
|
||||
super().__init__()
|
||||
self.ignored_layers = ignored_layers
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config):
|
||||
return cls()
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 80
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "mxfp4"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return []
|
||||
|
||||
# TODO (zyongye) This is only temporaty fallback.
|
||||
# We should have `Mxfp4MoEMethod` after this migration is complete.
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
if isinstance(layer, LinearBase):
|
||||
if self.ignored_layers and is_layer_skipped(
|
||||
prefix=prefix,
|
||||
ignored_layers=self.ignored_layers,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedLinearMethod()
|
||||
logger.debug_once(
|
||||
"MXFP4 linear layer is not implemented - falling back to "
|
||||
"UnquantizedLinearMethod.",
|
||||
)
|
||||
return UnquantizedLinearMethod()
|
||||
elif isinstance(layer, RoutedExperts):
|
||||
return GptOssMxfp4MoEMethod(layer.moe_config)
|
||||
elif isinstance(layer, Attention):
|
||||
logger.debug_once(
|
||||
"MXFP4 attention layer is not implemented. "
|
||||
"Skipping quantization for this layer.",
|
||||
)
|
||||
return None
|
||||
|
||||
def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool:
|
||||
"""MXFP4 config always uses MXFP4 quantization."""
|
||||
return True
|
||||
|
||||
|
||||
class GptOssMxfp4Config(Mxfp4Config):
|
||||
"""MXFP4 config for GPT-OSS checkpoints.
|
||||
|
||||
Checkpoints carry ``"quant_method": "mxfp4"`` in their JSON config.
|
||||
override_quantization_method() maps that to the canonical internal name
|
||||
so that the rest of the loading path uses "gpt_oss_mxfp4" consistently.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "gpt_oss_mxfp4"
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(
|
||||
cls, hf_quant_cfg, user_quant, hf_config=None
|
||||
) -> QuantizationMethods | None:
|
||||
# Match both "mxfp4" (original checkpoint value) and "gpt_oss_mxfp4"
|
||||
# (already normalized by verify_and_update_model_config) so that
|
||||
# explicit --quantization mxfp4 from the user doesn't cause a mismatch.
|
||||
if not (
|
||||
isinstance(hf_quant_cfg, dict)
|
||||
and hf_quant_cfg.get("quant_method") in ("mxfp4", "gpt_oss_mxfp4")
|
||||
):
|
||||
return None
|
||||
# Require explicit confirmation that this is a GPT-OSS model.
|
||||
# Do NOT fall back to returning the override when hf_config is None,
|
||||
# as that would silently claim all mxfp4 checkpoints.
|
||||
model_type = getattr(hf_config, "model_type", None)
|
||||
if model_type != "gpt_oss":
|
||||
return None
|
||||
return "gpt_oss_mxfp4"
|
||||
|
||||
|
||||
class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
|
||||
"""MXFP4 MoE quantization method."""
|
||||
|
||||
def __init__(self, moe: FusedMoEConfig):
|
||||
super().__init__(moe)
|
||||
self.weight_dtype = "gpt_oss_mxfp4"
|
||||
self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
|
||||
|
||||
self.max_capture_size = moe.max_capture_size
|
||||
|
||||
self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
|
||||
self.moe_kernel: mk.FusedMoEKernel | None = None
|
||||
|
||||
# Used for triton kernel precision configs
|
||||
self.w13_precision_config = None
|
||||
self.w2_precision_config = None
|
||||
|
||||
@property
|
||||
def skip_forward_padding(self) -> bool:
|
||||
# SM100_FI_MXFP4_MXFP8_TRTLLM supports padding with mxfp8 quant
|
||||
# so can skip the padding in the forward before applying the moe method
|
||||
return self.mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8
|
||||
|
||||
# TODO(bnell): move to MK/expert_class?
|
||||
@property
|
||||
def has_unpadded_output(self) -> bool:
|
||||
return self.mxfp4_backend in [
|
||||
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
|
||||
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
|
||||
]
|
||||
|
||||
def maybe_roundup_sizes(
|
||||
self,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
act_dtype: torch.dtype,
|
||||
moe_parallel_config: FusedMoEParallelConfig,
|
||||
) -> tuple[int, int]:
|
||||
hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size_per_partition=intermediate_size_per_partition,
|
||||
act_dtype=act_dtype,
|
||||
moe_parallel_config=moe_parallel_config,
|
||||
)
|
||||
return mxfp4_round_up_hidden_size_and_intermediate_size(
|
||||
self.mxfp4_backend, hidden_size, intermediate_size_per_partition
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
self.num_experts = num_experts
|
||||
weight_dtype = torch.uint8
|
||||
scale_dtype = torch.uint8
|
||||
mxfp4_block = 32
|
||||
|
||||
layer.params_dtype = params_dtype
|
||||
layer.num_experts = num_experts
|
||||
self.intermediate_size = intermediate_size_per_partition
|
||||
self.hidden_size = hidden_size
|
||||
|
||||
# Fused gate_up_proj (column parallel)
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // 2,
|
||||
dtype=weight_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // mxfp4_block,
|
||||
dtype=scale_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
w13_weight_scale.quant_method = "block"
|
||||
|
||||
# down_proj (row parallel)
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // 2,
|
||||
dtype=weight_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // mxfp4_block,
|
||||
dtype=scale_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
w2_weight_scale.quant_method = "block"
|
||||
|
||||
if self.moe.has_bias:
|
||||
w13_bias = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
dtype=torch.bfloat16,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_bias", w13_bias)
|
||||
set_weight_attrs(w13_bias, extra_weight_attrs)
|
||||
|
||||
w2_bias = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
dtype=torch.bfloat16,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_bias", w2_bias)
|
||||
set_weight_attrs(w2_bias, extra_weight_attrs)
|
||||
|
||||
def _setup_kernel(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
w13_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
w13_bias: torch.Tensor | None = None,
|
||||
w2_bias: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
num_experts = self.num_experts
|
||||
intermediate_size = self.intermediate_size
|
||||
hidden_size = self.hidden_size
|
||||
sf_block_size = 32
|
||||
|
||||
# Shape assertions
|
||||
assert (
|
||||
w13.dim() == 3
|
||||
and w13.shape[0] == num_experts
|
||||
and w13.shape[1] == intermediate_size * 2
|
||||
and w13.shape[2] == hidden_size // 2
|
||||
)
|
||||
assert (
|
||||
w13_scale.dim() == 3
|
||||
and w13_scale.shape[0] == num_experts
|
||||
and w13_scale.shape[1] == intermediate_size * 2
|
||||
and w13_scale.shape[2] == hidden_size // sf_block_size
|
||||
)
|
||||
assert (
|
||||
w2.dim() == 3
|
||||
and w2.shape[0] == num_experts
|
||||
and w2.shape[1] == hidden_size
|
||||
and w2.shape[2] == intermediate_size // 2
|
||||
)
|
||||
assert (
|
||||
w2_scale.dim() == 3
|
||||
and w2_scale.shape[1] == hidden_size
|
||||
and w2_scale.shape[2] == intermediate_size // sf_block_size
|
||||
)
|
||||
if w13_bias is not None:
|
||||
assert (
|
||||
w13_bias.dim() == 2
|
||||
and w13_bias.shape[0] == num_experts
|
||||
and w13_bias.shape[1] == intermediate_size * 2
|
||||
)
|
||||
if w2_bias is not None:
|
||||
assert (
|
||||
w2_bias.dim() == 2
|
||||
and w2_bias.shape[0] == num_experts
|
||||
and w2_bias.shape[1] == hidden_size
|
||||
)
|
||||
|
||||
# Convert weights to kernel format
|
||||
w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (
|
||||
convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
|
||||
mxfp4_backend=self.mxfp4_backend,
|
||||
layer=layer,
|
||||
w13_weight=w13,
|
||||
w2_weight=w2,
|
||||
w13_weight_scale=w13_scale,
|
||||
w2_weight_scale=w2_scale,
|
||||
w13_bias=w13_bias,
|
||||
w2_bias=w2_bias,
|
||||
_cache_permute_indices=self._cache_permute_indices,
|
||||
)
|
||||
)
|
||||
|
||||
# For TRITON backends, weights are wrapped tensors from triton_kernels
|
||||
# that don't support .detach(). Manually assign parameters.
|
||||
if self.mxfp4_backend not in TRITON_BACKENDS:
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
replace_parameter(layer, "w13_weight_scale", w13_scale)
|
||||
replace_parameter(layer, "w2_weight_scale", w2_scale)
|
||||
else:
|
||||
layer.w13_weight = w13
|
||||
layer.w2_weight = w2
|
||||
self.w13_precision_config = w13_scale
|
||||
self.w2_precision_config = w2_scale
|
||||
|
||||
# AITER backend requires weights to be marked as shuffled.
|
||||
if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
|
||||
layer.w13_weight.is_shuffled = True
|
||||
layer.w2_weight.is_shuffled = True
|
||||
|
||||
if w13_bias is not None and w2_bias is not None:
|
||||
replace_parameter(layer, "w13_bias", w13_bias)
|
||||
replace_parameter(layer, "w2_bias", w2_bias)
|
||||
|
||||
# Build quant config
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
|
||||
# Build kernel (modular or monolithic)
|
||||
if self.moe_quant_config is not None and self.experts_cls is not None:
|
||||
self.moe_kernel = make_mxfp4_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
mxfp4_backend=self.mxfp4_backend,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
w13 = layer.w13_weight
|
||||
w2 = layer.w2_weight
|
||||
w13_scale = layer.w13_weight_scale
|
||||
w2_scale = layer.w2_weight_scale
|
||||
w13_bias = getattr(layer, "w13_bias", None)
|
||||
w2_bias = getattr(layer, "w2_bias", None)
|
||||
|
||||
if self.mxfp4_backend == Mxfp4MoeBackend.NONE:
|
||||
return
|
||||
|
||||
self._setup_kernel(layer, w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: RoutedExperts
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
w1_bias = getattr(layer, "w13_bias", None)
|
||||
w2_bias = getattr(layer, "w2_bias", None)
|
||||
|
||||
if self.mxfp4_backend in TRITON_BACKENDS:
|
||||
# TRITON backends free w13/w2_weight_scale after swizzling; the
|
||||
# swizzled scales live inside the precision configs instead.
|
||||
assert self.w13_precision_config is not None
|
||||
assert self.w2_precision_config is not None
|
||||
w1_scale = self.w13_precision_config
|
||||
w2_scale = self.w2_precision_config
|
||||
else:
|
||||
w1_scale = layer.w13_weight_scale
|
||||
w2_scale = layer.w2_weight_scale
|
||||
|
||||
return make_mxfp4_moe_quant_config(
|
||||
mxfp4_backend=self.mxfp4_backend,
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
w1_bias=w1_bias,
|
||||
w2_bias=w2_bias,
|
||||
gemm1_alpha=1.702,
|
||||
gemm1_beta=1.0,
|
||||
swiglu_limit=7.0,
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def select_gemm_impl(
|
||||
self,
|
||||
prepare_finalize: mk.FusedMoEPrepareAndFinalize,
|
||||
layer: RoutedExperts,
|
||||
) -> mk.FusedMoEExpertsModular:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} uses the new modular kernel "
|
||||
"initialization logic. This function should not be called."
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
hidden_states=x,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
expert_map=layer.expert_map,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
hidden_states=x,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
router_logits=router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
)
|
||||
|
||||
|
||||
class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
"""MXFP4 MoE quantization method."""
|
||||
|
||||
def __init__(self, moe: FusedMoEConfig):
|
||||
super().__init__(moe)
|
||||
self.weight_dtype = "mxfp4"
|
||||
self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe)
|
||||
|
||||
self.max_capture_size = moe.max_capture_size
|
||||
|
||||
self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
|
||||
self.moe_kernel: mk.FusedMoEKernel | None = None
|
||||
|
||||
# Used for triton kernel precision configs
|
||||
self.w13_precision_config = None
|
||||
self.w2_precision_config = None
|
||||
|
||||
@property
|
||||
def supports_eplb(self) -> bool:
|
||||
return True
|
||||
|
||||
@property
|
||||
def skip_forward_padding(self) -> bool:
|
||||
# SM100_FI_MXFP4_MXFP8_TRTLLM supports padding with mxfp8 quant
|
||||
# so can skip the padding in the forward before applying the moe method
|
||||
return self.mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8
|
||||
|
||||
# TODO(bnell): move to MK/expert_class?
|
||||
@property
|
||||
def has_unpadded_output(self) -> bool:
|
||||
return self.mxfp4_backend in [
|
||||
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
|
||||
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
|
||||
]
|
||||
|
||||
def maybe_roundup_sizes(
|
||||
self,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
act_dtype: torch.dtype,
|
||||
moe_parallel_config: FusedMoEParallelConfig,
|
||||
) -> tuple[int, int]:
|
||||
hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size_per_partition=intermediate_size_per_partition,
|
||||
act_dtype=act_dtype,
|
||||
moe_parallel_config=moe_parallel_config,
|
||||
)
|
||||
return mxfp4_round_up_hidden_size_and_intermediate_size(
|
||||
self.mxfp4_backend, hidden_size, intermediate_size_per_partition
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
self.num_experts = num_experts
|
||||
weight_dtype = torch.uint8
|
||||
scale_dtype = torch.uint8
|
||||
mxfp4_block = 32
|
||||
|
||||
layer.params_dtype = params_dtype
|
||||
layer.num_experts = num_experts
|
||||
self.intermediate_size = intermediate_size_per_partition
|
||||
self.hidden_size = hidden_size
|
||||
|
||||
# Fused gate_up_proj (column parallel)
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // 2,
|
||||
dtype=weight_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // mxfp4_block,
|
||||
dtype=scale_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
w13_weight_scale.quant_method = "block"
|
||||
|
||||
# down_proj (row parallel)
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // 2,
|
||||
dtype=weight_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // mxfp4_block,
|
||||
dtype=scale_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
w2_weight_scale.quant_method = "block"
|
||||
|
||||
if self.moe.has_bias:
|
||||
w13_bias = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
dtype=torch.bfloat16,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_bias", w13_bias)
|
||||
set_weight_attrs(w13_bias, extra_weight_attrs)
|
||||
|
||||
w2_bias = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
dtype=torch.bfloat16,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_bias", w2_bias)
|
||||
set_weight_attrs(w2_bias, extra_weight_attrs)
|
||||
|
||||
def _setup_kernel(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
w13_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
w13_bias: torch.Tensor | None = None,
|
||||
w2_bias: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
num_experts = self.num_experts
|
||||
intermediate_size = self.intermediate_size
|
||||
hidden_size = self.hidden_size
|
||||
sf_block_size = 32
|
||||
|
||||
# Shape assertions
|
||||
assert (
|
||||
w13.dim() == 3
|
||||
and w13.shape[0] == num_experts
|
||||
and w13.shape[1] == intermediate_size * 2
|
||||
and w13.shape[2] == hidden_size // 2
|
||||
)
|
||||
assert (
|
||||
w13_scale.dim() == 3
|
||||
and w13_scale.shape[0] == num_experts
|
||||
and w13_scale.shape[1] == intermediate_size * 2
|
||||
and w13_scale.shape[2] == hidden_size // sf_block_size
|
||||
)
|
||||
assert (
|
||||
w2.dim() == 3
|
||||
and w2.shape[0] == num_experts
|
||||
and w2.shape[1] == hidden_size
|
||||
and w2.shape[2] == intermediate_size // 2
|
||||
)
|
||||
assert (
|
||||
w2_scale.dim() == 3
|
||||
and w2_scale.shape[1] == hidden_size
|
||||
and w2_scale.shape[2] == intermediate_size // sf_block_size
|
||||
)
|
||||
if w13_bias is not None:
|
||||
assert (
|
||||
w13_bias.dim() == 2
|
||||
and w13_bias.shape[0] == num_experts
|
||||
and w13_bias.shape[1] == intermediate_size * 2
|
||||
)
|
||||
if w2_bias is not None:
|
||||
assert (
|
||||
w2_bias.dim() == 2
|
||||
and w2_bias.shape[0] == num_experts
|
||||
and w2_bias.shape[1] == hidden_size
|
||||
)
|
||||
|
||||
# Convert weights to kernel format
|
||||
w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (
|
||||
convert_weight_to_mxfp4_moe_kernel_format(
|
||||
mxfp4_backend=self.mxfp4_backend,
|
||||
layer=layer,
|
||||
w13_weight=w13,
|
||||
w2_weight=w2,
|
||||
w13_weight_scale=w13_scale,
|
||||
w2_weight_scale=w2_scale,
|
||||
w13_bias=w13_bias,
|
||||
w2_bias=w2_bias,
|
||||
_cache_permute_indices=self._cache_permute_indices,
|
||||
)
|
||||
)
|
||||
|
||||
# For TRITON backends, weights are wrapped tensors from triton_kernels
|
||||
# that don't support .detach(). Manually assign parameters.
|
||||
if self.mxfp4_backend not in TRITON_BACKENDS:
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
replace_parameter(layer, "w13_weight_scale", w13_scale)
|
||||
replace_parameter(layer, "w2_weight_scale", w2_scale)
|
||||
else:
|
||||
layer.w13_weight = w13
|
||||
layer.w2_weight = w2
|
||||
self.w13_precision_config = w13_scale
|
||||
self.w2_precision_config = w2_scale
|
||||
|
||||
# AITER backend requires weights to be marked as shuffled.
|
||||
if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
|
||||
layer.w13_weight.is_shuffled = True
|
||||
layer.w2_weight.is_shuffled = True
|
||||
|
||||
if w13_bias is not None and w2_bias is not None:
|
||||
replace_parameter(layer, "w13_bias", w13_bias)
|
||||
replace_parameter(layer, "w2_bias", w2_bias)
|
||||
|
||||
# Build quant config
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
|
||||
# Build kernel (modular or monolithic)
|
||||
if self.moe_quant_config is not None and self.experts_cls is not None:
|
||||
self.moe_kernel = make_mxfp4_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
mxfp4_backend=self.mxfp4_backend,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer):
|
||||
w13 = layer.w13_weight
|
||||
w2 = layer.w2_weight
|
||||
w13_scale = layer.w13_weight_scale
|
||||
w2_scale = layer.w2_weight_scale
|
||||
w13_bias = getattr(layer, "w13_bias", None)
|
||||
w2_bias = getattr(layer, "w2_bias", None)
|
||||
|
||||
if self.mxfp4_backend == Mxfp4MoeBackend.NONE:
|
||||
return
|
||||
|
||||
self._setup_kernel(layer, w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
w1_bias = getattr(layer, "w13_bias", None)
|
||||
w2_bias = getattr(layer, "w2_bias", None)
|
||||
swiglu_limit = getattr(layer, "swiglu_limit", None)
|
||||
|
||||
if self.mxfp4_backend in TRITON_BACKENDS:
|
||||
# TRITON backends free w13/w2_weight_scale after swizzling; the
|
||||
# swizzled scales live inside the precision configs instead.
|
||||
assert self.w13_precision_config is not None
|
||||
assert self.w2_precision_config is not None
|
||||
w1_scale = self.w13_precision_config
|
||||
w2_scale = self.w2_precision_config
|
||||
else:
|
||||
w1_scale = layer.w13_weight_scale
|
||||
w2_scale = layer.w2_weight_scale
|
||||
|
||||
return make_mxfp4_moe_quant_config(
|
||||
mxfp4_backend=self.mxfp4_backend,
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
w1_bias=w1_bias,
|
||||
w2_bias=w2_bias,
|
||||
swiglu_limit=swiglu_limit,
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def select_gemm_impl(
|
||||
self,
|
||||
prepare_finalize: mk.FusedMoEPrepareAndFinalize,
|
||||
layer: RoutedExperts,
|
||||
) -> mk.FusedMoEExpertsModular:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} uses the new modular kernel "
|
||||
"initialization logic. This function should not be called."
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
hidden_states=x,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
expert_map=layer.expert_map,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
hidden_states=x,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
router_logits=router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
)
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,170 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
RoutedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
|
||||
UnquantizedFusedMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import (
|
||||
LinearBase,
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
|
||||
should_ignore_layer,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.online.fp8 import (
|
||||
Fp8PerBlockOnlineLinearMethod,
|
||||
Fp8PerBlockOnlineMoEMethod,
|
||||
Fp8PerTensorOnlineLinearMethod,
|
||||
Fp8PerTensorOnlineMoEMethod,
|
||||
Fp8PtpcOnlineLinearMethod,
|
||||
Fp8PtpcOnlineMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.online.int8 import (
|
||||
Int8OnlineMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.online.mxfp8 import (
|
||||
Mxfp8OnlineLinearMethod,
|
||||
Mxfp8OnlineMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
QuantKey,
|
||||
kFp8Static128BlockSym,
|
||||
kFp8StaticChannelSym,
|
||||
kFp8StaticTensorSym,
|
||||
kInt8StaticChannelSym,
|
||||
kMxfp8Dynamic,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
# Online dispatch tables, keyed by the QuantSpec.weight QuantKey. The
|
||||
# corresponding method class handles the activation choice via its
|
||||
# `supported_activation_quant` set.
|
||||
_ONLINE_LINEAR_METHODS: dict[QuantKey, type] = {
|
||||
kFp8StaticTensorSym: Fp8PerTensorOnlineLinearMethod,
|
||||
kFp8Static128BlockSym: Fp8PerBlockOnlineLinearMethod,
|
||||
kFp8StaticChannelSym: Fp8PtpcOnlineLinearMethod,
|
||||
kMxfp8Dynamic: Mxfp8OnlineLinearMethod,
|
||||
}
|
||||
|
||||
_ONLINE_MOE_METHODS: dict[QuantKey, type] = {
|
||||
kFp8StaticTensorSym: Fp8PerTensorOnlineMoEMethod,
|
||||
kFp8Static128BlockSym: Fp8PerBlockOnlineMoEMethod,
|
||||
kFp8StaticChannelSym: Fp8PtpcOnlineMoEMethod,
|
||||
kMxfp8Dynamic: Mxfp8OnlineMoEMethod,
|
||||
kInt8StaticChannelSym: Int8OnlineMoEMethod,
|
||||
}
|
||||
|
||||
|
||||
class OnlineQuantizationConfig(QuantizationConfig):
|
||||
"""Model-level config for online quantization (quantize fp16/bf16 weights
|
||||
during model loading, without requiring a pre-quantized checkpoint)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
args: QuantizationConfigArgs,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if args.linear is None and args.moe is None:
|
||||
raise ValueError(
|
||||
"OnlineQuantizationConfig requires at least one of "
|
||||
"quantization_config.linear or quantization_config.moe "
|
||||
"to be set."
|
||||
)
|
||||
self.args = args
|
||||
self.ignored_layers: list[str] = args.ignore
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "online"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.bfloat16, torch.half]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
# Note: as more online quant schemes will be added, this
|
||||
# value will become the minimum across all supported schemes.
|
||||
return 75
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "OnlineQuantizationConfig":
|
||||
raise NotImplementedError(
|
||||
"OnlineQuantizationConfig does not support loading from a "
|
||||
"checkpoint config. Use quantization_config or "
|
||||
"quantization='fp8_per_tensor'/'fp8_per_block' instead."
|
||||
)
|
||||
|
||||
def _dispatch(
|
||||
self,
|
||||
spec: QuantSpec | None,
|
||||
table: dict[QuantKey, type],
|
||||
layer: torch.nn.Module,
|
||||
) -> "QuantizeMethodBase | None":
|
||||
if spec is None or spec.weight is None:
|
||||
return None
|
||||
cls = table.get(spec.weight)
|
||||
if cls is None:
|
||||
raise ValueError(
|
||||
f"online quantization for {type(layer).__name__} with "
|
||||
f"weight={spec.weight} is not supported; supported weight "
|
||||
f"keys: {sorted(str(k) for k in table)}"
|
||||
)
|
||||
# Online method classes pick their own activation format internally.
|
||||
# Per-class activation overrides are not yet wired through; reject
|
||||
# explicit overrides until the relevant method class opts in.
|
||||
if spec.activation is not None:
|
||||
raise ValueError(
|
||||
f"activation override (activation={spec.activation}) is not "
|
||||
f"yet supported for online {cls.__name__}"
|
||||
)
|
||||
if isinstance(layer, RoutedExperts):
|
||||
return cls(layer=layer)
|
||||
return cls()
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
if isinstance(layer, LinearBase):
|
||||
if should_ignore_layer(
|
||||
prefix,
|
||||
ignore=self.ignored_layers,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedLinearMethod()
|
||||
method = self._dispatch(self.args.linear, _ONLINE_LINEAR_METHODS, layer)
|
||||
return method if method is not None else UnquantizedLinearMethod()
|
||||
elif isinstance(layer, RoutedExperts):
|
||||
if should_ignore_layer(
|
||||
prefix,
|
||||
ignore=self.ignored_layers,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
method = self._dispatch(self.args.moe, _ONLINE_MOE_METHODS, layer)
|
||||
return (
|
||||
method
|
||||
if method is not None
|
||||
else UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
)
|
||||
return None
|
||||
@@ -0,0 +1,760 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch.nn import Module
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
|
||||
from vllm.model_executor.kernels.linear.scaled_mm import (
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
MarlinFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe import RoutedExperts
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
select_fp8_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import (
|
||||
LinearMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.online.moe_base import (
|
||||
OnlineMoEMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
create_fp8_quant_key,
|
||||
kFp8Dynamic128Sym,
|
||||
kFp8DynamicTensorSym,
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8Static128BlockSym,
|
||||
kFp8StaticChannelSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
cutlass_fp8_supported,
|
||||
)
|
||||
from vllm.model_executor.model_loader.reload.layerwise import (
|
||||
initialize_online_processing,
|
||||
)
|
||||
from vllm.model_executor.parameter import ModelWeightParameter
|
||||
from vllm.model_executor.utils import replace_parameter
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.deep_gemm import per_block_cast_to_fp8
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Online FP8 Linear Methods
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _Fp8OnlineLinearBase(LinearMethodBase):
|
||||
"""Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint
|
||||
weights onto meta device and materializes them just-in-time."""
|
||||
|
||||
uses_meta_device: bool = True
|
||||
|
||||
def __init__(self):
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
self.input_dtype = get_current_vllm_config().model_config.dtype
|
||||
|
||||
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,
|
||||
):
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
layer.logical_widths = output_partition_sizes
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition,
|
||||
device="meta", # materialized and processed during loading
|
||||
dtype=params_dtype,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
initialize_online_processing(layer)
|
||||
|
||||
|
||||
class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
|
||||
"""Online tensorwise FP8 linear quantization.
|
||||
Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
self.block_quant = False
|
||||
self.use_deep_gemm = False
|
||||
self.use_marlin = False
|
||||
self.marlin_input_dtype = None
|
||||
self.weight_quant_key = kFp8StaticTensorSym
|
||||
# Use per-token quantization for better perf if dynamic and cutlass
|
||||
if cutlass_fp8_supported():
|
||||
self.activation_quant_key = kFp8DynamicTokenSym
|
||||
else:
|
||||
self.activation_quant_key = kFp8DynamicTensorSym
|
||||
|
||||
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,
|
||||
):
|
||||
super().create_weights(
|
||||
layer,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
input_size,
|
||||
output_size,
|
||||
params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
|
||||
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__,
|
||||
)
|
||||
self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
layer.input_scale = None
|
||||
qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
|
||||
|
||||
# Update layer with new values.
|
||||
replace_parameter(layer, "weight", qweight.t().data)
|
||||
replace_parameter(layer, "weight_scale", weight_scale.data)
|
||||
|
||||
if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"):
|
||||
self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
|
||||
# Prevent duplicate processing (e.g., during weight reload)
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
# if batch invariant mode is enabled, use BF16 dequant
|
||||
if envs.VLLM_BATCH_INVARIANT:
|
||||
if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel):
|
||||
return self.fp8_linear.apply_weights(layer, x, bias)
|
||||
|
||||
weight_fp8 = layer.weight.to(torch.bfloat16)
|
||||
weight_scale = layer.weight_scale.to(torch.bfloat16)
|
||||
if weight_scale.numel() == 1:
|
||||
# Per-tensor: simple scalar multiplication
|
||||
weight_bf16 = weight_fp8 * weight_scale
|
||||
else:
|
||||
# Multiple scales (fused modules like QKV)
|
||||
if (
|
||||
weight_scale.dim() == 1
|
||||
and weight_scale.shape[0] == weight_fp8.shape[0]
|
||||
):
|
||||
# Per-row scaling
|
||||
weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1)
|
||||
else:
|
||||
# Fallback
|
||||
weight_bf16 = weight_fp8 * weight_scale
|
||||
return torch.nn.functional.linear(x, weight_bf16.t(), bias)
|
||||
|
||||
return self.fp8_linear.apply_weights(layer, x, bias)
|
||||
|
||||
|
||||
class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
|
||||
"""Online blockwise FP8 linear quantization.
|
||||
Loads fp16/bf16 weights and quantizes them per-block during loading."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.weight_block_size = [128, 128]
|
||||
self.activation_quant_key = create_fp8_quant_key(
|
||||
static=False,
|
||||
group_shape=GroupShape(1, self.weight_block_size[0]),
|
||||
)
|
||||
self.weight_quant_key = create_fp8_quant_key(
|
||||
static=True, group_shape=GroupShape(*self.weight_block_size)
|
||||
)
|
||||
|
||||
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,
|
||||
):
|
||||
super().create_weights(
|
||||
layer,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
input_size,
|
||||
output_size,
|
||||
params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
layer.weight_block_size = self.weight_block_size
|
||||
|
||||
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 process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
layer.input_scale = None
|
||||
block_size = self.weight_block_size
|
||||
|
||||
qweight, weight_scale_inv = per_block_cast_to_fp8(
|
||||
layer.weight, block_size=block_size, use_ue8m0=False
|
||||
)
|
||||
|
||||
replace_parameter(layer, "weight", qweight.data)
|
||||
replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
|
||||
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
|
||||
# Prevent duplicate processing (e.g., during weight reload)
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.weight_block_size is not None
|
||||
|
||||
# Note: batch invariance already handled in the function below
|
||||
return self.fp8_linear.apply_weights(
|
||||
layer,
|
||||
x,
|
||||
bias,
|
||||
)
|
||||
|
||||
|
||||
class Fp8PtpcOnlineLinearMethod(_Fp8OnlineLinearBase):
|
||||
"""Online PTPC FP8 linear quantization.
|
||||
|
||||
Per-output-channel weight scale + dynamic per-token activation scale. The
|
||||
layout matches the llmcompressor's FP8_DYNAMIC recipe, so accuracy
|
||||
is comparable but no pre-quantized checkpoint is required.
|
||||
"""
|
||||
|
||||
weight_quant_key = kFp8StaticChannelSym
|
||||
activation_quant_key = kFp8DynamicTokenSym
|
||||
|
||||
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,
|
||||
):
|
||||
super().create_weights(
|
||||
layer,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
input_size,
|
||||
output_size,
|
||||
params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
|
||||
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__,
|
||||
)
|
||||
# PTPC requires per-token activation FP8; MarlinFP8 is W8A16 and
|
||||
# would silently produce a weight-only fp8 model.
|
||||
if isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel):
|
||||
raise ValueError(
|
||||
"FP8 PTPC online quant requires a kernel that honors "
|
||||
"per-token activation quantization; MarlinFP8 is W8A16 "
|
||||
"weight-only. Requires SM89+ for Cutlass FP8 or ROCm MI3xx "
|
||||
"for rowwise scaled_mm."
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
layer.input_scale = None
|
||||
qweight, weight_scale = ops.scaled_fp8_quant(
|
||||
layer.weight, scale=None, use_per_token_if_dynamic=True
|
||||
)
|
||||
|
||||
replace_parameter(layer, "weight", qweight.t())
|
||||
replace_parameter(layer, "weight_scale", weight_scale)
|
||||
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
# if batch invariant mode is enabled dequant
|
||||
if envs.VLLM_BATCH_INVARIANT and not isinstance(
|
||||
self.fp8_linear, CutlassFP8ScaledMMLinearKernel
|
||||
):
|
||||
weight_dequant = (
|
||||
layer.weight.to(x.dtype) * layer.weight_scale.to(x.dtype).t()
|
||||
)
|
||||
return torch.nn.functional.linear(x, weight_dequant.t(), bias)
|
||||
|
||||
return self.fp8_linear.apply_weights(layer, x, bias)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Online FP8 MoE Methods
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
|
||||
"""Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint
|
||||
weights onto meta device and materializes them just-in-time."""
|
||||
|
||||
# Declared here for mypy; actual values are set in __init__.
|
||||
fp8_backend: "Fp8MoeBackend"
|
||||
experts_cls: "type[mk.FusedMoEExperts] | None"
|
||||
weight_scale_name: str
|
||||
weight_block_size: list[int] | None
|
||||
per_act_token_quant: bool = False
|
||||
per_out_ch_quant: bool = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
weight_block_size: list[int] | None,
|
||||
layer: torch.nn.Module,
|
||||
weight_key: "QuantKey | None" = None,
|
||||
activation_key: "QuantKey | None" = None,
|
||||
allow_vllm_cutlass: bool = False,
|
||||
):
|
||||
super().__init__(layer.moe_config)
|
||||
self.weight_block_size = weight_block_size
|
||||
self.block_quant: bool = self.weight_block_size is not None
|
||||
self.weight_scale_name = (
|
||||
"weight_scale_inv" if self.block_quant else "weight_scale"
|
||||
)
|
||||
|
||||
# Subclasses may pass explicit kernel keys (PTPC needs channelwise +
|
||||
# per-token).
|
||||
if weight_key is None or activation_key is None:
|
||||
if self.block_quant:
|
||||
weight_key = kFp8Static128BlockSym
|
||||
activation_key = kFp8Dynamic128Sym
|
||||
else:
|
||||
weight_key = kFp8StaticTensorSym
|
||||
activation_key = kFp8DynamicTensorSym
|
||||
|
||||
# Select Fp8 MoE backend
|
||||
self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
|
||||
config=self.moe,
|
||||
weight_key=weight_key,
|
||||
activation_key=activation_key,
|
||||
allow_vllm_cutlass=allow_vllm_cutlass,
|
||||
)
|
||||
|
||||
def _setup_kernel(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
w13_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
w13_input_scale: torch.Tensor | None,
|
||||
w2_input_scale: torch.Tensor | None,
|
||||
) -> None:
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
convert_to_fp8_moe_kernel_format,
|
||||
make_fp8_moe_kernel,
|
||||
)
|
||||
|
||||
# Shuffle weights to runtime format.
|
||||
w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
|
||||
fp8_backend=self.fp8_backend,
|
||||
layer=layer,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
w13_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
w13_input_scale=w13_input_scale,
|
||||
w2_input_scale=w2_input_scale,
|
||||
)
|
||||
|
||||
# Replace parameters with updated versions. Note that this helper
|
||||
# function ensures the replacement is compatible with RL weight reloads.
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
|
||||
replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
if self.moe_quant_config:
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_fp8_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
fp8_backend=self.fp8_backend,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: torch.nn.Module
|
||||
) -> "FusedMoEQuantConfig":
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
make_fp8_moe_quant_config,
|
||||
)
|
||||
|
||||
w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
|
||||
w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
|
||||
a1_scale = layer.w13_input_scale
|
||||
a2_scale = layer.w2_input_scale
|
||||
|
||||
return make_fp8_moe_quant_config(
|
||||
fp8_backend=self.fp8_backend,
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
a1_scale=a1_scale,
|
||||
a2_scale=a2_scale,
|
||||
w1_bias=getattr(layer, "w13_bias", None),
|
||||
w2_bias=getattr(layer, "w2_bias", None),
|
||||
block_shape=self.weight_block_size,
|
||||
per_act_token_quant=self.per_act_token_quant,
|
||||
per_out_ch_quant=self.per_out_ch_quant,
|
||||
swiglu_limit=getattr(layer, "swiglu_limit", None),
|
||||
gemm1_alpha=getattr(layer, "swiglu_alpha", None),
|
||||
gemm1_beta=getattr(layer, "swiglu_beta", None),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
|
||||
class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
|
||||
"""Online tensorwise FP8 MoE quantization.
|
||||
Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
layer: torch.nn.Module,
|
||||
):
|
||||
super().__init__(
|
||||
weight_block_size=None,
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
# TODO(@ksayers): inplace fp8 quant kernel, initialize scales with ones
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
# If checkpoint is fp16, quantize in place.
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
|
||||
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
|
||||
w13_scale = torch.ones(
|
||||
layer.num_experts, device=w13.device, dtype=torch.float32
|
||||
)
|
||||
w2_scale = torch.ones(layer.num_experts, device=w2.device, dtype=torch.float32)
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
for expert in range(layer.local_num_experts):
|
||||
w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
|
||||
layer.w13_weight[expert, :, :]
|
||||
)
|
||||
w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant(
|
||||
layer.w2_weight[expert, :, :]
|
||||
)
|
||||
|
||||
# Shuffle weights to runtime format and setup kernel.
|
||||
self._setup_kernel(
|
||||
layer,
|
||||
w13,
|
||||
w2,
|
||||
w13_scale,
|
||||
w2_scale,
|
||||
w13_input_scale=layer.w13_input_scale,
|
||||
w2_input_scale=layer.w2_input_scale,
|
||||
)
|
||||
|
||||
# Prevent duplicate processing (e.g., during weight reload)
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
|
||||
|
||||
class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
|
||||
"""Online blockwise FP8 MoE quantization.
|
||||
Loads fp16/bf16 weights and quantizes them per-block during loading."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
layer: torch.nn.Module,
|
||||
):
|
||||
super().__init__(
|
||||
weight_block_size=[128, 128],
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def maybe_roundup_sizes(
|
||||
self,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
act_dtype: torch.dtype,
|
||||
moe_parallel_config,
|
||||
) -> tuple[int, int]:
|
||||
hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size_per_partition=intermediate_size_per_partition,
|
||||
act_dtype=act_dtype,
|
||||
moe_parallel_config=moe_parallel_config,
|
||||
)
|
||||
assert self.weight_block_size is not None
|
||||
block_size = self.weight_block_size[0]
|
||||
return (
|
||||
round_up(hidden_size, block_size),
|
||||
round_up(intermediate_size_per_partition, block_size),
|
||||
)
|
||||
|
||||
def _zero_padding(self, layer: Module) -> None:
|
||||
hidden_size = layer.moe_config.hidden_dim_unpadded
|
||||
intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded
|
||||
|
||||
w13_half_size = layer.w13_weight.shape[1] // 2
|
||||
if w13_half_size > intermediate_size:
|
||||
layer.w13_weight[:, intermediate_size:w13_half_size, :] = 0
|
||||
layer.w13_weight[
|
||||
:, w13_half_size + intermediate_size : 2 * w13_half_size, :
|
||||
] = 0
|
||||
if layer.w13_weight.shape[2] > hidden_size:
|
||||
layer.w13_weight[:, :, hidden_size:] = 0
|
||||
|
||||
if layer.w2_weight.shape[1] > hidden_size:
|
||||
layer.w2_weight[:, hidden_size:, :] = 0
|
||||
if layer.w2_weight.shape[2] > intermediate_size:
|
||||
layer.w2_weight[:, :, intermediate_size:] = 0
|
||||
|
||||
if getattr(layer, "w13_bias", None) is not None:
|
||||
w13_bias_half_size = layer.w13_bias.shape[1] // 2
|
||||
if w13_bias_half_size > intermediate_size:
|
||||
layer.w13_bias[:, intermediate_size:w13_bias_half_size] = 0
|
||||
layer.w13_bias[
|
||||
:, w13_bias_half_size + intermediate_size : 2 * w13_bias_half_size
|
||||
] = 0
|
||||
|
||||
if (
|
||||
getattr(layer, "w2_bias", None) is not None
|
||||
and layer.w2_bias.shape[1] > hidden_size
|
||||
):
|
||||
layer.w2_bias[:, hidden_size:] = 0
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
self._zero_padding(layer)
|
||||
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
|
||||
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
|
||||
|
||||
block_size = self.weight_block_size
|
||||
assert block_size is not None
|
||||
block_n, block_k = block_size
|
||||
|
||||
# Create block-shaped scales (computed here rather than in
|
||||
# create_weights because online quant doesn't need them until now).
|
||||
num_experts = layer.local_num_experts
|
||||
_, w13_out, w13_in = layer.w13_weight.shape
|
||||
_, w2_out, w2_in = layer.w2_weight.shape
|
||||
|
||||
w13_scale = torch.ones(
|
||||
num_experts,
|
||||
(w13_out + block_n - 1) // block_n,
|
||||
(w13_in + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
device=w13.device,
|
||||
)
|
||||
w2_scale = torch.ones(
|
||||
num_experts,
|
||||
(w2_out + block_n - 1) // block_n,
|
||||
(w2_in + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
device=w2.device,
|
||||
)
|
||||
|
||||
for expert in range(num_experts):
|
||||
w13[expert], w13_scale[expert] = per_block_cast_to_fp8(
|
||||
layer.w13_weight[expert],
|
||||
block_size=block_size,
|
||||
use_ue8m0=False,
|
||||
)
|
||||
w2[expert], w2_scale[expert] = per_block_cast_to_fp8(
|
||||
layer.w2_weight[expert],
|
||||
block_size=block_size,
|
||||
use_ue8m0=False,
|
||||
)
|
||||
|
||||
layer.weight_block_size = block_size
|
||||
|
||||
# Shuffle weights to runtime format and setup kernel.
|
||||
self._setup_kernel(
|
||||
layer,
|
||||
w13,
|
||||
w2,
|
||||
w13_scale,
|
||||
w2_scale,
|
||||
layer.w13_input_scale,
|
||||
layer.w2_input_scale,
|
||||
)
|
||||
|
||||
# Prevent duplicate processing (e.g., during weight reload)
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
|
||||
|
||||
class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase):
|
||||
"""Online PTPC FP8 MoE quantization.
|
||||
|
||||
Quantizes each expert's weights per output channel during loading.
|
||||
Activations are quantized dynamically per token at runtime.
|
||||
"""
|
||||
|
||||
per_act_token_quant: bool = True
|
||||
per_out_ch_quant: bool = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
layer: torch.nn.Module,
|
||||
):
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
|
||||
|
||||
super().__init__(
|
||||
weight_block_size=None,
|
||||
layer=layer,
|
||||
weight_key=kFp8StaticChannelSym,
|
||||
activation_key=kFp8DynamicTokenSym,
|
||||
allow_vllm_cutlass=True,
|
||||
)
|
||||
# Reject backends whose make_fp8_moe_quant_config branch silently
|
||||
# drops per_act_token_quant / per_out_ch_quant or collapses scales:
|
||||
# MARLIN / CPU route through fp8_w8a16_moe_quant_config; FLASHINFER_*
|
||||
# fold scales into a per-tensor alpha (oracle/fp8.py).
|
||||
if self.fp8_backend in (
|
||||
Fp8MoeBackend.MARLIN,
|
||||
Fp8MoeBackend.CPU,
|
||||
Fp8MoeBackend.FLASHINFER_CUTLASS,
|
||||
Fp8MoeBackend.FLASHINFER_TRTLLM,
|
||||
):
|
||||
raise ValueError(
|
||||
f"FP8 PTPC online MoE quant is not supported with the "
|
||||
f"{self.fp8_backend.value} backend, which does not implement "
|
||||
"per-output-channel weight scales."
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
|
||||
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
|
||||
# Scale's leading dim is taken from the fp8 weight tensor by
|
||||
# construction, so it cannot drift from the weight's expert count
|
||||
# under EP / padded MoE.
|
||||
n_w13 = layer.w13_weight.shape[1]
|
||||
n_w2 = layer.w2_weight.shape[1]
|
||||
w13_scale = torch.ones(
|
||||
w13.shape[0], n_w13, 1, device=w13.device, dtype=torch.float32
|
||||
)
|
||||
w2_scale = torch.ones(
|
||||
w2.shape[0], n_w2, 1, device=w2.device, dtype=torch.float32
|
||||
)
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
for expert in range(layer.local_num_experts):
|
||||
w13[expert], w13_scale[expert] = ops.scaled_fp8_quant(
|
||||
layer.w13_weight[expert],
|
||||
scale=None,
|
||||
use_per_token_if_dynamic=True,
|
||||
)
|
||||
w2[expert], w2_scale[expert] = ops.scaled_fp8_quant(
|
||||
layer.w2_weight[expert],
|
||||
scale=None,
|
||||
use_per_token_if_dynamic=True,
|
||||
)
|
||||
|
||||
self._setup_kernel(
|
||||
layer,
|
||||
w13,
|
||||
w2,
|
||||
w13_scale,
|
||||
w2_scale,
|
||||
w13_input_scale=None,
|
||||
w2_input_scale=None,
|
||||
)
|
||||
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
@@ -0,0 +1,128 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch.nn import Module
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
|
||||
from vllm.model_executor.layers.fused_moe import RoutedExperts
|
||||
from vllm.model_executor.layers.fused_moe.oracle.int8 import (
|
||||
convert_to_int8_moe_kernel_format,
|
||||
make_int8_moe_kernel,
|
||||
make_int8_moe_quant_config,
|
||||
select_int8_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.online.moe_base import (
|
||||
OnlineMoEMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kInt8DynamicTokenSym,
|
||||
kInt8StaticChannelSym,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter
|
||||
|
||||
|
||||
class Int8OnlineMoEMethod(OnlineMoEMethodBase):
|
||||
"""Online per-channel INT8 MoE quantization.
|
||||
Loads fp16/bf16 weights and quantizes them per-row to int8 during loading.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
layer: torch.nn.Module,
|
||||
):
|
||||
super().__init__(layer.moe_config)
|
||||
self.int8_backend, self.experts_cls = select_int8_moe_backend(
|
||||
config=self.moe,
|
||||
weight_key=kInt8StaticChannelSym,
|
||||
activation_key=kInt8DynamicTokenSym,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
self._quantize_weights(layer)
|
||||
self._setup_kernel(layer)
|
||||
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
|
||||
def _quantize_weights(self, layer: Module) -> None:
|
||||
vmax = torch.iinfo(torch.int8).max
|
||||
|
||||
w13 = torch.empty_like(layer.w13_weight, dtype=torch.int8)
|
||||
w2 = torch.empty_like(layer.w2_weight, dtype=torch.int8)
|
||||
w13_scale = torch.zeros(
|
||||
layer.num_experts,
|
||||
layer.w13_weight.shape[1],
|
||||
device=w13.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
w2_scale = torch.zeros(
|
||||
layer.num_experts,
|
||||
layer.w2_weight.shape[1],
|
||||
device=w2.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
for expert in range(layer.local_num_experts):
|
||||
# w13: per-row quantization over hidden_size dim
|
||||
w = layer.w13_weight[expert, :, :]
|
||||
scales = w.abs().amax(dim=1) / vmax
|
||||
q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
|
||||
w13[expert, :, :] = q.to(torch.int8)
|
||||
w13_scale[expert, :] = scales
|
||||
|
||||
# w2: per-row quantization over intermediate_size dim
|
||||
w = layer.w2_weight[expert, :, :]
|
||||
scales = w.abs().amax(dim=1) / vmax
|
||||
q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
|
||||
w2[expert, :, :] = q.to(torch.int8)
|
||||
w2_scale[expert, :] = scales
|
||||
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
replace_parameter(layer, "w13_scale", w13_scale)
|
||||
replace_parameter(layer, "w2_scale", w2_scale)
|
||||
|
||||
def _setup_kernel(self, layer: RoutedExperts) -> None:
|
||||
w13, w2 = convert_to_int8_moe_kernel_format(
|
||||
int8_backend=self.int8_backend,
|
||||
w13=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
layer=layer,
|
||||
w13_scale=layer.w13_scale,
|
||||
)
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert self.moe_quant_config is not None
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_int8_moe_kernel(
|
||||
int8_backend=self.int8_backend,
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: torch.nn.Module
|
||||
) -> "FusedMoEQuantConfig | None":
|
||||
return make_int8_moe_quant_config(
|
||||
int8_backend=self.int8_backend,
|
||||
w1_scale=getattr(layer, "w13_scale", None),
|
||||
w2_scale=getattr(layer, "w2_scale", None),
|
||||
w1_bias=getattr(layer, "w13_bias", None),
|
||||
w2_bias=getattr(layer, "w2_bias", None),
|
||||
layer=layer,
|
||||
)
|
||||
@@ -0,0 +1,163 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from abc import abstractmethod
|
||||
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEMethodBase,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.model_loader.reload.layerwise import (
|
||||
initialize_online_processing,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
|
||||
|
||||
class OnlineMoEMethodBase(FusedMoEMethodBase):
|
||||
"""Base for MoE methods that load full-precision weights on meta device
|
||||
and quantize them after loading via the QeRL layerwise processing system.
|
||||
"""
|
||||
|
||||
uses_meta_device: bool = True
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.num_experts = num_experts
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
# Fused gate_up_proj (column parallel) — full precision on meta device
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size,
|
||||
device="meta",
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
# down_proj (row parallel) — full precision on meta device
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
device="meta",
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
# BIASES (for models like GPT-OSS that have biased MoE)
|
||||
if self.moe.has_bias:
|
||||
w13_bias = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
device="meta",
|
||||
dtype=layer.orig_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_bias", w13_bias)
|
||||
set_weight_attrs(w13_bias, extra_weight_attrs)
|
||||
|
||||
w2_bias = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
device="meta",
|
||||
dtype=layer.orig_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_bias", w2_bias)
|
||||
set_weight_attrs(w2_bias, extra_weight_attrs)
|
||||
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
initialize_online_processing(layer)
|
||||
|
||||
@abstractmethod
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
pass
|
||||
|
||||
def maybe_make_prepare_finalize(
|
||||
self,
|
||||
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> mk.FusedMoEPrepareAndFinalizeModular | None:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} uses the new modular kernel "
|
||||
"initialization logic. This function should not be called."
|
||||
)
|
||||
|
||||
@property
|
||||
def supports_eplb(self) -> bool:
|
||||
return True
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
num_expert_group=layer.num_expert_group,
|
||||
topk_group=layer.topk_group,
|
||||
e_score_correction_bias=layer.e_score_correction_bias,
|
||||
routed_scaling_factor=layer.routed_scaling_factor,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
@@ -0,0 +1,256 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Online MXFP8 (microscaling FP8, block-32) quantization methods."""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch.nn import Module
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEQuantConfig,
|
||||
RoutedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
|
||||
|
||||
from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel
|
||||
from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
|
||||
select_mxfp8_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.online.fp8 import (
|
||||
_Fp8OnlineLinearBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.online.moe_base import (
|
||||
OnlineMoEMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
|
||||
MXFP8_BLOCK_SIZE,
|
||||
mxfp8_e4m3_quantize,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
class Mxfp8OnlineLinearMethod(_Fp8OnlineLinearBase):
|
||||
"""Online MXFP8 linear method.
|
||||
Loads bf16/fp16 checkpoints and quantizes weights to MXFP8 (microscaling
|
||||
FP8 with block-32 scales) during weight loading.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.kernel = init_mxfp8_linear_kernel()
|
||||
|
||||
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,
|
||||
):
|
||||
if input_size_per_partition % MXFP8_BLOCK_SIZE != 0:
|
||||
raise ValueError(
|
||||
f"MXFP8 requires input_size_per_partition "
|
||||
f"({input_size_per_partition}) to be divisible by "
|
||||
f"{MXFP8_BLOCK_SIZE}."
|
||||
)
|
||||
|
||||
super().create_weights(
|
||||
layer,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
input_size,
|
||||
output_size,
|
||||
params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
weight_fp8, weight_scale = mxfp8_e4m3_quantize(layer.weight.contiguous())
|
||||
|
||||
layer.input_scale = None
|
||||
replace_parameter(layer, "weight", weight_fp8.data)
|
||||
replace_parameter(layer, "weight_scale", weight_scale.data)
|
||||
|
||||
self.kernel.process_weights_after_loading(layer)
|
||||
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return self.kernel.apply_weights(layer, x, bias)
|
||||
|
||||
|
||||
class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
|
||||
"""MoE method for online MXFP8 (block) quantization."""
|
||||
|
||||
fp8_backend: "Fp8MoeBackend"
|
||||
experts_cls: "type[mk.FusedMoEExperts] | None"
|
||||
|
||||
def __init__(self, *, layer: torch.nn.Module):
|
||||
super().__init__(layer.moe_config)
|
||||
self.weight_block_size: list[int] = [1, MXFP8_BLOCK_SIZE]
|
||||
self.weight_scale_name = "weight_scale"
|
||||
|
||||
self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
if (
|
||||
hidden_size % MXFP8_BLOCK_SIZE != 0
|
||||
or intermediate_size_per_partition % MXFP8_BLOCK_SIZE != 0
|
||||
):
|
||||
raise ValueError(
|
||||
"Online MXFP8 MoE requires hidden/intermediate sizes divisible "
|
||||
f"by {MXFP8_BLOCK_SIZE}."
|
||||
)
|
||||
|
||||
super().create_weights(
|
||||
layer=layer,
|
||||
num_experts=num_experts,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size_per_partition=intermediate_size_per_partition,
|
||||
params_dtype=params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
|
||||
layer.weight_block_size = [1, MXFP8_BLOCK_SIZE]
|
||||
|
||||
def _quantize_mxfp8_moe_weight(
|
||||
self, weight: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Batch quantization: bf16/fp16 weights -> MXFP8 (fp8 + uint8 scales)."""
|
||||
E = weight.size(0)
|
||||
first_q, first_s = mxfp8_e4m3_quantize(weight[0], is_sf_swizzled_layout=False)
|
||||
# Pre-allocate the output tensors rather than stacking.
|
||||
# This is important for consistent memory layout.
|
||||
w_quant = torch.empty(
|
||||
(E, *first_q.shape), dtype=first_q.dtype, device=weight.device
|
||||
)
|
||||
w_scales = torch.empty(
|
||||
(E, *first_s.shape), dtype=first_s.dtype, device=weight.device
|
||||
)
|
||||
w_quant[0] = first_q
|
||||
w_scales[0] = first_s
|
||||
for i in range(1, E):
|
||||
w_quant[i], w_scales[i] = mxfp8_e4m3_quantize(
|
||||
weight[i], is_sf_swizzled_layout=False
|
||||
)
|
||||
|
||||
return w_quant, w_scales
|
||||
|
||||
def _setup_kernel(
|
||||
self,
|
||||
layer: "RoutedExperts",
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
w13_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
w13_input_scale: torch.Tensor | None,
|
||||
w2_input_scale: torch.Tensor | None,
|
||||
) -> None:
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
convert_to_fp8_moe_kernel_format,
|
||||
make_fp8_moe_kernel,
|
||||
)
|
||||
|
||||
# Shuffle weights to runtime format.
|
||||
w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
|
||||
fp8_backend=self.fp8_backend,
|
||||
layer=layer,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
w13_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
w13_input_scale=w13_input_scale,
|
||||
w2_input_scale=w2_input_scale,
|
||||
)
|
||||
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
|
||||
replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
if self.moe_quant_config:
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_fp8_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
fp8_backend=self.fp8_backend,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: torch.nn.Module
|
||||
) -> "FusedMoEQuantConfig":
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
make_fp8_moe_quant_config,
|
||||
)
|
||||
|
||||
w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
|
||||
w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
|
||||
a1_scale = layer.w13_input_scale
|
||||
a2_scale = layer.w2_input_scale
|
||||
|
||||
return make_fp8_moe_quant_config(
|
||||
fp8_backend=self.fp8_backend,
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
a1_scale=a1_scale,
|
||||
a2_scale=a2_scale,
|
||||
w1_bias=getattr(layer, "w13_bias", None),
|
||||
w2_bias=getattr(layer, "w2_bias", None),
|
||||
block_shape=self.weight_block_size,
|
||||
swiglu_limit=getattr(layer, "swiglu_limit", None),
|
||||
gemm1_alpha=getattr(layer, "swiglu_alpha", None),
|
||||
gemm1_beta=getattr(layer, "swiglu_beta", None),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
|
||||
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
w13, w13_scale = self._quantize_mxfp8_moe_weight(layer.w13_weight)
|
||||
w2, w2_scale = self._quantize_mxfp8_moe_weight(layer.w2_weight)
|
||||
|
||||
self._setup_kernel(
|
||||
layer,
|
||||
w13,
|
||||
w2,
|
||||
w13_scale,
|
||||
w2_scale,
|
||||
layer.w13_input_scale,
|
||||
layer.w2_input_scale,
|
||||
)
|
||||
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
@@ -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)
|
||||
@@ -0,0 +1,182 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Modified by Roberto L. Castro (Roberto.LopezCastro@ist.ac.at).
|
||||
#
|
||||
# Copied from https://github.com/pytorch/ao/tree/main/torchao/prototype/mx_formats
|
||||
#
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from typing import Literal
|
||||
|
||||
import torch
|
||||
from torch.library import wrap_triton
|
||||
|
||||
from vllm.triton_utils import tl, triton
|
||||
from vllm.utils.math_utils import cdiv
|
||||
|
||||
|
||||
@triton.jit
|
||||
def triton_scale_swizzle(
|
||||
scale_ptr: torch.Tensor,
|
||||
scale_rows: int,
|
||||
scale_cols: int,
|
||||
output_ptr: torch.Tensor,
|
||||
input_row_stride: int,
|
||||
output_block_stride: int,
|
||||
BLOCK_ROWS: tl.constexpr,
|
||||
BLOCK_COLS: tl.constexpr,
|
||||
):
|
||||
"""
|
||||
Rearranges tensor data from row-major to block-scaled swizzle format.
|
||||
|
||||
Args:
|
||||
scale_ptr: Pointer to the input scale tensor
|
||||
scale_rows: Number of rows in the scale tensor
|
||||
scale_cols: Number of columns in the scale tensor
|
||||
output_ptr: Pointer to the output tensor
|
||||
input_row_stride: Stride between rows in the input tensor
|
||||
output_block_stride: Stride between blocks in the output tensor
|
||||
BLOCK_ROWS: Number of rows in a tile (compile-time constant)
|
||||
BLOCK_COLS: Number of columns in a tile (compile-time constant)
|
||||
"""
|
||||
pid_row = tl.program_id(0)
|
||||
pid_col = tl.program_id(1)
|
||||
|
||||
rows = tl.arange(0, BLOCK_ROWS)[:, None]
|
||||
cols = tl.arange(0, BLOCK_COLS)[None, :]
|
||||
|
||||
# Calculate starting row and column for this tile
|
||||
start_row = pid_row * BLOCK_ROWS
|
||||
start_col = pid_col * BLOCK_COLS
|
||||
global_rows = start_row + rows
|
||||
global_cols = start_col + cols
|
||||
|
||||
mask = (global_rows < scale_rows) & (global_cols < scale_cols)
|
||||
|
||||
input_scales = tl.load(
|
||||
scale_ptr + global_rows * input_row_stride + global_cols,
|
||||
mask=mask,
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
r_div_32 = rows // 32
|
||||
r_mod_32 = rows % 32
|
||||
|
||||
# 2) Rearrange to (32, 4, 4) then to final (32, 16) coordinates
|
||||
dest_indices = r_mod_32 * 16 + r_div_32 * 4 + cols
|
||||
|
||||
# Flatten
|
||||
dest_indices_flat = tl.reshape(dest_indices, (BLOCK_ROWS * BLOCK_COLS))
|
||||
scales_flat = tl.reshape(input_scales, (BLOCK_ROWS * BLOCK_COLS))
|
||||
|
||||
# Calculate block offset using provided output block stride
|
||||
LOCAL_NUMEL = BLOCK_ROWS * BLOCK_COLS
|
||||
block_offset = pid_col * LOCAL_NUMEL + (pid_row * output_block_stride)
|
||||
|
||||
tl.store(
|
||||
output_ptr + block_offset + dest_indices_flat,
|
||||
scales_flat,
|
||||
)
|
||||
|
||||
|
||||
def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Rearranges an E8M0 tensor scale from row-major format to
|
||||
block-scaled swizzle format.
|
||||
|
||||
This format is suitable for Tmem as described in NVIDIA documentation:
|
||||
https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
|
||||
|
||||
Args:
|
||||
scale_tensor: Input tensor in row-major format with 8-bit elements
|
||||
|
||||
Returns:
|
||||
Rearranged tensor in block-scaled swizzle format
|
||||
"""
|
||||
assert scale_tensor.element_size() == 1, (
|
||||
"Expected element size to be 1 byte (8 bits)"
|
||||
)
|
||||
assert scale_tensor.is_contiguous(), "Input tensor must be contiguous"
|
||||
|
||||
rows, cols = scale_tensor.shape
|
||||
|
||||
# Calculate blocks needed
|
||||
n_row_blocks = triton.cdiv(rows, 128)
|
||||
n_col_blocks = triton.cdiv(cols, 4)
|
||||
padded_rows = n_row_blocks * 128
|
||||
padded_cols = n_col_blocks * 4
|
||||
|
||||
out = scale_tensor.new_empty((padded_rows, padded_cols))
|
||||
|
||||
# Input stride (for row-major format)
|
||||
input_row_stride = cols
|
||||
|
||||
# We probably want handle multiple blocks per tile but
|
||||
# for now keep it simple
|
||||
BLOCK_ROWS, BLOCK_COLS = 128, 4
|
||||
|
||||
# Output block stride for the rearranged format
|
||||
output_block_stride = BLOCK_ROWS * BLOCK_COLS * (padded_cols // BLOCK_COLS)
|
||||
|
||||
grid = lambda META: (
|
||||
triton.cdiv(padded_rows, BLOCK_ROWS),
|
||||
triton.cdiv(padded_cols, BLOCK_COLS),
|
||||
)
|
||||
|
||||
wrap_triton(triton_scale_swizzle)[grid](
|
||||
scale_tensor.view(torch.uint8),
|
||||
rows,
|
||||
cols,
|
||||
out.view(torch.uint8),
|
||||
input_row_stride,
|
||||
output_block_stride,
|
||||
BLOCK_ROWS=BLOCK_ROWS,
|
||||
BLOCK_COLS=BLOCK_COLS,
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def to_blocked(
|
||||
input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = "triton"
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Rearrange a large matrix by breaking it into blocks and applying
|
||||
the rearrangement pattern.
|
||||
|
||||
See:
|
||||
https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
|
||||
|
||||
Args:
|
||||
input_matrix: Input tensor of shape (H, W)
|
||||
backend: "torch" (PyTorch path) or "triton" (Triton kernel)
|
||||
|
||||
Returns:
|
||||
Rearranged tensor of shape (32*cdiv(H,128), 16*cdiv(W,4))
|
||||
"""
|
||||
if backend == "triton":
|
||||
return triton_mx_block_rearrange(input_matrix).flatten()
|
||||
elif backend != "torch":
|
||||
raise ValueError(f'backend must be "torch" or "triton", got {backend!r}')
|
||||
|
||||
rows, cols = input_matrix.shape
|
||||
n_row_blocks = cdiv(rows, 128)
|
||||
n_col_blocks = cdiv(cols, 4)
|
||||
|
||||
# Calculate the padded shape
|
||||
padded_rows = n_row_blocks * 128
|
||||
padded_cols = n_col_blocks * 4
|
||||
|
||||
padded = input_matrix
|
||||
assert (rows, cols) == (padded_rows, padded_cols)
|
||||
|
||||
# Rearrange the blocks
|
||||
blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
|
||||
rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
|
||||
|
||||
return rearranged.flatten()
|
||||
@@ -0,0 +1,399 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import importlib
|
||||
import json
|
||||
import types
|
||||
from importlib.util import find_spec
|
||||
from typing import Any
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from packaging import version
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from vllm.logger import init_logger
|
||||
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 (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _bond_method_to_cls(func, obj):
|
||||
if hasattr(func, "__self__") or not callable(func):
|
||||
# If the function is already bound to an instance, return it as is
|
||||
return func
|
||||
else:
|
||||
return types.MethodType(func, obj)
|
||||
|
||||
|
||||
def _get_weight_attrs(param):
|
||||
# record attributes attached to the weight, so we can
|
||||
# recover later
|
||||
recorded_weight_attr = {}
|
||||
for key in param.__dict__:
|
||||
if hasattr(param, key):
|
||||
attr = getattr(param, key)
|
||||
if not callable(attr):
|
||||
recorded_weight_attr[key] = attr
|
||||
elif hasattr(attr, "__self__") and param is attr.__self__:
|
||||
# if attr is a bonded method for an instance, and
|
||||
# attr.__self__ points to the instance (param)
|
||||
# we'll record the underlying function object
|
||||
recorded_weight_attr[key] = attr.__func__
|
||||
else:
|
||||
recorded_weight_attr[key] = attr
|
||||
return recorded_weight_attr
|
||||
|
||||
|
||||
def _restore_weight_attrs(param, recorded_weight_attr):
|
||||
for attr_name, attr in recorded_weight_attr.items():
|
||||
if not hasattr(param, attr_name):
|
||||
setattr(param, attr_name, _bond_method_to_cls(attr, param))
|
||||
|
||||
|
||||
def torchao_version_at_least(torchao_version: str) -> bool:
|
||||
if find_spec("torchao"):
|
||||
try:
|
||||
if version.parse(importlib.metadata.version("torchao")) >= version.parse(
|
||||
torchao_version
|
||||
):
|
||||
return True
|
||||
except (ImportError, version.InvalidVersion):
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
def should_skip(prefix: str, skip_modules: list[str]) -> bool:
|
||||
"""
|
||||
Robust skipping logic:
|
||||
should_skip("model.model.layers.1.q_proj",
|
||||
["model.model.layers.1.q_proj"]) # True
|
||||
should_skip("model.model.layers.10.o_proj", ["o_proj"]) -> True
|
||||
should_skip("visual.model.layers.1.q_proj", ["visual"]) -> True
|
||||
should_skip("model.model.layers.1.q_proj", ["layers.1"]) -> True
|
||||
should_skip("model.model.layers.11.q_proj", ["layers.1"]) -> False
|
||||
"""
|
||||
for s in skip_modules:
|
||||
if prefix == s:
|
||||
return True
|
||||
if f".{s}." in f".{prefix}.":
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
if torchao_version_at_least("0.15.0"):
|
||||
from torchao.prototype.tensor_conversion.api import (
|
||||
convert_to_packed_tensor_based_on_current_hardware,
|
||||
)
|
||||
else:
|
||||
convert_to_packed_tensor_based_on_current_hardware = lambda t: t
|
||||
|
||||
|
||||
def _check_torchao_fp8_activation_capability(torchao_config) -> None:
|
||||
"""Check if the current GPU supports FP8 activation quantization.
|
||||
|
||||
FP8 activation configs (e.g., Float8DynamicActivationFloat8WeightConfig)
|
||||
require GPU compute capability >= 8.9 (Ada Lovelace / Hopper) on NVIDIA,
|
||||
or MI300+ on AMD. This check provides a clear error message before
|
||||
torchao's internal assertion fires with a confusing message.
|
||||
"""
|
||||
config_name = type(torchao_config).__name__
|
||||
if "Float8" not in config_name or "Activation" not in config_name:
|
||||
return
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.supports_fp8():
|
||||
return
|
||||
|
||||
capability = current_platform.get_device_capability()
|
||||
capability_str = (
|
||||
f" (current GPU compute capability: {capability.major}.{capability.minor})"
|
||||
if capability is not None
|
||||
else ""
|
||||
)
|
||||
raise ValueError(
|
||||
f"torchao FP8 activation quantization config '{config_name}' "
|
||||
f"requires GPU compute capability >= 8.9 (e.g., NVIDIA Ada Lovelace "
|
||||
f"/ Hopper or AMD MI300+){capability_str}. "
|
||||
f"For older GPUs, consider using a non-FP8 config such as "
|
||||
f"Int8WeightOnlyConfig or Int4WeightOnlyConfig."
|
||||
)
|
||||
|
||||
|
||||
class TorchAOConfig(QuantizationConfig):
|
||||
"""Config class for torchao."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
torchao_config,
|
||||
skip_modules: list[str] | None = None,
|
||||
is_checkpoint_torchao_serialized: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.torchao_config = torchao_config
|
||||
self.skip_modules = skip_modules or []
|
||||
self.is_checkpoint_torchao_serialized = is_checkpoint_torchao_serialized
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"TorchAOConfig({self.torchao_config=}, {self.skip_modules=}, "
|
||||
f"{self.is_checkpoint_torchao_serialized=})"
|
||||
)
|
||||
|
||||
def get_name(self) -> QuantizationMethods:
|
||||
return "torchao"
|
||||
|
||||
def get_supported_act_dtypes(self) -> list[torch.dtype]:
|
||||
return [torch.float32, torch.float16, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 75
|
||||
|
||||
@staticmethod
|
||||
def get_config_filenames() -> list[str]:
|
||||
"""torchao doesn't require additional config files, we use
|
||||
`config.json` from huggingface: `model_config.hf_config`
|
||||
"""
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "TorchAOConfig":
|
||||
"""Create the quant config from an hf model config"""
|
||||
try:
|
||||
from torchao.core.config import config_from_dict
|
||||
except ImportError as err:
|
||||
raise ImportError(
|
||||
"Please install torchao>=0.10.0 via "
|
||||
"`pip install torchao>=0.10.0` to use torchao quantization."
|
||||
) from err
|
||||
|
||||
quant_method = cls.get_from_keys_or(config, ["quant_method"], None)
|
||||
is_checkpoint_torchao_serialized = (
|
||||
quant_method is not None and "torchao" in quant_method
|
||||
)
|
||||
|
||||
hf_config = cls.get_from_keys_or(config, ["quant_type"], None)
|
||||
assert hf_config is not None, "quant_type must be specified"
|
||||
assert len(hf_config) == 1 and "default" in hf_config, (
|
||||
"Expected only one key 'default' in quant_type dictionary"
|
||||
)
|
||||
quant_type = hf_config["default"]
|
||||
ao_config = config_from_dict(quant_type)
|
||||
|
||||
# Adds skipped modules defined in "modules_to_not_convert"
|
||||
skip_modules = config.get("modules_to_not_convert", []) or []
|
||||
|
||||
# Adds skipped modules defined in "module_fqn_to_config"
|
||||
_data = quant_type.get("_data", {})
|
||||
if not isinstance(_data, dict):
|
||||
_data = {}
|
||||
|
||||
module_fqn = _data.get("module_fqn_to_config", {})
|
||||
if not isinstance(module_fqn, dict):
|
||||
module_fqn = {}
|
||||
|
||||
for layer, layer_cfg in module_fqn.items():
|
||||
if layer_cfg is None:
|
||||
skip_modules.append(layer)
|
||||
|
||||
return cls(ao_config, skip_modules, is_checkpoint_torchao_serialized)
|
||||
|
||||
@classmethod
|
||||
def from_config_file(cls, config_file: str) -> "TorchAOConfig":
|
||||
"""Initialize class from a config file. Example:
|
||||
```
|
||||
config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
|
||||
fn = "torchao_config.json"
|
||||
|
||||
with open(fn, "w") as f:
|
||||
f.write(json.dumps(config_to_dict(config)))
|
||||
```
|
||||
"""
|
||||
with open(config_file) as f:
|
||||
f.seek(0)
|
||||
f_read = f.read()
|
||||
config_dict = json.loads(f_read)
|
||||
|
||||
hf_config = {"quant_type": {"default": config_dict}}
|
||||
return cls.from_config(hf_config)
|
||||
|
||||
@classmethod
|
||||
def from_config_dict_json(cls, config_dict_json: str) -> "TorchAOConfig":
|
||||
"""Initialize class from a config_dict json string, got from
|
||||
torchao_config_object = some AOBaseConfig object
|
||||
json.dumps(config_to_dict(torchao_config_object))
|
||||
"""
|
||||
config_dict = json.loads(config_dict_json)
|
||||
hf_config = {"quant_type": {"default": config_dict}}
|
||||
return cls.from_config(hf_config)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
if not isinstance(layer, LinearBase):
|
||||
return None
|
||||
|
||||
from torchao.quantization import ModuleFqnToConfig
|
||||
|
||||
if should_skip(prefix, self.skip_modules):
|
||||
return UnquantizedLinearMethod()
|
||||
|
||||
module_fqn = prefix
|
||||
if isinstance(self.torchao_config, ModuleFqnToConfig):
|
||||
module_fqn_to_config = self.torchao_config.module_fqn_to_config
|
||||
c = None
|
||||
if module_fqn in module_fqn_to_config:
|
||||
assert not module_fqn.startswith("re:"), (
|
||||
"module fqn should not start with"
|
||||
"`re:`, which is used for specifying regex"
|
||||
)
|
||||
c = module_fqn_to_config[module_fqn]
|
||||
else:
|
||||
for maybe_module_fqn_pattern in module_fqn_to_config:
|
||||
if not maybe_module_fqn_pattern.startswith("re:"):
|
||||
continue
|
||||
elif re.fullmatch(maybe_module_fqn_pattern[3:], module_fqn):
|
||||
# we'll apply the config for first fully matched pattern
|
||||
c = module_fqn_to_config[maybe_module_fqn_pattern]
|
||||
break
|
||||
else:
|
||||
# fallback to use default if no module specific
|
||||
# config is provided
|
||||
c = module_fqn_to_config.get("_default", None)
|
||||
|
||||
if c is not None:
|
||||
current_torchao_config = TorchAOConfig(
|
||||
c, self.skip_modules, self.is_checkpoint_torchao_serialized
|
||||
)
|
||||
return TorchAOLinearMethod(current_torchao_config)
|
||||
else:
|
||||
return UnquantizedLinearMethod()
|
||||
|
||||
return TorchAOLinearMethod(self)
|
||||
|
||||
def get_scaled_act_names(self) -> list[str]:
|
||||
return []
|
||||
|
||||
|
||||
def torchao_quantize_param_data(
|
||||
param: torch.Tensor, torchao_config: Any
|
||||
) -> torch.nn.Parameter:
|
||||
"""Quantize a Tensor with torchao quantization specified by torchao_config
|
||||
|
||||
Args:
|
||||
param: weight parameter of the linear module
|
||||
torchao_config: type of quantization and their arguments we want to
|
||||
use to quantize the Tensor
|
||||
"""
|
||||
from torchao.core.config import AOBaseConfig
|
||||
from torchao.quantization import quantize_
|
||||
|
||||
assert isinstance(torchao_config, AOBaseConfig), f"{torchao_config}"
|
||||
_check_torchao_fp8_activation_capability(torchao_config)
|
||||
"""
|
||||
Avoid real weight allocation for faster load, since we will
|
||||
end up setting it to param.
|
||||
"""
|
||||
with torch.device("meta"):
|
||||
# linear can't be top level module since quantize_ is inplace
|
||||
# while some of our configs need to do module swap, and only non-top
|
||||
# level modules support module swap
|
||||
dummy_linear = torch.nn.Sequential(
|
||||
torch.nn.Linear(param.shape[1], param.shape[0], bias=False)
|
||||
)
|
||||
|
||||
dummy_linear[0].weight = param
|
||||
quantize_(dummy_linear, torchao_config)
|
||||
return dummy_linear[0].weight
|
||||
|
||||
|
||||
class TorchAOLinearMethod(LinearMethodBase):
|
||||
"""Linear method for torchao.
|
||||
|
||||
Args:
|
||||
quant_config: The torchao quantization config, a string that encodes
|
||||
the type of quantization and all relevant arguments.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: TorchAOConfig):
|
||||
self.quant_config = quant_config
|
||||
|
||||
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,
|
||||
):
|
||||
weight = Parameter(
|
||||
torch.empty(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
if self.quant_config.is_checkpoint_torchao_serialized:
|
||||
weight = torchao_quantize_param_data(
|
||||
weight, self.quant_config.torchao_config
|
||||
)
|
||||
|
||||
set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
|
||||
|
||||
layer.register_parameter("weight", weight)
|
||||
set_weight_attrs(weight, extra_weight_attrs)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return F.linear(x, layer.weight, bias)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
if self.quant_config.is_checkpoint_torchao_serialized:
|
||||
if not hasattr(layer, "weight"):
|
||||
return
|
||||
|
||||
# record attributes attached to the weight, so we can
|
||||
# recover later
|
||||
recorded_weight_attr = _get_weight_attrs(layer.weight)
|
||||
|
||||
layer.weight = Parameter(
|
||||
convert_to_packed_tensor_based_on_current_hardware(layer.weight),
|
||||
requires_grad=layer.weight.requires_grad,
|
||||
)
|
||||
|
||||
_restore_weight_attrs(layer.weight, recorded_weight_attr)
|
||||
return
|
||||
|
||||
# online quantize the weight if the checkpoint is not already
|
||||
# quantized by torchao
|
||||
recorded_weight_attr = _get_weight_attrs(layer.weight)
|
||||
|
||||
weight = torchao_quantize_param_data(
|
||||
layer.weight, self.quant_config.torchao_config
|
||||
)
|
||||
weight = torch.nn.Parameter(
|
||||
convert_to_packed_tensor_based_on_current_hardware(weight),
|
||||
weight.requires_grad,
|
||||
)
|
||||
|
||||
_restore_weight_attrs(weight, recorded_weight_attr)
|
||||
layer.register_parameter("weight", weight)
|
||||
@@ -0,0 +1,26 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""TurboQuant: KV-cache quantization for vLLM.
|
||||
|
||||
Hadamard rotation + per-coordinate Lloyd-Max scalar quantization for
|
||||
keys, uniform quantization for values.
|
||||
|
||||
The core algorithmic pattern implemented for key quantization (Hadamard
|
||||
rotation followed by deterministic scalar quantization and
|
||||
re-normalization) was originally established in DRIVE (Vargaftik et al.,
|
||||
NeurIPS 2021) and EDEN (Vargaftik et al., ICML 2022). This formulation is
|
||||
also mathematically equivalent to the scalar case of the HIGGS
|
||||
quantization method (Malinovskii et al., "Pushing the Limits of Large
|
||||
Language Model Quantization via the Linearity Theorem", NAACL 2025;
|
||||
preprint arXiv:2411.17525), which subsequently generalized these concepts.
|
||||
|
||||
A first application of this approach to KV-cache compression is in "Cache
|
||||
Me If You Must: Adaptive Key-Value Quantization for Large Language Models"
|
||||
(Shutova et al., ICML 2025; preprint arXiv:2501.19392). All of these
|
||||
foundational and application references pre-date the TurboQuant paper
|
||||
(Zandieh et al., ICLR 2026).
|
||||
"""
|
||||
|
||||
from vllm.model_executor.layers.quantization.turboquant.config import TurboQuantConfig
|
||||
|
||||
__all__ = ["TurboQuantConfig"]
|
||||
@@ -0,0 +1,86 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Lloyd-Max optimal scalar quantizer for TurboQuant.
|
||||
|
||||
After rotating a d-dimensional unit vector by a random orthogonal matrix,
|
||||
each coordinate approximately follows N(0, 1/d) for d >= 64.
|
||||
We solve the Lloyd-Max conditions to find optimal centroids.
|
||||
|
||||
Based on: turboquant-pytorch/lloyd_max.py (Zandieh et al.)
|
||||
"""
|
||||
|
||||
import math
|
||||
from functools import lru_cache
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _gaussian_pdf(x: float, sigma2: float) -> float:
|
||||
return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp(-x * x / (2 * sigma2))
|
||||
|
||||
|
||||
def _trapz(f, a: float, b: float, n: int = 200) -> float:
|
||||
"""Trapezoidal numerical integration (replaces scipy.integrate.quad)."""
|
||||
h = (b - a) / n
|
||||
result = 0.5 * (f(a) + f(b))
|
||||
for i in range(1, n):
|
||||
result += f(a + i * h)
|
||||
return result * h
|
||||
|
||||
|
||||
def solve_lloyd_max(
|
||||
d: int,
|
||||
bits: int,
|
||||
max_iter: int = 200,
|
||||
tol: float = 1e-10,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Solve Lloyd-Max optimal quantizer for N(0, 1/d) distribution.
|
||||
|
||||
Args:
|
||||
d: Vector dimension (determines variance = 1/d).
|
||||
bits: Number of quantization bits.
|
||||
max_iter: Maximum Lloyd-Max iterations.
|
||||
tol: Convergence tolerance.
|
||||
|
||||
Returns:
|
||||
centroids: Sorted tensor of 2^bits optimal centroids.
|
||||
boundaries: Sorted tensor of 2^bits - 1 decision boundaries.
|
||||
"""
|
||||
n_levels = 2**bits
|
||||
sigma2 = 1.0 / d
|
||||
sigma = math.sqrt(sigma2)
|
||||
|
||||
def pdf(x):
|
||||
return _gaussian_pdf(x, sigma2)
|
||||
|
||||
lo, hi = -3.5 * sigma, 3.5 * sigma
|
||||
centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
|
||||
|
||||
for _ in range(max_iter):
|
||||
boundaries = [
|
||||
(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)
|
||||
]
|
||||
edges = [lo * 3] + boundaries + [hi * 3]
|
||||
new_centroids = []
|
||||
for i in range(n_levels):
|
||||
a, b = edges[i], edges[i + 1]
|
||||
num = _trapz(lambda x: x * pdf(x), a, b)
|
||||
den = _trapz(pdf, a, b)
|
||||
new_centroids.append(num / den if den > 1e-15 else centroids[i])
|
||||
|
||||
if max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels)) < tol:
|
||||
break
|
||||
centroids = new_centroids
|
||||
|
||||
boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)]
|
||||
return (
|
||||
torch.tensor(centroids, dtype=torch.float32),
|
||||
torch.tensor(boundaries, dtype=torch.float32),
|
||||
)
|
||||
|
||||
|
||||
@lru_cache(maxsize=32)
|
||||
def get_centroids(d: int, bits: int) -> torch.Tensor:
|
||||
"""Get precomputed Lloyd-Max centroids (cached)."""
|
||||
centroids, _ = solve_lloyd_max(d, bits)
|
||||
return centroids
|
||||
@@ -0,0 +1,260 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""TurboQuant configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.config import ModelConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Named TQ presets: each maps to frozen config parameters.
|
||||
# key_quant_bits: 8 = FP8 keys, 3-4 = MSE (Lloyd-Max) quantized keys.
|
||||
# value_quant_bits: 3-4 = uniform quantized values.
|
||||
TQ_PRESETS: dict[str, dict] = {
|
||||
"turboquant_k8v4": {
|
||||
"key_quant_bits": 8,
|
||||
"value_quant_bits": 4,
|
||||
"norm_correction": False,
|
||||
},
|
||||
"turboquant_4bit_nc": {
|
||||
"key_quant_bits": 4,
|
||||
"value_quant_bits": 4,
|
||||
"norm_correction": True,
|
||||
},
|
||||
"turboquant_k3v4_nc": {
|
||||
"key_quant_bits": 3,
|
||||
"value_quant_bits": 4,
|
||||
"norm_correction": True,
|
||||
},
|
||||
"turboquant_3bit_nc": {
|
||||
"key_quant_bits": 3,
|
||||
"value_quant_bits": 3,
|
||||
"norm_correction": True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class TurboQuantConfig:
|
||||
"""Configuration for TurboQuant KV-cache quantization.
|
||||
|
||||
Applies Hadamard rotation followed by per-coordinate Lloyd-Max scalar
|
||||
quantization for keys, and uniform quantization for values.
|
||||
|
||||
Historical note: the core algorithmic pattern implemented for key
|
||||
quantization (Hadamard rotation followed by deterministic scalar
|
||||
quantization and re-normalization) was originally established in DRIVE
|
||||
(Vargaftik et al., NeurIPS 2021) and EDEN (Vargaftik et al., ICML
|
||||
2022). This formulation is also mathematically equivalent to the
|
||||
scalar case of the HIGGS quantization method (Malinovskii et al.,
|
||||
"Pushing the Limits of Large Language Model Quantization via the
|
||||
Linearity Theorem", NAACL 2025; preprint arXiv:2411.17525), which
|
||||
subsequently generalized these concepts.
|
||||
|
||||
A first application of this approach to KV-cache compression is in
|
||||
"Cache Me If You Must: Adaptive Key-Value Quantization for Large
|
||||
Language Models" (Shutova et al., ICML 2025; preprint
|
||||
arXiv:2501.19392). All of these foundational and application
|
||||
references pre-date the TurboQuant paper (Zandieh et al., ICLR 2026).
|
||||
|
||||
QJL is intentionally omitted: community consensus (5+ independent
|
||||
groups) found it hurts attention quality by amplifying variance
|
||||
through softmax.
|
||||
|
||||
Named presets (use via --kv-cache-dtype):
|
||||
turboquant_k8v4: FP8 keys + 4-bit values, 2.6x, +1.17% PPL
|
||||
turboquant_4bit_nc: 4-bit MSE keys + 4-bit values + NC, 3.8x, +2.71%
|
||||
turboquant_k3v4_nc: 3-bit MSE keys + 4-bit values + NC, ~3.5x, +10.63%
|
||||
turboquant_3bit_nc: 3-bit MSE keys + 3-bit values + NC, 4.9x, +20.59%
|
||||
|
||||
Args:
|
||||
head_dim: Attention head dimension (e.g. 64, 96, 128).
|
||||
key_quant_bits: Bits for key quantization. 8 = FP8 keys (no
|
||||
rotation/MSE). 3-4 = Lloyd-Max MSE quantized keys.
|
||||
value_quant_bits: Bits per value dimension for uniform quantization.
|
||||
3 = 8 levels, 4 = 16 levels (default).
|
||||
norm_correction: Re-normalize centroid vectors to unit norm before
|
||||
inverse rotation during dequant. Fixes quantization-induced norm
|
||||
distortion, improving PPL by ~0.8% at 4-bit.
|
||||
"""
|
||||
|
||||
head_dim: int = 128
|
||||
key_quant_bits: int = 3 # 3-4 = MSE keys, 8 = FP8 keys
|
||||
value_quant_bits: int = 4 # 3-4 = uniform quantized values
|
||||
seed: int = 42 # kept for backward compatibility; no longer used internally
|
||||
norm_correction: bool = False
|
||||
|
||||
@property
|
||||
def key_fp8(self) -> bool:
|
||||
"""Whether keys are stored as FP8 — no rotation/quantization needed."""
|
||||
return self.key_quant_bits == 8
|
||||
|
||||
@property
|
||||
def mse_bits(self) -> int:
|
||||
"""MSE quantizer bit-width (determines centroid count: 2^mse_bits).
|
||||
|
||||
For MSE key modes, equals key_quant_bits.
|
||||
For FP8 key mode, falls back to value_quant_bits (centroids are still
|
||||
needed for continuation-prefill dequant and decode kernel params).
|
||||
"""
|
||||
if self.key_fp8:
|
||||
return self.value_quant_bits
|
||||
return self.key_quant_bits
|
||||
|
||||
@property
|
||||
def key_mse_bits(self) -> int:
|
||||
"""MSE bits actually used for key quantization (0 if FP8 keys)."""
|
||||
if self.key_fp8:
|
||||
return 0
|
||||
return self.key_quant_bits
|
||||
|
||||
@property
|
||||
def centroid_bits(self) -> int:
|
||||
"""Bits for centroid generation — always non-zero."""
|
||||
return self.mse_bits
|
||||
|
||||
@property
|
||||
def n_centroids(self) -> int:
|
||||
return 2**self.mse_bits
|
||||
|
||||
@property
|
||||
def key_packed_size(self) -> int:
|
||||
"""Packed bytes for a single KEY vector.
|
||||
|
||||
FP8 mode (key_quant_bits=8):
|
||||
head_dim bytes (1 byte per element, no overhead).
|
||||
|
||||
TQ mode:
|
||||
- MSE indices: ceil(head_dim * key_mse_bits / 8) bytes
|
||||
- vec_norm: 2 bytes (float16)
|
||||
"""
|
||||
if self.key_fp8:
|
||||
return self.head_dim # 1 byte per element
|
||||
mse_bytes = math.ceil(self.head_dim * self.key_mse_bits / 8)
|
||||
norm_bytes = 2 # vec_norm fp16
|
||||
return mse_bytes + norm_bytes
|
||||
|
||||
@property
|
||||
def effective_value_quant_bits(self) -> int:
|
||||
"""Actual bits used for value storage."""
|
||||
return self.value_quant_bits
|
||||
|
||||
@property
|
||||
def value_packed_size(self) -> int:
|
||||
"""Packed bytes for a single VALUE vector.
|
||||
|
||||
Uniform quantization: ceil(head_dim * bits / 8) + 4 bytes (scale + zero fp16).
|
||||
"""
|
||||
data_bytes = math.ceil(self.head_dim * self.value_quant_bits / 8)
|
||||
return data_bytes + 4 # +2 scale(fp16) +2 zero(fp16)
|
||||
|
||||
@property
|
||||
def slot_size(self) -> int:
|
||||
"""Total packed bytes per head per position (key + value combined).
|
||||
|
||||
Layout: [key_packed | value_packed]
|
||||
"""
|
||||
return self.key_packed_size + self.value_packed_size
|
||||
|
||||
@property
|
||||
def slot_size_aligned(self) -> int:
|
||||
"""Slot size rounded up to next even number.
|
||||
|
||||
Even-number is required so effective_head_size = slot_size_aligned // 2
|
||||
is integral.
|
||||
"""
|
||||
s = self.slot_size
|
||||
return s + (s % 2) # round up to even
|
||||
|
||||
@staticmethod
|
||||
def get_boundary_skip_layers(
|
||||
model_config: ModelConfig,
|
||||
n: int = 2,
|
||||
) -> list[str]:
|
||||
"""Layer indices to skip TQ compression (boundary protection).
|
||||
|
||||
For hybrid models (attention + Mamba/linear-attention), boundary
|
||||
protection is disabled — hybrids typically have only 8-12
|
||||
full-attention layers and a hard n=2 on each side would cover
|
||||
~40 % of them. The dense GSM8K baselines that motivate n=2
|
||||
don't apply to hybrids.
|
||||
|
||||
For dense models, skips first N and last N attention layers.
|
||||
Empirically required for aggressive presets (k3v4_nc, 3bit_nc)
|
||||
— without it GSM8K drops ~30 points on Qwen3-4B.
|
||||
"""
|
||||
if model_config.is_hybrid:
|
||||
attn_indices = _get_full_attention_layer_indices(model_config)
|
||||
if not attn_indices:
|
||||
raise NotImplementedError(
|
||||
"TurboQuant KV cache requires identifiable "
|
||||
"full-attention layers, but none were found in "
|
||||
"the hybrid model config."
|
||||
)
|
||||
logger.info("TQ hybrid: full-attention layers %s", attn_indices)
|
||||
return []
|
||||
|
||||
num_layers = model_config.hf_text_config.num_hidden_layers
|
||||
if n <= 0 or num_layers <= 0:
|
||||
return []
|
||||
n = min(n, num_layers // 2) # don't skip more than half
|
||||
first = list(range(n))
|
||||
last = list(range(num_layers - n, num_layers))
|
||||
# Deduplicate (if num_layers <= 2*n)
|
||||
indices = sorted(set(first + last))
|
||||
return [str(i) for i in indices]
|
||||
|
||||
@staticmethod
|
||||
def from_cache_dtype(cache_dtype: str, head_dim: int) -> TurboQuantConfig:
|
||||
"""Create config from a named preset.
|
||||
|
||||
Valid presets: turboquant_k8v4, turboquant_4bit_nc, etc.
|
||||
"""
|
||||
if cache_dtype not in TQ_PRESETS:
|
||||
valid = ", ".join(TQ_PRESETS.keys())
|
||||
raise ValueError(
|
||||
f"Unknown TurboQuant cache dtype: {cache_dtype!r}. "
|
||||
f"Valid presets: {valid}"
|
||||
)
|
||||
preset = TQ_PRESETS[cache_dtype]
|
||||
return TurboQuantConfig(
|
||||
head_dim=head_dim,
|
||||
key_quant_bits=preset["key_quant_bits"],
|
||||
value_quant_bits=preset["value_quant_bits"],
|
||||
norm_correction=preset["norm_correction"],
|
||||
)
|
||||
|
||||
|
||||
def _get_full_attention_layer_indices(model_config: ModelConfig) -> list[int]:
|
||||
"""Global indices of full-attention layers in a hybrid model.
|
||||
|
||||
Covers the conventions used across vLLM: ``layer_types`` (Qwen3.5/Next),
|
||||
``layers_block_type`` (Jamba/Zamba2), ``attn_type_list`` (Minimax).
|
||||
"""
|
||||
text_cfg = model_config.hf_text_config
|
||||
hf_cfg = model_config.hf_config
|
||||
|
||||
layer_types = getattr(text_cfg, "layer_types", None)
|
||||
if layer_types is not None:
|
||||
return [
|
||||
i for i, t in enumerate(layer_types) if t in ("full_attention", "attention")
|
||||
]
|
||||
|
||||
layers_block_type = getattr(text_cfg, "layers_block_type", None)
|
||||
if layers_block_type is not None:
|
||||
return [
|
||||
i for i, t in enumerate(layers_block_type) if t in ("attention", "hybrid")
|
||||
]
|
||||
|
||||
attn_type_list = getattr(hf_cfg, "attn_type_list", None)
|
||||
if attn_type_list is not None:
|
||||
return [i for i, t in enumerate(attn_type_list) if t == 1]
|
||||
|
||||
return []
|
||||
@@ -0,0 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from .layer_utils import replace_parameter, update_tensor_inplace
|
||||
|
||||
__all__ = ["update_tensor_inplace", "replace_parameter"]
|
||||
@@ -0,0 +1,67 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import ScalarType, scalar_types
|
||||
|
||||
ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD = 1024
|
||||
ALLSPARK_SUPPORTED_QUANT_TYPES = [scalar_types.uint8b128]
|
||||
ALLSPARK_AMPERE_N_ALIGN = 16
|
||||
ALLSPARK_AMPERE_K_ALIGN = 16
|
||||
|
||||
|
||||
def check_allspark_supported_dtype_shape(
|
||||
input_size_per_partition: int,
|
||||
output_size_per_partition: int,
|
||||
group_size: int,
|
||||
weight_dtype: ScalarType,
|
||||
act_dtype: torch.dtype,
|
||||
):
|
||||
capability_tuple = current_platform.get_device_capability()
|
||||
device_capability = -1 if capability_tuple is None else capability_tuple.to_int()
|
||||
|
||||
# For Ampere GPU
|
||||
if device_capability >= 80 and device_capability < 90:
|
||||
if group_size != -1:
|
||||
return (
|
||||
False,
|
||||
"For Ampere GPU, AllSpark does not support group_size "
|
||||
f"= {group_size}. Only group_size = -1 are supported.",
|
||||
)
|
||||
|
||||
if weight_dtype not in ALLSPARK_SUPPORTED_QUANT_TYPES:
|
||||
return (
|
||||
False,
|
||||
"For Ampere GPU, AllSpark does not support "
|
||||
f"quant type ({weight_dtype}). Only quant type "
|
||||
f"({ALLSPARK_SUPPORTED_QUANT_TYPES}) are supported.",
|
||||
)
|
||||
|
||||
if (
|
||||
input_size_per_partition % ALLSPARK_AMPERE_K_ALIGN != 0
|
||||
or output_size_per_partition % ALLSPARK_AMPERE_N_ALIGN != 0
|
||||
):
|
||||
return (
|
||||
False,
|
||||
"AllSpark needs input_size_per_partition % "
|
||||
f"{ALLSPARK_AMPERE_K_ALIGN} = 0 and "
|
||||
f"output_size_per_partition % {ALLSPARK_AMPERE_N_ALIGN} = 0 "
|
||||
"for Ampere GPU optimized kernels.",
|
||||
)
|
||||
|
||||
if act_dtype != torch.float16 and act_dtype != torch.bfloat16:
|
||||
return (
|
||||
False,
|
||||
"AllSpark only supports act_dtype = float16 or bfloat16,"
|
||||
f"for Ampere GPU, but got act_dtype = {act_dtype}.",
|
||||
)
|
||||
else:
|
||||
return (
|
||||
False,
|
||||
"AllSpark currently does not support "
|
||||
f"device_capability = {device_capability}.",
|
||||
)
|
||||
|
||||
return True, None
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 8,
|
||||
"num_stages": 2
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
}
|
||||
}
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 8,
|
||||
"num_stages": 5
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 8,
|
||||
"num_stages": 4
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 8,
|
||||
"num_stages": 4
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 8,
|
||||
"num_stages": 4
|
||||
}
|
||||
}
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 8,
|
||||
"num_stages": 4
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 8,
|
||||
"num_stages": 2
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_M": 256,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 64,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 64,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 64,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
}
|
||||
}
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
}
|
||||
}
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 8,
|
||||
"num_stages": 4
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
}
|
||||
}
|
||||
+164
@@ -0,0 +1,164 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
}
|
||||
}
|
||||
+164
@@ -0,0 +1,164 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
}
|
||||
}
|
||||
+164
@@ -0,0 +1,164 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"kpack": 1,
|
||||
"matrix_instr_nonkdim": 16,
|
||||
"num_warps": 4
|
||||
}
|
||||
}
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 8,
|
||||
"num_stages": 5
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 256,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 8,
|
||||
"num_stages": 4
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
}
|
||||
}
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 64,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 64,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 8,
|
||||
"num_stages": 4
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 64,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_M": 16,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_M": 32,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
}
|
||||
}
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
}
|
||||
}
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
{
|
||||
"1": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 8,
|
||||
"num_stages": 4
|
||||
},
|
||||
"2": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"4": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"8": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"16": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"24": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"32": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 8,
|
||||
"num_stages": 5
|
||||
},
|
||||
"48": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"64": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 4
|
||||
},
|
||||
"96": {
|
||||
"BLOCK_SIZE_M": 128,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 8,
|
||||
"num_stages": 3
|
||||
},
|
||||
"128": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 5
|
||||
},
|
||||
"256": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"512": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 32,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 2
|
||||
},
|
||||
"1024": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 16,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"1536": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"2048": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 1,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
}
|
||||
}
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
{
|
||||
"2048": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 64,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"3072": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 64,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
"4096": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 128,
|
||||
"GROUP_SIZE_M": 32,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user