实现 TurboQuant (第5部分)

这是一个将 Google Research 的 TurboQuant (arXiv:2504.19874) 论文用 PyTorch 实现的项目。

在 LLM 推理时，最主要的内存瓶颈是 KV Cache，它在不增加量化开销的情况下被压缩。

TurboQuant = PolarQuant (b-1 bit) + QJL (1 bit)
             主信号压缩              残差校正 (无偏)
项目 传统方法 (KIVI 等) TurboQuant
量化常数开销 +1 ~ +2 bit 0 bit
3-bit 压缩时实际位数 4 ~ 5 bit 3 bit
内积估计偏差 存在 不存在
KV Cache 压缩率 ~4x ~6x

需求

硬件

项目 最低 推荐
GPU 支持 CUDA 的 NVIDIA GPU RTX 3090 24GB
RAM 16GB 32GB
磁盘 20GB（模型缓存） 50GB

软件

项目 版本
OS Ubuntu 20.04+（包括 WSL2）
Python 3.10 ~ 3.12
NVIDIA 驱动 595+
CUDA 13.0+（驱动 13.2 向下兼容）
PyTorch 2.10+（cu130）
Transformers 4.40+
vLLM 0.8+（serve_vllm.py 专用，可选）
gcc 必需（用于 Unsloth/Triton 构建）

在 Windows 上通过 WSL2 Ubuntu 运行。 PyTorch 自带 CUDA 运行时，因此不需要单独安装 CUDA Toolkit。 WSL 必需：使用 Unsloth 时需要 C 编译器：sudo apt update && sudo apt install -y gcc

DeepSeek 模型指南 (RTX 3090 24GB)

模型 参数 VRAM 24GB 可用
DeepSeek-R1-Distill-Qwen-1.5B 1.5B ~3GB O
DeepSeek-R1-Distill-Qwen-7B 7B ~14GB O (默认)
DeepSeek-R1-Distill-Qwen-14B 14B ~9GB（4-bit） O (--load-in-4bit)
DeepSeek-R1-Distill-Llama-70B 70B ~35GB+ X

项目结构

[Github] https://github.com/javafa/turboquant 

turboquant/
├── requirements.txt             # Python 依赖
│
├── turboquant.py                # 核心库 (PolarQuant + QJL + TurboQuant)
├── compare.py                   # 原始 vs TurboQuant 比较（质量 + 时间 + 内存）
├── benchmark.py                 # 综合基准测试（6种组合 x 内存/速度比较）
├── demo.py                      # HuggingFace 模型 KV cache 提取 + 压缩演示
├── serve_vllm.py                # vLLM OpenAI-compatible 服务器 + 分析
│
└── scripts/
    ├── setup_wsl.sh             # WSL Ubuntu 环境设置
    └── setup_conda.bat          # Windows Conda 环境设置

文件说明

文件 角色
turboquant.py PolarQuant, QJL, TurboQuant, TurboQuantKVCache 类实现。支持批处理。
compare.py 主脚本。 模型推理后比较原始 KV cache 与 TurboQuant 压缩结果。一次性输出 Attention score 质量、处理时间、内存节省。
benchmark.py 综合基准测试。 Baseline/BnB 4-bit/TurboQuant/Unsloth 等 6种组合的 VRAM，KV cache，速度在短/长提示中比较。
demo.py 从 HuggingFace 模型中提取 KV cache，并测量 TurboQuant 压缩率/准确度。
serve_vllm.py 基于 vLLM 的 OpenAI-compatible API 服务器 + TurboQuant KV cache 分析。

2. 综合基准测试（6种组合比较）

比较组合：

# 配置 权重 KV Cache
1 Baseline (FP16) FP16 FP16
2 BnB 4-bit 4-bit FP16
3 TurboQuant 3-bit FP16 3-bit
4 BnB 4-bit + TurboQuant 4-bit 3-bit
5 Unsloth 4-bit 4-bit FP16
6 Unsloth 4-bit + TurboQuant 4-bit 3-bit

基准测试结果（DeepSeek-R1-Distill-Qwen-1.5B，RTX 3090 24GB）：

短提示（4 个 token 输入 → 64 个 token 生成）：

Config                       | Model VRAM |  Peak VRAM |   KV Cache |      Speed |     Time
-----------------------------|------------|------------|------------|------------|----------
Baseline (FP16)              |    3022MB  |    3396MB  |    1.5 MB  |   39.4t/s  |    1.6s
BnB 4-bit                    |    1078MB  |    1578MB  |    1.5 MB  |   28.1t/s  |    2.3s
TurboQuant 3-bit             |    3022MB  |    3396MB  |  0.9MB(40%↓) |   48.5t/s  |    1.3s
BnB 4-bit + TurboQuant       |    1078MB  |    1578MB  |  0.9MB(40%↓) |   27.8t/s  |    2.3s
Unsloth 4-bit                |    1754MB  |    1974MB  |    1.5 MB  |   23.2t/s  |    2.8s
Unsloth 4-bit + TurboQuant   |    1754MB  |    1974MB  |  0.9MB(40%↓) |   37.2t/s  |    1.7s

长提示（191 个 token 输入 → 256 个 token 生成）：

Config                       | Model VRAM |  Peak VRAM |   KV Cache |      Speed |     Time
-----------------------------|------------|------------|------------|------------|----------
Baseline (FP16)              |   3391MB  |   3418MB  |   12.2 MB  |   27.0t/s  |    9.5s
BnB 4-bit                    |   1555MB  |   1595MB  |   12.2 MB  |   17.0t/s  |   15.1s
TurboQuant 3-bit             |   3399MB  |   3418MB  |  7.3MB(40%↓) |   28.5t/s  |    9.0s
BnB 4-bit + TurboQuant       |   1555MB  |   1595MB  |  7.3MB(40%↓) |   16.6t/s  |   15.4s
Unsloth 4-bit                |   1754MB  |   1979MB  |   12.2 MB  |   22.3t/s  |   11.5s
Unsloth 4-bit + TurboQuant   |   1754MB  |   1979MB  |  7.3MB(40%↓) |   23.6t/s  |   10.8s

与基线(FP16)的比较:

配置 显存 KV Cache 速度 (短) 速度 (长)
BnB 4-bit -53.5% 相同 -28.6% -37.3%
TurboQuant 3-bit 相同 -39.9% +23.1% +5.3%
BnB 4-bit + TurboQuant -53.5% -39.9% -29.4% -38.5%
Unsloth 4-bit -42% 相同 -41.1% -17.5%
Unsloth 4-bit + TurboQuant -42% -39.9% -5.5% -12.7%
  • TurboQuant是唯一比基线更快的配置 (短 +23%, 长 +5%)
  • 内存优化: BnB 4-bit + TurboQuant (显存 -53%, KV -40%)
  • 平衡: Unsloth + TurboQuant (显存 -42%, KV -40%, 速度小幅下降)
  • BnB 4-bit在长序列中由于去量化成本累积导致速度下降严重
  • Unsloth短结果包含Triton内核初始编译开销

输出项目:

部分 内容
1. 加载模型 模型加载时间, 显存使用量
2. 运行推理 原始生成结果, 速度 (tok/s)
3. 注意力得分比较 余弦相似度, Top-1/5 匹配, KL散度 (2/3/4-bit)
4. 时间基准 原始 vs TurboQuant 每个查询的处理时间 (ms)
5. 生成级别 注意力模式 Top-1 一致率
6. 内存节省 FP16 vs TurboQuant 内存比较
7. 缩放预测 长上下文(1k~128k)内存预测

CLI 选项

benchmark.py

--models             all | 1.5b | 7b (默认: all)
--skip-baseline      跳过FP16基线
--short-tokens       短提示生成的令牌数 (默认: 64)
--long-tokens        长提示生成的令牌数 (默认: 256)
--turboquant-bits    TurboQuant比特数 (默认: 3)

compare.py

--model              模型ID (默认: deepseek-ai/DeepSeek-R1-Distill-Qwen-7B)
--bits               TurboQuant比特数, 2-8 (默认: 3)
--max-new-tokens     生成的令牌数 (默认: 64)
--prompt             输入提示
--dtype              float16 | bfloat16 (默认: float16)
--load-in-4bit       4-bit量化加载 (14B模型用)
--load-in-8bit       8-bit量化加载
--layers-to-compare  比较的层数 (默认: 4)

demo.py

--model              模型ID
--bits               TurboQuant比特数 (默认: 3)
--max-new-tokens     生成的令牌数 (默认: 64)
--prompt             输入提示
--dtype              float16 | bfloat16
--load-in-4bit       4-bit量化加载
--load-in-8bit       8-bit量化加载

serve_vllm.py

--model              模型ID
--host               服务器主机 (默认: 0.0.0.0)
--port               服务器端口 (默认: 8000)
--max-model-len      最大序列长度 (默认: 4096)
--gpu-memory-utilization  GPU内存利用率 (默认: 0.90)
--quantization       vLLM量化: awq, gptq 等
--analyze-only       仅执行分析，不启动服务器
--turboquant-bits    TurboQuant比特数 (默认: 3)

算法概述

[压缩]
  k (原始键向量)
    ├── PolarQuant (b-1 bits): 随机旋转 → 极坐标 → 均匀量化
    │     存储: (r, codes)         开销: ~0.25 bit (一个r)
    │
    ├── 残差: residual = k - decompress(r, codes)
    │
    └── QJL (1 bit): sign(S @ residual)
          存储: sign_bits            开销: 0 bit

[估计 <q, k>]
  <q, k> ≈ <q, k_hat> + sqrt(pi/2) * (S@q)^T * sign_bits
            PolarQuant    QJL 残差修正 (无偏)

参考文献