量化层加载:AWQ / GPTQ / Marlin / fp8

源码版本v0.25.1

职责

vLLM 不把量化当成"训完再压缩"的事后步骤,而是把每种量化方案实现成一对 QuantizationConfig + LinearMethodBase(以及 MoE 的 FusedMoEMethodBase)。Config 负责读 Hugging Face checkpoint 里的 quantization_config 字段,决定用哪个 LinearMethod;LinearMethod 负责三件事:create_weights 在层构造时建好 qweight / scales / qzeros 这些参数对象,process_weights_after_loading 在权重全部灌完后做一次性重排,apply 在前向时调用真正的量化 GEMM kernel。这套抽象让同一个 ColumnParallelLinear 既能跑未量化、也能跑 AWQ / GPTQ / fp8 / Marlin,只是 quant_method 不同。

AWQ 和 GPTQ 在 vLLM 里都走 Marlin 内核家族:AutoAWQMarlinLinearMethod(auto_awq.py:414-426)和 AutoGPTQLinearMethod(auto_gptq.py:306-324),通过 choose_mp_linear_kernel 在 Conch / Exllama / Marlin / Machete 里挑最优。fp8 则是另一条路:Fp8Config(fp8.py:95-101)支持 per-tensor / per-block / 静态 / 动态激活量化,通过 init_fp8_linear_kernel 选 cutlass / Marlin / triton 后端。

设计动机

为什么把量化拆成 Config + Method + Kernel 三层?

• 统一接口:LinearMethodBase.create_weights(layer, input_size_per_partition, output_partition_sizes, input_size, output_size, ...)(linear.py:144-168)签名对各种量化一致,ColumnParallelLinear 不需要知道是 AWQ 还是 fp8,只要 quant_config.get_quant_method(layer) 返回一个 LinearMethodBase 实例就行。
• kernel 选择延迟到运行时:AWQ / GPTQ 的 create_weights 会先构造 MPLinearLayerConfig 描述 weight_type / group_size / zero_points / has_g_idx,再 choose_mp_linear_kernel(...) 挑一个具体 kernel(auto_gptq.py:341-354),kernel 自己再决定要不要 process_weights_after_loading 重排。
• AWQ → GPTQ 格式转换:AWQ checkpoint 用非标准的 packing 顺序且沿输出维打包,Marlin kernel 只认 GPTQ 风格(沿输入维打包、标准 bit 顺序),所以 process_weights_after_loading 先调 _convert_awq_to_standard_format 再交给 kernel(auto_awq.py:528-536)。
• GPTQ 的 desc_act / g_idx:AutoGPTQLinearMethod.create_weights 根据 desc_act 决定是否建 g_idx 参数,根据 marlin_repeat_scales_on_all_ranks 决定 scales 在 TP 下是复制还是分片(auto_gptq.py:366-440)。
• fp8 checkpoint 不是 fp8 时在线量化:如果 checkpoint 是 bf16,process_weights_after_loading 调 process_fp8_weight_tensor_strategy 把它在线量化成 fp8,顺便处理 per-tensor 重整(fp8.py:416-437)。
• Marlin 重排:fp8 选到 MarlinFP8ScaledMMLinearKernel 时,process_weights_after_loading 会把 weight 转置、按 Marlin 期望的排布重写,并设 marlin_input_dtype(fp8.py:398-408)。
• 参数走 v2 加载器:AWQ / GPTQ 用 PackedvLLMParameter / GroupQuantScaleParameter / PackedColumnParameter 等 v2 参数类,自带 packed_dim / packed_factor / input_dim / output_dim 元信息,weight_loader_v2 据此正确切分(auto_awq.py:478-519)。

关键文件

• QuantizationConfig:87-101 — 抽象基类,定义 get_quant_method、get_cache_method 等接口。
• LinearMethodBase:141-180 — 量化方法基类,create_weights / apply / process_weights_after_loading。
• AutoAWQConfig:171-180 — 解析 AWQ checkpoint 的 quantization_config。
• AutoAWQMarlinLinearMethod.create_weights:439-526 — 建 qweight / qzeros / scales,挑 Marlin kernel。
• AutoAWQMarlinLinearMethod.process_weights_after_loading:528-536 — _convert_awq_to_standard_format + kernel 自己重排。
• _convert_awq_to_standard_format:93-100 — AWQ packing → GPTQ packing 的转换函数。
• AutoGPTQConfig:97-103 — 解析 GPTQ checkpoint,处理 bits / group_size / desc_act。
• AutoGPTQLinearMethod.create_weights:326-453 — 建 qweight / g_idx / scales / qzeros,根据 desc_act 选 ChannelQuantScaleParameter 或 GroupQuantScaleParameter。
• AutoGPTQLinearMethod.process+apply:455-464 — kernel.process_weights_after_loading(layer) + kernel.apply_weights(layer, x, bias)。
• Fp8Config:95-101 — fp8 配置,区分 per-tensor / block / 静态-动态激活。
• Fp8LinearMethod.create_weights:322-394 — 建 weight / weight_scale / 可选 input_scale,通过 init_fp8_linear_kernel 选 kernel。
• Fp8LinearMethod.process_weights_after_loading:398-441 — Marlin 时转置 + kernel 重排,否则在线量化非 fp8 checkpoint。
• choose_mp_linear_kernel:685-687 — 按平台 / 形状挑 Conch / Exllama / Marlin / Machete 等 kernel。
• MarlinLinearKernel.process_weights_after_loading:88-137 — ops.gptq_marlin_repack 把 qweight 重排成 Marlin 期望布局。

数据流

以 AWQ Marlin 为例,ColumnParallelLinear.__init__ 调 self.quant_method.create_weights(...),AutoAWQMarlinLinearMethod.create_weights 先构 MPLinearLayerConfig 再 choose_mp_linear_kernel,挑出 MarlinLinearKernel 实例,然后建三个 v2 参数对象:

qweight = PackedvLLMParameter(
    data=torch.empty(
        input_size_per_partition,
        output_size_per_partition // self.quant_config.pack_factor,
        dtype=torch.int32,
    ),
    input_dim=0,
    output_dim=1,
    packed_dim=1,
    packed_factor=self.quant_config.pack_factor,
    weight_loader=weight_loader,
)

num_groups = input_size_per_partition // group_size

qzeros = PackedvLLMParameter(
    data=torch.empty(
        num_groups,
        output_size_per_partition // self.quant_config.pack_factor,
        dtype=torch.int32,
    ),
    input_dim=0,
    output_dim=1,
    packed_dim=1,
    packed_factor=self.quant_config.pack_factor,
    weight_loader=weight_loader,
)

scales = GroupQuantScaleParameter(
    data=torch.empty(
        num_groups,
        output_size_per_partition,
        dtype=params_dtype,
    ),
    input_dim=0,
    output_dim=1,
    weight_loader=weight_loader,
)

(auto_awq.py:478-515)之后 DefaultModelLoader 通过 model.load_weights 把磁盘上的 qweight / scales / qzeros 按 v2 加载器切分写入。全部灌完后,ColumnParallelLinear 调 quant_method.process_weights_after_loading,AutoAWQMarlinLinearMethod 先把 AWQ packing 转成 GPTQ 风格再交给 MarlinLinearKernel.process_weights_after_loading,后者用 ops.gptq_marlin_repack 把权重重排成 Marlin 内核期望的布局(marlin.py:124-137)。前向时 apply 调 self.kernel.apply_weights(layer, x, bias)(auto_awq.py:538-545),最终落到 Marlin 的 CUDA / Triton kernel。

边界与失败

• 平台不支持 Marlin:AutoAWQMarlinLinearMethod.__init__ 在非 CPU 平台调 verify_marlin_supported(quant_type, group_size, has_zp=...),不满足就抛错(auto_awq.py:431-437)。
• group_size=-1:create_weights 把 group_size 退化成 input_size,即 per-channel 量化(auto_awq.py:452-455)。
• fp8 block_quant + act_q_static:块量化强制要求动态激活(assert not self.act_q_static),否则 raise(fp8.py:367-368)。
• fp8 非 fp8 checkpoint:process_weights_after_loading 走 process_fp8_weight_tensor_tensor_strategy 把 bf16 在线量化,per-tensor 还要对 fused module 重整(fp8.py:416-429)。
• GPTQ row-parallel 的 scales 复制:marlin_repeat_scales_on_all_ranks 决定 scales 是 ChannelQuantScaleParameter(scale_dim=None 在所有 rank 复制)还是 GroupQuantScaleParameter(scale_dim=0 分片)(auto_gptq.py:367-440)。
• Phi-3 融合 QKV:磁盘上 QKV 已融合时没有 shard id,QKVParallelLinear._load_fused_module_from_checkpoint 自己切三段(linear.py:1049-1097),量化时要处理 packed_dim 的 offset 调整。
• Marlin act_type 是 fp8:MarlinLinearKernel.process_weights_after_loading 在 c.act_type == torch.float8_e4m3fn 时调 ops.marlin_int4_fp8_preprocess 并把 scales 乘 512(marlin.py:98-103)。

小结

量化这套抽象把 Config(读 checkpoint 元信息) → Method(建参数 + 后处理 + apply 委托) → Kernel(真正算)三层分开,AWQ / GPTQ / Marlin 共用 MPLinearLayerConfig + choose_mp_linear_kernel,fp8 走自己的 init_fp8_linear_kernel。它依赖 张量并行线性层 在 create_weights 时把 v2 加载器塞进参数,加载流程由 DefaultModelLoader 驱动;后处理完的权重最终在前向被 attention 层和 块表 描述的 KV cache 上消费。

对照官方资料:vLLM 文档 · README。