ds4.c 学习笔记

Appearance

Multi-head Latent Attention——DeepSeek 的压缩注意力机制，将 KV Cache 压缩到低维潜在空间，大幅减少缓存内存。

为什么需要 MLA ​

标准 Multi-Head Attention 中，每个 token 需要缓存所有头的 K 和 V 向量。DeepSeek V4 有 64 个 Q 头，head_dim = 512。MLA 的思路：不缓存完整的 K/V，而是缓存一个低维"潜在表示"，需要时再投影回高维。

核心原理 ​

DeepSeek V4 的 KV Cache 节省来自两个机制叠加：

  1. GQA（Grouped-Query Attention）：n_head_kv = 1，64 个 Q 头共享 1 组 KV，将标准 KV 从 64×512×2 = 65536 压缩到 1×512×2 = 1024 floats
  2. MLA 低秩压缩：KV 投影到 512 维潜在空间，每个位置存 512 floats（而非 1024）
  3. 序列压缩：部分层按 ratio-4 或 ratio-128 压缩序列长度（多个 token 共享一个压缩位置）
标准 Multi-Head KV Cache:  64 × 512 × 2 = 65536 floats/token/layer
仅 GQA:                    1 × 512 × 2  = 1024 floats/token/layer
MLA (GQA + 低秩压缩):      1 × 512      = 512 floats/token/layer
序列压缩后 (ratio=128):    512 / 128    ≈ 4 floats/token/layer

DeepSeek V4 Flash 使用 512 维的潜在表示，配合序列压缩，使得百万级上下文在本地成为可能。

在 ds4.c 中的实现 ​

KV Cache 存储压缩后的潜在向量（attn_comp_kv），以及压缩器的状态（attn_state_kv）。推理时通过投影矩阵将潜在表示解压为完整的 K/V 向量。

相关概念 ​

  • kv-cache — MLA 是 KV Cache 的压缩变体
  • softmax — 解压后的 KV 用于注意力计算
  • moe — 注意力层和 MoE 层在每层前向中串联

详见 Part 3 — 注意力机制。

出链