如何通过 INT4 KV Cache 量化大幅提升移动端 LLM 的多轮对话上限

在大模型（LLM）落地移动端的过程中，内存占用是最大的瓶颈。除了模型权重（Weights）外，KV Cache 的增长直接决定了多轮对话的上下文长度上限。本文将教你如何通过 INT4 量化技术对 KV Cache 进行压缩，将内存占用降低约 75%。

1. 为什么要对 KV Cache 进行 INT4 量化？

在 Transformer 推理过程中，KV Cache 会随着 Token 数量线性增长。以 FP16 精度的 Llama-7B 模型为例，当上下文达到 4096 个 Token 时，KV Cache 约占用 4GB 内存。对于大多数 8GB 或 12GB 内存的手机来说，这极易导致 OOM（内存溢出）。通过 INT4 量化，我们可以将每个元素从 16bit 压缩到 4bit，显著释放内存空间。

2. 核心技术原理

KV Cache 量化通常采用 Per-Group 对称量化。我们将 Key/Value 向量按通道分组（例如每 64 个元素一组），计算每组的最大绝对值作为缩放因子（Scale）。
– 量化公式: $q = clamp(round(x / scale), -8, 7)$
– 打包存储: 由于 4bit 数据不足一个字节，我们通常将两个 4bit 数据打包进一个 INT8 类型的字节中进行存储。

3. 实战代码：PyTorch 实现 INT4 量化压缩

以下代码展示了如何对 Tensor 进行量化、打包存储以及在推理时还原（反量化）的逻辑：

import torch

def pack_int4_kv(tensor, group_size=64):
    # 1. 重塑形状以进行分组量化
    orig_shape = tensor.shape
    t = tensor.view(-1, group_size)

    # 2. 计算 Scale (对称量化)
    max_val = t.abs().max(dim=-1, keepdim=True)[0]
    scale = max_val / 7.0
    scale = scale.clamp(min=1e-5)

    # 3. 量化并截断到 [-8, 7]
    q = torch.round(t / scale).clamp(-8, 7).to(torch.int8)

    # 4. 打包：将两个 int4 组合到一个 int8 字节中
    q_flat = q.view(-1)
    # 取低 4 位并偏移打包
    packed = ((q_flat[0::2] & 0x0F) << 4) | (q_flat[1::2] & 0x0F)

    return packed, scale, orig_shape

def unpack_int4_kv(packed, scale, orig_shape, group_size=64):
    # 1. 解包：提取高 4 位和低 4 位
    high = (packed >> 4).to(torch.int8)
    low = (packed & 0x0F).to(torch.int8)

    # 2. 符号位恢复 (针对补码逻辑)
    high = torch.where(high > 7, high - 16, high)
    low = torch.where(low > 7, low - 16, low)

    # 3. 反量化
    q_unpacked = torch.stack([high, low], dim=1).view(-1, group_size)
    recovered = q_unpacked * scale

    return recovered.view(orig_shape)

# --- 测试运行 ---
kv_sample = torch.randn(1, 32, 128) # [heads, seq_len, dim]
packed_data, scales, shape = pack_int4_kv(kv_sample)
recovered_data = unpack_int4_kv(packed_data, scales, shape)

print(f'原始 Tensor 大小: {kv_sample.nelement() * 2} 字节 (FP16)')
print(f'量化打包后大小: {packed_data.nelement() + scales.nelement() * 4} 字节')
print(f'平均绝对误差: {(kv_sample - recovered_data).abs().mean().item():.4f}')

4. 适配建议与优化建议

1. 误差补偿: Key 对精度更敏感，建议对 Key 使用 Per-Channel 量化，而 Value 使用 Per-Token 量化。
2. 硬件对齐: 在端侧（如 Android/iOS）使用 NCNN 或 MNN 时，确保打包逻辑与底层 C++ 的 reinterpret_cast 逻辑一致。
3. 算子融合: 实际部署时，应将反量化逻辑融合进 Attention 算子中，避免产生额外的访存开销。

通过 INT4 KV Cache，你可以将移动端 LLM 的对话上下文长度提升 3-4 倍，这对于构建实用的端侧 AI 助手至关重要。