量化与矩阵乘内核

这是 colibrì “内核层”的核心：所有权重都以极低精度存储（int8/int4/int2），推理时用到哪一行才反量化（dequant-on-use），从不把整个反量化矩阵物化到内存。decode 单 token 时还有一条把激活也量化成 int8 的整数点积（IDOT）快路径，量化矩阵乘快 2–3 倍。

1. QT：量化张量类型

引擎里所有二维权重都用 QT 结构（quantized tensor）表示：

typedef struct {
    int fmt; float *qf; int8_t *q8; uint8_t *q4; float *s; int O, I;
    ... // CUDA 字段
} QT;
  • fmt 是格式选择器：0=F32、1=INT8、2=INT4（2 个/字节）、3=INT2（4 个/字节）。
  • qf/q8/q4 分别是三种存储指针；注释说明 q4 同时容纳 int4 和 int2 的 packed 缓冲。
  • s 是每行一个 f32 缩放系数（长度 O）——这是贯穿全引擎的 “per-row scale” 方案。
  • O, I 是输出/输入维度，逻辑形状 [O, I]，矩阵乘算 x @ W^T。

qt_bytes（76 行）按格式返回常驻字节数，是内存预算的基础。qt_alloc（569 行）根据 bits 选格式：≥16→F32，≥5→int8，≥3→int4，否则 int2。qt_fill（576 行）分派到对应的量化器。

2. 量化器：per-row 对称量化

三个量化器共用同一套 “per-row absmax” 方案，且与 Python 转换器 bit-identical（详见 §6）。

2.1 quantize_rows（int8）

quantize_rows（493 行）对每行独立：qmax = (1<<(bits-1))-1，算出该行绝对值最大 amax，缩放 s = max(amax/qmax, 1e-8)，然后 v = lrintf(w/s) 截断到 [-qmax-1, qmax]。按行 OpenMP 并行。

2.2 pack_int4（两个 nibble 一字节）

pack_int4（505 行）每行 rb = (I+1)/2 字节。同样的 per-row 缩放，然后两个 4-bit 打包进一个字节：v0、v1 各 lrintf(w/s) 截断到 [-8, 7]，存储时偏置 +8：qr[i>>1] = (v0+8) | ((v1+8)<<4)。低 nibble 是偶数下标，高 nibble 是奇数下标，奇数尾部有 guard。

2.3 pack_int2（四个 2-bit 一字节）

pack_int2（521 行）每行 rb = (I+3)/4 字节，值域 [-2, 1]，四个 2-bit 值打包，偏置 +2 后移位 (v+2)<<(k*2)。

g_nopack（535 行，NOPACK=1）可以把 ≤4-bit 的值留在 int8 容器里，用来验证 packing 的正确性。

3. dequant-on-use 矩阵乘

这组内核的共同点：把原始（整数或反量化后的）点积累加起来，最后每个输出元素只乘一次 scale[o]——反量化只在用到时发生，从不物化整个反量化权重矩阵。

内核 行 说明
matmul 201 F32 权重，朴素 x @ W^T，OpenMP over O
matmul_q 207 int8：内层现场把 int8 权重展宽转 float 再 FMA（AVX2/NEON/标量）
matmul_i4 226 int4：现场拆 nibble 并减去偏置 8，再 FMA
matmul_i2 261 int2：现场拆 2-bit 并减去偏置 2

以 matmul_i4 为例，AVX2 路径会掩码取低/高 nibble、用 _mm_unpacklo_epi8 恢复顺序、cvtepu8_epi32 后减 8，一次处理 16 个值；NEON 用 vzip_u8 交错、vsub_s8 去偏置；标量尾部逐字节 (byte&0xF)-8、(byte>>4)-8。

辅助的横向求和：hsum256（41 行）把 8 个 f32 求和，hsum256_i32（327 行）把 8 个 i32 求和，都在 #ifdef __AVX2__ 内。

4. IDOT：整数点积快路径

decode 时朴素的 dequant-on-use 要为每个权重做一次整数→浮点转换。IDOT（integer dot）走另一条路：把激活也量化成 int8，用 SIMD 的整数乘加指令直接算整数点积，最后乘回权重和激活两个缩放系数。设计注释（298–301 行）：量化矩阵乘快约 2–3 倍，只带来约 0.3% 的 RMS 误差，IDOT=0 可退回精确 f32 路径。

4.1 编译期选择 SIMD 内核

IDOT_KERNEL 宏（302–310 行）在编译期决定走哪条：

__AVX512VNNI__ && __AVX512BW__ → "avx512-vnni"
__AVX2__                       → "avx2"
__ARM_NEON                     → "neon"
否则                            → "scalar"

4.2 激活量化与阈值

qrow_i8（320 行）把一行激活按 absmax 量化成 int8（s = amax/127，Q8_0 风格），返回缩放。

两个阈值控制何时启用 IDOT：

  • g_idot（311 行）：IDOT 总开关（env IDOT）。
  • g_i4s：int4 IDOT 的最小 S 阈值。在有 ARM SDOT 的机器上是 1（int4 IDOT 即便 S=1 的 decode 也划算，Apple M 系测得 +14%）；否则是 2（作者原始 AVX2 阈值下 S=1 不划算）。这是”按 shape 逐个测量决定路由”哲学的体现。

4.3 整数点积内核

两个核心内核都用了一个 “sign-trick”，因为 SIMD 没有直接的无符号×有符号指令：

  • dot_i8i8（335 行）：int8·int8。AVX512-VNNI 用 _mm512_dpbusd_epi32 作用于 abs(w) 和符号折叠后的 x；AVX2 用 _mm256_maddubs_epi16(sign(w,w), sign(x,w))——sign(w,w)=|w|；NEON 优先用原生 vdotq_s32，无 DOTPROD 时用 vmull_s8/vpadalq_s16 回退。溢出安全性注释（333–334 行）说明每对乘积 ≤ 32512，到 I=16384 都安全。
  • dot_i4i8（380 行）：int4(packed)·int8，边拆 nibble 边做同样的 sign-trick。AVX512 路径巧妙地对 x 做 permute 而非重排 w（点积对顺序不敏感），每 64 个值一条 vpdpbusd。

对应的矩阵乘包装：matmul_q_idot（442 行）和 matmul_i4_idot（448 行），都算 y = dot(...) * sc * sx[s]——权重行缩放 sc 乘激活行缩放 sx[s]。

5. matmul_qt：中央分派器

matmul_qt（456 行）是所有矩阵乘的入口，按顺序决定走哪条路：

1. CUDA 路径（457–470）：仅当 cuda_eligible && !omp_in_parallel()，常驻张量走 GPU
2. fmt0 → matmul（f32）
3. IDOT 快路径（478–486）：g_idot 且 (fmt==1 或 (fmt==2 且 S>=g_i4s))
     → qrow_i8 量化激活 → matmul_q_idot / matmul_i4_idot
4. 回退 dequant-on-fly：fmt1→matmul_q, fmt3→matmul_i2, else→matmul_i4

注意 fmt3（int2）没有 IDOT 路径，总是走 matmul_i2。

6. FP8 → int4 离线转换器

引擎运行时只读预量化的容器；把官方 GLM-5.2 的 FP8 权重转成这个容器的是 c/tools/convert_fp8_to_int4.py（离线，纯 Python）。它的磁盘友好流式设计（5–6 行注释）：一次下载一个 ~5 GB 分片、转成 int4、删掉、下一片——所以 756 GB 的 FP8 权重从不需要同时存在于磁盘。

6.1 量化与 C 完全一致

三个量化器镜像 glm.c，注释说 “identica al C”：quant_int8、quant_int4、quant_int2——np.rint 对应 C 的 lrintf，偏置、nibble 顺序、值域全部一致，docstring 明确说这保证 “token identici”。

6.2 FP8 反量化

dequant（99–108 行）：FP8（E4M3）权重带 128×128 的 block scale（*_scale_inv），通过 repeat_interleave(128) 展开到 [O,I] 后还原成 f32，再重新量化成 colibrì 的 per-row 方案。

6.3 张量分类：谁留 f32、谁 int4

classify（75–96 行）决定每个张量的去向：

6.4 为什么 MTP 头必须 int8

默认 ebits（148–151 行）：--mtp 或 --indexer 时是 8，否则 4。注释给出了原因：

testa MTP a int4 = acceptance ~0-4%（misurato, issue #8）：draft 总是猜错，投机永远不启动。int8：39-59%，2.2-2.8 token/forward。

也就是说 MTP draft 头对量化极其敏感：int4 下 draft 几乎总错（接受率 0–4%），投机解码永不触发；int8 下接受率升到 39–59%，一次前向多产出 2.2–2.8 token。转换器默认就把这个头转成 int8——这与 glm.c 的注释 “layer MTP int8 = 2x i layer int4” 一致。详见 08-MTP投机解码与预取。

7. 小结

技术 出处 收益 / 代价
per-row 对称量化 quantize_rows:493 每行一个 scale，简单、和转换器 bit-identical
int4 packing（2/字节，偏置 8） pack_int4:505 内存/磁盘减到 1/4
dequant-on-use matmul_i4:226 从不物化反量化矩阵
IDOT 整数点积 dot_i8i8:335, dot_i4i8:380 量化矩阵乘快 2–3×，~0.3% RMS 误差
按 shape 测量决定路由 matmul_qt:478, g_i4s:313 int4 单行 decode 在 AVX2 上反而走 f32
MTP 头保 int8 converter:149 保住投机接受率，否则崩到 0–4%
This site uses Just the Docs, a documentation theme for Jekyll.