vLLM 架构拆解

繁體中文

简体中文
English
日本語
Español
Français
Deutsch

繁體中文

简体中文
English
日本語
Español
Français
Deutsch

Appearance

量化層載入: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)。

關鍵檔案 ​

資料流 ​

以 AWQ Marlin 為例,ColumnParallelLinear.__init__ 調 self.quant_method.create_weights(...),AutoAWQMarlinLinearMethod.create_weights 先構 MPLinearLayerConfig 再 choose_mp_linear_kernel,挑出 MarlinLinearKernel 實例,然後建三個 v2 參數物件:

python
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