企业官网建设流程全解析

上个月接了个活儿，要把一个权重全部收敛到 {-1, 0, +1} 的三元量化模型塞进 vLLM 里跑起来。拿到手才发现，事情远没有“改个 dtype、换个 weight_loader”那么简单：checkpoint 里的权重是 2bit 打包的紧凑格式，vLLM 原生算子完全不认识，直接加载就是一堆形状错误和乱码输出。我只好把整条路重新走了一遍，从逆向权重格式开始，到逐层对账验证输出，最后写 CUDA kernel 做性能优化，折腾了将近两周才跑稳。这篇就把完整的踩坑和实现过程记录下来，给同样要接自定义低比特量化的同学做个参考。

1. 项目背景与整体思路拆解

1.1 三元量化模型是什么，为什么值得折腾

三元量化（Ternary Quantization）就是把神经网络的权重强制约束到三个值：-1、0、+1。和常见的 int8 / int4 量化不同，它不需要乘法器，矩阵乘法的本质从“乘累加”退化成了“符号加减”，0 的位置还能直接跳过。从理论上讲，这是极端低比特量化里仅次于二值量化的一种方案，内存占用能压得非常低。

但代价也很明显：表达能力受限，训练不好精度掉得惨。我手上这个模型是内部同学用蒸馏方式训出来的，参数量约 3B，权重在三值化之后仍保留了 per-channel 的 scale 因子，用于缓解数值范围丢失的问题。也就是说，它不是简单地把一个 FP16 模型“硬截断”成三值，而是专门训练的量化模型，精度还算能看，跑评测比 FP16 原版只低了两三个点。

这里稍微强调一下：标题里的“逆向”不是破解什么闭源算法，而是对 checkpoint 序列化格式做反向解析。模型权重在磁盘上是字节流，用什么编码存、怎么解包、scale 放在哪、打包维度是哪个轴，这些信息不会写在使用文档里，需要一点点抠出来。类似的过程在做端侧模型移植时很常见，把模型文件当黑盒打开，从字节层面搞懂它内部的规定。

1.2 为什么选 vLLM 而不是 HuggingFace 直接跑

答案是吞吐量和显存。单纯用 HuggingFace transformers 能跑，但三值权重即使算得再快，也绕不过 PyTorch 动态图的调度开销和缺乏连续批处理的问题。vLLM 的 scheduler 会把多个请求拼成连续批（continuous batching），配合 PagedAttention 做 KV cache 分页管理，在服务场景下吞吐差距非常大。

vLLM 的问题在于“支持的模型列表”是写死的。它对常见模型的架构、权重切分方式、量化格式做了高度定制，一旦遇到自定义量化格式，内部的 Linear 层、attention 实现、weight_loader 全都对不上。把三元权重塞进去，相当于要把vLLM 的模型层做一次整容：保留它的调度和显存管理优势，替换掉我们的算子路径。

所以整个项目天然分成三段：逆向搞清楚“权重到底是什么布局”，对账解决“换了实现之后输出还对不对”，kernel 优化解决“就算对了但跑得够不够快”。三者有严格的先后关系，跳过任何一步，后面都很难定位问题。

1.3 三步走的全流程设计

我的计划是这样的：

• 逆向：扫描 checkpoint 里的 tensor 名称、形状、dtype 和值域，推导出打包格式和 scale 对应关系。先写一个 Python 解包函数，把打包权重还原成逻辑上的三值矩阵。
• 对账：不直接上 vLLM，先用 PyTorch 搭一个参考实现，加载解包后的权重跑前向，把每一层中间 tensor 保存下来。之后再把同一份输入喂给 vLLM 自定义建模路径，逐层比较差异。
• kernel 优化：有了正确的输出基线，再写 CUDA kernel。先写朴素的、性能不行的版本，保证输出和参考实现一致；再优化内存读取和解包逻辑，最后接进 vLLM 的算子调用点。

为什么要这么绕？因为 GPU kernel 一旦写错，表现往往不是报错，而是“结果悄悄偏了一点点”或者干脆“有时候对有时候错”，比如位序反了、scale 没乘、维度没对齐，这些用肉眼很难看出来。先建立对账基线，再动 kernel，出了问题能快速定位是哪个环节引入的。

2. 逆向：把权重布局和模型结构“抠”出来

2.1 先读配置，再扫权重，不急着写加载

拿到 checkpoint 之后第一件事不是写加载代码，而是把结构摸清楚。我习惯先打开 config.json 和模型目录里的文件清单，确认这几个信息：

• 模型的 model_type 是什么，对应 transformers 里哪个类；
• 是否有 quantization_config 字段，量化框架是自己写的还是现成的；
• snapshot 文件是 safetensors 还是 pickle（安全起见一定要用 safetensors）；
• 每个权重 tensor 的 shape、dtype、key 名称规律。

最快捷的排查方式是用一个小脚本把所有 tensor 的值域打出来，找“只有几个唯一值”的权重。这一步基本能确认哪些层做了三元量化、哪些层还保留高精度。

from safetensors import safe_open import torch fname = "model-00001-of-00002.safetensors" with safe_open(fname, framework="pt", device="cpu") as f: for key in f.keys(): tensor = f.get_tensor(key) uniq = torch.unique(tensor) if uniq.numel() <= 3: print(key, tensor.shape, tensor.dtype, uniq.tolist())

我这边扫出来的结果非常典型：attention 里的 q_proj、k_proj、v_proj、o_proj，还有 MLP 里的 gate_proj、up_proj、down_proj，权重文件都是 uint8 类型，唯一值只有 0、1、2 三个编码数字；而 layer_norm 的 weight、bias、position_ids 以及最后的 lm_head 仍是 float16 或 float32。这说明三值化只作用在线性层的权重矩阵上，没有动归一化层和输出层。

顺带说一下，为什么走“值域扫描”而不是直接看代码：很多自定义量化框架会自己定义权重参数名，比如weight_ternary，或者在保存前做了 pack，导致 PyTorch 加载后的 shape 和逻辑 shape 对不上。靠配置文档摸不出来，直接看唯一值是最不容易误判的方式。

2.2 解开打包格式：从 uint8 到 -1/0/+1

扫描结果里唯一值是 0、1、2，并不是最终的三值语义。再结合 shape 特征判断，比如一个逻辑 shape 为 (1024, 2048) 的矩阵，在文件里实际存成了 (1024, 512) 的 uint8，那基本可以断定是按“每个字节存 4 个 2bit 权重”的方式打包的。

我需要先确定编码表。常见的两种编码方式：

• 0 表示 0，1 表示 +1，2 表示 -1，3 表示无效/未使用；
• 0 表示 -1，1 表示 0，2 表示 +1，3 表示填充。

这两种表反了会对账时立刻暴露出来——输出完全乱掉。我习惯写一个“暴力解包 + 常量检查”的脚本，把解包后的权重传给一个简单的线性层前向，先看输出数值范围是不是正常，再进逐层对账流程。

下面是我用的解包函数，按字节拆分 4 个 2bit 编码：

import torch def unpack_ternary_packed( packed: torch.Tensor, out_dim: int, in_dim: int ) -> torch.Tensor: """把 (out_dim, in_dim // 4) 的 uint8 解包成 (out_dim, in_dim) 的 int8.""" packed = packed.reshape(out_dim, in_dim // 4).to(torch.uint8) out = torch.empty(out_dim, in_dim, dtype=torch.int8) # 编码：0 -> 0, 1 -> +1, 2 -> -1, 3 -> 当作0 for i in range(4): code = (packed >> (2 * i)) & 0x3 values = torch.where( code == 1, torch.tensor(1, dtype=torch.int8), torch.where( code == 2, torch.tensor(-1, dtype=torch.int8), torch.tensor(0, dtype=torch.int8), ), ) out[:, i::4] = values return out

解包之后有一个关键动作：统计每个权重矩阵中 0 的比例。三值模型里 0 的比例往往不低，这个数值直接影响后面 kernel 的设计方向。如果 0 的比例超过 50%，可以做稀疏跳过；如果只有百分之几，那么“跳过 0 收益”的意义就不大，反而白白增加分支判断。我这边统计下来，大部分线性层的 0 占比在 20%-30% 之间，偏分布均匀，这个比例更适合把重点放在“内存翻倍拆包 + 符号选择”，而不是走稀疏索引。

2.3 用 CPU 参考实现验证逆向结果

逆向做到“解包没错”还不够，必须验证解出来的权重能和模型结构对上。在写 vLLM 接入代码之前，我先做一件省事的事情：把 transformers 里对应模型的 modeling 代码复制一份出来，把里面的 nn.Linear 全部替换成“解包后权重 + torch.matmul 或 einops”。

参考实现的核心代码很简单，关键是检查每个层的输入输出维度：

class TernaryLinear(nn.Module): def __init__(self, packed_weight, scale, out_dim, in_dim): super().__init__() self.weight = unpack_ternary_packed(packed_weight, out_dim, in_dim) self.scale = scale # (out_dim,) 或 (out_dim, in_dim // block) def forward(self, x): # x: (batch, seq, in_dim) w = self.weight.float() y = torch.einsum("bsi,oi->bso", x, w) return y * self.scale

跑通前向之后，把每一层的输出都保存成 pt 文件，后续所有实现都以这份输出为“上帝视角”。这里有一个很重要的取舍：参考实现里的 scale 位置、样本 tokenizer 输出、logits 的 softmax 前数值，都必须固定下来，否则之后对账没有稳定基准。我直接把输入文本固定成 4 条业务数据，截断到 64 token，避免 batch 和序列长度带来的额外变量。

2.4 逆向阶段容易踩的坑和检查清单

逆向阶段最常见的坑有三个：

第一个是打包位序。上面解包代码里用>> (2 * i)取低位，但有的框架会按高位字节序存，比如>> (6 - 2 * i)。如果一个模型是 A 框架训的、B 框架存的，极易搞混。我的经验是先取前几个字节手动算一遍编码，肉眼确认解包结果与语义匹配。

第二个是 scale 的归属。三元权重矩阵本身只给方向，不给大小，幅度信息全靠 scale。scale 可能是 per-channel（每个输出维度一个值）、per-token（每个 token 一个值）或者 per-block（一个 block 共享一个值）。我这边模型用的是 per-channel，处理起来还算简单；如果是 per-block，对账时必须把 block 大小搞准确，否则输出会呈周期性偏差。

第三个是权重切分。vLLM 的模型代码里经常把一个大矩阵切成 q/k/v 三份，或者把 up/gate 两份合并。从 checkpoint 原始 key 看到的是拆开的，但 vLLM 内部可能又有自己的一套 merge 逻辑。我的习惯是在逆向阶段就列一张“原始权重 key -> vLLM 模型参数名 -> 是否需要切分/合并/转置”的表格，一行一行比对清楚，后面写 weight_loader 时照着这张表直接搬，省去反复试错的时间。

逆向完成后，我应该得到这样一组信息：每个三值权重矩阵的逻辑形状、打包后的实际形状、scale 形状、编码表、非零比例统计、以及与 vLLM 参数名的映射关系。这部分看似不起眼，实际上决定了后面所有 kernel 和 loader 是否顺利。

3. 对账：让 vLLM 输出和参考实现达成一致

3.1 对账的本质：给所有中间结果一个“统一口径”

对账是我觉得整个项目里容易卡住的地方，因为它不是简单比两个精度数值。三值量化模型的“正确输出”是相对概念。拿它和 FP16 原模型比，任何一层都会有肉眼可见的误差；真正有意义的是——同一套量化权重，我们用 PyTorch 参考实现算出来的结果，和使用 vLLM 自定义算子算出来的结果，它们之间是否一致。

所以我把对账基准定成：

• 参考实现：加载解包后的权重，使用理论上“最准确”的高精度计算（例如 fp32 中间累加）。
• vLLM 实现：加载同一份打包权重，走我准备嵌入的 CUDA kernel 或 initial naive kernel。
• 判定方法：逐层比较输出 tensor，要求最大绝对误差和相对误差都达到可接受阈值。

这个基准在数学上看就是一个“控制变量”实验：权重相同、输入相同、框架不同，差异若能被控制到 1e-3 以内，就说明算子语义是对的。

3.2 搭建可复现的对账环境

对账环境的关键是“固化一切变量”。我做了一个独立目录，专门放参考输出、输入样本和对比脚本。步骤固定如下：

1. 用固定的随机种子生成输入 hidden_states，或者直接固定输入文本后走 tokenizer 拿到 input_ids；
2. 用参考实现跑一次前向，把以下中间结果全部保存：

• 每个 q/k/v 线性层输出；
• attention 输出（经过 softmax 和 value 加权之后）；
• MLP gate/up 投影输出；
• 每个 decoder layer 的最终输出；
• 最后一层 logits 与 top-10 索引。

保存中间结果时有几个细节：tensor 要 clone 成独立张量，不能只是视图引用；文件命名建议带层级名称，比如layer0_attn_q.pt、layer0_mlp_out.pt。

3. 在 vLLM 自定义模型里同样插入 hook，把同名中间 tensor dump 出来，和参考输出做torch.allclose，同时打印每个层的最大绝对误差。如果是大规模张量，建议用(a - b).abs().max().item()这种精确指标，而不是看 mean，因为均值会把零星大误差“洗白”。

3.3 逐层对账：算子级 -> 模块级 -> 模型级

我强烈建议不要一上来就比整个模型输出，那样分不清误差来自哪个层。我在实践里按三层漏斗方式做：

第一层是“算子级对账”。只比单个线性层的输出。做法是把同一份输入张量分别喂给参考实现里的 TernaryLinear，以及 vLLM 里对应的自定义 Linear 类。看输出的最大误差。

第二层是“模块级对账”。把整个 decoder layer（含 attention、MLP、残差、RMSNorm）作为 unit，比较整个 transformer block 的输出。这一层最容易暴露出 scale 切分、维度排布、残差路径等拼接问题。

第三层才是“模型级对账”。整模型前向跑一遍，比较最终 logits。模型级都对齐到 1e-3 以内，才说明可以进入 kernel 优化。如果模型级有明显偏差，比如 top-10 乱掉，说问题出现在早期层并且一直累积放大，要靠模块级对账往回追。

对账的容差阈值不是越严越好。考虑到 GPU 上 fp16 和 fp32 的浮点误差，以及 softmax 等非线性放大的情况，我一般把rtol=1e-3, atol=1e-3设为默认，遇到 attention 层可以放宽到 1e-2。重要的不是单层完全相等，而是“每一层误差都没有跳跃性放大”。如果第 1 层误差 1e-6，到第 12 层变成max diff = 2.0，这就不正常了——大概率是某个残差或归一化处理环节没对齐，而不是单纯浮点误差累积。

3.4 对账遇到的高频问题速查

对账时常见的 diff 现象和原因我整理成了表格，每次遇到直接对号入座：

对账脚本一旦通过，我会把参考输出文件和比对代码一起提交到仓库，后面每次改 kernel 都重新跑一遍这个回归。不保留回归基线，后面 kernel 优化很容易犯“改一处、坏一处”的错误。

4. Kernel 优化：把三元权重转化为带宽收益

4.1 先写功能正确的朴素 kernel，再谈优化

很多人一想到自定义量化就直奔高性能 kernel，实际最容易翻车。我习惯先写一个“能跑但性能一般”的朴素 CUDA kernel——一个线程算输出矩阵的一个元素，边解包边计算。这样的 kernel 即使性能不行，但逻辑清晰，对它做性能剖析和错误排查都会容易得多。

朴素 kernel 的伪码大致这样：

// y[m][n] = sum_k x[m][k] * w[n][k] * scale[n] // w_packed shape: (N, K / 4) 的 uint8，每个字节存 4 个 2bit 权重 __global__ void ternary_linear_naive( const float* __restrict__ x, // (M, K) const uint8_t* __restrict__ w, // (N, K / 4) const float* __restrict__ scale, // (N,) float* __restrict__ y, // (M, N) int M, int N, int K ) { int row = blockIdx.y * blockDim.y + threadIdx.y; int col = blockIdx.x * blockDim.x + threadIdx.x; if (row >= M || col >= N) return; float acc = 0.0f; const float* x_row = x + row * K; const uint8_t* w_row = w + col * (K / 4); for (int k = 0; k < K / 4; ++k) { uint8_t code = w_row[k]; #pragma unroll for (int i = 0; i < 4; ++i) { int c = (code >> (2 * i)) & 3; float wval = (c == 1) ? 1.0f : ((c == 2) ? -1.0f : 0.0f); acc += wval * x_row[k * 4 + i]; } } y[row * N + col] = acc * (scale ? scale[col] : 1.0f); }

这段代码能跑通，性能大约只有 cuBLAS FP16 的 1/10 甚至更低。因为每个线程的 for 循环里既有位运算又有条件分支，写入全局内存也很粗糙。但它的存在意义是产出“可信输出”，用于对账验证。

4.2 位打包与掩码技巧：把三值计算变成加减法

写完朴素 kernel，对账通过后，开始优化。优化的核心思路是：三元权重带来的优势不在 FLOPs，而在内存带宽和“乘法变加减”。

把权重拆成两个 bitmask 是常见技巧：

• pos_mask = 1表示该位置是 +1；
• neg_mask = 1表示该位置是 -1；
• 两者都为 0 表示该位置是 0。

那么矩阵乘法可以变换成：

y = sum(x * pos_mask) - sum(x * neg_mask)

这意味着原先每个权重需要判断三态（0/+1/-1），现在只需要做两次“与 mask 后的求和”。如果配合 popcount 指令，甚至可以进一步加速：sum(x & mask)在支持向量化 popcount 的硬件上可以一次处理多个通道，但 popcount 用的是整数运算，如何保持浮点累加精度需要仔细设计。我这边当前版本没有走到纯 popcount 路线，而是先做“按位解包 + 符号选择”，已经比朴素版快好几倍。popcount 方式适合激活值本身也能被阈值化的极端场景，普通线性层直接上会引入额外的比特级拆分复杂度。

2bit 打包之后的显存收益非常直观：FP16 一个权重占 2 字节，2bit 打包后每个权重只占 0.25 字节，对应约 8 倍的权重内存减少。一个 3B 模型，原本 FP16 的权重文件接近 6GB，三值打包后权重部分大约 0.75GB——在显存紧张的推理场景里这个优势极其明显。

4.3 CUDA kernel 优化要点与示例

实际优化时，我按下面几个方向逐个推进：

第一是向量化内存读取。权重按 uint8 读很亏，最好一次读一个 uint32，然后拆成 4 个 uint8 处理。激活侧用float4读取，一次处理 4 个 float，能够明显提升内存吞吐。调整前后的代码结构对比很直观，但要注意对齐要求：权重按 uint32 读取时，指针地址必须 4 字节对齐；激活按 float4 读取时，K 维度和指针要对齐到 16 字节，否则会直接 illegal memory access。

第二是减少重复读取。每个线程如果自己从全局内存取同一份激活向量，会导致大量冗余访存。经典做法是把激活 tile 搬进 shared memory，权重矩阵做 block 分块。虽然三元权重解包仍需要计算，但全局内存访存少了，整体带宽瓶颈就缓解了。

第三是控制分支。一开始的朴素 kernel 每读一个字节做 4 次条件判断，分支开销非常大。优化后可以把“抽取 2bit 编码”和“根据编码选择符号”拆成更紧凑的算术运算，例如用(int8_t)(code << 6) >> 6之类的位操作直接把 2bit 映射到 -1/0/1，避开分支。这个技巧在不同的 GPU 架构上表现略有差异，实测下来在 A100 上收益明显。

下面是我优化后的一个 kernel 片段，重点是 uint32 批量读取和符号展开：

__device__ __forceinline__ float ternary_relu(int mask_sign) { // 把 2bit 编码映射为 -1/0/+1 的一种紧凑做法 // 具体映射关系与打包格式强绑定，这里只展示思路 } __global__ void ternary_linear_vec( const float4* __restrict__ x, // (M, K/4) const uint32_t* __restrict__ wpack, // (N, K/16) const float* __restrict__ scale, float* __restrict__ y, int M, int N, int K ) { int row = blockIdx.y * blockDim.y + threadIdx.y; int col = blockIdx.x * blockDim.x + threadIdx.x; if (row >= M || col >= N) return; const float4* x_row = x + row * (K / 4); const uint32_t* w_row = wpack + col * (K / 16); float acc = 0.0f; for (int i = 0; i < K / 4; i += 4) { float4 xv = x_row[i / 4]; uint32_t wp = w_row[i / 4]; // wp 拆成 16 个 2bit 权重 #pragma unroll for (int j = 0; j < 4; ++j) { int c0 = (wp >> (2 * j)) & 3; int c1 = (wp >> (8 + 2 * j)) & 3; acc += wval(c0) * xv.x; acc += wval(c1) * xv.y; } } y[row * N + col] = acc * (scale ? scale[col] : 1.0f); }

当然，上面的写法还有进一步调优空间。实际落地时我又加了 shared memory tile、每个线程计算多个输出列以减少索引开销、限制 block 内的线程数量以避免占用率过高。

优化的核心原则是：不要迷信某一个“神优化”，每个改动都要用 ncu profile 看一遍访存比例和计算吞吐，针对瓶颈再动刀。

4.4 把自定义 kernel 挂进 vLLM

kernel 本身跑通之后，下一步是把算子接到 vLLM 的模型层上。vLLM 在 0.6.x 版本里主要通过LinearMethodBase和模型类的weight_loader来做量化扩展，流程大致如下：

1. 定义TernaryLinearMethod(LinearMethodBase)，在create_weights里注册参数packed_weight和scale；
2. 实现process_weights，把 HuggingFace 原始权重转成 pack 后的 uint8 格式；
3. 实现apply，在前向时直接调用torch.ops.ternary_linear这个自定义 op；
4. 在模型类的__init__里，把原来的QKVParallelLinear替换成使用新 method 的自定义 Linear；
5. 用torch.library把 CUDA kernel 封装成自定义 PyTorch op，并用opcheck做一遍 shape 和 striding 验证。

vLLM 版本之间 API 差异很大。0.6.x 和 0.7.x、0.8.x 的LinearMethodBase方法签名都不一样。我这边锁的是当时项目仓库里的 0.6.3 镜像，后续如果要升级，必须先看新版本的linear_method接口再迁。这里也建议所有接 vLLM 定制量化的人都把镜像 tag 和 vLLM 版本固写在文档里，避免同事后期无脑升级打碎一切。

还有一个很隐蔽的问题：vLLM 对模型输出的 shape 有强假设，尤其是 hidden_states 的 batch 维度和 attention mask 的写法。自定义 kernel 只需要保证输入输出 shape 和原来的 Linear 层一致，其他交给上层的 transformer block 处理。我踩过的一个坑是：我一开始在 apply 方法里直接对二维 hidden_states 写 matmul，但 vLLM 在某些情况下会传 3D 张量（比如处理 prefix caching 或 cross attention），导致 shape 对不上。解决办法很简单——在 apply 里把输入reshape(-1, K)，算完再恢复原形状，别让 kernel 层假设什么 3D/2D 结构。

真正接入 vLLM 之后，vLLM 的 scheduler 和 CUDA graph 不会因为你的 kernel 而改变行为。scheduler 只负责管理 token 序列的启停、KV cache 分配和连续性；自定义 kernel 只要按照标准 Linear 的输入输出契约来，就不会干扰调度。但有一个例外：如果你的 kernel 内部有 host 端的 if/else 或者动态分配显存，在开启 CUDA graph 捕获时可能直接报错或产生不可预期行为。所以自定义 op 的函数体里尽量只写固定 shape、固定 launch 配置，不要搞“运行时判断 batch 大小然后切不同 kernel”——这种逻辑放到 apply 方法里，用 Python 层判断，避免破坏 graph capture。

4.5 实测性能记录

写完优化版 kernel 之后，我做了完整的性能测试。测试环境是单卡 A100 80G，模型为内部蒸馏的三元量化模型（约 3B 参数，28 层 transformer，hidden_size 2048）。

单项数据：

• 权重显存：FP16 6GB → 三元打包 0.8GB，下降约 86%；
• 单层 QKV 线性层平均耗时：朴素 kernel 1.4ms → 向量化优化后 0.4ms → cuBLAS FP16 0.7ms；
• 端到端吞吐（batch=32，seq_len=512，最大生成长度 256）：相比把权重解包回 FP16 后跑 cuBLAS，三元 kernel 大约快 30%-40%；相比直接跑 FP16 模型，显存和长文本场景收益更明显；
• 小 batch（batch=1，seq_len=32）场景，三元 kernel 反而比 cuBLAS 慢，主要因为 kernel launch 和位解包开销占比太大。

所以最终在 vLLM 的 apply 方法里做了分支：batch*seq_len 较小时走 cuBLAS/标准高精度路径，只有长序列、大 batch 场景才切换到三元 kernel。不要觉得“自定义 kernel 一定更快”，实际是否要用、什么阈值切，必须拿数据说话。

5. 踩坑实录与排查技巧

5.1 vLLM 集成时的版本与接口问题

vLLM 的版本接口变动是接自定义量化的第一大坑。我刚开始照着某个 0.7.x 版本博客写代码，结果在 0.6.3 里create_weights的参数签名完全不一样，力白费了一半。后面学乖了，直接从本地 vLLM 源码里找LinearMethodBase的抽象方法，对着当前源码的签名写实现。

还有一个小坑是 vLLM 的weight_loader是按原始 HuggingFace 参数名来的。我自定义了packed_weight这个参数名之后，默认 loader 根本不会去读 checkpoint 里的q_proj.weight_ternary这种 key，必须在模型类里覆写load_weights或者使用自定义 load 函数，把 torch 的 state_dict key 和 vLLM 权重一一映射起来。这段代码最容易写成“隐藏逻辑”，建议所有 mapping 都写成显式字典，留着注释，不然过两周自己都看不懂。

5.2 CUDA kernel 的常见崩溃与静默错误

kernel 阶段的报错分两类：一类是直接崩溃，另一类更可怕——没有任何 crash，但输出悄悄不对。

直接崩溃最常见的是 illegal memory access。通常来自三种情况：指针未对齐（float4 读了一个奇数偏移地址）、网格边界没写 full check、以及 pack 维度算错导致越界读。崩溃之后先用 compute-sanitizer 跑一遍，它能精确定位越界的内核和访问地址。

静默错误最阴险的是“编码表只错一半”。比如我把 2bit 编码 2 当成 -1，而实际应该当成 0，覆盖 vol 时会发现大张量整体误差很小，但某一列会周期性出现两倍偏差。我最后定位是靠对账脚本的逐层 dump，看到第 3 层的输出里每个 hidden 维的误差呈现 block 规律，才反推回编码表问题。所以再次强调，对账回归不能省。

5.3 量化模型特有的对账问题

三值量化模型和普通 FP16 模型的对账有一个额外难点：scale 本身可能是低精度存储的。如果训练阶段把 scale 存成 fp16，加载时再转 fp32，和原始 fp32 scale 会有细微差别，多层层累加后也可能造成可见误差。解决方式是在对账脚本里显式指定 scale 加载 dtype，然后用同一份 scale 喂给参考实现和 vLLM 实现，避免把“模型本身的精度损失”误判成“算子误差”。

另一个问题是 RMSNorm 的 epsilon。vLLM 和 transformers 的 RMSNorm 实现细节不同，epsilon 默认值也可能不同，三值权重 + 低精度计算会把 1e-6 级别的 epsilon 差异放大到不可忽略的程度。对账时遇到“越往后偏得越多”的情况，先检查 layer_norm 部分的 epsilon 和计算 dtype 是否完全一致。

5.4 一些工程习惯建议

项目做完之后，最大的体会是工程规范很重要。我把下面这几件事固化成了仓库里的 checklists：

• 权重格式文档必须在项目第一天建：编码表、pack 维度、scale 规则、维度映射关系，全部白纸黑字写下来；
• 对账脚本必须可一键运行，且纳入 CI：每次 kernel 改动、vLLM 版本升级都要重新跑；
• 保留“失败现场”：曾经跑出过偏离的中间输出文件不要覆盖，留一个 snapshot 目录，方便事后对比定位回归是哪次改动引入的；
• 小步快跑：kernel 改动一次只动一个优化点，提交信息里写清楚改了什么以及对应的 profiling 数据，不要一次堆三四个 optimization 再一起调。

这些习惯不复杂，但能让你在踩坑之后快速复原，不用重头再来一遍。

6. 这个项目后续还能怎么扩展

如果后续继续深入，我觉得有三条可以走的方向。第一条是把 2bit 打包升级成更彻底的角色分离：权重完全转成pos_mask/neg_mask两个 bitmask，激活也按 batch 做一次量化，利用 popcount 类指令把核心计算变成纯位运算，理论上还有一到两倍的性能空间。第二条是加一层离线 kernel 缓存：因为三元模型权重固定，calibration 完的解包/索引表可以离线生成，运行时不再重复做 unpack，对服务场景的冷启动也有优化。第三条是把这套流程抽象成更通用的“任意低比特量化模型接入 vLLM 脚手架”，不只限于三元，int2、混合精度的线性层也可以用同一套对账和 loader 体系。

我个人最深的体会是：把量化模型接进 vLLM 这类推理框架时，写 kernel 往往是最前端的技能，但真正决定项目成败的反而是逆向和对账这两个容易被低估的环节。权重格式没吃透，kernel 写再好也是白搭；对账基线没留好，后面每次改算子都战战兢兢。把这两个步骤做扎实，后续接任何自定义量化模型都会顺畅很多。