《Born》第17章：KV-Cache、Flash Attention 与推理优化

LLM 推理的成本主要来自两个瓶颈：重复计算历史 token 的键值（KV-Cache 解决），以及注意力机制的二次方复杂度（Flash Attention 解决）。

KV-Cache

Transformer 的自回归生成每次都要重新计算所有历史 token 的 K 和 V。KV-Cache 把它们存起来，只计算新 token：

// 没有 KV-Cache：每次 forward 都重新计算所有历史
// O(n²) 复杂度

// 有 KV-Cache：只计算新 token 的 K/V
// O(n) 复杂度（对当前 token）
type KVCache struct {
    K []Tensor // [layer, head, seq, headDim]
    V []Tensor
}

func (cache *KVCache) Update(layer int, newK, newV Tensor) {
    cache.K[layer] = cache.K[layer].Cat(newK, 2) // 沿 seq 维度拼接
    cache.V[layer] = cache.V[layer].Cat(newV, 2)
}

Flash Attention

标准注意力的内存访问模式对 GPU 不友好：

标准 Attention：
  1. 计算 S = Q @ K^T      → 加载 Q, K，写出 S
  2. 计算 P = softmax(S)   → 加载 S，写出 P
  3. 计算 O = P @ V        → 加载 P, V，写出 O
  
Flash Attention（分块 + 在线 softmax）：
  1. 把 Q, K, V 分成小块
  2. 每块在 SRAM 内完成 softmax + 输出
  3. 不需要写出中间矩阵 S 和 P

Born 的 flashAttentionShader 实现了此算法，WGSL 代码约 200 行。

推理优化总结

📘 《Born》连载技术书，第 17/22 章。