大模型注意力机制深度解析：从 FlashAttention 到工业级最优组合

给定输入 X ∈ R^(n×d)，其中 n 为序列长度，d 为模型维度

1. 线性投影：Q = X·W_Q, K = X·W_K, V = X·W_V
2. 多头拆分：将 Q, K, V 拆分为 h 个头，每个头维度 d_k = d/h
3. 注意力计算：Attention(Q, K, V) = softmax(Q·K^T / √d_k)·V
4. 多头拼接：将 h 个头的输出拼接后做线性投影

HBM (40-80GB, 带宽 ~2TB/s)  ← 大但慢
    ↕ 数据搬运（瓶颈所在！）
SRAM (每 SM ~192KB, 带宽 ~19TB/s)  ← 小但快

分块 j 计算：
  S_ij = Q_i · K_j^T           // 分块矩阵乘
  m_ij = rowmax(S_ij)           // 当前块行最大值
  P_ij = exp(S_ij - m_ij)       // 数值稳定
  l_ij = rowsum(P_ij)           // 当前块行和

与之前的结果合并：
  m_new = max(m_prev, m_ij)     // 更新全局最大值
  l_new = exp(m_prev - m_new)·l_prev + exp(m_ij - m_new)·l_ij  // 修正累加和
  O_new = (l_prev·exp(m_prev - m_new)·O_prev + P_ij·V_j) / l_new  // 修正输出

# 标准 Attention（伪代码）—— 需要 HBM 中间结果
S = Q @ K.T                    # [n, n] 写入 HBM
P = softmax(S)                 # [n, n] 写入 HBM
O = P @ V                      # [n, d] 写入 HBM

# FlashAttention（伪代码）—— SRAM 内完成
for block_i in Q:              # 分块遍历 Q
    O_i = zeros(...)
    m_i = -inf
    l_i = 0
    for block_j in K, V:       # 分块遍历 K, V
        S_ij = block_i @ block_j.T    # SRAM 内计算
        m_new = max(m_i, rowmax(S_ij))
        P_ij = exp(S_ij - m_new)      # SRAM 内 Softmax
        # 在线修正 O_i
        O_i = (exp(m_i - m_new) * l_i * O_i + P_ij @ block_j_v) / l_new
        m_i, l_i = m_new, l_new
    write O_i to HBM          # 只写回最终结果

MHA (Multi-Head Attention):
  Q: h 个头    K: h 个头    V: h 个头
  每个 Q 头独立对应一个 K/V 头
  → 推理时 KV Cache 最大，但质量最高

MQA (Multi-Query Attention, 2019):
  Q: h 个头    K: 1 个头    V: 1 个头
  所有 Q 头共享同一组 K/V
  → 推理时 KV Cache 最小，但质量有损

GQA (Grouped Query Attention, 2023):
  Q: h 个头    K: g 个头    V: g 个头    (1 < g < h)
  每 (h/g) 个 Q 头共享一组 K/V
  → 在 MHA 和 MQA 之间取得平衡

Q_heads:  [Q0, Q1, ..., Q7]  [Q8, Q9, ..., Q15]  ...  [Q56, ..., Q63]
              ↕ 共享               ↕ 共享                   ↕ 共享
KV_heads:      K0, V0              K1, V1                   K7, V7

标准 MHA:
  K = X · W_K    → K ∈ R^(n×h×d_k)     // 直接存储完整 KV
  V = X · W_V    → V ∈ R^(n×h×d_k)

MLA:
  c_KV = X · W_DKV    → c_KV ∈ R^(n×d_c)    // 压缩到低维隐向量！
  K = c_KV · W_UK     → K ∈ R^(n×h×d_k)      // 上投影恢复
  V = c_KV · W_UV     → V ∈ R^(n×h×d_k)      // 上投影恢复

Q = [c_Q · W_UQ | x · W_QR]       // Q 分为 content 和 position 两部分
K = [c_KV · W_UK | x · W_KR]      // K 同样分为 content 和 position 两部分

其中：
  c_Q · W_UQ 部分：不携带位置信息，可以与 c_KV · W_UK 吸收合并
  x · W_QR / W_KR 部分：携带 RoPE 位置信息，单独计算

问题 1：预分配浪费
  为最大序列长度预分配连续显存 → 大量浪费（实际序列通常远短于最大长度）

问题 2：显存碎片
  不同请求的序列长度不同，频繁分配/释放导致碎片化
  → 即使总显存足够，也无法为新请求分配连续空间

问题 3：无法共享
  同一 prompt 的不同生成请求，prefix 部分的 KV Cache 完全相同
  → 但在连续分配模式下无法共享

传统方式：
  KV Cache 连续存储：[Request1_KV | Request2_KV | Request3_KV | ... 空闲 ...]
  → 碎片化，无法共享，预分配浪费

Paged Attention：
  将 KV Cache 划分为固定大小的 Page（块）：
  Physical Pages:  [P0][P1][P2][P3][P4][P5][P6][P7]...
  
  Request1 的 Page Table: 逻辑页0→P0, 逻辑页1→P3, 逻辑页2→P5  (非连续！)
  Request2 的 Page Table: 逻辑页0→P0, 逻辑页1→P3  (共享 prefix！)
  Request3 的 Page Table: 逻辑页0→P1, 逻辑页1→P2

传统注意力：Q_i 关注 K_0...K_{i-1}（连续内存）
Paged Attention：Q_i 关注其 Page Table 中所有物理块内的 K/V（非连续内存）

每个 Block 存储固定数量（如 16）的 token 的 KV：
  Block b: K_b ∈ R^(block_size × h × d_k), V_b ∈ R^(block_size × h × d_k)

注意力计算时，按 Block 遍历 Page Table 中的物理块

当多个请求共享 prefix 时：
  共享的 Block 只存一份，引用计数 > 1
  当某个请求需要修改某个 Block 时：
    → 复制该 Block（Copy-on-Write）
    → 修改副本，原 Block 引用计数 -1

热门 system prompt 的 KV Cache 可以常驻显存：
  /v1/chat/completions 中 system 字段相同 → 复用 KV Cache
  → 首 token 延迟大幅降低

1. 请求到达 → 分配逻辑 Page Table
2. Prefill 阶段 → 计算 KV Cache，按 Block 写入物理页
3. Decode 阶段 → 每生成一个 token，写入当前 Block
   → Block 满 → 分配新物理页，更新 Page Table
4. 请求完成 → 释放物理页（引用计数 -1，为 0 则回收）

对于二维特征 [x0, x1]，位置 m 的旋转编码：

RoPE(x, m) = [x0·cos(mθ) - x1·sin(mθ), x0·sin(mθ) + x1·cos(mθ)]

即：将特征向量在二维平面上旋转 mθ 角度

Q_m · K_n^T = (RoPE(Q, m)) · (RoPE(K, n))^T
            = f(Q·K^T, m-n)     // 内积只依赖相对位置 m-n！

原始 RoPE：训练长度 L，推理长度 > L 时性能急剧下降

NTK-aware Scaling：
  修改 base 频率：base_new = base × (scale_factor)^(d/(d-2))
  → 高频分量保持不变，低频分量按比例缩放
  → 实现近乎无损的长度外推

┌─────────────────────────────────────────────────────────────┐
│                    大模型推理技术栈                           │
│                                                             │
│  ┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────────┐   │
│  │   RoPE   │  │   GQA    │  │   FA     │  │  Paged   │   │
│  │ 位置编码  │  │ KV头复用  │  │ 计算加速  │  │ 显存管理  │   │
│  └──────────┘  └──────────┘  └──────────┘  └──────────┘   │
│       ↓             ↓             ↓             ↓          │
│  相对位置信息    KV Cache压缩   IO感知加速    碎片消除       │
│  长度外推能力    推理延迟降低   精确注意力     按需分配       │
│                               无损加速       跨请求共享     │
└─────────────────────────────────────────────────────────────┘

1. 请求到达 → Paged Attention 分配 Page Table

2. Prefill 阶段（处理 prompt）：
   a. Token Embedding → 线性投影得到 Q, K, V
   b. RoPE：对 Q, K 施加旋转位置编码
      Q = apply_rope(Q, positions)
      K = apply_rope(K, positions)
   c. GQA：将多个 Q 头映射到共享的 KV 组
      K_grouped = repeat_kv(K, n_rep=h//g)
      V_grouped = repeat_kv(V, n_rep=h//g)
   d. FlashAttention：分块计算注意力
      O = flash_attention(Q, K_grouped, V_grouped)
      → 不在 HBM 中实例化注意力矩阵
   e. 将 KV Cache 写入 Paged Attention 的物理页

3. Decode 阶段（逐 token 生成）：
   a. 新 token → Q_new, K_new, V_new
   b. RoPE：对新 Q, K 施加位置编码
   c. GQA：KV 头复用
   d. FlashAttention：Q_new 与全部 KV Cache 计算注意力
      → 从 Paged Attention 的物理页中读取 KV Cache
   e. 新 KV 写入当前物理页，页满则分配新页

4. 请求完成 → Paged Attention 释放物理页

2023: MHA + 标准 Attention + 连续 KV Cache
      ↓
2024: GQA + FlashAttention-2 + Paged Attention (vLLM)
      ↓
2025: GQA + FlashAttention-3 + Paged Attention + Prefix Caching
      ↓
未来?: MLA-like 压缩 + FlashAttention-3 + 更智能的 KV Cache 管理