首页
/ vLLM W8A8 量化 GEMM 的 CUTLASS Epilogue 设计：反量化融合的数学推导与源码实现

vLLM W8A8 量化 GEMM 的 CUTLASS Epilogue 设计：反量化融合的数学推导与源码实现

2026-09-05 14:23:36作者：戚魁泉Nursing
vllm
A high-throughput and memory-efficient inference and serving engine for LLMs
项目地址：https://gitcode.com/GitHub_Trending/vl/vllm

导读

本文基于 vLLM 仓库中 Epilogues.md 的设计文档，系统讲解 vLLM 的 W8A8（int8/fp8 权重量化 + int8/fp8 激活量化）GEMM 如何通过 4 种 CUTLASS Epilogue 把反量化操作融合进矩阵乘法，避免 GEMM 输出后再做一遍完整的反量化 kernel。读完本文，你将掌握：epilogue 融合反量化的数学原理（含零点对齐项 azp_adj/azp_with_adj 的推导）、4 种 epilogue 的适用条件与参数约定、对应 C++ 源码（EVT 访问者树）的实现细节、GPU 架构分派逻辑，以及 Python 侧如何离线预计算零点修正项。

一、背景：为什么要在 Epilogue 中融合反量化

W8A8 量化推理中，权重 B 和激活 A 都以 8 位整数（int8）或浮点（fp8_e4m3）存储，真正的数值是

  • 权重：B = s_b * B_hat（B_hat 为量化后的整数矩阵，s_b 为权重 scale）
  • 激活：A = s_a * (A_hat - J_a * z_a)（A_hat 为量化后的激活，z_a 为零点，J_a 为与 A 同维的全 1 矩阵）

若 GEMM 直接对 A_hat * B_hat 做整型矩阵乘，输出是"未反量化"的累加结果 D_hat，必须再乘上 scale 并修正零点误差才能得到最终输出。若把这一步单独作为一个 kernel，就需要把整个中间结果写回再读回，带来额外显存带宽开销。CUTLASS 的 GEMM 在写出主循环结果前有一个"尾声"（epilogue）阶段，vLLM 正是利用它把 scale 乘法和零点修正融合进 GEMM，使整条链路只产生一次最终输出。

vLLM 当前支持的量化形式（来自设计文档的约束）：

  • 权重只支持对称量化（零点恒为 0），可以是 per-tensor 或 per-channel；
  • 激活支持对称与非对称量化，可以是 per-tensor 或 per-token。

据此实现了 4 种 epilogue：

Epilogue 激活量化 bias 额外参数
ScaledEpilogue 对称 无（C = 0） 无
ScaledEpilogueBias 对称 支持（per-channel） bias
ScaledEpilogueAzp 非对称 per-tensor 支持 azp_with_adj
ScaledEpilogueAzpPerToken 非对称 per-token 支持 azp_adj、azp

文档明确说明：为了减小最终二进制体积，没有为"非对称激活 + 无 bias"单独提供 epilogue。当调用方不传 bias 时，内核使用 0 作为 bias，这会引入一次冗余加法（以及一次运行时判空检查），但性能影响很小。这一设计在源码中对应 RowOrZeroLoad 访问者（见下文第三节）。

二、底层线性代数推导

设 X_hat 为量化后的 X，则矩阵关系为：

A = s_a (A_hat - J_a z_a)
B = s_b B_hat
D = A B + C
D = s_a s_b D_hat + C

其中 D 是 GEMM 输出，C 是 bias；A（激活）支持非对称量化，B（权重）仅支持对称量化；s_a、s_b 分别为激活与权重的 scale，z_a 为激活零点，J_a 是与 A 同维的全 1 矩阵。要支持权重的非对称量化还需要另外的 epilogue（当前未实现）。

展开 AB 项：

AB      = s_a (A_hat - J_a z_a) s_b B_hat
AB      = s_a s_b (A_hat B_hat - J_a z_a B_hat)
D_hat   = A_hat B_hat - z_a J_a B_hat

两个关键观察：

  1. A_hat B_hat 就是整型 GEMM 的原始输出（主循环直接产出）；
  2. J_a B_hat 可以提前算出来——它的每一行都等于 1 B_hat，即 B_hat 的列和组成的行向量。

这就是后文 azp_adj（per-token 情形）与 azp_with_adj（per-tensor 情形）的数学来源：零点修正项中凡是只依赖权重 B_hat 的部分，都可以离线预计算为一个 per-channel 的行向量，运行时只需一次广播乘/减，避免在 epilogue 里做额外的矩阵级运算。

2.1 ScaledEpilogue（对称激活、无 bias）

此时 C = 0 且 z_a = 0：

D_hat = A_hat B_hat
D     = s_a s_b D_hat = s_a s_b A_hat B_hat

参数：

  • scale_a：激活 scale，可为 per-tensor（标量）或 per-token（列向量）；
  • scale_b：权重 scale，可为 per-tensor（标量）或 per-channel（行向量）。

2.2 ScaledEpilogueBias（对称激活、带 bias）

此时 z_a = 0：

D_hat = A_hat B_hat
D     = s_a s_b D_hat + C = s_a s_b A_hat B_hat + C

参数：

  • scale_a：per-tensor（标量）或 per-token（列向量）；
  • scale_b：per-tensor（标量）或 per-channel（行向量）；
  • bias：恒为 per-channel（行向量），长度等于输出通道数 n。

2.3 ScaledEpilogueAzp（非对称 per-tensor 激活、带 bias）

输出为：

D_hat = A_hat B_hat - z_a J_a B_hat
D     = s_a s_b D_hat + C
D     = s_a s_b (A_hat B_hat - z_a J_a B_hat) + C

由于 z_a 是标量，零点修正项 z_a J_a B_hat 的每一行都等于 z_a 1 B_hat，即一个行向量。该项在加载权重阶段就预先计算并存储在 azp_with_adj 中。

参数：

  • scale_a：per-tensor（标量）或 per-token（列向量）。实际上由于零点是 per-tensor 的，这里一般就是 per-tensor；
  • scale_b：per-tensor（标量）或 per-channel（行向量）；
  • azp_with_adj：预计算的零点修正项（z_a J_a B_hat），per-channel（行向量）；
  • bias：per-channel（行向量）。

使用前提：调用方必须离线预计算 azp_with_adj 并传给内核。

2.4 ScaledEpilogueAzpPerToken（非对称 per-token 激活、带 bias）

计算形式与 2.3 相同，但此时 z_a 是列向量（每个 token 一个零点），零点修正项 z_a J_a B_hat 变成 z_a 与 1 B_hat 的外积：

out = scale_a * scale_b * (Dq - azp_adj * azp) + bias

其中 Dq 是 GEMM 的原始量化输出。参数：

  • scale_a：per-tensor 或 per-token；由于零点是 per-token 的，这里一般为 per-token；
  • scale_b：per-tensor 或 per-channel；
  • azp_adj：预计算的修正项（1 B_hat，即 B_hat 的列和），per-channel（行向量）；
  • azp：零点 z_a，per-token（列向量）；
  • bias：per-channel（行向量）。

同样，调用方必须离线预计算 azp_adj 并传入内核。

三、源码实现：EVT 访问者树与四个 Epilogue 类

上述 4 种 epilogue 的 CUTLASS 2.x 实现位于 scaled_mm_epilogues_c2x.hpp（命名空间 vllm::c2x，"c2x" 即 CUTLASS 2.x）。该文件头注释说明：epilogue 必须包含一个名为 EVTCompute、类型为 Sm80EVT 的公开类型，以及一个构造 EVTCompute::Arguments 的静态 prepare_args 函数。

先对照命名：设计文档中的 ScaledEpilogueAzp / ScaledEpilogueAzpPerToken，在源码中的实际类名是 ScaledEpilogueBiasAzp / ScaledEpilogueBiasAzpToken（都带 "Bias" 字样，因为如前所述非对称激活路径统一走"bias 可选、缺省为 0"的模板）。

3.1 公共基类：广播加载访问者

ScaledEpilogueBase 为所有 epilogue 提供统一的加载描述符模板：

template <typename T>
using ColOrScalarLoad =  // per-token（列向量）或标量
    cutlass::epilogue::threadblock::VisitorColOrScalarBroadcast<...>;
template <typename T>
using RowOrScalarLoad =  // per-channel（行向量）或标量
    cutlass::epilogue::threadblock::VisitorRowOrScalarBroadcast<...>;
template <typename T>
using RowLoad = ...;   // 恒为行向量（bias、azp 修正项）
template <typename T>
using RowOrZeroLoad = ...;  // 行向量或 0（bias 可选时使用）

配套的 args_from_tensor 工具函数从 PyTorch 张量构造加载参数：对 ColOrScalar/RowOrScalar 类描述符，它会检查 tensor.numel() != 1 来决定按"向量广播"还是"标量广播"处理——这正是文档中"scale 可以是 per-tensor（标量）或 per-token/per-channel（向量）"的机制来源。针对可选 bias，重载版本接受 std::optional<Tensor>：为空时传 nullptr，内核端按常量 0 处理，对应文档所述"不传 bias 时用 0 代替，产生一次冗余加法"。

3.2 ScaledEpilogue / ScaledEpilogueBias

ScaledEpilogue 的 EVT 结构是一棵两层的乘法树：

EVTCompute = Compute(multiplies, ElementD)
                ├── ScaleA (ColOrScalarLoad<float>)
                └── EVTCompute0 = Compute(multiplies, float)
                                     ├── ScaleB (RowOrScalarLoad<float>)
                                     └── Accum (GEMM 累加器)

即先做 scale_b * Accum，再与 scale_a 相乘，最后转成输出类型 ElementD（bf16/fp16）。ScaledEpilogueBias 在此基础上把 Accum 换成 Bias 行向量加载，并把顶层运算换为 homogeneous_multiply_add（scale_a * (scale_b * accum) + bias），bias 类型为 RowLoad<ElementD>，即与输出同类型的 per-channel 行向量。

3.3 ScaledEpilogueBiasAzp（对应文档的 ScaledEpilogueAzp）

ScaledEpilogueBiasAzp 的类注释写明：它直接支持 int32 形式的 per-tensor 零点，azp_adj 参数是形状 (1, n) 的一维张量，等于 azp * (J @ B)——即文档中已乘以标量零点的 azp_with_adj。其 EVT 自内向外为：

  1. ComputeAzp = minus<float, int32_t>：float(accum) - azp_with_adj，两个操作数一个是 int32 累加器、一个是 int32 修正项；
  2. ComputeScaleB = multiplies<float>：乘 scale_b；
  3. ComputeScaleBiasA = homogeneous_multiply_add<ElementD>：乘 scale_a 并加 bias。

注意 bias 用的是 RowOrZeroLoad，因此该模板同时覆盖"有/无 bias"两种调用，与文档"无非对称无 bias 专用 epilogue"的设计一致。prepare_args 的四个张量参数顺序为 a_scales, b_scales, azp_adj, bias(optional)。

3.4 ScaledEpilogueBiasAzpToken（对应文档的 ScaledEpilogueAzpPerToken）

ScaledEpilogueBiasAzpToken 处理 per-token 零点。类注释给出了一个重要的工程动机：若把外积 azp * azp_adj 物化成矩阵需要 O(m·n) 空间；用秩-1 更新方式只在 epilogue 里做逐元素相乘，空间复杂度仅 O(m + n)——azp 是形状 (m, 1) 的列向量（ColLoad<int32_t>），azp_adj 是 (1, n) 的行向量（RowLoad<int32_t>）。其 EVT 计算链为：

azp * azp_adj            (multiplies, int32_t)
→ accum - (azp * azp_adj) (minus, float, int32_t)
→ × scale_b              (multiplies, float)
→ × scale_a + bias       (homogeneous_multiply_add, ElementD)

最终公式即文档给出的 out = scale_a * scale_b * (Dq - azp_adj * azp) + bias。

四、内核入口：形状、内存布局与 GPU 架构分派

这 4 种 epilogue 由统一的 cutlass_scaled_mm / cutlass_scaled_mm_azp 入口调度，入口实现见 scaled_mm_entry.cu。

4.1 接口约束（STD_TORCH_CHECK 逐项校验）

cutlass_scaled_mm（L197-L270）与 cutlass_scaled_mm_azp（L381-L457）在分派前做了以下校验，这些约束是使用该内核的硬性前提：

  • a、b、c 均为 2D，且维度满足 c[m,n] = a[m,k] @ b[k,n]；
  • a 与 c 行主序（stride(1) == 1），b 列主序（stride(0) == 1）——调用方需先把权重转置；
  • 16 字节对齐：c.stride(0) % 16 == 0 且 b.stride(1) % 16 == 0；
  • a_scales 的 numel 为 1（per-tensor）或 m（per-token）；b_scales 的 numel 为 1（per-tensor）或 n（per-channel）；二者必须连续；
  • azp 接口：azp_adj 必须为 int32、1D 连续、长度 n；azp（per-token 时）必须为 int32、1D 连续、长度 m；bias 若提供则长度 n、1D 连续，且 dtype 必须与输出 dtype 一致（见错误信息 "currently bias dtype must match output dtype"）。

按硬件分派：SM ≥ 120 走 sm120 内核，SM 100–119 走 sm100（Blackwell，CUTLASS 3.x），SM 90 走 sm90（Hopper），SM 89 走 Ada 专用路径，SM 80–88 与 SM 75 分别走 Ampere / Turing 的 2.x 路径（L222-L263）。cutlass_scaled_mm_azp 则支持 SM 90/89/80/75（L426-L449）。FP8 支持还有版本门槛：cutlass_scaled_mm_supports_fp8 要求 SM90 上 CUDA ≥ 12.0、SM89 上 CUDA ≥ 12.4（L145-L159）；block 量化 FP8 要求 SM100 上 CUDA ≥ 12.8、SM90 上 CUDA ≥ 12.0（L161-L174）。

4.2 2.x 路径中的 epilogue 选择逻辑

以 scaled_mm_c2x.cu 为例，epilogue 的选择完全由参数有无决定：

// cutlass_scaled_mm_sm80：对称路径（仅 int8，输出 fp16/bf16）
if (bias) {
    cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBias>(out, a, b,
        a_scales, b_scales, *bias);
} else {
    cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogue>(out, a, b,
        a_scales, b_scales);
}

// cutlass_scaled_mm_azp_sm80：非对称路径
if (azp) {   // 传了 per-token 零点列向量
    cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBiasAzpToken>(
        out, a, b, a_scales, b_scales, azp_adj, *azp, bias);
} else {     // 仅 azp_adj，即 per-tensor 情形
    cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBiasAzp>(
        out, a, b, a_scales, b_scales, azp_adj, bias);
}

（见 L102-L142。）也就是说，同一个 azp 入口函数用"是否传入 azp 张量"区分 per-tensor 与 per-token 两种非对称情形，与文档的两个 ScaledEpilogueAzp* 一一对应。sm89 路径（Ada，同时支持 int8 与 fp8_e4m3）还额外区分：int8 输入走 cutlass_gemm_sm89_int8_dispatch，fp8 输入走 cutlass_gemm_sm89_fp8_dispatch；当处于 batch invariance 模式（vllm_is_batch_invariant()）时会改走 fp8 batch-invariant 特化。sm75（Turing）仅支持 int8，且 azp 路径覆盖 75/80/89/90 各档。

Python 侧绑定见 vllm/_custom_ops.py 中的 cutlass_scaled_mm（L820 附近）与 cutlass_scaled_mm_azp（L871 附近）。

五、azp_adj 的离线预计算：Python 侧如何满足"先算好"的前提

设计文档强调"用户必须离线预计算 azp_with_adj / azp_adj 并传给内核"。vLLM 主引擎中这一前提由量化线性层内核在权重加载完成后（而非每次前向）自动满足。CutlassInt8ScaledMMLinearKernel.process_weights_after_loading 的关键步骤：

  1. 转置权重：Cutlass 内核要求 b 列主序，因此把量化权重 weight.t() 后放回层参数；
  2. scale 规整：Cutlass 内核只支持 per-tensor 与 per-channel 两种权重 scale。对融合模块（如 QKV、MLP，len(layer.logical_widths) > 1）且配置为 per-tensor 时，用 convert_to_channelwise 把 per-tensor scale 展开为 per-channel；
  3. azp_adj 预计算（对应文档"离线预计算"要求，代码注释也直接引用了本设计文档）：
# azp_adj is the AZP adjustment term, used to account for weights.
# It does not depend on scales or azp, so it is the same for
# static and dynamic quantization.
# For more details, see csrc/quantization/w8a8/cutlass/Epilogues.md
if not config.input_symmetric:
    weight = getattr(layer, w_q_name)
    azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32)  # 1 @ B_hat
    if config.is_static_input_scheme:
        # cutlass_w8a8 requires azp to be folded into azp_adj
        # in the per-tensor case
        azp_adj = getattr(layer, i_zp_name) * azp_adj             # z_a * (1 @ B_hat)
    setattr(layer, azp_adj_name,
           torch.nn.Parameter(azp_adj, requires_grad=False))

这段代码恰好实现了文档的两条数学约定：per-token 情形传的是 1 B_hat（权重列和，int32）；per-tensor 情形则把标量零点折进去，传 z_a * 1 B_hat（即 azp_with_adj）。注释还点明一个重要性质：azp_adj 只依赖权重，与 scale 和 azp 无关，因此静态/动态激活量化方案共用同一个预计算结果，只需在加载阶段算一次。静态非对称输入时，输入 scale 还会按 int8 数值范围与零点重新构造（range_max/range_min 公式，L83-L101），零点转成 int32 保存以匹配内核对 azp 的 dtype 要求。

六、测试验证

上述接口的行为由 tests/kernels/quantization/test_cutlass_scaled_mm.py 覆盖，该测试针对 cutlass_scaled_mm / cutlass_scaled_mm_azp 的对称与非对称路径做数值正确性校验，可作为修改 epilogue 或分派逻辑后的回归依据。相关的 FP8 量化工具函数见 csrc/quantization/w8a8/fp8/common.cuh 及 nvidia/quant_utils.cuh、amd/quant_utils.cuh。

七、小结

  • vLLM 的 W8A8 CUTLASS GEMM 通过 4 种 epilogue 把"乘 scale + 加 bias + 修零点"全部融合进矩阵乘法尾声，权重仅对称、激活可对称/非对称，粒度上支持 per-tensor/per-channel（权重）与 per-tensor/per-token（激活）的任意组合；
  • 非对称路径的核心技巧是把只依赖权重的零点修正项（J_a B_hat 的列和）离线预计算成 per-channel 行向量，per-token 情形进一步用秩-1 外积把空间开销控制在 O(m + n)；
  • 为控制二进制体积，非对称激活不提供"无 bias"专用 epilogue，而是统一走 bias 可选（缺省为 0）的 *BiasAzp* 模板；
  • 使用时的硬性前提是：激活行主序、权重复列主序、16 字节对齐、azp/azp_adj 为 int32、bias 与输出同 dtype，且按 GPU 架构（SM 75/80/89/90/100/120）与 CUDA 版本满足各分派路径的编译门槛。
vllm
A high-throughput and memory-efficient inference and serving engine for LLMs
项目地址：https://gitcode.com/GitHub_Trending/vl/vllm
登录后查看全文
热门项目推荐
相关项目推荐

项目优选

收起
kernelkernel
deepin linux kernel
C
33
18
ops-transformerops-transformer
本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。
C++
1.13 K
2.75 K
pytorchpytorch
作为 Ascend for PyTorch 社区的核心组件，TorchNPU 是昇腾专为 PyTorch 打造的深度学习适配插件，使 PyTorch 框架能够直接调用昇腾 NPU，为开发者提供昇腾 AI 处理器的超强算力。
Python
857
1.35 K
docsdocs
暂无描述
Markdown
897
5.8 K
kernelkernel
openEuler内核是openEuler操作系统的核心，既是系统性能与稳定性的基石，也是连接处理器、设备与服务的桥梁。
C
529
593
ops-nnops-nn
本项目是CANN提供的神经网络类计算算子库，实现网络在NPU上加速计算。
C++
915
1.83 K
jiuwenswarmjiuwenswarm
JiuwenSwarm 是一款基于openJiuwen开发的智能AI Agent，它能够将大语言模型的强大能力，通过你日常使用的各类通讯应用，直接延伸至你的指尖。
Python
3.58 K
1.01 K
ops-mathops-math
本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。
C++
1.35 K
1.46 K
cann-learning-hubcann-learning-hub
CANN 学习中心仓，支持在线互动运行、边学边练，提供教程、示例与优化方案，一站式助力昇腾开发者快速上手。
Jupyter Notebook
1.01 K
515
AscendNPU-IRAscendNPU-IR
AscendNPU-IR是基于MLIR（Multi-Level Intermediate Representation）构建的，面向昇腾亲和算子编译时使用的中间表示，提供昇腾完备表达能力，通过编译优化提升昇腾AI处理器计算效率，支持通过生态框架使能昇腾AI处理器与深度调优
C++
547
388