KV 缓存与 Flash Attention:推理优化

KV 缓存与 Flash Attention:推理优化

本节摘要:训练是并行且受 FLOPs 限制的,推理是串行且受内存带宽限制的——瓶颈不同,技巧也不同。一个朴素的解码器生成 N 个 token 要做 O(N²) 次注意力:每步都重算整个前缀。一个 4K token 的回复就是 1600 万次注意力操作,绝大多数是冗余的——前缀 token 的隐状态一旦算出来就是确定的,你只需用新 token 的查询对前面所有 token 缓存的键值做一次查询。更糟的是,标准注意力要把 N×N 分数矩阵物化到显存,N≥2K 时注意力先撞内存墙再撞算力墙,经典 kernel 浪费现代 GPU 410 倍。Dao 等人的两个优化把前沿推理从「慢」推到「快」:KV 缓存把每 token 的 K/V 存起来,每 token 注意力变成一次查询对缓存键,推理从 O(N²) 降到每步 O(N);Flash Attention 把注意力分块计算,N×N 矩阵永不落地到 HBM,softmax 和矩阵乘全在 SRAM 里做,A100 上 24 倍墙钟加速、H100 配 FP8 达 5~10 倍。本节带你纯标准库实现 KV 缓存、分块 softmax(模拟 Flash 的 running-max),数清朴素 vs 缓存解码的操作数差距(5050 vs 100),并打通 PagedAttention、连续批处理、投机解码这些 2026 年生产推理栈的标配。

学习目标

阅读完本节,你应当能够:

1. 说清训练并行 FLOPs-bound、推理串行 memory-bound的不对称,以及为什么朴素自回归解码是 O(N²)。
2. 实现 KV 缓存:存每个前缀 token 的 K/V,新 token 注意力变成一次查询对缓存键。
3. 算出 KV 缓存的显存占用:2·d_head·dtype 每层每头每 token`,并说清 GQA/MLA 为何是缓存大杀器。
4. 解释 Flash Attention 的分块技巧:把 N×N 矩阵挡在 HBM 外,用 running (max,sum) 做精确 softmax,内存足迹从 O(N²) 降到 O(N),输出与标准注意力 bit 一致。
5. 区分 2026 生产推理的四件套:PagedAttention(KV 当虚拟内存)、连续批处理(不排空就换入换出)、前缀缓存(跨请求复用系统提示)、投机解码(草稿+验证)。

一、问题与直觉

朴素自回归解码每步重算整个前缀,O(N²)。而且标准注意力把 N×N 分数矩阵物化到 HBM,读写了太多次——N≥2K 时注意力先变成 memory-bound。两个 Dao 等人的优化改写了局面。

KV 缓存的数学

每个解码器层、每个 token、每个头:

bytes_per_token_per_layer = 2 * d_head * dtype_size
                          ^
                          K 和 V 各一份
​

7B 模型(32 层、32 头、d_head=128、fp16):

每 token 每层 = 2 * 128 * 2 = 512 字节
每 token(32 层)= 16 KB
每 32K 上下文 = 512 MB
​

Llama 3 70B(80 层、d_head=128、GQA 8 个 KV 头):

每 token 每层 = 2 * 8 * 128 * 2 = 4096 字节(4 KB)
每 32K 上下文 = 10.4 GB
​

这 10 GB 就是为什么 Llama 3 70B 在 128K 上下文下,批大小 1 就吃掉 40 GB A100 的大半——光 KV 缓存。

💡 GQA 是 KV 缓存的胜利:MHA(64 头)会是 32 GB,GQA 把 KV 头压到 8,直接砍到四分之一;MLA(DeepSeek)压得更狠。架构选择直接决定能服务多长的上下文。

Flash Attention——分块技巧

标准注意力三次 HBM 往返:

S = Q @ Kᵀ     (HBM 读, N×N, HBM 写)
P = softmax(S) (HBM 读, HBM 写)
O = P @ V      (HBM 读, HBM 写)
​

H100 上 HBM 带宽 3 TB/s,SRAM 30 TB/s——每次 HBM 往返都比片上慢 10 倍。

Flash Attention 分块:

对每个 Q 块(tile ~128×128):
    把 Q_tile 载入 SRAM
    对每个 K、V 块:
        把 K_tile、V_tile 载入 SRAM
        算 S_tile = Q_tile @ K_tileᵀ     (SRAM)
        running softmax 聚合              (SRAM)
        累加进 O_tile                      (SRAM)
    把 O_tile 写回 HBM
​

每 tile 一次 HBM 往返,总内存足迹从 O(N²) 降到 O(N)。反向传播从前向重算部分值而非存储——又一个内存胜利。

数值技巧:running softmax 跨 tile 维护 (max, sum),最终归一化是精确的。Flash Attention 不是近似——它的输出与标准注意力 bit 一致(除 fp16 非结合性)。

版本演进:

Flash 4 发布时仅前向,训练仍用 Flash 3;GQA 和 varlen 支持在 2026 年中补齐。

投机解码——另一个延迟胜利

便宜模型提议 N 个 token,大模型并行验证全部 N 个。若验证接受 k 个,你为 k 个生成只付 1 次大模型前向。代码和散文上典型 k=35。2026 默认:EAGLE 2/Medusa(集成草稿头,共享验证者隐状态,23 倍加速无质量损失)、草稿模型投机解码(消费硬件 2~4 倍)、lookahead 解码(Jacobi 迭代,无需草稿模型,小众但免费)。(第 16 节详述)

连续批处理

经典批处理:等最慢的序列完成再开新批,短回复早完成时浪费 GPU。连续批处理(最早在 Orca,现 vLLM、TensorRT-LLM、SGLang 都有):旧序列一完成就把新请求换入,典型聊天负载吞吐 5~10 倍。

PagedAttention——KV 缓存当虚拟内存

vLLM 的招牌功能。KV 缓存按 16 token 块分配,页表把逻辑位置映射到物理块。于是可以跨并行样本(beam search、并行采样)共享 KV、为提示缓存热替换前缀、给内存去碎片。比朴素连续分配吞吐高 4 倍。

二、从零实现

完整代码见原课程 phases/07-transformers-deep-dive/12-kv-cache-flash-attention/code/main.py。实现:① 朴素 O(N²) 增量解码器、② O(N) KV 缓存解码器、③ 模拟 Flash Attention running-max 的分块 softmax。

Step 1:KV 缓存

class KVCache:
    def __init__(self, n_layers, n_heads, d_head):
        self.K = [[[] for _ in range(n_heads)] for _ in range(n_layers)]
        self.V = [[[] for _ in range(n_heads)] for _ in range(n_layers)]

    def append(self, layer, head, k, v):
        self.K[layer][head].append(k)
        self.V[layer][head].append(v)

    def read(self, layer, head):
        return self.K[layer][head], self.V[layer][head]
​

简单:按层、按头维护不断增长的每 token K、V 列表。

Step 2:分块 softmax

def tiled_softmax_dot(q, K, V, tile=4):
    """Flash 风格的 softmax(qKᵀ)V,带 running max/sum。"""
    m = float("-inf"); s = 0.0
    out = [0.0] * len(V[0])
    for start in range(0, len(K), tile):
        k_block = K[start:start + tile]
        v_block = V[start:start + tile]
        scores = [sum(qi * ki for qi, ki in zip(q, k)) for k in k_block]
        new_m = max(m, *scores)
        exp_old = math.exp(m - new_m) if m != float("-inf") else 0.0
        exp_new = [math.exp(sc - new_m) for sc in scores]
        s = s * exp_old + sum(exp_new)
        for j in range(len(out)):
            out[j] = out[j] * exp_old + sum(e * v[j] for e, v in zip(exp_new, v_block))
        m = new_m
    return [o / s for o in out]
​

输出与一次性的 softmax(qK)V bit 一致,但任意时刻工作集只是 tile × d_head 块,而非完整 N × d_head。

设计要点:running (m, s) 是 Flash 的精髓——softmax 在数学上需要全局 max 来归一化,但通过维护跨 tile 的运行 max 和加权和,可以在只看一块时逐步逼近正确的归一化因子。这是「精确而非近似」的关键。

Step 3:对比朴素 vs 缓存解码

数注意力操作:朴素 O(N²) = 5050,缓存 O(N) = 100。代码会打印两者。

三、框架对比:HuggingFace 与 vLLM

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B",
    attn_implementation="flash_attention_2",   # Hopper 用 FA3
    torch_dtype="bfloat16",
)
# generate() 自动用 KV 缓存
​

vLLM 生产:

vllm serve meta-llama/Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4 \
    --max-model-len 32768 \
    --enable-prefix-caching \
    --kv-cache-dtype fp8
​

跨请求的前缀缓存是 2026 的大赢——同一个系统提示、few-shot 示例、长上下文文档可跨调用复用 KV。对有重复工具提示的代理工作负载,前缀缓存常带来 5 倍吞吐。

四、可复用产物

原课程产出 outputs/skill-inference-optimizer.md:一个优化器 Skill,为新推理部署选注意力实现、KV 缓存策略、量化、投机解码。

五、练习

1. (Easy) 跑 code/main.py,确认朴素与缓存解码器输出相同,记录操作数差距(5050 vs 100)。
2. (Medium) 实现前缀缓存:给定提示 P 和多个补全,对 P 跑一次前向填 KV 缓存,再按补全分支。测相对每次重编码 P 的加速。
3. (Hard) 实现玩具 PagedAttention:KV 缓存按固定 16 token 块 + 自由链表;序列完成时把块还回池子。模拟 1000 条变长聊天补全,对比连续分配的内存碎片。

本节要点回顾

1. 训练并行 FLOPs-bound、推理串行 memory-bound:瓶颈不同,技巧也不同——朴素自回归 O(N²) 撞的是显存墙不是算力墙。
2. KV 缓存把推理降到 O(N):存每 token 的 K/V,新 token 注意力是一次查询对缓存键,100 token 生成从 5050 次操作降到 100 次。
3. KV 缓存显存公式:2·d_head·dtype 每层每头每 token;Llama 3 70B 32K 上下文光 KV 就 10.4 GB。
4. GQA/MLA 是缓存大杀器:KV 头从 64 压到 8(GQA)或更狠(MLA),直接决定能服务多长上下文。
5. Flash Attention 分块:N×N 不落地 HBM,softmax+matmul 全在 SRAM,内存足迹 O(N²)→O(N),输出 bit 一致(非近似)。
6. running (max,sum) 是 Flash 精髓:跨 tile 维护运行 max 和加权和,逐步逼近正确归一化因子——精确而非近似。
7. Flash 版本:Flash 1(2022 分块)→2(并行)→3(Hopper FP8)→4(2026 Blackwell,初期仅前向)。
8. 2026 生产四件套:PagedAttention(KV 当虚拟内存,4× 吞吐)、连续批处理(5~10×)、前缀缓存(代理负载 5×)、投机解码(草稿+验证)。

下一节,我们从「怎么快」转向「该训多大」——Chinchilla 缩放定律如何告诉你给定算力预算下,参数、数据、模型规模的最优配比。