GQA MQA

以 Llama 2 7B 的一层 decoder 为例，batch = 1，输入 prompt 长度为 T 个 token。

MHA 固定参数：hidden = 4096，head 数 = 32，head_dim = 128。

GQA（n_kv = 8）的 prefill

• 第 0 步：Embedding

  token ids: [T]
  查表 E: [32000, 4096]
  X = E[ids]                                  → X: [T, 4096]
  

• 第 1 步：RMSNorm

  X̂ = X / sqrt(mean(X², 按行) + ε) * γ        → X̂: [T, 4096]
  

• 第 2 步：K、V 的投影矩阵变窄

  Q = X̂ @ W_Q     [T,4096] @ [4096,4096]     → [T,4096] → [32, T, 128]   （不变）
  K = X̂ @ W_K     [T,4096] @ [4096,1024]     → [T,1024] → [8,  T, 128]
  V = X̂ @ W_V     [T,4096] @ [4096,1024]     → [T,1024] → [8,  T, 128]
  

  Q 还是 32 个 head，K 和 V 只有 8 个 head。

• 第 3 步：RoPE

  对 Q 和 K 的每个 head、每个位置 i，把 128 维两两配对，按位置 i 旋转对应角度。形状不变。V 不做 RoPE。

• 第 4 步：KV cache 只存 8 个 head

  cache_K[layer]: [8, T, 128]      （原来是 [32, T, 128]，缩小到 1/4）
  

• 第 5 步：每 4 个 query head 共用同一个 K

  Q 的 head 0、1、2、3   →  都用 K 的 head 0
  Q 的 head 4、5、6、7   →  都用 K 的 head 1
  ...
  Q 的 head 28..31       →  都用 K 的 head 7
  
  即 S_h = Q_h @ K_{h//4}ᵀ / √128              → 仍然是 [32, T, T]
  

• 第 6 步：causal mask

  S_h[i][j] 表示"第 i 个 token 看第 j 个 token"的分数，不允许看未来，所以 j > i 的位置设为 −∞。以 T = 4 为例：

         j=0   j=1   j=2   j=3
  i=0  [ s00   -∞    -∞    -∞  ]
  i=1  [ s10   s11   -∞    -∞  ]
  i=2  [ s20   s21   s22   -∞  ]
  i=3  [ s30   s31   s32   s33 ]
  

• 第 7 步：按行 softmax

  P_h = softmax(S_h, 每一行)                  → P: [32, T, T]
  

  每一行加起来等于 1，−∞ 的位置变成 0。

• 第 8 步：同样的对应关系用于 V

  O_h = P_h @ V_{h//4}                         → 仍然是 [32, T, 128]
  P 的 head 0  1  2  3  乘上 V 的 head 0
  P 的 head 4  5  6  7  乘上 V 的 head 1
  ...
  P 的 head 28 29 30 31 乘上 V 的 head 7
  

• 第 9 步：输出投影 + 残差

  A = O @ W_O     [T,4096] @ [4096,4096]     → [T, 4096]
  X = X + A
  

• 后面的 FFN

  之后进入 FFN（RMSNorm → gate/up 投影 → SiLU 相乘 → down 投影 → 残差），和 attention 的类型无关，这里略过。32 层都按这个流程走一遍。

MQA（n_kv = 1） 的 prefill

和 GQA 一样，只是共享得更彻底：

第 2 步：W_K、W_V 变成 [4096, 128]，K、V 只有 1 个 head：[1, T, 128]
第 4 步：cache 存 [1, T, 128]
第 5 步：S_h = Q_h @ K_0ᵀ / √128              32 个 query head 全部用同一个 K
第 8 步：O_h = P_h @ V_0