[Feature]Support and update INT8/INT4 quantization inference (#299)

Signed-off-by: Li Wei <liwei.109@outlook.com>

Author: liwei109

ci/scripts/env/install_env.sh

@@ -54,6 +54,10 @@ docker exec "${DOCKER_NAME}" bash -lc "
   # Patch torch dynamo eval_frame
   cp vllm_kunlun/patches/eval_frame.py \
     /root/miniconda/envs/${CONDA_ENV}/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py
+  
+  # Patch quantization __init__.py
+  cp vllm_kunlun/quantization/__init__.py \
+    /root/miniconda/envs/${CONDA_ENV}/lib/python3.10/site-packages/vllm/model_executor/layers/quantization/__init__.py
 
   ########################################
   # Kunlun runtime dependencies

docs/source/installation.md

@@ -76,6 +76,12 @@ Copy the eval_frame.py patch:
 
 ```
 cp vllm_kunlun/patches/eval_frame.py "${CONDA_PREFIX:-$VIRTUAL_ENV}"/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py
+```
+
+ ### Replace quantization __init__.py
+
+```
+cp vllm_kunlun/quantization/__init__.py "${CONDA_PREFIX:-$VIRTUAL_ENV}"/lib/python3.10/site-packages/vllm/model_executor/layers/quantization/__init__.py
 ```
 
 ## Choose to download customized xpytorch

vllm_kunlun/ops/__init__.py

@@ -29,13 +29,3 @@
 
 # TODO @xyDong0223 remove v0.16.0
 # import vllm_kunlun.ops.mla
-
-# quantization
-# TODO @liwei109 enable quantization in v0.16.0
-# import vllm_kunlun.ops.quantization.awq
-# import vllm_kunlun.ops.quantization.compressed_tensors.compressed_tensors
-# import vllm_kunlun.ops.quantization.compressed_tensors.compressed_tensors_moe
-# import vllm_kunlun.ops.quantization.gptq
-# import vllm_kunlun.ops.quantization.kernels.kunlun_exllama_linear
-# import vllm_kunlun.ops.quantization.kernels.kunlun_scale_mm
-# import vllm_kunlun.ops.quantization.moe_wna16

vllm_kunlun/ops/linear.py

@@ -336,22 +336,47 @@ def _load_fused_module_from_checkpoint(
         self.weight_loader_v2(param, loaded_weight_shard, shard_id)
 
 
+def validate_shard_id(self, loaded_shard_id: int | tuple[int, ...] | None):
+    if loaded_shard_id is None:
+        return
+    if isinstance(loaded_shard_id, tuple):
+        for idx in loaded_shard_id:
+            if not (0 <= idx < len(self.output_sizes)):
+                raise ValueError(
+                    f"Shard id index {idx} should be between 0 and "
+                    f"{len(self.output_sizes) - 1}. Got shard id {loaded_shard_id}."
+                )
+        if len(loaded_shard_id) > 1 and any(
+            b - a != 1 for a, b in zip(loaded_shard_id[:-1], loaded_shard_id[1:])
+        ):
+            raise ValueError(
+                "Shard id with multiple indices should be consecutive. "
+                f"Got shard id {loaded_shard_id}."
+            )
+        return
+    elif isinstance(loaded_shard_id, int):
+        if loaded_shard_id < 0 or loaded_shard_id >= len(self.output_sizes):
+            raise ValueError(
+                f"Shard id should be between 0 and {len(self.output_sizes) - 1}. "
+                f"Got shard id {loaded_shard_id}."
+            )
+        return
+
+
 def weight_loader_v2(
     self,
     param: BasevLLMParameter,
     loaded_weight: torch.Tensor,
     loaded_shard_id: tuple[int, ...] | int | None = None,
 ):
+    self.validate_shard_id(loaded_shard_id)
     if loaded_shard_id is None or isinstance(loaded_shard_id, tuple):
         if isinstance(param, PerTensorScaleParameter):
             param.load_merged_column_weight(loaded_weight=loaded_weight, shard_id=0)
             return
         elif type(param) in (RowvLLMParameter, BasevLLMParameter):
             param.load_merged_column_weight(loaded_weight=loaded_weight)
             return
-        # TODO: @dsikka - move to parameter.py
-        self._load_fused_module_from_checkpoint(param, loaded_weight)
-        return
         output_sizes = (
             [self.output_sizes[idx] for idx in loaded_shard_id]
             if loaded_shard_id
@@ -363,6 +388,7 @@ def weight_loader_v2(
                 adjust_block_scale_shard(weight_block_size, size, 0)[0]
                 for size in (output_sizes or self.output_sizes)
             ]
+        # TODO: @dsikka - move to parameter.py
         self._load_fused_module_from_checkpoint(
             param, loaded_weight, output_sizes=output_sizes
         )
@@ -394,6 +420,7 @@ def weight_loader_v2(
 MergedColumnParallelLinear._load_fused_module_from_checkpoint = (
     _load_fused_module_from_checkpoint
 )
+MergedColumnParallelLinear.validate_shard_id = validate_shard_id
 MergedColumnParallelLinear.weight_loader_v2 = weight_loader_v2
 
 

vllm_kunlun/ops/quantization/__init__.py

@@ -7,6 +7,7 @@
 import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.platforms.interface import DeviceCapability, Platform, PlatformEnum
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
 if TYPE_CHECKING:
@@ -375,3 +376,15 @@ def support_hybrid_kv_cache(cls) -> bool:
     @classmethod
     def support_static_graph_mode(cls) -> bool:
         return True
+
+    @classmethod
+    def pre_register_and_update(
+        cls, parser: FlexibleArgumentParser | None = None
+    ) -> None:
+        from vllm_kunlun.quantization.awq import KunlunAWQConfig  # noqa
+        from vllm_kunlun.quantization.compressed_tensors import (  # noqa
+            KunlunCompressedTensorsConfig,
+        )
+        from vllm_kunlun.quantization.gptq import KunlunGPTQConfig  # noqa
+        from vllm_kunlun.quantization.kernels import _POSSIBLE_INT8_KERNELS  # noqa
+        from vllm_kunlun.quantization.kernels import _POSSIBLE_KERNELS  # noqa

vllm_kunlun/ops/quantization/compressed_tensors/__init__.py

@@ -0,0 +1,187 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# patched by vLLM-Kunlun
+
+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",
+    "fp8",
+    "ptpc_fp8",
+    "fbgemm_fp8",
+    # "fp_quant",
+    "modelopt",
+    "modelopt_fp4",
+    "bitblas",
+    "gguf",
+    "gptq_marlin_24",
+    "gptq_marlin",
+    "gptq_bitblas",
+    "awq_marlin",
+    "gptq",
+    "compressed-tensors",
+    "bitsandbytes",
+    "experts_int8",
+    "ipex",
+    "quark",
+    "moe_wna16",
+    "torchao",
+    "inc",
+    "mxfp4",
+    "petit_nvfp4",
+    "cpu_awq",
+]
+QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
+
+DEPRECATED_QUANTIZATION_METHODS = [
+    "tpu_int8",
+    "ptpc_fp8",
+    "fbgemm_fp8",
+    # "fp_quant",
+    "bitblas",
+    "gptq_marlin_24",
+    "gptq_bitblas",
+    "experts_int8",
+    "ipex",
+    "petit_nvfp4",
+]
+
+# 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.warning(
+                "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.model_executor.layers.quantization.awq import AWQConfig
+    from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
+    from vllm.model_executor.layers.quantization.bitblas import BitBLASConfig
+    from vllm.model_executor.layers.quantization.bitsandbytes import BitsAndBytesConfig
+    from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
+        CompressedTensorsConfig,
+    )
+    from vllm.model_executor.layers.quantization.cpu_wna16 import CPUAWQConfig
+    from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config
+    from vllm.model_executor.layers.quantization.fbgemm_fp8 import FBGEMMFp8Config
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    # from vllm.model_executor.layers.quantization.fp_quant import FPQuantConfig
+    from vllm.model_executor.layers.quantization.gguf import GGUFConfig
+    from vllm.model_executor.layers.quantization.gptq import GPTQConfig
+    from vllm.model_executor.layers.quantization.gptq_bitblas import GPTQBitBLASConfig
+    from vllm.model_executor.layers.quantization.gptq_marlin import GPTQMarlinConfig
+    from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
+        GPTQMarlin24Config,
+    )
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+    from vllm.model_executor.layers.quantization.ipex_quant import IPEXConfig
+    from vllm.model_executor.layers.quantization.modelopt import (
+        ModelOptFp8Config,
+        ModelOptNvFp4Config,
+    )
+    from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
+    from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4Config
+    from vllm.model_executor.layers.quantization.petit import PetitNvFp4Config
+    from vllm.model_executor.layers.quantization.ptpc_fp8 import PTPCFp8Config
+    from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
+    from vllm.model_executor.layers.quantization.torchao import TorchAOConfig
+
+    method_to_config: dict[str, type[QuantizationConfig]] = {
+        "awq": AWQConfig,
+        "fp8": Fp8Config,
+        "fbgemm_fp8": FBGEMMFp8Config,
+        # "fp_quant": FPQuantConfig,
+        "modelopt": ModelOptFp8Config,
+        "modelopt_fp4": ModelOptNvFp4Config,
+        "bitblas": BitBLASConfig,
+        "gguf": GGUFConfig,
+        "gptq_marlin_24": GPTQMarlin24Config,
+        "gptq_marlin": GPTQMarlinConfig,
+        "gptq_bitblas": GPTQBitBLASConfig,
+        "awq_marlin": AWQMarlinConfig,
+        "gptq": GPTQConfig,
+        "compressed-tensors": CompressedTensorsConfig,
+        "bitsandbytes": BitsAndBytesConfig,
+        "ptpc_fp8": PTPCFp8Config,
+        "experts_int8": ExpertsInt8Config,
+        "ipex": IPEXConfig,
+        "quark": QuarkConfig,
+        "moe_wna16": MoeWNA16Config,
+        "torchao": TorchAOConfig,
+        "auto-round": INCConfig,
+        "inc": INCConfig,
+        "mxfp4": Mxfp4Config,
+        "petit_nvfp4": PetitNvFp4Config,
+        "cpu_awq": CPUAWQConfig,
+    }
+    # 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",
+]