KV Cache 与 PagedAttention：显存复用与长上下文推理

KV Cache 是什么：
□ 解码时每个 token 的注意力要复用历史 Key/Value
□ 缓存下来 → 避免每步重复计算前文注意力
□ 每层、每个注意力头各存一份

显存公式：
KV = 2 × layers × kv_heads × head_dim × seq_len × bytes
□ 2：Key 与 Value 各一份
□ kv_heads：GQA 后的 KV 头数
□ seq_len：上下文长度；bytes：精度（FP16=2，INT8=1）

def kv_cache_size(layers, kv_heads, head_dim, seq_len, bytes_per=2):
    return 2 * layers * kv_heads * head_dim * seq_len * bytes_per

# Llama-2-70B：80 层，GQA 8 个 KV 头，head_dim=128
for seq in (2048, 8192, 32768):
    gb = kv_cache_size(80, 8, 128, seq) / 1024**3
    print(f"seq_len={seq:>6} → {gb:.2f} GiB")
# 2048 → 0.63 GiB；8192 → 2.50 GiB；32768 → 10.00 GiB
# batch=8、seq=32K → 80 GiB，直接吃掉整张 A100

传统方案的问题：
□ 预分配 (max_batch, max_seq_len) 的连续 KV Cache
□ 100 token 与 4000 token 的请求占一样大的张量
  → 内部碎片吃一半以上显存
□ batch 容量由「显存」决定，而不是由「算力」决定

PagedAttention 的设计：
□ KV Cache 切成固定块（默认 16 个 token）
□ 逻辑连续的序列 → 物理上存放在不连续的块里
□ 每请求一张块表（block table）：逻辑位置 → 物理块
□ 空闲块进全局池，任何请求按需分配
□ 内部碎片只剩「最后一个块」，块池共享消除外部碎片

物理布局示意：
请求 A（35 token）→ 块 7、块 12、块 3（末块只用 3/16）
请求 B（10 token） → 块 5
请求 C（70 token） → 块 8、块 1、块 2、块 0、块 9
块表（A）：[逻辑 0-15]→块 7、[16-31]→块 12、[32-34]→块 3

三大运行时组件：
□ Block Allocator：显存池按块管理，支持换出到 CPU RAM
□ Block Table：每请求一张，decode 时按块加载 KV
□ Reference Count + 写时复制（COW）：
  - 物理块可被多请求共享 → 计数递增
  - 写共享块时先复制一份私有块

共享场景的价值：
□ 同一 prompt 并行采样 4 条答案 → 前缀 KV 完全共享
□ beam search：分叉前共享，分叉时 COW
□ RAG/多轮：system prompt 被大量请求复用

class Block:
    def __init__(self, bid):
        self.bid, self.refcount = bid, 1

def copy_on_write(seq, logical_idx):
    old = seq.block_table[logical_idx]
    if old.refcount > 1:        # 有人共享 → 复制一份
        new = alloc_block()
        new.data.copy_(old.data)
        old.refcount -= 1
        seq.block_table[logical_idx] = new.bid

量化原理：
□ FP16 → INT8：显存减半（2 字节 → 1 字节），主流
□ FP16 → FP8（E4M3）：减半，精度略好于 INT8
□ 常见方案：per-channel / per-token 缩放
□ 长上下文下累积误差放大 → 敏感层保留 FP16（混合精度）

收益评估：
□ 单请求 seq=32K、80 层：FP16≈10 GiB → INT8≈5 GiB
□ 省下的显存直接换成更大 batch 或更长 max_model_len
□ 代价：1%-2% 困惑度上升（长上下文更明显）

# vLLM 开启 KV Cache FP8 量化，配合长上下文
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Meta-Llama-3-8B-Instruct \
    --kv-cache-dtype fp8 \
    --max-model-len 32768

为什么可以剪：
□ 注意力天然稀疏：多数 token 只 attend 到少数关键位置
□ 长上下文中大量是噪音（重复文本、无关段落）

两类经典方法：
□ H2O（Heavy Hitter Oracle）：
  - 保留累计注意力分数最高的 token
□ StreamingLLM：
  - 保留初始 token（attention sink）保证数值稳定
  - 滑动窗口保留最近 token，中间直接丢弃
  - 显存占用与上下文长度「脱钩」，O(seq) → O(budget)

组合：初始 token + 最近窗口 + 高注意力 token

sink_size, window_size = 4, 256
h2o_size = 512 - sink_size - window_size

def should_cache(pos, scores, is_recent):
    if pos < sink_size or is_recent:
        return True                  # 初始/最近：必缓存
    return pos in scores.argsort()[-h2o_size:]  # 高注意力

共享场景：
□ System prompt：几乎所有请求共享同一段指令
□ RAG 上下文：同一文档被多次检索提问
□ 多轮对话 / 并行采样：历史 KV 天然可复用

两种形态：
□ vLLM Prefix Caching（COW 物理共享）：
  - 请求到达时按 token 前缀匹配缓存块，命中则跳过 prefill
□ SGLang RadixAttention（结构共享）：
  - 前缀组织成前缀树，任意前缀细粒度共享

收益：
□ 共享前缀场景 TTFT 降 50%-90%
□ 省掉的是 prefill 计算 → 相当于白赚一批算力

# vLLM 开启前缀缓存（RAG/Agent 场景强烈建议）
python -m vllm.entrypoints.openai.api_server \
    --model ... \
    --enable-prefix-caching \
    --max-model-len 32768

命中率优化：
□ 把高复用前缀放在 prompt 开头（缓存按前缀匹配）
□ 保持 system prompt 完全一致，避免拼接顺序抖动

显存侧组合拳：
□ KV 量化：INT8/FP8 减半
□ 剪枝/淘汰：StreamingLLM 封顶缓存
□ 前缀共享：跨请求省显存
□ 分级存储：热 KV 在 GPU，冷 KV 在 CPU/NVMe（offload）

精度侧：混合精度（敏感层 FP16）、RoPE 大角度数值稳定性

算法侧：
□ 稀疏/近似注意力：局部窗口 + 全局 token
□ 上下文压缩：先摘要、再推理（长文档分段问答）

工程侧：
□ max_model_len 与 batch 容量互为代价
□ 长请求与短请求混合调度（chunked prefill 交错）

容量规划示例（单卡 A100 80GB，FP8 KV）：
权重（8B FP16）≈16 GiB + 激活/工作区 ≈10 GiB
剩余 ≈54 GiB 全给 KV → seq=32K 单请求 KV≈0.5 GiB
→ 理论并发 ~100 条，实际受调度/共享/碎片影响需压测

vLLM 核心参数：
□ --gpu-memory-utilization：0.85-0.95
  - 太高：kernel 工作区不足，长序列 OOM
  - 太低：KV 池小，并发受限
□ --max-model-len：单请求上下文上限，与并发互为代价
□ --max-num-seqs：最大并发请求数（防 OOM 保险丝）
□ --enable-prefix-caching：共享前缀场景必开
□ --kv-cache-dtype：fp8/int8，长上下文推荐
□ --enable-chunked-prefill：长 prompt 与短请求混合时开

监控：/metrics 暴露 free/used blocks、前缀命中率、抢占次数

# 生产级启动：长上下文 + 前缀缓存 + FP8 KV
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Meta-Llama-3-8B-Instruct \
    --tensor-parallel-size 4 \
    --gpu-memory-utilization 0.92 \
    --max-model-len 65536 \
    --kv-cache-dtype fp8 \
    --enable-prefix-caching \
    --enable-chunked-prefill \
    --max-num-seqs 256

局限性：
□ 页式管理：极短序列（seq<64）时块表/分配开销可能抵消收益
□ KV 量化：长上下文 + 低精度 → 质量明显回退的案例常见
□ 剪枝：高信息密度任务（代码、数学）剪枝损失大
□ 前缀共享：请求前缀不一致（用户 prompt 在前）时命中率极低

高频踩坑：
□ gpu-memory-utilization 设 0.98 → 长序列请求偶发 OOM
□ 量化后没做端到端质量回归 → 生成悄悄变差
□ 开启 prefix caching 但拼接顺序抖动 → 命中率为零
□ 长上下文 + batch=1 单条大请求占满显存 → 无并发能力

选型建议：
□ 在线服务：PagedAttention + FP8 KV + prefix caching
□ 超长上下文：加剪枝/StreamingLLM 或摘要压缩
□ 极短交互：评估页式开销，短序列直接预分配