DeepSeek V4 源码剖析

第4章 Indexer：稀疏注意力的可学路由

作者 杨艺韬 · 7,833 字 · 发布于 · 更新于

"index_topk": 1024、"index_n_heads": 64、"index_head_dim": 128 —— DeepSeek-V4-Pro config.json

Indexer 是模型自己的"选择性阅读"——它决定了 1M context 中哪 1024 个压缩位置值得这一步 attention 真正去看。

4.1 引子：稀疏注意力的"选择题"

V4 的稀疏注意力面对一个工程"选择题"：

已经有了 1M 长度的压缩 KV 序列（通过 Compressor 压成 n_tok / ratio 组），attention 该从中选哪 top-k 组真正参与计算？

这个"选哪些"的问题，就是 Indexer 要解决的核心。它有几个看似可行但都不对的方案：

方案 A：随机选——理论上可行，但稀疏选取必须能"针对当前 query 选最相关的 KV"。随机选会丢掉关键 token。

方案 B：用主 attention 的 q · k 直接打分——但主 attention 的 q 维度是 512，dot product 成本仍然太高（1M / ratio 个位置都要算一遍）。

方案 C：用一个小型独立的 attention 给 KV 打分——这就是 Indexer 的方案。它有自己的小维度（head_dim=128，比主 attention 的 512 小很多）、自己的 KV 投影、自己的 head 数（64，比主 attention 的 128 少一半），把"打分"成本降到主 attention 的一小部分。

flowchart LR
  Q["主 attention 的 query<br/>[B, S, 128 heads, 512]"] --> MainAttn["主 attention 计算"]
  XInput["输入 x"] --> Indexer
  Indexer["Indexer (轻量)<br/>[B, S, 64 heads, 128]"] --> Score["score 打分"]
  Score --> TopK["topk(1024) 选取"]
  TopK --> SparseIdx["topk_idxs"]
  SparseIdx --> SparseAttn["sparse_attn(q, kv, sink, idxs)"]
  MainAttn -.数据流入.-> SparseAttn

Indexer 是个"小而专"的 attention——专门服务"选 top-1024"这一件事，不参与最终输出。

4.2 Indexer 类的源码结构

V4 的 Indexer 类（inference/model.py）的 __init__：

class Indexer(torch.nn.Module):
    def __init__(self, args: ModelArgs, compress_ratio: int = 4):
        super().__init__()
        self.dim = args.dim
        self.n_heads = args.index_n_heads               # 64
        self.n_local_heads = args.index_n_heads // world_size
        self.head_dim = args.index_head_dim             # 128
        self.rope_head_dim = args.rope_head_dim         # 64
        self.index_topk = args.index_topk               # 1024（04-24 首发版是 512，首发当天稍后即改成 1024）
        self.q_lora_rank = args.q_lora_rank             # 1536 (与主 attention 共享)
        self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim)
        self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16)
        self.softmax_scale = self.head_dim ** -0.5
        self.compress_ratio = compress_ratio
        self.compressor = Compressor(args, compress_ratio, self.head_dim, True)
        self.register_buffer("kv_cache",
            torch.zeros(args.max_batch_size, args.max_seq_len // compress_ratio, self.head_dim),
            persistent=False)
        self.freqs_cis = None

注意几个关键点：

1. q 的来源是主 attention 的 qr：

# Attention.forward 里
qr = q = self.q_norm(self.wq_a(x))         # qr: [B, S, q_lora_rank=1536]
q = self.wq_b(q).unflatten(-1, ...)         # 主 attention 用的 q
...
compress_topk_idxs = self.indexer(x, qr, start_pos, offset)  # qr 传给 Indexer

Indexer 的 wq_b 输入是 qr（来自主 attention 的 LoRA 中间表示），输出是 [B, S, 64, 128] 的 query。Indexer 与主 attention 共享 q_lora_rank=1536 的 wq_a 投影——这一层 LoRA 投影不重复算，节省 FLOPs。

2. Indexer 有自己的 Compressor（rotate=True）：

Indexer 的压缩 KV 是从 x 直接经过自己的 Compressor 算出来的——这个 Compressor 与主 attention 的 Compressor 是不同的实例，参数独立训练，量化策略也不同（FP4 + Hadamard 旋转）。

3. weights_proj：每 head 一个标量权重：

self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16)

输出形状 [B, S, n_heads=64]——每个 token 的每个 head 一个 BF16 权重。这些权重在 forward 中作用是"head-level 的打分加权"——某些 head 在某个 token 上"更可信"，权重就大。

4. softmax_scale 用 head_dim=128 而非主 attention 的 512：

softmax_scale = 128 ** -0.5 ≈ 0.088——就是 dot product 的标准 1/d1/\sqrt{d} 缩放。注意 Indexer 里并没有 softmax：这个系数在 forward 里被折进了 weights（weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)），作用只是把 q · k 的量级摁回 O(1)O(1)，让跨 head 求和之后的总分不随 head_dim / n_heads 漂移。

4.3 forward 流程与 score 计算

Indexer 的 forward 是 V4 稀疏注意力的核心算法：

# inference/model.py:402
def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, offset: int):
    bsz, seqlen, _ = x.size()
    freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
    ratio = self.compress_ratio
    rd = self.rope_head_dim
    end_pos = start_pos + seqlen

    # 第一次调用时建立 Indexer 自己的 Compressor 状态
    if self.compressor.kv_cache is None:
        self.compressor.kv_cache = self.kv_cache
        self.compressor.freqs_cis = self.freqs_cis

    # q: 用主 attention 的 qr 经过 Indexer 自己的 wq_b
    q = self.wq_b(qr)
    q = q.unflatten(-1, (self.n_local_heads, self.head_dim))   # [B, S, n_heads, head_dim]
    apply_rotary_emb(q[..., -rd:], freqs_cis)
    q = rotate_activation(q)                                    # Hadamard 旋转
    fp4_act_quant(q, fp4_block_size, True)                       # FP4 量化

    # KV: Indexer 的 Compressor 把 x 压成 KV
    self.compressor(x, start_pos)

    # weights: 每 head 一个权重 + 缩放
    weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)

    # 打分：q 与所有压缩 KV 做 dot product，再用 weights 加权 sum 跨 head
    index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, :end_pos // ratio])
    index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)

    if world_size > 1:
        dist.all_reduce(index_score)

    # mask 当前 token 之后的位置
    if start_pos == 0:
        mask = torch.arange(seqlen // ratio).repeat(seqlen, 1) >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio
        index_score += torch.where(mask, float("-inf"), 0)

    # topk 选取
    topk_idxs = index_score.topk(min(self.index_topk, end_pos // ratio), dim=-1)[1]

    # 加 offset 把 topk_idxs 从"压缩 KV 索引"变成"绝对 KV cache 索引"
    if start_pos == 0:
        mask = topk_idxs >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio
        topk_idxs = torch.where(mask, -1, topk_idxs + offset)
    else:
        topk_idxs += offset

    return topk_idxs

把这段代码拆成 6 个步骤：

步骤 1：q 投影 + 旋转

q = self.wq_b(qr)
q = q.unflatten(-1, (self.n_local_heads, self.head_dim))
apply_rotary_emb(q[..., -rd:], freqs_cis)
q = rotate_activation(q)
fp4_act_quant(q, fp4_block_size, True)

qr 来自主 attention 的 q_norm 输出 [B, S, 1536]。wq_b 是 Indexer 独立的 q 投影矩阵——把 qr 投到 n_heads × head_dim = 64 × 128 = 8192 维。然后 RoPE 处理最后 64 维（rope_head_dim）。再过 Hadamard 旋转 + FP4 量化。

为什么 q 也走 FP4？因为 score 计算只是个"打分排序"，不参与最终输出——精度损失对 topk 选取的影响可以接受。注意这里的 fp4_act_quant(q, fp4_block_size, True) 最后一个参数是 inplace=True（inference/kernel.py:186）：走的是"量化再反量化写回"的路径，q 出来还是 BF16，只是数值落到了 FP4 e2m1 的格点上。源码注释把这件事讲得很明白——# use fp4 simulation for q and kv in indexer（inference/model.py:415）。也就是说参考实现里那个 einsum 仍然是 BF16 算的，FP4 只是模拟。按技术报告第 3.4 节，推理部署和 RL rollout 时 Indexer 的 QK 路径才真正以 FP4 缓存、读取和相乘，用来加速长上下文下的打分；参考实现没有走这条路。下文说的"Indexer 打分成本只有主 attention 的 1/16"是按乘加数算的，靠的是 head 数减半（64 vs 128）与 head_dim 减到 1/4（128 vs 512），跟 dtype 无关。

步骤 2：KV 压缩

self.compressor(x, start_pos)

Indexer 的 Compressor 把当前 token 段 x 压成压缩 KV，存到 self.kv_cache 里。这个 Compressor 与主 attention 的 Compressor 是同一个 Compressor 类（head_dim 为 128、rotate=True），但参数独立——参数独立意味着 Indexer 学到的"应该压什么"与主 attention 学到的"应该存什么"是不一样的。

步骤 3：每 head 权重投影

weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)

weights_proj 输出 [B, S, 64]——每 head 一个标量。乘以缩放系数 softmax_scale × n_heads^(-0.5)：前者是 1/dhead1/\sqrt{d_\text{head}}，把单个 head 的内积摁回 O(1)O(1)；后者是 1/nheads1/\sqrt{n_\text{heads}}，抵消"64 个 head 的分数相加"带来的 64\sqrt{64} 量级放大。两个因子都由维度直接算出，不是调出来的超参。

步骤 4：打分

index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, :end_pos // ratio])
index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)

第一行：q 形状 [B, S, 64 heads, 128]，kv_cache 形状 [B, end_pos/ratio, 128]。einsum 把每个 head 的 q 与所有压缩 KV 做内积，得到 [B, S, 64 heads, end_pos/ratio]。

第二行：relu_() 砍掉负相关分（in-place）；乘以 head-level 权重；最后跨 head sum——得到每个压缩 KV 位置的"总分" [B, S, end_pos/ratio]。

注意 relu_ 是 in-place 操作——节省显存。为什么用 ReLU 而非 softmax？DeepSeek-V3.2 论文给出的理由是吞吐（§4.4 细说）；从选取效果看，ReLU 让"完全不相关"的 KV 直接得分 0，topk 选取更稳定；softmax 会强行把概率分到所有位置，topk 选取的边界容易模糊。

步骤 5：mask + topk

mask = ...
index_score += torch.where(mask, float("-inf"), 0)
topk_idxs = index_score.topk(min(self.index_topk, end_pos // ratio), dim=-1)[1]

mask 把"当前 token 之后" 的位置打成 -inf——causal mask 的稀疏版本。topk 选 1024 个分最高的位置（如果总数不足 1024，全选）。

步骤 6：offset 修正

if start_pos == 0:
    mask = topk_idxs >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio
    topk_idxs = torch.where(mask, -1, topk_idxs + offset)
else:
    topk_idxs += offset

topk_idxs 出来时是"压缩 KV 数组的索引"（0 到 end_pos / ratio）。需要加 offset 转换成"sparse_attn 所读 KV 数组里的绝对索引"——offset 是压缩 KV 之前那一段的长度（prefill 时是本次 seqlen 条逐 token KV，decode 时是 win 个滑窗槽位，见 §4.12 陷阱四）。-1 表示"无效位置"。

最终返回的 topk_idxs 形状 [B, S, 1024]（上下文足够长时）——对每个 query token，返回它应该关注的 1024 个 KV 位置的绝对索引。这个数组随后被 cat 到滑窗 idxs，传给 sparse_attn kernel。

4.4 ReLU vs Softmax：score net 的代数选择

V4 的 Indexer 在打分阶段用 relu_() 而不是 softmax。这是个看起来违反直觉的选择——传统 attention 用 softmax 是因为它是"概率分布"，符合 attention weights 的语义。但 score net 不是 attention，它的输出是"排序"，用 ReLU 反而更合适。

Softmax 的特性：

  • 输出是 [0, 1] 区间内的概率分布，所有位置和为 1
  • 数值上"温和"——再差的位置也会分到一点概率
  • topk 选取依赖"概率最大的几个"，但概率之间的差异可能很小
  • 容易受温度（temperature）影响——同样的 logit 在不同温度下选出来的 topk 可能完全不同

ReLU 的特性：

  • 每个 head 的输出是 [0, ∞) 区间内的非归一化值，负值直接砍成 0
  • 数值上"果断"——在单个 head 上完全不相关的位置得分 0，不再贡献分数
  • topk 选取直接基于"分数最大的几个"，分数差异越大、选取越稳定
  • 没有温度参数，行为可预测

V4 选 ReLU 的具体好处：

  1. 稀疏选取的稳定性：ReLU 让分数"二值化倾向"更强——要么得分不为 0、要么得分为 0。topk 选取在这种分布上更不容易被噪声扰动。
  2. 训练梯度的清晰性：当一个 KV 位置完全不相关（softmax 仍会给 1e-10 量级概率），它在反向传播里仍会有微弱梯度——这些梯度累积起来会污染 score 网络。ReLU 直接砍掉负值，反向传播 mask 掉这部分干扰。
  3. 吞吐：softmax 要在 1M / ratio 个候选位置上算 exp，再做求最大值、求和两轮跨位置归约；ReLU 是逐元素操作，不需要跨位置归约。

但 ReLU 也有代价：

  • 输出不是概率分布——不能直接当作"注意力权重"用。这就是为什么 V4 把 Indexer 的输出仅用于 topk 选取，不直接参与最终 attention 输出。最终输出由主 attention 用 softmax 算。
  • ReLU 本身不带"head 间归一化"——V4 通过 weights_proj 给每 head 一个独立权重，再跨 head sum，间接实现 head 间的相对重要性归一。

这种"score net 用 ReLU、主 attention 用 softmax"的分工沿用自 DeepSeek-V3.2 的 DSA：V4 技术报告第 2.3.1 节写明 CSA "applies the DSA strategy" 做 top-k 选取，其式 (16) 与 DeepSeek-V3.2 论文（arXiv:2512.02556）的 lightning indexer 打分式同形。V3.2 论文对选 ReLU 只给了一句理由："for throughput consideration"（出于吞吐考虑）；上面第 1、2 条是本专栏的解读。

4.5 Indexer 与主 attention 的资源对比

把 Indexer 与主 attention 的资源占用摆出来：

这里要先说清楚跟谁比，否则很容易得出相反的结论。

Indexer 要给全部压缩 KV 打分（源码里是 self.kv_cache[:bsz, :end_pos // ratio]，没有任何提前裁剪），所以它的成本是 O(n/ratio)O(n/\text{ratio})——随上下文线性增长。而主 attention 因为吃的是 Indexer 选出来的 1024 条加 128 条滑窗，成本被钉死在 O(window+topk)O(\text{window} + \text{topk})，与上下文无关。

所以正确的对照物是"假如不做稀疏、让主 attention 看完所有压缩 KV"：

维度 dense 主 attention（看完所有压缩 KV） Indexer 比率
n_heads 128 64 1/2
head_dim 512 128 1/4
需要算几次矩阵乘 QK + PV 两次 只有 QK 一次 1/2
q / k 数据类型 BF16（QAT 到 FP8） FP4（参考实现为模拟） —
每 token 乘加数 128×512×2=131072128\times512\times2 = 131072 64×128=819264\times128 = 8192 1/16

也就是说：Indexer 用 1/16 的算力，把主 attention 从"看 262144 条"降到"看 1152 条"（1M context、ratio=4 的层）。这才是 V4 稀疏注意力的 ROI 公式。

反过来的推论同样重要：在 1M 这种长度上，ratio=4 的层里 Indexer 本身才是算力大头——它是这条路径上唯一还随 nn 线性增长的部分。§4.6 会把这笔账逐项算出来。

4.6 Indexer 输出维度的张量代数

把 Indexer 的张量形状变化按维度逐步追踪一遍。下面所有"计算量"都统一记成每个 query token、每层的乘加数（MAC），在 1M context、ratio=4 的层上取值：

输入:
  x:  [B, S, dim=7168]
  qr: [B, S, q_lora_rank=1536]   ← 来自主 attention 的 q_norm 输出
  start_pos, offset

────────────────────
步骤 1：q 投影
────────────────────
q = wq_b(qr)   # [B, S, n_heads × head_dim] = [B, S, 64 × 128] = [B, S, 8192]
              # 乘加数：1536 × 8192 = 12.6M / token
q = q.unflatten(-1, (64, 128))  # [B, S, 64, 128]
apply_rotary_emb(q[..., -64:], freqs_cis)  # 仅最后 64 维
q = rotate_activation(q)         # Hadamard 不改变形状
fp4_act_quant(q, 32, True)        # FP4 量化（每 32 个元素一个 scale）

────────────────────
步骤 2：KV 压缩
────────────────────
self.compressor(x, start_pos)
# 内部：wkv(x): [B, S, 2*128=256] (overlap=True 翻倍维度)
#       softmax(score) 加权求和后输出 [B, S/4, 128]
#       存到 self.kv_cache[:, :end_pos/4]

────────────────────
步骤 3：weights 投影
────────────────────
weights = self.weights_proj(x)  # [B, S, 64]
                               # 乘加数：7168 × 64 = 459K / token
weights *= self.softmax_scale * self.n_heads ** -0.5
                               # softmax_scale = 128**-0.5 ≈ 0.0884
                               # n_heads**-0.5 = 64**-0.5 = 0.125
                               # 总缩放系数 ≈ 0.011

────────────────────
步骤 4：score 计算
────────────────────
index_score = einsum("bshd,btd->bsht", q, kv_cache[:bsz, :end_pos/4])
            # q: [B, S, 64, 128], kv_cache: [B, end_pos/4, 128]
            # 输出: [B, S, 64, end_pos/4]
            # 乘加数：64 × (end_pos/4) × 128 / token
            # 在 1M context 下 = 64 × 262144 × 128 ≈ 2147M / token  ← 绝对大头
index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)
            # 输出: [B, S, end_pos/4]

────────────────────
步骤 5：mask + topk
────────────────────
index_score += causal_mask                  # 维度同上
topk_idxs = index_score.topk(min(1024, end_pos/4), dim=-1)[1]
            # 输出: [B, S, 1024]

────────────────────
步骤 6：offset 修正
────────────────────
topk_idxs += offset
            # 输出: [B, S, 1024]，整数索引

把一个 ratio=4 层里所有矩阵乘按"每 query token 的乘加数"列全（1M context）：

归属 算子 乘加数 / token
Indexer wq_b（1536 → 8192） 12.6M
Indexer weights_proj（7168 → 64） 0.46M
Indexer 自带 Compressor 的 wkv + wgate（→ 2×128） 3.67M
Indexer score einsum（64 × 262144 × 128） 2147M
主 attention wq_a + wq_b 111.7M
主 attention wkv + 自带 Compressor 18.4M
主 attention sparse_attn 的 QK + PV（128 × 1152 × 512 × 2） 151.0M
主 attention wo_a + wo_b 184.5M

两边一加：Indexer ≈ 2164M，主 attention ≈ 466M。在 1M 这个长度上，Indexer 的 score einsum 一项就吃掉了这一层 attention 侧七成以上的算力。

那省在哪里？对照"不做稀疏、让主 attention 看完全部 262272 条压缩 KV"的假想版本：

128 heads × 262272 positions × 512 dims × 2 (QK + PV) ≈ 34400M 乘加 / token

34400M → 2630M，省下约 13 倍。

访存那一侧要分开算，容易被算漂：主 attention 这一段确实从"读 262272 条 × 512 维"降到"读 1152 条 × 512 维"，是 228 倍；但 Indexer 自己要把 262144 条 × 128 维的压缩 KV 完整读一遍，这笔账甩不掉。两边加起来，decode 时每 token 在 attention 侧的 KV 访存从 262272×512≈1.34×108262272\times512 \approx 1.34\times10^8 个元素降到 1152×512+262144×128≈3.41×1071152\times512 + 262144\times128 \approx 3.41\times10^7 个元素——约 3.9 倍，而不是 200 倍。稀疏路径真正的数量级红利在算力（13 倍）和 KV cache 体积（§3.10 的 8.9 倍）上，访存这一项被 Indexer 的线性打分吃掉了大半。

这笔账也解释了为什么 index_topk 从首发的 512 调到 1024 是"便宜的"：topk 只影响主 attention 那 151M 的一项，多出来的 512 条位置也就是 128×512×512×2≈67128\times512\times512\times2 \approx 67M，相对 2147M 的打分开销可以忽略。

4.7 Hadamard 旋转：FP4 量化前的"信息打散"

Indexer 的 q 在量化前过了一次 rotate_activation(q)。这个函数实际是 Hadamard 变换：

def rotate_activation(x: torch.Tensor) -> torch.Tensor:
    """Applies randomized Hadamard rotation to spread information across dims before FP8 quant."""
    assert x.dtype == torch.bfloat16
    from fast_hadamard_transform import hadamard_transform
    return hadamard_transform(x, scale=x.size(-1) ** -0.5)

Hadamard 变换是一种正交变换——它不改变向量长度，但把每个维度的信息"打散"到所有维度。具体地，Hadamard 矩阵 H 满足 H × H^T = n × I，对一个向量 v 做 H × v 后，原本集中在某些维度的能量被均匀分布到所有维度。

为什么量化前要打散？因为 FP4 e2m1 的精度损失对"维度间能量集中"特别敏感——如果某些维度数值远大于其他维度，FP4 量化会把小数值维度"舍入到 0"，丢失信息。Hadamard 打散后所有维度数值大致同量级，FP4 量化的精度损失被均摊。

这种"先正交旋转再量化"的技巧不是 V4 首创——QuaRot 论文（arXiv:2404.00456）系统提出了这个方法。但 V4 把它工业化用到了 1.6T 模型的 attention 路径上——在 score net 这种"对精度可以放宽"的位置。

第 12 章讲 FP4 e2m1 的格点几何，第 14 章讲 Indexer 这一路"Hadamard + FP4 + block 32"的量化配置。

4.8 Hadamard 旋转的几何直觉

Hadamard 旋转看起来神秘，但它的几何直觉非常简单——把一个向量的"信息"从某些集中维度均匀打散到所有维度。

考虑一个简化例子：四维向量 v = [10, 0.1, 0.1, 0.1]。

直接 FP4 量化（为了看清机制，这里先不考虑 block scale）：

  • FP4 e2m1 能表示的正值只有 0.5、1、1.5、2、3、4、6 这几个格点（mantissa 1 bit + exponent 2 bit）
  • 元素 10 量化后变成 6（饱和）；元素 0.1 量化后变成 0（下溢）
  • 量化后 v ≈ [6, 0, 0, 0]——三个小分量全部丢失，总平方误差约 16.0

先 Hadamard 旋转再 FP4 量化：

  • Hadamard 矩阵 H_4 = [[1,1,1,1], [1,-1,1,-1], [1,1,-1,-1], [1,-1,-1,1]] / 2
  • H_4 × v = [5.15, 4.95, 4.95, 4.95]（每个维度都接近 5）
  • FP4 格点里没有 5，就近落到 [6, 4, 4, 4]
  • 反 Hadamard 后 = [9, 1, 1, 1]——峰值误差从 4 降到 1，小分量误差从 0.1 升到 0.9，总平方误差降到约 3.4

误差没有消失，而是从一个维度摊到了所有维度。若按 V4 实际的做法给每 32 个元素配一个 2 的幂 scale，直接量化那一路不再饱和（得 [8, 0, 0, 0]，总平方误差约 4.0），Hadamard 那一路仍是约 3.4——block scale 解决饱和，Hadamard 解决离群值把同组小数值挤到 0 的问题。

Hadamard 的关键性质：输入越"集中"（高峰值），输出越"均匀"。但它不是对任何输入都有效——一个本来就平坦的向量反而可能被变成尖峰（[1, 1, 1, 1] 旋转后是 [2, 0, 0, 0]）。它对量化友好，是因为激活里的离群值通常集中在少数维度上，这正是它擅长摊平的情形。

flowchart LR
  subgraph 量化前["FP4 量化挑战"]
    direction TB
    Concentrated["v = [10, 0.1, 0.1, 0.1]<br/>能量集中于 1 维"]
  end
  subgraph Hadamard["Hadamard 旋转"]
    Spread["H × v = [5.15, 4.95, 4.95, 4.95]<br/>能量均匀分布"]
  end
  subgraph 量化后["FP4 量化"]
    Quantized["[6, 4, 4, 4]<br/>误差在所有维度均摊"]
  end
  subgraph 反向["反 Hadamard"]
    Recovered["[9, 1, 1, 1]<br/>总平方误差 16.0 → 3.4"]
  end
  Concentrated --> Spread --> Quantized --> Recovered

V4 的 rotate_activation 函数用 fast_hadamard_transform 库实现，scale 设为 x.size(-1) ** -0.5——这个 scale 让 Hadamard 矩阵正交（H × H^T = I），保证旋转不改变向量长度。

为什么只有 Indexer 用 Hadamard 而主 attention 不用？源码与技术报告都没有给出理由。一个合理的解释是：主 attention 的 KV 走 FP8（且 RoPE 维留在 BF16），量化损失已经够小；Indexer 走 FP4，格点稀疏得多，更需要 Hadamard 先把离群值摊平。

4.9 Indexer 与 attn_sink 的协同

第 2 章讲过 attn_sink 是稀疏注意力的"兜底参数"——Indexer 选错时，允许这一步的注意力总和小于 1、甚至接近 0。Indexer 与 attn_sink 形成一个明确的协同分工：

  • Indexer 努力选对：用学到的 score 函数尽可能选出最相关的 1024 个 KV
  • attn_sink 兜住选错：当所有 1024 个 KV 都不相关时，sink 接住注意力质量

这个分工的工程意义：Indexer 不需要选得"完美"，只需要选得"够好"。如果它偶尔选错，attn_sink 让这一步注意力可以少输出甚至不输出，而不是被迫把权重分给不相关的 KV。

需要说明，这是本专栏的解读：技术报告（第 2.3.3 节）对 sink 的表述只是让每个 query head 的注意力总和可以不等于 1、甚至接近 0，并没有把它与 Indexer 的精度要求挂钩；sink 对没有 Indexer 的 ratio=128 层同样生效。另外每个 query 还固定带着 128 条滑窗 KV，Indexer 选偏时主 attention 也不至于无处可看。

4.10 Indexer 在训练时的辅助监督

V4 的 Indexer 在推理时只输出 topk_idxs，不参与最终输出。但在训练时，Indexer 必须有自己的"学习目标"——否则它的参数无从更新。

虽然 V4 的训练源码不在公开仓库内，但从架构设计可以反推 Indexer 训练的几个标准做法（V3.2-Exp 论文已经描述过类似机制）：

做法一：用主 attention 的 attention weights 作为监督信号

主 attention 在训练时会算完整的 dense softmax(Q · K^T)——这给每个 query token 一组"理想"的 KV 重要性分布。Indexer 的 score 应该与这组分布正相关——不需要完全相等，但 top-k 重合度要高。这可以用 KL 散度或 listwise ranking loss 作为监督。DeepSeek-V3.2 论文（arXiv:2512.02556）用的就是这种做法：以主 attention 分布与 Softmax(I) 之间的 KL 散度作为 Indexer 的损失，先在 dense attention 下冻结其余参数单独预热 Indexer，进入稀疏训练后再把 Indexer 的输入从计算图上 detach，让它只受这一项损失驱动。

做法二：稀疏 attention 输出与 dense attention 输出的距离

更直接的做法是：用 Indexer 的 topk 选取做一次稀疏 attention，再用主 attention 做一次 dense attention（仅训练时），对比两者的输出差异——这个差异就是 Indexer 的损失。

做法三：辅助 next-token prediction

让 Indexer 在主 lm_head 之外有一个"小 head"，直接预测下一个 token。如果稀疏选取错了关键 KV，next-token loss 会上升——这是另一种间接监督。

V4 实际用了哪种、或几种的组合，技术报告没有写明：它只说引入稀疏注意力时"先用一个短阶段预热 lightning indexer，再以稀疏注意力训练大部分步数"（技术报告第 4.2.2 节），没有给出 Indexer 的损失形式。但架构上保留了所有这些训练接口——Indexer 类的输入 qr 是 q_norm 之后、q_b 之前的中间表示，意味着 Indexer 可以在训练时被独立 supervise。

4.11 一段动手实验：Indexer 的最小可运行版本

import torch
import torch.nn as nn
import torch.nn.functional as F

class MiniIndexer(nn.Module):
    """简化版 Indexer：去掉 Hadamard / FP4 / 多卡，保留核心逻辑"""
    def __init__(self, dim=512, n_heads=8, head_dim=64, ratio=4, topk=16):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.ratio = ratio
        self.topk = topk
        self.wq = nn.Linear(dim, n_heads * head_dim)
        self.weights_proj = nn.Linear(dim, n_heads)
        # 简化：这里用一个 Linear 模拟 Compressor 的输出（实际上应该是 Compressor）
        self.kv_proj = nn.Linear(dim, head_dim)
        self.scale = head_dim ** -0.5

    def forward(self, x):
        B, S, _ = x.shape
        # q
        q = self.wq(x).view(B, S, self.n_heads, self.head_dim)
        # 模拟压缩 KV（每 ratio 个 token 取平均）
        kv = x.view(B, S // self.ratio, self.ratio, -1).mean(dim=2)
        kv = self.kv_proj(kv)                   # [B, S/ratio, head_dim]
        # weights
        weights = self.weights_proj(x) * self.scale * self.n_heads ** -0.5  # [B, S, n_heads]
        # score
        score = torch.einsum("bshd,btd->bsht", q, kv)
        score = (score.relu_() * weights.unsqueeze(-1)).sum(dim=2)  # [B, S, S/ratio]
        # mask
        mask = torch.arange(S // self.ratio).repeat(S, 1) >= torch.arange(1, S+1).unsqueeze(1) // self.ratio
        score += torch.where(mask, float("-inf"), 0)
        # topk
        topk_idxs = score.topk(min(self.topk, S // self.ratio), dim=-1)[1]
        return topk_idxs


# 测试
mi = MiniIndexer()
x = torch.randn(2, 32, 512)
idxs = mi(x)
print(idxs.shape)   # 应该是 [2, 32, 8] (topk 在不足时被 clip)

注意 mini 版省掉了源码在 topk 之后把不可见位置改成 -1 的那一步（§4.3 步骤 6）：凡是可见组不足 k 个的 token，返回的索引里都会混进得分为 -inf 的未来组（本例 k=8 等于总组数，除最后一个 token 外都是如此），拿去做 attention 就会泄漏未来信息。

跑通这个 mini 版本后再看 V4 源码，会觉得"原来 V4 的 Indexer 不过是把 mini 版本加上了 Compressor、Hadamard、FP4 三层工业化壳"。核心算法没那么神秘。

4.12 Indexer 的工程实现陷阱清单

读 V4 的 Indexer 源码时，至少有四个细节如果没注意会在自己的实现里翻车：

陷阱一：weights_proj 的 dtype

self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16)

注意明确指定 dtype=torch.bfloat16。对照 Transformer.__init__：default_dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16——默认权重 dtype 是 FP8 e4m3，不是 FP4（FP4 只用在 expert_dtype="fp4" 的 MoE 专家权重上）。所以 Indexer 里没写 dtype 的 wq_b 走 FP8，而 weights_proj 被显式拉回 BF16：它输出的是 64 个 head-level 标量，直接乘在每个候选位置的总分上，量化误差会成比例地扰动 topk 的排序边界。

陷阱二：rotate_activation 的 dtype 断言

def rotate_activation(x: torch.Tensor) -> torch.Tensor:
    assert x.dtype == torch.bfloat16
    ...

Hadamard 变换强制要求输入是 BF16。如果你给它 FP8 或 FP4，会直接 assert 失败。Indexer 的 wq_b 没写 dtype，权重走默认的 FP8 e4m3，而 fp8_gemm 的输出 dtype 是 torch.get_default_dtype()（inference/kernel.py:270）也就是 BF16，可以直接传给 rotate_activation；如果你的实现里 q 已经被提前量化到 FP4，需要先反量化回 BF16 再做 Hadamard。

陷阱三：causal mask 的边界

mask = torch.arange(seqlen // ratio).repeat(seqlen, 1) >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio

注意是 >= 不是 >，且右边是 (arange + 1) // ratio 而不是 arange // ratio。这两个细节确保了"压缩组的右边界"被正确处理——一个压缩组 g 覆盖到 token (g+1)*ratio - 1 为止，所以 token i 能看见它的条件是 i >= (g+1)*ratio - 1，也就是 (g+1)*ratio <= i+1，整数上恰好等价于 g < (i+1) // ratio。源码的 >= 把"不可见"的位置打上 -inf。注意 i 恰好是某一组最后一个 token 时那一组是可见的（它只含 i 及更早的 token），如果把条件写成 i >= (g+1)*ratio 就会多屏蔽掉一组。

陷阱四：offset 在 prefill / decode 上的不同语义

# Attention.forward 里
offset = kv.size(1) if start_pos == 0 else win

prefill 时 offset = kv.size(1)，即 seqlen：这一步 sparse_attn 读的不是 cache，而是本次拼出来的 torch.cat([kv, kv_compress], dim=1)，前 seqlen 条是逐 token 的 KV，压缩 KV 排在它们后面；decode 时 offset = win：sparse_attn 读的是 self.kv_cache，前 win 个槽位是滑窗，压缩 KV 从第 win 个槽位开始。如果你写 reuse 一个常量 offset，两条 codepath 会出现"压缩 KV 索引与 cache 实际位置错位" 的 bug——稀疏 attention 读到的就是错位的 KV。

除了陷阱二会直接 assert 失败，其余三个都不会报错，只会让模型输出"看起来对、实际错"——是稀疏 attention 实现里最难调试的一类问题。

4.13 Indexer 的"输入信号" 完整性分析

Indexer 的 forward 输入是 (x, qr, start_pos, offset)——四个参数提供完整的"上下文信号"。把每个信号的作用讲清楚。

信号 1：x（hidden state）

x 是当前 token 的 hidden state。Indexer 用 x 算两件事：

  • 算 weights_proj 输出（每 head 权重）
  • 通过自己的 Compressor 把 x 压成新的压缩 KV

x 提供"这个 token 是什么内容"的信息——决定它对哪些历史 KV 感兴趣。

信号 2：qr（query 中间表示）

qr 是主 attention 的 q_norm 输出 [B, S, q_lora_rank=1536]。Indexer 用 qr 经过自己的 wq_b 算 query。

qr 提供"主 attention 视角下的 query 信号"——让 Indexer 的 query 与主 attention 对齐。两者共享 q_lora 让 Indexer 不需要重学这一层投影。

信号 3：start_pos

当前 token 在序列中的位置。Indexer 用这个：

  • 决定 freqs_cis 的切片（RoPE 旋转角度）
  • 决定 causal mask（不能 attend 到未来）
  • 决定 prefill / decode codepath

信号 4：offset

把 topk_idxs 从"压缩 KV 索引"换算到"主 KV cache 绝对索引"的偏移量。prefill 与 decode 的 offset 含义不同（详见 §4.12 陷阱清单）。

这 4 个信号合起来让 Indexer 有"决定 top-k KV 选择" 所需的全部输入。少任何一个都会让 Indexer 工作不正确。

和主 attention 的"信号共享"：

Indexer 与主 attention 共享：

  • q_lora_rank 投影（qr）
  • 共享 freqs_cis（RoPE）的源
  • 共享 input x

但不共享：

  • wq_b（独立的 query 投影）
  • Compressor 实例（独立的 KV 压缩）
  • 量化策略（FP4 vs FP8）

这种"共享 + 独立"的设计让 Indexer 又快又准——快是因为复用了主 attention 的中间表示，准是因为有独立的精修空间。

4.14 Indexer 的运行时性能特征

Indexer 在生产中的运行时性能对最终延迟影响显著。把它的几个特征列出来：

特征 1：每层只有 ratio=4 的层有 Indexer

V4 Pro 61 层中恰好 30 层是 ratio=4——意味着 Indexer 只在这 30 层运行。其余 31 层 ratio=128，不建 Indexer，直接用 get_compress_topk_idxs 取全部可见的压缩 KV（compress_ratios 末尾那个 0 属于 MTP 层）。

特征 2：Indexer 与主 attention 部分并行

Indexer 的 wq_b 与主 attention 的 wq_b 输入都是 qr，两者之间没有依赖，理论上可以并行或合并成一次 GEMM launch。官方参考实现里是顺序调用的，是否合并取决于下游引擎。

特征 3：Indexer 的输出在主 attention 之前完成

Indexer 必须先完成（输出 topk_idxs），主 attention 才能开始 sparse_attn。这是串行依赖——不能并行。这是 Indexer 在 critical path 上的体现。

特征 4：上下文越长，Indexer 在这一层里的占比越高

Indexer 的打分成本随上下文线性增长，主 attention 的成本却是常数（§4.6 算过：1M 下 Indexer 约 2164M 乘加、主 attention 约 466M）。短上下文时 Indexer 只是零头，长上下文时它反过来成为这一层的主要开销。这不是 bug，是稀疏路径的定价方式：把 O(n2)O(n^2) 换成 O(n)O(n) 的打分 + O(1)O(1) 的注意力，代价就是那个 O(n)O(n) 的打分永远甩不掉。真要继续往下压，只能做"分层打分"（先粗筛再细筛），而 V4 Preview 没有这么做。

特征 5：Indexer 不能跨层并发

Indexer 的输入 x、qr 来自本层的输入，而本层输入又是上一层的输出，所以同一个 token 不同层的 Indexer 之间有严格的数据依赖，没法跨层并行或流水。

理解这些性能特征让你在调优 V4 时知道"Indexer 是否是瓶颈"：短上下文下通常不是；长上下文下它本来就是 ratio=4 层 attention 侧的主要开销（§4.6）。这时该看的是打分 kernel 的实现——技术报告第 3.4 节的做法是 QK 路径全程 FP4，并把 index score 从 FP32 降到 BF16，让 top-k 选取提速 2 倍、KV 召回率保持 99.7%。

4.15 延伸阅读

  • DeepSeek-V3.2-Exp：DSA score net 的最早期实现
  • QuaRot 论文（arXiv:2404.00456）：Hadamard 旋转 + INT4 量化的代表性工作
  • Native Sparse Attention 论文（arXiv:2502.11089）：稀疏 attention 训练理论基础
  • 本专栏第 5 章：inference/kernel.py 里的 sparse_attn_kernel 把 Indexer 输出的 topk_idxs 真正"用起来"
  • 本专栏第 3 章：Indexer 的 Compressor 与 Attention 的 Compressor 是同源结构——回看 §3.8

4.16 本章小结

  • Indexer 是 V4 稀疏注意力的"score net"——专门给压缩 KV 打分并选 top-1024
  • 它与主 attention 共享 q_lora_rank 投影，但有独立的 wq_b、weights_proj、Compressor
  • 它的 head_dim=128（而非主 attention 的 512）、n_heads=64（而非 128）、q/k 走 FP4 模拟——每个候选位置的打分成本只有 dense attention 的 1/16
  • ReLU + 加权 sum 替代 softmax + sum 是 V4 的工程取舍——topk 边界更稳定
  • Hadamard 旋转保证 FP4 量化的精度损失被均摊到所有维度
  • attn_sink 与 Indexer 形成"努力选对 + 兜住选错"的工程冗余

第 5 章我们进入 V4 注意力革命的最后一站：sparse_attn kernel——官方推理仓库里那个 TileLang 写的 sparse_attn_kernel，到底怎么把 topk 索引变成一次 online-softmax 的稀疏 attention。