vLLM 架构拆解

English

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

English

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

Appearance

Quantized layer loading: AWQ / GPTQ / Marlin / fp8 ​

源码版本v0.25.1

Responsibilities ​

vLLM does not treat quantization as a post-training "compress after training" step. Instead, each quantization scheme is implemented as a pair of QuantizationConfig + LinearMethodBase (and, for MoE, FusedMoEMethodBase). The Config reads the quantization_config field from the Hugging Face checkpoint and decides which LinearMethod to use; the LinearMethod is responsible for three things: create_weights builds the qweight / scales / qzeros parameter objects at layer construction time, process_weights_after_loading performs a one-shot reorder once all weights have been loaded, and apply invokes the actual quantized GEMM kernel during the forward pass. This abstraction lets the same ColumnParallelLinear run unquantized as well as AWQ / GPTQ / fp8 / Marlin, differing only in quant_method.

Both AWQ and GPTQ go through the Marlin kernel family in vLLM: AutoAWQMarlinLinearMethod(auto_awq.py:414-426) and AutoGPTQLinearMethod(auto_gptq.py:306-324), selecting the best among Conch / Exllama / Marlin / Machete via choose_mp_linear_kernel. fp8 takes a different route: Fp8Config(fp8.py:95-101) supports per-tensor / per-block / static / dynamic activation quantization, and picks a cutlass / Marlin / triton backend via init_fp8_linear_kernel.

Design motivation ​

Why split quantization into three layers: Config + Method + Kernel?

  • Uniform interface: LinearMethodBase.create_weights(layer, input_size_per_partition, output_partition_sizes, input_size, output_size, ...)(linear.py:144-168) has the same signature across quantization schemes, so ColumnParallelLinear does not need to know whether it is AWQ or fp8; it only needs quant_config.get_quant_method(layer) to return a LinearMethodBase instance.
  • Kernel selection is deferred to runtime: create_weights for AWQ / GPTQ first builds an MPLinearLayerConfig describing weight_type / group_size / zero_points / has_g_idx, then calls choose_mp_linear_kernel(...) to pick a concrete kernel(auto_gptq.py:341-354), and the kernel itself decides whether to reorder via process_weights_after_loading.
  • AWQ -> GPTQ format conversion: AWQ checkpoints use a non-standard packing order packed along the output dimension, while the Marlin kernel only understands the GPTQ style (packed along the input dimension, standard bit order), so process_weights_after_loading first calls _convert_awq_to_standard_format before handing off to the kernel(auto_awq.py:528-536).
  • GPTQ desc_act / g_idx: AutoGPTQLinearMethod.create_weights decides whether to build the g_idx parameter based on desc_act, and decides whether scales are replicated or sharded under TP based on marlin_repeat_scales_on_all_ranks(auto_gptq.py:366-440).
  • fp8 checkpoints that are not fp8 are quantized online: if the checkpoint is bf16, process_weights_after_loading calls process_fp8_weight_tensor_strategy to quantize it to fp8 online, also handling per-tensor reshaping(fp8.py:416-437).
  • Marlin reorder: when fp8 picks MarlinFP8ScaledMMLinearKernel, process_weights_after_loading transposes the weight, rewrites it in the layout Marlin expects, and sets marlin_input_dtype(fp8.py:398-408).
  • Parameters go through the v2 loader: AWQ / GPTQ use v2 parameter classes like PackedvLLMParameter / GroupQuantScaleParameter / PackedColumnParameter, which carry packed_dim / packed_factor / input_dim / output_dim metadata, and weight_loader_v2 splits them correctly based on this(auto_awq.py:478-519).

Key files ​

Data flow ​

Taking AWQ Marlin as an example, ColumnParallelLinear.__init__ calls self.quant_method.create_weights(...), AutoAWQMarlinLinearMethod.create_weights first builds an MPLinearLayerConfig, then calls choose_mp_linear_kernel to pick a MarlinLinearKernel instance, then builds three v2 parameter objects:

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)Afterwards DefaultModelLoader drives model.load_weights to split and write qweight / scales / qzeros from disk through the v2 loader. Once everything is loaded, ColumnParallelLinear calls quant_method.process_weights_after_loading, AutoAWQMarlinLinearMethod first converts AWQ packing to GPTQ style, then hands off to MarlinLinearKernel.process_weights_after_loading, which uses ops.gptq_marlin_repack to repack the weight into the layout expected by the Marlin kernel(marlin.py:124-137). During forward, apply calls self.kernel.apply_weights(layer, x, bias)(auto_awq.py:538-545), ultimately landing in Marlin's CUDA / Triton kernels.

Boundaries and failure modes ​

  • Platform does not support Marlin: AutoAWQMarlinLinearMethod.__init__ calls verify_marlin_supported(quant_type, group_size, has_zp=...) on non-CPU platforms and raises if the check fails(auto_awq.py:431-437).
  • group_size=-1: create_weights degenerates group_size to input_size, i.e. per-channel quantization(auto_awq.py:452-455).
  • fp8 block_quant + act_q_static: block quantization strictly requires dynamic activation (assert not self.act_q_static), otherwise it raises(fp8.py:367-368).
  • fp8 non-fp8 checkpoint: process_weights_after_loading goes through process_fp8_weight_tensor_tensor_strategy to quantize bf16 online, and per-tensor also has to reshape fused modules(fp8.py:416-429).
  • GPTQ row-parallel scale replication: marlin_repeat_scales_on_all_ranks decides whether scales are ChannelQuantScaleParameter (scale_dim=None, replicated on all ranks) or GroupQuantScaleParameter (scale_dim=0, sharded)(auto_gptq.py:367-440).
  • Phi-3 fused QKV: when QKV is already fused on disk there is no shard id, so QKVParallelLinear._load_fused_module_from_checkpoint splits it into three segments itself(linear.py:1049-1097), and quantization has to handle the packed_dim offset adjustment.
  • Marlin act_type is fp8: MarlinLinearKernel.process_weights_after_loading calls ops.marlin_int4_fp8_preprocess and multiplies scales by 512 when c.act_type == torch.float8_e4m3fn(marlin.py:98-103).

Summary ​

The quantization abstraction splits Config (reading checkpoint metadata) -> Method (building params + post-processing + delegating apply) -> Kernel (the actual computation) into three layers. AWQ / GPTQ / Marlin share MPLinearLayerConfig + choose_mp_linear_kernel, while fp8 goes through its own init_fp8_linear_kernel. It relies on tensor-parallel linear layers to inject the v2 loader into parameters at create_weights time, and the load flow is driven by DefaultModelLoader; the post-processed weights are ultimately consumed on the forward pass by the attention layers and the KV cache described by block tables.

See official docs: vLLM docs · README