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 は 3 つのことを担当します:create_weights はレイヤー構築時に qweight / scales / qzeros などのパラメータオブジェクトを作り、process_weights_after_loading は重みのロード完了後に 1 回だけ再配置を行い、apply は前向き計算時に実際の量子化 GEMM kernel を呼び出します。この抽象により、同じ ColumnParallelLinear で未量子化にも AWQ / GPTQ / fp8 / Marlin にも対応できます。quant_method が違うだけです。

AWQ と GPTQ は vLLM 内ではどちらも Marlin kernel ファミリーを使います: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 の 3 層に分けるのか?

  • 統一インターフェース: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 順序で出力次元に沿って pack されますが、Marlin kernel は GPTQ スタイル (入力次元で pack、標準 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 インスタンスを選び、3 つの 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 kernel が期待する配置に再配置します (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 が自力で 3 つに分割します (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 (実際に計算) の 3 層に分けます。AWQ / GPTQ / Marlin は MPLinearLayerConfig + choose_mp_linear_kernel を共有し、fp8 は独自の init_fp8_linear_kernel を使います。テンサル並列線形レイヤー に依存して create_weights 時に v2 ローダーをパラメータに仕込み、ロードフローは DefaultModelLoader が駆動します。後処理済みの重みは最終的に前向き計算時に attention レイヤーと ブロック表 が記述する KV cache 上で消費されます。

公式資料: vLLM 文档 · README