mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-26 23:16:32 +08:00
25 lines
778 B
Python
25 lines
778 B
Python
from typing import Type
|
|
|
|
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
|
from vllm.model_executor.layers.quantization.awq import AWQConfig
|
|
from vllm.model_executor.layers.quantization.gptq import GPTQConfig
|
|
from vllm.model_executor.layers.quantization.squeezellm import SqueezeLLMConfig
|
|
|
|
_QUANTIZATION_CONFIG_REGISTRY = {
|
|
"awq": AWQConfig,
|
|
"gptq": GPTQConfig,
|
|
"squeezellm": SqueezeLLMConfig,
|
|
}
|
|
|
|
|
|
def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
|
|
if quantization not in _QUANTIZATION_CONFIG_REGISTRY:
|
|
raise ValueError(f"Invalid quantization method: {quantization}")
|
|
return _QUANTIZATION_CONFIG_REGISTRY[quantization]
|
|
|
|
|
|
__all__ = [
|
|
"QuantizationConfig",
|
|
"get_quantization_config",
|
|
]
|