Quantization¶

import torch
from vllm.model_executor.layers.quantization import (
    register_quantization_config,
)
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig,
    QuantizeMethodBase,
)
from vllm.model_executor.layers.linear import LinearBase
from vllm.model_executor.layers.fused_moe import FusedMoE

@register_quantization_config("my_quant")
class MyQuantConfig(QuantizationConfig):
    """Custom quantization config."""

    def get_name(self) -> str:
        return "my_quant"

    def get_supported_act_dtypes(self) -> list:
        return [torch.float16, torch.bfloat16]

    @classmethod
    def get_min_capability(cls) -> int:
        # Minimum GPU compute capability, -1 for no restriction
        return -1

    @staticmethod
    def get_config_filenames() -> list[str]:
        # Config files to search for in model directory
        return []

    @classmethod
    def from_config(cls, config: dict) -> "MyQuantConfig":
        # Create config from model's quantization config
        return cls()

    def get_quant_method(
        self, layer: torch.nn.Module, prefix: str
    ) -> QuantizeMethodBase | None:
        # Dispatch based on layer type
        # NOTE: you only need to implement methods you care about
        if isinstance(layer, LinearBase):
            return MyQuantLinearMethod()
        elif isinstance(layer, FusedMoE):
            return MyQuantMoEMethod(layer.moe_config)
        return None

from vllm.model_executor.layers.linear import UnquantizedLinearMethod

class MyQuantLinearMethod(UnquantizedLinearMethod):
    """Custom quantization method for linear layers."""

    def create_weights(
        self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
    ):
        # Create quantized weights for the layer
        ...

    def apply(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        # Apply custom quantization logic here
        ...

from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
    FusedMoEMethodBase,
)
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

class MyQuantMoEMethod(FusedMoEMethodBase):
    """Custom quantization method for MoE layers."""

    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,
    ):
        # Create quantized weights for the MoE layer
        ...

```markdown
---  
title: "使用 vLLM 插件系统开发自定义量化方法"  
description: "了解如何使用 vLLM 插件系统创建自定义量化方法。"  
linkTitle: "自定义量化方法"  
weight: 30  
---

这个指南展示了如何通过 vLLM 的插件系统开发自定义量化方法。  
vLLM 目前支持 FP8、AWQ 和 GPTQ 等多种量化方法，您可以创建自己的量化插件来支持新的量化格式。

## 插件结构

自定义量化插件必须实现以下接口：

```python
from vllm.model_executor.layers.quantization import (
    QuantizationConfig,
    QuantizedModule,
    FusedMoEMethod,
    FusedMoEQuantConfig,
)
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig as BaseQuantizationConfig,
)


class MyQuantConfig(QuantizationConfig):
    # 实现量化配置的具体细节
    ...


class MyQuantizedModule(QuantizedModule):
    # 实现量化模块的具体细节
    ...


class MyFusedMoEMethod(FusedMoEMethod):
    def apply(
        self,
        layer: torch.nn.Module,
        router: "FusedMoERouter",
        x: torch.Tensor,
        router_logits: torch.Tensor,
    ) -> torch.Tensor:
        # 应用量化权重的 MoE 计算
        ...

    def get_fused_moe_quant_config(
        self, layer: torch.nn.Module
    ) -> FusedMoEQuantConfig | None:
        # 返回 MoE 量化配置
        ...

# 注册您的量化方法（导入包含配置的模块）
import my_quant_plugin

from vllm import LLM

# 使用自定义量化方法
llm = LLM(model="your-model", quantization="my_quant")