Class QuantConfig

tensorrt_llm/models/modeling_utils.py:131–271

Serializable quantization configuration class, part of the PretrainedConfig. Args: quant_algo (tensorrt_llm.quantization.mode.QuantAlgo, optional): Quantization algorithm. Defaults to None. kv_cache_quant_algo (tensorrt_llm.quantization.mode.QuantAlgo, optional): KV cache q

@dataclasses.dataclass
class QuantConfig:
    """
    Serializable quantization configuration class, part of the PretrainedConfig.

    Args:
        quant_algo (tensorrt_llm.quantization.mode.QuantAlgo, optional): Quantization algorithm. Defaults to None.
        kv_cache_quant_algo (tensorrt_llm.quantization.mode.QuantAlgo, optional): KV cache quantization algorithm. Defaults to None.
        group_size (int): The group size for group-wise quantization. Defaults to 128.
        smoothquant_val (float): The smoothing parameter alpha used in smooth quant. Defaults to 0.5.
        clamp_val (List[float], optional): The clamp values used in FP8 rowwise quantization. Defaults to None.
        use_meta_recipe (bool): Whether to use Meta's recipe for FP8 rowwise quantization. Defaults to False.
        has_zero_point (bool): Whether to use zero point for quantization. Defaults to False.
        pre_quant_scale (bool): Whether to use pre-quant scale for quantization. Defaults to False.
        exclude_modules (List[str], optional): The module name patterns that are skipped in quantization. Defaults to None.
        mamba_ssm_cache_dtype (str, optional): The data type for mamba SSM cache. Defaults to None.
    """
    quant_algo: Optional[QuantAlgo] = None
    kv_cache_quant_algo: Optional[QuantAlgo] = None
    group_size: int = 128
    smoothquant_val: float = 0.5
    clamp_val: Optional[List[float]] = None
    use_meta_recipe: bool = False
    has_zero_point: bool = False
    pre_quant_scale: bool = False
    exclude_modules: Optional[List[str]] = None
    mamba_ssm_cache_dtype: Optional[str] = None

    @cached_property
    def quant_mode(self) -> QuantModeWrapper:
        quant_mode_list = [
            QuantMode.from_quant_algo(
                self.quant_algo,
                self.kv_cache_quant_algo,
            )
        ]
        return QuantModeWrapper(quant_mode_list)

    @cached_property
    def layer_quant_mode(self) -> QuantMode:
        return QuantMode.from_quant_algo(
            self.quant_algo,
            self.kv_cache_quant_algo,
        )

    @property
    def _use_plugin_sq(self):
        return self.quant_algo in W8A8_SQ_PLUGIN_LIST

    @property
    def _requires_calibration(self):
        return self.quant_algo in (set(QUANT_ALGO_LIST) - {
            QuantAlgo.W8A16, QuantAlgo.W4A16,
            QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN
        }) or self.kv_cache_quant_algo in KV_CACHE_QUANT_ALGO_LIST

    @property
    def _requires_modelopt_quantization(self):
        if self.quant_algo in [