Skip to content

针对项目分析 #1

@GGzili

Description

@GGzili

mac-code 项目 TurboQuant 技术分析

概述

Google 的 TurboQuant 论文提出了一种极端 KV cache 压缩技术。本项目受其启发,在 MLX 后端实现了 KV cache 的 per-group N-bit 量化,目的是大幅缩小推理上下文在磁盘上的存储体积,实现跨会话/跨设备的上下文持久化。


一、背景知识

1.1 什么是 KV Cache

大模型生成文字时,每产出一个字,都需要"回顾"之前所有的对话内容。如果每次都从头算,极其浪费。

所以模型会把之前算过的中间结果缓存起来,这就是 KV Cache(Key-Value Cache)。

  • K(Key) = "这段话在说什么"的特征向量
  • V(Value) = "这段话具体内容"的特征向量

模型有很多层(比如 32 层),每一层都有自己的 K 和 V。所以 KV Cache 就是一个很大的数字矩阵集合

类比:你在写一篇作文,KV Cache 就像你写在草稿纸上的提纲和笔记。有了它你不用重新构思,直接接着写。

1.2 为什么要压缩 KV Cache

KV Cache 里每个数字默认用 16 bit(2 字节) 的浮点数存储,比如:

0.0234375, -0.1523438, 0.8710938, ...

聊天聊久了,这些数字堆起来就有 26MB 甚至更多。

如果你想:

  • 保存到硬盘(下次秒恢复对话)
  • 上传到云端(换台电脑继续聊)

26MB 太大了,需要压缩。


二、项目具体实现

核心代码在 mlx/turboquant.py,主要做了以下几件事:

2.1 Per-group 非对称量化(turboquant.py:40-89

tensor → 分组(group_size) → 每组算 min/max → 映射到 [0, 2^bits-1] → 存储量化值 + scale + zero
  • 支持 2/3/4-bit 量化
  • 默认 group_size=64,bits=3
  • 采用非对称 min-max 量化:每组独立计算 scale 和 zero point,精度优于全局量化

什么是量化

把高精度的小数,映射成低精度的整数

原始(16-bit 浮点):  0.0234, -0.1523, 0.8711, 0.3301, ...
量化后(4-bit 整数):  2,       0,       15,     7,      ...

16-bit 能表示 65536 种值,4-bit 只能表示 16 种值(0~15)。体积直接缩小 4 倍,但精度会有损失。

什么是"分组"(per-group)

不对整个矩阵用同一套映射规则,而是每 64 个数字一组,每组有自己的映射规则。

为什么?因为不同位置的数字范围差异很大:

  • 第一组数字可能在 [-0.2, 0.3] 之间
  • 第二组数字可能在 [-5.0, 8.0] 之间

如果用同一套规则,小数字全被压成 0,精度惨不忍睹。分组后每组独立处理,精度大幅提高。

什么是"非对称 min-max"

每一组的映射规则是这样算的:

一组数字:[0.1, 0.5, 0.3, 0.9, 0.2, ...]

min = 0.1(最小值)
max = 0.9(最大值)
scale = (max - min) / 15 = 0.0533   (15 是 4-bit 的最大整数)
zero = min = 0.1

量化公式:整数值 = round((原始值 - zero) / scale)
  0.1 → round((0.1 - 0.1) / 0.0533) = 0
  0.5 → round((0.5 - 0.1) / 0.0533) = 8
  0.9 → round((0.9 - 0.1) / 0.0533) = 15

反量化(恢复):原始值 ≈ 整数值 × scale + zero
  0  → 0 × 0.0533 + 0.1 = 0.1      ✅ 完美恢复
  8  → 8 × 0.0533 + 0.1 = 0.5264   ≈ 0.5(略有误差)
  15 → 15 × 0.0533 + 0.1 = 0.8995  ≈ 0.9 ✅

"非对称"是指 min 和 max 可以不对称(不一定以 0 为中心),这比对称量化更灵活。每组需要额外存一个 scale 和一个 zero,这就是代码里的 scaleszeros

对称量化 vs 非对称量化

对称量化:假设数据以 0 为中心,zero 固定为 0,只存 scale。

假设数据范围是 [-0.8, 0.8]
scale = 0.8 / 7 = 0.114    (4-bit 有符号:-8 到 +7)
zero = 0(固定,不用存)
量化:整数值 = round(原始值 / scale)

非对称量化:不假设以 0 为中心,每组存 scale + zero。

假设数据范围是 [0.1, 0.9]
scale = (0.9 - 0.1) / 15 = 0.0533
zero = 0.1(必须存)

为什么不能全局固定 scale/zero?因为每组数字的范围差异很大

第 1 组:[0.1, 0.3, 0.2, 0.15, ...]   范围 [0.1, 0.3]
第 2 组:[-5.0, 3.0, 8.0, -2.0, ...]   范围 [-5.0, 8.0]
第 3 组:[0.001, 0.003, 0.002, ...]     范围 [0.001, 0.003]

如果用全局统一的 scale/zero(按最大范围 [-5.0, 8.0] 算):

全局 scale = (8.0 - (-5.0)) / 15 = 0.867

第 3 组的数字 0.001, 0.003:
  round((0.001 - (-5.0)) / 0.867) = round(5.77) = 6
  round((0.003 - (-5.0)) / 0.867) = round(5.77) = 6
  → 0.001 和 0.003 量化成了同一个整数!区别完全丢失了

所以必须每组独立算 scale/zero。这就是"非对称"的代价——更灵活,但要额外存储。

存储开销的权衡

存储量 = 量化后的整数 + 每组的 scale 和 zero

省了:每个数从 16-bit 变成 4-bit(省 4 倍)
多了:每 64 个数多存 2 个 float(scale + zero)

每 64 个数多存 2 个 float16(4 字节),相当于每个数多 0.5 bit 的开销。从 16-bit 压到 4.5-bit,还是赚了很多。

但如果要压到 3-bit 甚至 2-bit,这 0.5 bit 的开销占比就变得很大了(2-bit + 0.5 = 2.5 bit,开销占了 20%)。这恰好就是 Google PolarQuant 要解决的问题(见后文)。

2.2 KV Cache 整体压缩/解压(turboquant.py:118-172

原始 KV Cache(32 层,每层有 K 和 V)
        │
        ▼
compress_kv_cache():逐层、逐 tensor 调用 quantize_tensor()
        │
        ▼
压缩后的数据:量化整数 + 每组的 scale/zero
        │
        ▼
decompress_kv_cache():逐层调用 dequantize_tensor() 恢复浮点数

compress_kv_cache() 遍历模型所有层的 KV 状态,逐层逐 tensor 进行量化,并统计压缩比。decompress_kv_cache() 做逆操作恢复原始精度。

2.3 质量度量(turboquant.py:175-205

measure_quality() 对压缩前后的 KV cache 计算两个指标:

MSE(均方误差):压缩前后每个数字差了多少

原始:[0.5, 0.9, 0.1]
恢复:[0.5264, 0.8995, 0.1]
误差:(0.0264² + 0.0005² + 0²) / 3 = 很小的数

余弦相似度:两组向量的方向有多接近

  • 1.0 = 方向完全一样
  • 0.0 = 完全无关
  • 项目实测 0.993 = 几乎一样

类比:把一张照片从 PNG 压缩成 JPEG,然后对比两张照片有多像。0.993 就像肉眼完全看不出区别。

2.4 序列化到磁盘(turboquant.py:208-276

serialize_compressed() / load_compressed() 将压缩后的 KV cache 保存为两个文件:

  • .npz:numpy 的压缩数组格式,存量化后的整数 + scale + zero
  • .meta.json:存元信息(原始形状、数据类型、bit 数、分组大小)

加载时读这两个文件,就能恢复出压缩后的 KV Cache,再反量化就得到近似原始数据。


三、效果

指标 数值
压缩比 26.6 MB → 6.7 MB(约 4x
余弦相似度 0.993(几乎无损)
SSD 加载时间 0.0003 秒(比重新计算快 6677x)

四、与 kv_cache.py 的区别

项目有两个压缩层:

turboquant.py kv_cache.py
压缩方式 量化(有损) gzip(无损)
压缩比 4x 约 1.5-2x
精度损失 有(0.7%)
用途 大幅缩小体积 通用文件压缩

实际代码中的使用情况

经过代码审查,这两者实际上并没有串联使用,而是各走各的路:

路径 1(实际在用):普通保存 → gzip → 上传 R2

KV cache → safetensors 保存到 SSD → gzip 压缩 → 上传 R2

这是 r2_store.pyupload_context() 走的路径(第 121 行调用 compress_cache(),就是 gzip)。

路径 2(仅 benchmark 测试):TurboQuant 量化 → 保存到 SSD

KV cache → turboquant 4-bit 量化 → serialize 成 .npz → 存到本地 SSD

这是 benchmark.py 里测试的路径。量化后直接存盘,没有再套一层 gzip。

两条路径从未合并:

  • agent.pymlx_engine.py完全没有引用 turboquant
  • r2_store.py 里的上传逻辑只用了 gzip,没有调用 turboquant
  • turboquant 只在 benchmark.pyagent_benchmark.py 里被调用过,纯粹是跑基准测试用的
实际在用的:  KV cache → gzip → R2 上传(kv_cache.py + r2_store.py)
只跑了测试:  KV cache → turboquant 4-bit 量化 → 存本地(benchmark.py)
从未实现的:  KV cache → turboquant → gzip → R2  ← 不存在

TurboQuant 在这个项目里目前还停留在 benchmark 验证阶段,没有集成到实际的保存/上传流程中。

理论上可行的串联方案

理论上完全可以将两者串联,进一步压缩上传体积:

KV cache → turboquant 4-bit 量化 → .npz → gzip 压缩 → 上传 R2

效果估算:

阶段 体积 做了什么
原始 KV cache 26.6 MB 16-bit 浮点数
turboquant 4-bit 量化 ~6.7 MB 有损压缩 4x
gzip 再压缩 ~4-5 MB 无损压缩(量化后的整数重复模式多,gzip 还能再压一些)

对比当前 R2 上传的路径:

路径 上传体积
当前:safetensors → gzip ~13-17 MB(gzip 对浮点数压缩比一般,约 1.5-2x)
串联:turboquant → gzip ~4-5 MB(先量化 4x,再无损压缩)

串联方案上传体积大约是当前方案的 1/3 到 1/4,上传下载速度也快了对应倍数。唯一的代价是 0.7% 的精度损失(余弦相似度 0.993),但这是在 turboquant 那一步已经产生的,gzip 是无损的不会再损失。

项目没有这样做的原因不是技术上做不到,而是开发顺序问题:

  1. 先把基础的 save/load + R2 上传跑通(用 safetensors + gzip,简单可靠)
  2. 再单独验证 turboquant 的压缩效果(benchmark)
  3. 最终集成到正式流程(还没做到这一步)

五、为什么 TurboQuant 等 KV Cache 操作只在 MLX 后端实现

35B 只能用 llama.cpp,9B 两个都能用

Mac mini M4 只有 16GB 内存,而 35B 模型量化后还有 10.6GB。加上系统占用、KV cache 等开销,内存根本不够把整个模型放进去。

llama.cpp 有一个关键能力:当模型放不进内存时,它会利用 macOS 的虚拟内存机制,通过 mmap 自动把模型的一部分从 SSD 分页加载(page in/out)。这就是 README 里说的 "SSD paging",也是 Apple 那篇 "LLM in a Flash" 论文的思路。

MLX 做不到这一点。 MLX 要求模型完整加载到统一内存里才能运行。如果模型放不进内存,MLX 直接跑不动。

9B 模型量化后只有 5.3GB,16GB 内存完全放得下,所以 llama.cpp 和 MLX 都能跑。

16GB Mac mini M4 内存分配:

35B 模型(10.6GB):塞不进去 → 只能用 llama.cpp(支持 SSD 分页)
9B 模型(5.3GB):  轻松放下 → llama.cpp 和 MLX 都能用

模型格式差异

两种后端加载的是不同格式的模型文件:

后端 35B 模型 9B 模型 格式
llama.cpp Qwen3.5-35B-A3B-UD-IQ2_M.gguf Qwen3.5-9B-Q4_K_M.gguf GGUF
MLX mlx-community/Qwen3.5-35B-A3B-4bit mlx-community/Qwen3.5-9B-MLX-4bit safetensors

GGUF 是 llama.cpp 专用的模型格式,由 Unsloth 等社区做好量化后发布。项目直接下载的就是已经量化好的 GGUF 文件,不需要自己量化。MLX 格式是 safetensors(HuggingFace 标准格式),由 mlx-community 转换发布。

为什么 MLX 更快

llama.cpp MLX
语言 C++ Python 原生(底层 C++/Metal)
定位 跨平台(Linux/Windows/Mac) Apple Silicon 专用
内存模型 通过 Metal API,有数据搬运开销 统一内存零拷贝
执行方式 即时执行 懒执行(Lazy Evaluation),操作融合
GPU 适配 Metal shader 通用适配 为 M 系列芯片专门优化的算子
llama.cpp:C++ 代码 → Metal API → GPU 执行 → 结果拷贝回 CPU 内存
MLX:      Python 代码 → 直接在统一内存上操作 → GPU 执行(零拷贝)

为什么 KV Cache 操作只在 MLX 实现

MLX(Python 原生):KV cache 的每个 tensor 都是 mx.array,和 numpy 一样可以随意切片、量化、保存、加载:

cache = make_prompt_cache(model)
logits = model(tokens, cache=cache)
for layer_cache in cache:
    state = layer_cache.state        # ← 直接拿到 KV tensor
    mx.save(state, "layer.safetensors")  # ← 可以保存

llama.cpp(C++ 服务器):KV cache 存在 C++ 进程的内存里,Python 只能通过 HTTP API 与之通信,根本碰不到内部的 tensor:

agent.py → HTTP 请求 → llama-server(C++ 进程) → 返回 JSON 文本
                        ↑ KV cache 在这里面,Python 摸不到
MLX 架构:
Python 代码 ←→ MLX tensor(统一内存)←→ GPU
                  ↑ 可以直接操作 KV cache

llama.cpp 架构:
Python 代码 → HTTP → C++ 服务器 → Metal → GPU
                      ↑ KV cache 在这里面,Python 摸不到

这就是为什么 TurboQuant、KV cache 持久化、R2 同步这些功能全都只在 MLX 后端实现——因为只有 MLX 让你在 Python 里直接拿到 KV cache 的原始数据。

llama.cpp MLX
模型超内存时 能跑(SSD 分页) 跑不了
模型在内存内 能跑 能跑,且更快 25%
KV cache 操作 无法直接访问 tensor 可以直接操作(Python 原生)

六、与 Google TurboQuant 论文的对比

论文实际做了什么

TurboQuant(发表在 ICLR 2026)是两种技术的组合:

技术 来源 做什么
PolarQuant AISTATS 2026 论文 笛卡尔→极坐标,消除量化常数开销
QJL AAAI 论文 1-bit 随机投影,修正残差误差

第一步:PolarQuant(极坐标量化)

传统量化是在笛卡尔坐标下做的——把每个数字直接映射到整数。但这有个问题:每一小组数字的范围不同,必须额外存储每组的 min/max(即 scale 和 zero),这些额外开销会占 1-2 bit,压缩到 3-bit 时开销占比就很大了。

PolarQuant 的做法:

  1. 先对向量做随机旋转(简化数据分布的几何形状)
  2. 把向量从笛卡尔坐标 (x, y, z, ...) 转换成极坐标 (半径, 角度1, 角度2, ...)
  3. 对每个分量独立量化

论文的比喻:导航时把"往东走 3 格、往北走 4 格"换成"朝 37° 方向走 5 步"。

核心洞察:转换成极坐标后,角度的分布是已知的、高度集中的,不需要再为每组计算 min/max 了。相当于数据自动落在一个"固定的圆形网格"上,边界是提前确定的。这样就消除了 per-group 量化常数的存储开销

类比:指南针告诉你"往东北走"(方向)和"走 500 米"(长度)。如果"走 500 米"变成"走 498 米"你不会迷路,但如果"东北"变成"正北"你就走错了。所以方向用高精度存,长度可以粗糙存。

第二步:QJL(1-bit 误差修正)

PolarQuant 量化后会有残差误差。QJL 用 Johnson-Lindenstrauss 随机投影,只用额外 1 bit 对误差做修正:

  • 把高维残差投影到低维空间
  • 每个数只保留正负号(+1 / -1)
  • 零额外内存开销

相当于一个"数学纠错器",消除第一步引入的偏差。

对比总结

Google TurboQuant 论文 mac-code 项目实现
坐标系 笛卡尔 → 极坐标转换 始终在笛卡尔坐标下操作
量化方式 极坐标分量独立量化,无需 per-group 常数 传统 per-group min-max 非对称量化
误差修正 QJL 1-bit 残差修正
额外开销 几乎为零(不需要存 scale/zero) 每组需要存 scale 和 zero
精度 3-bit 零精度损失 4-bit 0.993 余弦相似度
压缩比 6x+ 4x

结论

Google TurboQuant 论文
├── PolarQuant(极坐标量化)     ← 未实现
└── QJL(1-bit 随机投影误差修正) ← 未实现

mac-code 项目实际实现的
└── per-group min-max 量化(传统方法)← 和论文无直接关系,只是受到启发

项目里的 turboquant.py 用的是 per-group min-max 非对称量化,这是一种非常通用的、早已存在的量化方法,不属于 TurboQuant 论文的任何一部分。它和论文的关系只是"受到启发"——看到论文说 KV cache 可以极端压缩,于是用最基础的量化方法先做了一版,拿到了 4x 压缩比。


七、尚未完成的部分

mlx/PROJECT.md:99-103 的路线图中,Phase 5 标记为未完成:

  • PolarQuant(笛卡尔→极坐标量化,消除 per-group 常数开销)
  • QJL(随机投影 1-bit 误差修正)
  • 3-bit 目标:95% payload 缩减

这两个如果实现了,理论上能做到 3-bit 甚至 2-bit 量化,压缩比从 4x 提升到 5-8x,同时保持精度。


参考

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions