tqlite

PyTorch 实现的 KV cache 量化原型, 基于 TurboQuant (Google, 2025) 的核心思想: 随机正交旋转 + MSE 最优标量量化 + 可选的 QJL 残差校正 + 可选的 outlier 混合精度.

定位: 算法正确性验证 + 真实 HF 模型挂载测试, 不是生产级推理框架.

目录结构

tqlite/
├── tqlite/                # 核心库
│   ├── rotation.py        # 随机旋转 Π (Hadamard × 随机符号)
│   ├── codebook.py        # Lloyd-Max 高斯码本
│   ├── mse.py             # TurboQuantMSE: 基础 MSE 量化
│   ├── prod.py            # TurboQuantProd: MSE + QJL 残差校正
│   ├── grouped.py         # OutlierGrouped: 混合精度分组
│   ├── factory.py         # build_quantizer 统一工厂
│   └── kv_adapter.py      # HF DynamicCache fake-quant 适配器
├── scripts/
│   ├── selftest.py        # 核心算法自检 (无需模型)
│   ├── generate.py        # fp vs 量化 生成对比
│   ├── eval_quality.py    # KL/top-k/首次分叉 (多 prompt 质量评估)
│   ├── eval_ppl.py        # PPL 对比
│   ├── eval_hellaswag.py  # HellaSwag 多选准确率
│   ├── eval_needle.py     # Needle-in-Haystack 长上下文检索
│   ├── eval_gsm8k.py      # GSM8K 数学推理准确率
│   └── run_all.sh         # 一键跑全套 (建议在 B200 上用)
└── tests/
    └── test_core.py       # pytest 单元测试

快速开始

pip install -r requirements.txt

# 不需要模型, 纯算法验证
python scripts/selftest.py

# 国内可选 HF 镜像
export HF_ENDPOINT=https://hf-mirror.com

# 生成对比
python scripts/generate.py --model Qwen/Qwen3-0.6B --mode mse --bits 8

# pytest
pytest tests/

三种量化模式

mode	原理	每 key bit 开销	适用
mse	旋转 + Lloyd-Max 码本	d * bits	baseline, bits≥6 质量较好
prod	MSE + QJL 1-bit 残差	d * bits + d + 32	低 bit (≤4) 的质量补偿
grouped	高方差通道分配更高 bit	d_out*b_out + d_n*b_n	异方差分布显著的层

评测脚本

脚本	测什么	典型样本数	备注
eval_quality.py	next-token 分布差异 (KL, top-5)	6 prompts × 30 steps	不需要数据集
eval_ppl.py	WikiText 风格 PPL	100-200 tokens	PPL 对量化敏感, 仅作相对对比
eval_hellaswag.py	常识推理多选准确率	100-500	下载 HellaSwag
eval_needle.py	长上下文检索能力	5 depths × 3 needles	KV cache 量化的关键测试
eval_gsm8k.py	数学推理准确率	50-100	要求大模型才有信号

通用参数

所有脚本都支持:

• --model HF_ID — HuggingFace 模型名
• --device cuda|mps|cpu — 模型运行设备 (自动选择)
• --work-device cuda|mps|cpu — 量化计算设备, 默认同 --device
• --dtype fp32|fp16|bf16 — 模型精度 (B200 推荐 bf16)
• --configs fp,mse8,mse6,mse4,prod4,prod6,grouped4 — 逗号分隔的配置集

B200 推荐用法

# 一键全套评测 (selftest + hellaswag + needle + gsm8k + ppl)
MODEL=Qwen/Qwen2.5-7B-Instruct DTYPE=bf16 bash scripts/run_all.sh

# 或单独跑长上下文 (B200 最适合)
python scripts/eval_needle.py \
    --model Qwen/Qwen2.5-7B-Instruct \
    --dtype bf16 --work-device cuda \
    --context-len 32768 \
    --configs fp,mse8,mse4,prod4,grouped4

设计要点

• Fake quantization: 当前实现是 quantize → dequantize, cache 存储的仍是 fp32. 用于验证算法对模型输出的影响, 不追求真实显存/时间收益.
• Per-row L2 归一化: 真实 K 向量的 scale 变化很大 (std 从 2 到 20+). KVCacheAdapter 在量化前把每个 key 缩到 ||·||=√d, 量化后恢复原 norm.
• V 不量化: V 仅参与 softmax 后的加权求和, 对量化的敏感度远低于 K.
• Per-layer quantizer: 每层用独立 seed, 旋转矩阵不同. 代价低, 隔离层间相关性.

已知限制

• 无 FWHT, 旋转是 O(d²) 稠密矩阵乘 (生产需要 Fast Walsh-Hadamard 降到 O(d log d))
• 无 int packing, cache 仍是 fp32 (真实 4×/8× 显存收益需改存储格式)
• 无 fused kernel, attention 走标准 PyTorch (没有论文宣称的 8× 加速)
• PPL 对 KV 量化极敏感, 仅作相对对比, 不是绝对质量指标
• GSM8K 对 <1B 模型基线过低, 建议至少 1.5B+ 模型才有评测信号

参考文献

• TurboQuant (arXiv:2504.19874) — Google, 2025
• QJL (arXiv:2406.03482) — JL + sign quantization
• KIVI (arXiv:2402.02750) — 对比 baseline

License

MIT