
[InferLLM大模型推理框架项目](22)kern模块中量化的实现(src/kern/naive/quantize.h)
InferLLM 框架中 quantize.h 的代码结构与功能分析
quantize.h
是 InferLLM 框架中实现量化和反量化操作的核心头文件,位于 kern/naive
目录下,提供了基础的量化计算实现。这个文件主要实现了模型权重和激活值的量化与反量化操作,以及量化数据之间的点积计算,是实现低精度推理的关键组件。
1. 文件结构概览
文件结构可以分为以下几个部分:
- 头文件引入
- 反量化函数(dequantize)
- 量化函数(quantize)
- 量化数据点积计算函数
2. 反量化函数
2.1 4位整数反量化
inline void dequantize_row_q4_0_reference(
const void* __restrict x, float* __restrict y, int k) {
const int nb = k / QK40;
const size_t bs = sizeof(float) + QK40 / 2;
const uint8_t* __restrict pd = ((const uint8_t*)x + 0 * bs);
const uint8_t* __restrict pb = ((const uint8_t*)x + 0 * bs + sizeof(float));
// scalar
for (int i = 0; i < nb; i++) {
const float d = *(const float*)(pd + i * bs);
const uint8_t* __restrict pp = pb + i * bs;
for (int l = 0; l < QK40; l += 2) {
const uint8_t vi = pp[l / 2];
const int8_t vi0 = vi & 0xf;
const int8_t vi1 = vi >> 4;
const float v0 = (vi0 - 8) * d;
const float v1 = (vi1 - 8) * d;
y[i * QK40 + l + 0] = v0;
y[i * QK40 + l + 1] = v1;
}
}
}
这个函数将4位整数量化的数据反量化为浮点数:
- 将输入数据分为多个块(每块32个元素)
- 每个块有一个缩放因子
d
- 每个字节存储两个4位整数(高4位和低4位)
- 将4位整数解码为 -8 到 7 的范围,然后乘以缩放因子得到浮点数
2.2 8位整数反量化
inline void dequantize_row_q8_0_reference(
const void* __restrict x, float* __restrict y, int k) {
const int nb = k / QK80;
const size_t bs = sizeof(float) + QK40 / 2;
const BlockQ80* xx = reinterpret_cast<const BlockQ80*>(x);
// scalar
for (int i = 0; i < nb; i++) {
const float d = xx[i].d;
const int8_t* __restrict pp = xx[i].qs;
for (int l = 0; l < QK80; l++) {
y[i * QK80 + l] = pp[l] * d;
}
}
}
这个函数将8位整数量化的数据反量化为浮点数:
- 将输入数据分为多个块(每块32个元素)
- 每个块有一个缩放因子
d
- 直接将8位整数乘以缩放因子得到浮点数
3. 量化函数
3.1 浮点数量化为4位整数
inline size_t quantize_row_q4_0_reference(const float* x, BlockQ40* y, int k) {
const int nb = k / QK40;
uint8_t pp[QK40 / 2];
for (int i = 0; i < nb; i++) {
float amax = 0.0f; // absolute max
for (int l = 0; l < QK40; l++) {
const float v = x[i * QK40 + l];
amax = std::max(amax, fabsf(v));
}
const float d = amax / ((1 << 3) - 1);
const float id = d ? 1.0f / d : 0.0f;
y[i].d = d;
for (int l = 0; l < QK40; l += 2) {
const float v0 = x[i * QK40 + l + 0] * id;
const float v1 = x[i * QK40 + l + 1] * id;
const uint8_t vi0 = (int8_t)roundf(v0) + 8;
const uint8_t vi1 = (int8_t)roundf(v1) + 8;
assert(vi0 < 16);
assert(vi1 < 16);
pp[l / 2] = vi0 | (vi1 << 4);
}
memcpy(y[i].qs, pp, sizeof(pp));
}
return nb * sizeof(BlockQ40);
}
这个函数将浮点数量化为4位整数:
- 将输入数据分为多个块(每块32个元素)
- 计算每个块中的最大绝对值
amax
- 计算缩放因子
d = amax / 7
(7 = 2^3 - 1) - 将浮点数除以缩放因子,四舍五入后加上偏移量8,得到0到15的整数
- 将两个4位整数打包到一个字节中(低4位和高4位)
3.2 浮点数量化为8位整数
inline void quantize_row_q8_0_reference(const float* x, BlockQ80* y, int k) {
assert(k % QK80 == 0);
const int nb = k / QK80;
for (int i = 0; i < nb; i++) {
float amax = 0.0f; // absolute max
for (int l = 0; l < QK80; l++) {
const float v = x[i * QK80 + l];
amax = std::max(amax, fabsf(v));
}
const float d = amax / ((1 << 7) - 1);
const float id = d ? 1.0f / d : 0.0f;
y[i].d = d;
for (int l = 0; l < QK80; ++l) {
const float v0 = x[i * QK80 + l] * id;
y[i].qs[l] = roundf(v0);
}
}
}
这个函数将浮点数量化为8位整数:
- 将输入数据分为多个块(每块32个元素)
- 计算每个块中的最大绝对值
amax
- 计算缩放因子
d = amax / 127
(127 = 2^7 - 1) - 将浮点数除以缩放因子,四舍五入后得到-127到127的整数
4. 量化数据点积计算函数
4.1 4位整数与8位整数点积
inline float vec_vec_dot_q40_with_q80_reference(
const int n, const void* vx, const void* vy) {
const int nb = n / QK80;
assert(n % QK80 == 0);
assert(nb % 2 == 0);
const BlockQ40* __restrict x = (BlockQ40*)(vx);
const BlockQ80* __restrict y = (BlockQ80*)(vy);
float sumf = 0.0;
for (int i = 0; i < nb; i++) {
const float d0 = x[i].d;
const float d1 = y[i].d;
const uint8_t* __restrict p0 = x[i].qs;
const int8_t* __restrict p1 = y[i].qs;
int sumi = 0;
for (int j = 0; j < QK80 / 2; j++) {
const uint8_t v0 = p0[j];
const int i0 = (int8_t)(v0 & 0x0F) - 8;
const int i1 = (int8_t)(v0 >> 4) - 8;
const int i2 = p1[2 * j + 0];
const int i3 = p1[2 * j + 1];
sumi += i0 * i2 + i1 * i3;
}
sumf += d0 * d1 * sumi;
}
return sumf;
}
这个函数计算4位整数向量与8位整数向量的点积:
- 将输入数据分为多个块
- 对于每个块,解码4位整数和8位整数
- 计算整数之间的点积
- 乘以两个缩放因子得到最终结果
4.2 打包的4位整数与8位整数点积
inline void vec_vec_dot_q40_with_q80_packed_reference(
const int n, const void* vx, const void* vy, float* dst, const float* bias) {
const int nb = n / QK80;
assert(n % QK80 == 0);
const BlockQ40X8* __restrict x = (BlockQ40X8*)(vx);
const BlockQ80* __restrict y = (BlockQ80*)(vy);
float sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0, sum4 = 0, sum5 = 0, sum6 = 0, sum7 = 0;
if (bias) {
sum0 = bias[0], sum1 = bias[1], sum2 = bias[2], sum3 = bias[3], sum4 = bias[4],
sum5 = bias[5], sum6 = bias[6], sum7 = bias[7];
}
// ... 计算过程 ...
dst[0] = sum0;
dst[1] = sum1;
dst[2] = sum2;
dst[3] = sum3;
dst[4] = sum4;
dst[5] = sum5;
dst[6] = sum6;
dst[7] = sum7;
}
这个函数计算8组打包的4位整数向量与一个8位整数向量的点积:
- 将输入数据分为多个块
- 对于每个块,解码8组4位整数和1组8位整数
- 计算8组点积
- 乘以相应的缩放因子得到最终结果
- 如果有偏置,加上偏置值
4.3 8位整数与8位整数点积
inline float vec_vec_dot_q80_with_q80_reference(
const int n, const void* vx, const void* vy) {
const int nb = n / QK80;
assert(n % QK80 == 0);
assert(nb % 2 == 0);
const BlockQ80* __restrict x = (BlockQ80*)(vx);
const BlockQ80* __restrict y = (BlockQ80*)(vy);
float sumf = 0.0;
for (int i = 0; i < nb; i++) {
const float d0 = x[i].d;
const float d1 = y[i].d;
const int8_t* __restrict p0 = x[i].qs;
const int8_t* __restrict p1 = y[i].qs;
int sumi = 0;
for (int j = 0; j < QK80; j++) {
sumi += p0[j] * p1[j];
}
sumf += d0 * d1 * sumi;
}
return sumf;
}
这个函数计算两个8位整数向量的点积:
- 将输入数据分为多个块
- 对于每个块,计算8位整数之间的点积
- 乘以两个缩放因子得到最终结果
4.4 浮点数与浮点数点积
inline float vec_vec_dot_float_with_float_reference(
const int n, const float* x, const float* y) {
float sumf = 0.0;
for (int i = 0; i < n; i++) {
sumf += x[i] * y[i];
}
return sumf;
}
这个函数计算两个浮点数向量的点积,直接进行浮点数乘法和累加。
5. 功能分析
通过分析 quantize.h
文件,可以看出它实现了以下主要功能:
5.1 数据量化
将浮点数量化为低精度整数(4位或8位),减少内存占用和计算量:
- 4位整数量化:每个浮点数用4位表示,每个字节存储两个值
- 8位整数量化:每个浮点数用8位表示
量化过程中使用块量化方法,每个块(通常是32个元素)共享一个缩放因子,这样可以在保持计算精度的同时减少内存占用。
5.2 数据反量化
将低精度整数(4位或8位)反量化为浮点数,用于模型的输出或中间结果:
- 4位整数反量化:将4位整数解码为-8到7的范围,然后乘以缩放因子
- 8位整数反量化:将8位整数乘以缩放因子
5.3 量化数据点积计算
实现不同精度数据之间的点积计算,是矩阵乘法的基础操作:
- 4位整数与8位整数点积
- 打包的4位整数与8位整数点积(一次计算8组)
- 8位整数与8位整数点积
- 浮点数与浮点数点积
这些点积计算函数是实现低精度矩阵乘法的核心,通过整数乘法和累加,再乘以缩放因子,可以在保持计算精度的同时提高计算效率。
6. 实现细节分析
6.1 块量化策略
quantize.h
采用块量化策略,将数据分为固定大小的块(通常是32个元素),每个块共享一个缩放因子。这种策略有以下优点:
- 减少存储开销:只需要为每个块存储一个缩放因子,而不是每个元素一个
- 保持局部精度:每个块独立计算缩放因子,可以适应数据的局部分布
- 计算效率高:可以批量处理一个块内的所有元素
6.2 对称量化
代码中使用的是对称量化方法,具体表现为:
- 4位整数量化:值域为 -8 到 7(通过减去偏移量8实现)
- 8位整数量化:值域为 -127 到 127
对称量化相比非对称量化更简单,计算效率更高,但可能在处理非对称分布的数据时精度略低。
6.3 内存布局优化
代码中的内存布局经过优化,以提高访问效率:
BlockQ40
结构:先存储缩放因子,再存储量化值BlockQ40X8
结构:将8组4位整数量化值连续存储,然后是8个缩放因子- 打包存储:每个字节存储两个4位整数,减少内存占用
6.4 SIMD 友好设计
虽然 quantize.h
中的实现是标量版本,但其内存布局和计算模式设计考虑了 SIMD 指令集的优化可能:
- 固定大小的块(32个元素)适合 SIMD 寄存器宽度
- 连续的内存访问模式有利于 SIMD 加载指令
- 整数乘法和累加操作适合 SIMD 指令优化
7. 性能优化分析
7.1 计算复杂度优化
通过量化,将浮点数运算转换为整数运算,显著降低了计算复杂度:
- 浮点数乘法 → 整数乘法(更快)
- 浮点数加法 → 整数加法(更快)
- 最后只需要少量的浮点数乘法(乘以缩放因子)
7.2 内存占用优化
量化显著减少了内存占用:
- 4位量化:比32位浮点数减少87.5%的内存占用
- 8位量化:比32位浮点数减少75%的内存占用
减少内存占用不仅节省存储空间,还能提高缓存命中率,进一步提升性能。
7.3 批量处理优化
vec_vec_dot_q40_with_q80_packed_reference
函数实现了批量处理优化,一次计算8组点积,减少了函数调用开销和循环控制开销。
7.4 编译器优化友好
代码中使用了 __restrict
关键字,告诉编译器指针之间没有别名,有助于编译器进行更激进的优化。
8. 与其他优化实现的关系
quantize.h
中的函数都带有 _reference
后缀,表明这些是参考实现,主要用于:
- 功能验证:验证量化和反量化算法的正确性
- 后备实现:当特定平台的优化实现不可用时作为后备
- 基准比较:用于与优化实现进行性能比较
实际应用中,会根据目标平台选择更高效的实现,如:
- x86平台:使用 AVX/AVX2 指令集优化
- ARM平台:使用 NEON 指令集优化
- GPU平台:使用 CUDA 优化
9. 应用场景分析
9.1 模型权重量化
大型语言模型的权重通常占用大量内存,通过量化可以显著减少内存占用:
- 原始 LLaMA 7B 模型(FP32):约28GB
- 4位量化后:约3.5GB
这使得在资源受限的设备上运行大型模型成为可能。
9.2 激活值量化
在模型推理过程中,中间激活值也可以量化,减少内存占用和计算量:
- 将注意力矩阵量化为8位整数
- 将前馈网络的中间结果量化为8位整数
9.3 矩阵乘法加速
矩阵乘法是大型语言模型中最耗时的操作,通过量化可以显著加速:
- 权重量化为4位整数
- 激活值量化为8位整数
- 使用整数矩阵乘法(比浮点数矩阵乘法快)
10. 未来优化方向
基于 quantize.h
的实现,可以考虑以下优化方向:
10.1 更高效的量化方法
- 非对称量化:处理非对称分布的数据
- 组量化:不同组使用不同的缩放因子,提高精度
- 混合精度量化:重要权重使用更高精度,不重要权重使用更低精度
10.2 更多硬件优化
- ARM SVE:支持可变长度向量的 ARM 指令集
- RISC-V V 扩展:RISC-V 的向量扩展
- 特定加速器:针对 NPU、TPU 等专用加速器的优化
10.3 算法优化
- 稀疏量化:结合稀疏性和量化,进一步减少计算量
- 量化感知训练:在训练阶段考虑量化误差,提高量化模型精度
- 自适应量化:根据数据分布动态调整量化参数
总结
quantize.h
实现了大语言模型推理中的核心量化操作,包括4位和8位整数量化、反量化以及量化数据之间的点积计算。通过块量化、对称量化、内存布局优化和批量处理等技术,在保持计算精度的同时显著减少了内存占用和计算量。这些实现为 InferLLM 框架提供了高效的低精度推理能力,使得大型语言模型可以在资源受限的设备上运行。
