SGLang的AWQ量化笔记

文章目录

- 一、总体流程
- 二、代码讲解
- - [1）SGLang 中 AWQ 的配置层实现](#1）SGLang 中 AWQ 的配置层实现)
  - [2）awq sglang路径](#2）awq sglang路径)
  - [3）awq_marlin sglang路径](#3）awq_marlin sglang路径)
  - 4）端到端流程讲解

一、总体流程

AWQ 的完整生命周期可以分为两个阶段，它们之间的职责边界比较清晰。这种分工，是 AWQ 能够实现工程化部署的重要前提之一：

阶段	主要负责什么	最终产物
离线 AWQ 量化阶段	完成激活统计、显著通道识别、缩放保护搜索和量化参数生成。	生成量化后的 Checkpoint，其中包含 `qweight`、`qzeros`、`scales`，以及可能的通道缩放保护结果。
SGLang 运行时推理阶段	加载量化 Checkpoint，整理参数格式，并按后端选择合适的执行路径。	输出可供后续层继续使用的浮点激活表示，常见格式为 FP16 或 BF16。

如果把离线 AWQ 量化阶段再拆开看，它通常包含以下几个步骤：

步骤	发生了什么
Step 1	使用校准数据运行前向传播，收集各层激活的 Channel-wise 幅值统计。
Step 2	根据统计结果识别显著通道。
Step 3	对这些通道应用 Scaling 保护。
Step 4	执行 Per-group INT4 量化，并将量化结果打包存储。

这一阶段的重点在于完成激活统计、显著通道识别、缩放搜索和量化参数生成。SGLang 当前实现并不负责这套离线 AWQ 校准流程，而是消费其最终产物。

在 SGLang 的实现中，运行时阶段主要分成两条执行路径：

执行路径	理解	SGLang 在做什么
普通 AWQ 路径	先恢复权重，再进行矩阵乘法。	基于 `qweight`、`qzeros` 和 `scales` 执行 `awq_dequantize`，恢复计算所需的浮点权重表示，再执行 `torch.matmul`。
AWQ Marlin 路径	先重新排列量化数据，再调用高性能量化 Kernel。	权重加载完成后，将 AWQ 格式参数重新排列为适合 Marlin Kernel 的布局；前向时直接调用 Marlin 量化线性 Kernel。

二、代码讲解

1）SGLang 中 AWQ 的配置层实现

在 SGLang 中，AWQ 配置层的核心文件是 awq.py。这个文件中定义了三个关键配置类：

配置类	职责
`AWQConfig`	负责普通 AWQ 路径的配置解析与量化方法分发。
`AWQCPUConfig`	CPU 上的 AWQ 配置类，继承自 `AWQConfig`。
`AWQMarlinConfig`	负责 AWQ Marlin 优化路径的配置，以及自动切换到更高性能后端的判断逻辑。

AWQConfig 是普通 AWQ 路径的入口配置类。它负责从 HF/AWQ 格式的量化配置中读取参数，并在运行时为每个 Linear 层分配合适的量化方法。它的构造函数接收四个核心参数：

读取字段	用途
`w_bit` 或 `bits`	解析权重量化位宽。
`q_group_size` 或 `group_size`	解析量化分组大小。
`zero_point`	解析是否启用零点。
`modules_to_not_convert`	解析需要跳过量化的模块列表。

AWQMarlinConfig 是 AWQ Marlin 优化路径的配置类。它的作用不是改变 AWQ checkpoint 的基本格式，而是在运行时把 AWQ 权重组织成更适合 Marlin kernel 执行的形式，从而获得更高的推理效率。它和 AWQConfig 的关键区别主要体现在以下几个方面：

AWQ Marlin 与普通 AWQ 的差异

差异点	理解
仍面向 AWQ Checkpoint	仍然处理 `qweight`、`qzeros` 和 `scales` 这组 AWQ 参数。
布局会进一步重组	在后续执行链路中，将这些参数整理为更适合 Marlin Kernel 的数据布局。
激活数据类型受限	支持的激活数据类型为 FP16 和 BF16。
硬件要求更高	要求更高的 GPU Compute Capability，最低为 8.0。

AWQ Marlin 配置字段

读取字段	用途
`bits`	解析量化位宽。
`group_size`	解析量化分组大小。
`zero_point`	解析是否启用零点。
`lm_head`	决定 `ParallelLMHead` 是否也纳入量化路径。
`modules_to_not_convert`	解析需要跳过量化的模块列表。

其中，lm_head 会决定 ParallelLMHead 是否也纳入量化路径。override_quantization_method 是这一路径里的关键接口。它的作用是：在满足兼容条件时，把原本的 AWQ 路径自动切换成 awq_marlin，从而使用更快的执行后端。

需要注意的是，AWQMarlinConfig 的 TYPE_MAP 当前包含 4 和 8 两种位宽支持；而普通 AWQConfig 则只接受 4-bit 权重量化。这也是两者在配置约束上的一个重要差异。

是否支持awq_marlin主要是看
①硬件支持
②SGLang版本支持
③权重格式支持

用户输入场景	框架行为
不显式指定 `--quantization`	读取模型 Hugging Face 量化配置，并结合运行环境判断是否可以从 `awq` 自动切换到 `awq_marlin`。
显式指定 `--quantization awq`	固定使用普通 AWQ 路径，不执行 Marlin 自动升级。
显式指定 `--quantization awq_marlin`	优先尝试 Marlin 路径；若个别层不满足条件，可能回退到兼容实现。

2）awq sglang路径

qweight 使用 32-bit 整数打包多个 INT4 值，qzeros 与 scales 则按 group_size 分组组织

apply 方法可以按三段理解：

准备输入与输出形状。读取 qweight、scales、qzeros，根据 pack_factor 计算 out_shape，并把输入 x reshape 为二维张量 reshaped_x。
反量化权重。调用 awq_dequantize(qweight, scales, qzeros)，将 packed INT4 权重恢复为矩阵乘法可直接使用的浮点密集权重。
矩阵乘法与输出整理。调用 torch.matmul(reshaped_x, out) 完成计算；如果存在 bias，则通过 out.add_(bias) 原位相加，最后 reshape 回目标输出形状。

out = awq_dequantize(qweight, scales, qzeros) # ① 反量化
out = torch.matmul(reshaped_x, out) # ② 普通矩阵乘法

代价：dequantize-then-matmul 带来的额外开销
普通 AWQ 路径采用典型的 dequantize-then-matmul 两步策略：先把量化权重恢复为浮点密集权重表示，再执行标准矩阵乘法。这个设计使实现更直观，但也引入了几类开销。

开销类型	来源	影响
中间张量开销	反量化后显式构造浮点密集权重张量。	密集浮点权重体积明显大于低比特量化表示，会带来额外的临时显存压力。
反量化计算开销	展开 Packed 权重，并结合 `zero_point` 与 `scale` 恢复权重值。	Forward 前增加一段额外计算，尤其在频繁调用线性层时更加明显。
融合优化受限	执行链路被拆成"反量化"和"矩阵乘法"两个阶段。	难以像专门的量化 GEMM 路径一样，将更多计算和数据重用融合到同一个高性能 Kernel 中。

3）awq_marlin sglang路径

SGLang 为 AWQ 提供了一条更高性能的替代路径：awq_marlin。它在不改变 AWQ checkpoint 存储格式的前提下，通过加载时的权重重排（repack）以及前向阶段的 Marlin 专用 kernel，提供了区别于普通 AWQ 的优化执行路径。

出现的原因
Marlin 路径要解决的问题很明确：普通 AWQ 每次 forward 都要走 dequantize + matmul 两步，而 Marlin 希望把权重提前整理成更适合量化 GEMM 的执行布局，让 forward 阶段直接进入专用 kernel。
普通 AWQ 路径实现直接，但它的执行方式天然会产生额外开销。核心流程可以拆成两步：

先反量化。每次 forward 都需要把 packed INT4 权重恢复为浮点密集权重表示。
再矩阵乘法。反量化完成后，再把密集权重交给标准矩阵乘法执行。

特别点

真正面向 Marlin kernel 的布局调整并不发生在 create_weights(...)，而是留到权重加载完成后的 process_weights_after_loading(...) 中再做 repack 和重排。
两者对比

对比维度	普通 AWQ	AWQ Marlin
执行策略	反量化后执行 Dense Matmul，即 Dequantize-then-Multiply。	加载时完成 Repack，Forward 时直接执行 INT4 GEMM。
Forward 主线	`awq_dequantize → torch.matmul`	`apply_awq_marlin_linear → gptq_marlin_gemm`
权重转换时机	每次 Forward 时执行反量化。	加载后一次性完成 Repack / Permute。
中间内存开销	会生成 FP16 Dense 权重张量。	不需要生成完整的中间反量化权重张量。

AWQ Marlin把一次性的权重布局转换放在加载阶段，并把运行时的反量化融合进INT4 GEMM，而不是把反量化做到Repack里。

4）端到端流程讲解

流程主线

阶段	核心动作	关注点
配置解析	从 CLI、Hugging Face Config 或 ModelSlim Config 中确定 `quantization`。	条件兼容时，`awq` 可能自动升级为 `awq_marlin`。
模型初始化	`DefaultModelLoader` 获取 `quant_config`，并将其传入模型层。	`LinearBase` 保存量化配置，并调用 `get_quant_method()`。
权重生命周期	Scheme 创建 `qweight`、`qzeros`、`scales`；Loader 写入 Checkpoint 权重；加载完成后执行必要的后处理。	普通 AWQ 和 AWQ Marlin 的主要差异从权重后处理阶段开始体现。
Forward 执行	`Linear.forward()` 调用 `quant_method.apply()`，随后进入 `scheme.apply_weights()` 和底层 Kernel。	最终差异体现为"反量化后 Matmul"或"直接执行 INT4 GEMM"。

两条路径的差异

路径	执行方式	工程含义
普通 AWQ	保持原始 Packed 布局；每次 Forward 先执行 `awq_dequantize`，再执行 `torch.matmul`。	实现直观、兼容性较好，但会产生中间浮点权重和额外反量化开销。
AWQ Marlin	权重加载后执行 Repack、Permute 和 Zero Point 转换；Forward 直接进入 Marlin GEMM。	推理路径更高效，但依赖硬件能力、Shape 约束和加载后处理成本。

量化 Kernel 计算中三层职责的区分

层级	代表对象	职责
Config	`AWQConfig` / `AWQMarlinConfig`	解析量化配置、选择适用的量化方法，并根据硬件与层 Shape 判断是否进入 `awq_marlin` 等优化路径。
Method	`AWQLinearMethod`	作为量化线性层的统一方法入口：向上承接线性层生命周期；向下将 `create_weights()`、`apply()`、`process_weights_after_loading()` 分发给具体 Scheme。
Scheme / Kernel	`AWQLinearScheme`、`AWQMarlinLinearScheme` 与相关 Kernels	定义量化参数的创建方式、权重加载后的整理与重排逻辑，以及 Forward 阶段实际调用的 Kernel 计算路径。

调用链

复制代码

CLI / HF Config / ModelSlim Config
    ↓
AWQConfig / AWQMarlinConfig
    ↓
DefaultModelLoader
    ↓
LinearBase.get_quant_method()
    ↓
AWQLinearMethod
    ├── create_weights()
    ├── process_weights_after_loading()
    └── apply()
          ↓
AWQLinearScheme / AWQMarlinLinearScheme
          ↓
普通 AWQ：awq_dequantize → torch.matmul
AWQ Marlin：Marlin INT4 GEMM