KV-Cache-Compression-Report / kv_compress_plan.md
ljsysfurry's picture
Upload kv_compress_plan.md with huggingface_hub
f0ad681 verified
|
Raw
History Blame Contribute Delete
8.7 kB
# 🧠 纯文本 LLM KV 缓存压缩 — 详细定制方案
> **版本**: v1.0
> **日期**: 2026-08-06
> **适用**: 纯文本大语言模型推理(对话/长文/代码/文档)
> **目标**: 显存 -50%+、吞吐 +2x、精度损失 < 0.5 PPL
---
## 一、方案总览
### 1.1 核心思路
纯文本 LLM 的 KV 缓存压缩采用**分层正交叠加**策略:
```
┌─────────────────────────────────────────────────────┐
│ Layer 1: INT8 量化 (per-channel) → 显存 -50% │
│ Layer 2: StreamingLLM (attention sink) → 长上下稳定 │
│ Layer 3: H2O 驱逐 (重击者+近邻) → 再省 40% │
│ Layer 4: K/V 分离存储 (可选) → 显存换质量 │
└─────────────────────────────────────────────────────┘
每层独立开关, 可按场景组合
```
### 1.2 三层正交性
| 维度 | 压缩对象 | 对应层 | 说明 |
|------|----------|--------|------|
| 特征维度 | 每个 token 的 K/V 精度 | Layer 1 量化 | 与长度无关 |
| 序列维度 | 不重要的历史 token | Layer 3 驱逐 | 与精度无关 |
| 存储维度 | 冷数据放慢速介质 | Layer 4 分离 | 与上面都正交 |
| 结构维度 | 首 token 保留 | Layer 2 保底 | 免费防护 |
---
## 二、Layer 1: INT8 KV 量化(必做)
### 2.1 原理
```python
# per-channel 量化 (K/V 每通道一个 scale)
scale_k = max_abs(K, dim=channel) / 127
K_int8 = round(K / scale_k)
# 反量化在注意力计算时做
K_deq = K_int8 * scale_k
```
### 2.2 技术选型
| 方案 | 精度 | 速度 | 复杂度 | 推荐 |
|------|------|------|--------|------|
| **per-channel INT8** | 掉点 <0.3 | 快 | 低 | ✅ 首选 |
| per-token INT8 | 更稳 | 快 | 中 | ✅ 精度敏感时 |
| FP8 (E4M3) | 几乎无损 | 最快 | 高 | ⚠️ 需 H100/L40S 支持 |
| INT4 | 掉点 1-2 | 最快 | 高 | ❌ 纯文本不值当 |
### 2.3 实现要点
```
1. 量化时机: 每 token 写入 KV 缓存时即时量化 (不重算)
2. scale 计算: per-channel (head 维度), 避免 per-token 额外开销
3. 反量化位置: 在 FlashAttention 内核内融合, 不额外访存
4. 溢出处理: 用动态 scale (max-abs), 比固定 scale 稳
5. 特殊 token: attention sink 的 K/V 保持 FP16 (防精度损失)
```
### 2.4 预期收益
- 显存: K/V 从 FP16 → INT8, **-50%**
- 速度: 缓存命中率提升, decode 吞吐 **+30~50%**
- 精度: PPL 变化 < 0.3, 实际对话几乎无感
---
## 三、Layer 2: StreamingLLM 滑窗保底(必做)
### 3.1 原理
```
保留集合 = [attention_sink(前4个token)] + [最近 N 个token]
旧 token 滑出窗口后被丢弃, 但 sink 永不动
```
### 3.2 参数推荐
| 参数 | 默认 | 说明 |
|------|------|------|
| sink_tokens | 4 | 首 4 个 token 永存 |
| window_size | 2048 | 滑动窗口长度 |
| 触发时机 | 序列 > window 时 | 之前不启用 |
### 3.3 与 Layer 1 叠加
```
Layer 1 (量化) 对窗口内所有 token 生效
Layer 2 (滑窗) 决定哪些 token 保留
→ 实际显存 = window_size × 量化后大小
→ 长对话显存从 O(seq_len) 降到 O(window)
```
---
## 四、Layer 3: H2O 驱逐(可选, 按场景)
### 4.1 适用判定
| 场景 | 是否启用 | 原因 |
|------|----------|------|
| 多轮对话 | ⚠️ 谨慎 | 回头翻旧账可能丢信息 |
| 单轮长文 | ✅ 启用 | 注意力集中, 效果好 |
| RAG 检索 | ❌ 禁用 | 检索需全量上下文 |
| 代码补全 | ✅ 启用 | 代码局部性强 |
### 4.2 H2O 算法
```python
def h2o_evict(kv_cache, budget, attention_scores):
"""
1. 统计每个 token 的累计注意力分数 (heavy hitter 分数)
2. 保留 top-k 重击者 token
3. 加上最近 budget 内 token
4. 其余驱逐
"""
scores = cumulative_attention(attention_scores)
heavy = topk(scores, k=budget * 0.3) # 30% 重击者
recent = last_n_tokens(kv_cache, budget * 0.7) # 70% 最近
return heavy ∪ recent
```
### 4.3 参数
| 参数 | 默认 | 说明 |
|------|------|------|
| budget | 1024 | 驱逐后保留的 token 数 |
| heavy_ratio | 0.3 | 重击者占比 |
| 驱逐频率 | 每 64 token | 避免频繁计算 |
### 4.4 与推理模型的区别(重要)
```
⚠️ 推理模型 (DeepSeek-R1) 禁用 H2O:
思考链长且信息密度高, 驱逐会破坏推理连贯性
纯文本模型 (对话/总结/代码) 可用:
注意力集中, 驱逐掉低分 token 影响极小
```
---
## 五、Layer 4: K/V 分离存储(高级定制)
### 5.1 原理(HCAttention 思路)
```
┌─────────────────────────────────────────┐
│ GPU 高速显存 │
│ K_int8 (量化后, 只用来打分) ← 全保留 │
│ 最近窗口的 V (FP16) │
├─────────────────────────────────────────┤
│ CPU 内存 │
│ 旧 token 的完整 V (FP16) ← 按需调取 │
└─────────────────────────────────────────┘
```
### 5.2 工作流
```python
# 打分阶段 (GPU, 只要 K)
scores = Q @ K_int8.T / sqrt(d) # 全局近似打分
# 取用阶段 (按需)
if token 在 GPU 窗口: V = V_gpu
else: V = V_cpu[token] # 从 CPU 拉取
# 输出
out = softmax(scores) @ V
```
### 5.3 适用场景
| 场景 | 是否启用 | 原因 |
|------|----------|------|
| 多轮对话(客户翻旧账) | ✅ 启用 | 旧信息可回溯, 不丢 |
| 单轮长文 | ❌ 不必 | 滑窗够用 |
| 显存紧张 (24G 卡) | ✅ 启用 | 把显存让给 batch |
### 5.4 代价
- 冷 V 从 CPU 拉取: 延迟 +5~15ms(首 token 感知)
- 实现复杂度: 中高(需要 CPU/GPU 内存管理)
---
## 六、场景定制矩阵
| 业务场景 | Layer1 | Layer2 | Layer3 | Layer4 | 预期压缩 |
|----------|--------|--------|--------|--------|----------|
| **客服对话** | ✅ | ✅ | ⚠️ | ✅ | 4~6x |
| **长文总结** | ✅ | ✅ | ✅ | ❌ | 6~8x |
| **代码生成** | ✅ | ✅ | ✅ | ❌ | 6~8x |
| **RAG 问答** | ✅ | ✅ | ❌ | ✅ | 3~4x |
| **推理模型** | ✅ | ✅ | ❌ | ✅ | 3~4x |
| **批量离线** | ✅ | ❌ | ✅ | ❌ | 8x+ |
---
## 七、实现路线图(L40S 实测)
### Phase 1: Baseline(1 天)
```
1. 搭 vLLM 或自研推理脚本 (支持 7B/14B)
2. 跑通 FlashAttention 2
3. 记录 baseline: PPL / 吞吐 / 显存
```
### Phase 2: INT8 量化(1-2 天)
```
1. 实现 per-channel KV 量化
2. 内核融合反量化
3. 对比 PPL / 显存 / 吞吐
```
### Phase 3: StreamingLLM + H2O(1-2 天)
```
1. 实现滑窗 + attention sink
2. 实现 H2O 驱逐
3. 长上下文压力测试 (4k/8k/16k/32k)
```
### Phase 4: K/V 分离(2 天, 可选)
```
1. CPU 内存池管理
2. 冷热 V 迁移
3. 多轮回溯测试
```
### Phase 5: 评测 + 交付(1 天)
```
1. 三张核心曲线: 压缩率 vs PPL / 压缩率 vs 吞吐 / 压缩率 vs 延迟
2. 生成对比报告 (HTML, 类似 kv_cache_report)
3. 封装成可复用的推理加速模块
```
---
## 八、评测方案
### 8.1 精度指标
| 指标 | 工具 | 目标 |
|------|------|------|
| PPL | wikitext-2/4 | 变化 < 0.5 |
| 任务准确率 | MMLU / GSM8K | 掉点 < 1% |
| 长文检索 | LongBench | 保持 |
| 对话一致性 | 人工抽测 | 无明显劣化 |
### 8.2 性能指标
| 指标 | 定义 | 目标 |
|------|------|------|
| 显存占用 | KV 缓存峰值 | -50% |
| decode 吞吐 | tokens/s | +50% |
| TTFT | 首 token 延迟 | 不变 |
| 长上下文极限 | 可处理最大序列 | 2x 提升 |
---
## 九、交付物
```
1. 推理加速模块 (可插拔: kv_compress.py)
2. 四层开关配置 (yaml)
3. 评测报告 (HTML 带曲线)
4. 场景接入示例 (对话/长文/代码/RAG)
5. 文档: 参数调优指南
```
---
## 十、风险与对策
| 风险 | 概率 | 对策 |
|------|------|------|
| INT8 掉点超预期 | 低 | 回退 per-token 量化 / 关键层保 FP16 |
| H2O 在特定任务崩 | 中 | 按场景开关, 提供安全默认值 |
| K/V 分离延迟超标 | 中 | 预取 / 双缓冲 / LRU 热 V 常驻 |
| 长上下文 OOM | 低 | 滑窗兜底 + 显存监控自动降级 |
---
*本方案为纯文本 LLM KV 缓存压缩定制版, 基于 H2O / StreamingLLM / HCAttention / INT8 量化等技术路线综合设计。*