MLA 注意力与 KV 缓存

GLM-5.2 有 64 个注意力头且不用 GQA，如果用朴素注意力，KV-cache 每 token 要存 32,768 个 float——长上下文根本装不进消费级内存。colibrì 忠实实现了 MLA（Multi-head Latent Attention），把每 token 的 KV 压缩到 576 个 float（57 倍更小），还叠加了 DeepSeek 的权重吸收技巧和 GLM-5.2 的 DSA 稀疏注意力。核心是 attention（990 行）。

---

1. 压缩 KV-cache 的核心思想

朴素多头注意力要为每个头、每个 token 缓存完整的 k 和 v。MLA 的做法是：只缓存一个低秩 latent，k 和 v 在用时现场重建。

朴素（GLM-5.2, 64 头无 GQA）:
  每 token 缓存 = 64 头 × (256 k + 256 v) = 32768 floats

MLA（colibrì）:
  每 token 缓存 = kv_lora (latent) + qk_rope (旋转 key)
              ≈ 512 + 64 = 576 floats            ← 57× 更小
  k_nope 和 value 在注意力时通过 kv_b 现场重建

Model 结构的注释（115–117 行）明确记录了 576 vs 32768 这个比例。缓存本体是 Lc（latent）和 Rc（旋转 key）（118 行）。

---

2. attention 的数据流

attention（990–1140 行）分几个阶段。先看维度设置（991–996 行）：qh=qk_head、kvb_dim = H*(qk_nope+v_head)、Tk=pos_base+S。

2.1 down/up 投影 + 部分 RoPE + 写缓存

999–1012 行：

Q 路径:
  q_a 下投影 → qresid → RMSNorm(q_a_ln) → q_b 上投影 → qfull[H*qh]
  对每个头的 qk_rope 尾部（偏移 qk_nope）做交错 RoPE

KV 路径:
  kv_a 下投影 → comp[kv_lora + qk_rope]
    ├─ 前 kv_lora 段 → RMSNorm(kv_a_ln) → 写入 Lc[layer][pos]   ← latent
    └─ 后 qk_rope 段 → RoPE（所有头共享）→ 写入 Rc[layer][pos]  ← 旋转 key

交错部分 RoPE 是 rope_interleave（600 行）：只对每个头的 qk_rope 维做旋转（qk_nope 维不旋转），把偶/奇对旋转成 v[j]=a*cs-b*sn、v[half+j]=b*cs+a*sn。注意 key 的旋转部分是所有头共享的一小段——这正是 MLA 能压缩的原因。

2.2 DSA lightning indexer（长上下文稀疏）

1013–1067 行：GLM-5.2 的 DSA（DeepSeek Sparse Attention）”lightning indexer”，只在 full 层（idx_type[layer]）且上下文超过 index_topk 时激活。它用一套独立的轻量权重（ix_wk/ix_wq/ix_wp）算出每个 query 对每个历史 key 的分数，然后选出 top-2048 个因果 key存进 dsa_sel，后续 shared 层复用这个选择。

• indexer key：ix_wk + layernorm(eps 1e-6) + RoPE，写入 Ic（1021–1027 行）。
• indexer query 与打分：ix_wq/ix_wp，ReLU-gated 分数（1040–1054 行）。
• top-keep 选择：qsort + 阈值扫描进 dsa_sel/dsa_nsel（1055–1064 行）。

DSA=0 禁用，DSA_TOPK 覆盖数量。README 说这个 indexer 权重约 189 MB，用 --indexer 转换模式从 FP8 仓库提取。验证过：把选择强制成”保留所有 key”能逐 token 复现稠密注意力。

2.3 两条注意力路径：吸收 vs 非吸收

g_absorb（984 行）决定走哪条：-1 自动（decode S<=4 时吸收），0 从不，1 总是。

吸收路径（1068–1106 行，decode 用）——DeepSeek 权重吸收技巧：

不重建每个 token 的 k/v，而是：
  1. 把 W_K（kv_b 的一部分）折叠进 query：qt_addrow → qabs[kv_lora]
  2. 直接用缓存的 latent Lt + 旋转 key 打分（避免 kv_b 重建）
  3. 注意力后，输出 latent clat → qt_matvec_rows(kv_b) 重建上下文

注释（1068–1071 行）说明成本从 O(T·H·(nope+vh)) 降到 O(T·kv_lora)——decode 时每 token 省下的 k/v 重建是关键。用到的助手是 qt_addrow（958 行） 和 qt_matvec_rows（970 行）。

非吸收路径（1107–1136 行，prefill 用）：一次 kv_b 矩阵乘重建所有 token 的 k_nope + value（计时进 t_kvb），然后做因果注意力 score = q_pass·k_nope + q_rot·k_rot、softmax、value 累加。

两条路都以 o 输出投影（1102/1137 行）收尾。

README 声称吸收路径经过精确验证：TF 32/32 和生成 20/20 在”处处强制吸收”下逐 token 一致。

---

3. KV 分配与绑定

• kv_alloc（1376 行）：释放旧缓冲；为每个 full 层分配 DSA 的 Ic[i]；为 NR = n_layers+1 个层分配 Lc[i]（max_t*kv_lora）和 Rc[i]（max_t*qk_rope）——多出的一行给 MTP 层（1386–1389 行）。
• kv_bind（1393 行）：把 m->kv 和所有别名指向某个 KVState——这是多 slot 服务（10）切换会话上下文的机制。

---

4. KV 磁盘持久化：对话”暖开机”

colibrì 的一个独特能力：对话在引擎重启后重新打开时是”暖”的——服务/聊天模式每回合把压缩的 MLA KV 追加到 .coli_kv，重启时恢复，零重新 prefill。设计注释（1856–1862 行）：append-only、约 182 KB/token、nrec 最后写以保证崩溃一致性。

• g_kvsave（1863 行）（KVSAVE=0 禁用）、KV_MAGIC "COLIKV1\0"（1864 行）。
• kv_hdr（1865 行）：8 个 int32 头（n_layers、kv_lora、qk_rope、DSA index_hd、nic、vocab、nrec、0）。
• kv_disk_append（1880 行）：每条记录 = token id + 每层 Lc+Rc（+ DSA Ic）；只追加新位置，先 flush 数据再改写 nrec（1899–1900 行）。
• kv_disk_load（1903 行）：校验 magic + 头字段，拒绝 nrec >= maxctx-8-g_draft 的情况（1914 行），把记录读回 Lc/Rc/Ic。

README 说这个持久化验证过 byte-identical：关掉聊天、明天重开，模型还记得整段对话。:reset 清除，KVSAVE=0 禁用。多 slot 时每个 slot 有自己的文件（.coli_kv、.coli_kv.1…，见 10）。

---

5. 小结

机制	出处	收益
MLA 压缩 KV（576 floats）	attention:990, Model 注释:115	KV-cache 小 57×，长上下文装得下
交错部分 RoPE	rope_interleave:600	只旋转 qk_rope 段，key 旋转部分头间共享
权重吸收（decode）	attention:1068, qt_addrow:958	避免每 token 重建 k/v
DSA lightning indexer	attention:1013	长上下文选 top-2048 因果 key
KV 磁盘持久化	kv_disk_append:1880	对话暖开机，零重新 prefill