第3章 Compressor：滑窗与 KV 几何压缩

“Memory is not what you remember. It’s what’s worth keeping.” —— 引自一位老 OS kernel 工程师的名言

V4 的 KV cache 同样如此。它不是把所有 token 都记下来，而是把”值得长期记的东西”压成一份精炼的远距离记忆。

3.1 引子：1M context 的 KV 几何崩塌

如果 V4 不做任何压缩，按 head_dim=512 + 61 layers + 1M tokens 算，每个序列的 KV cache 是：

单序列 64 GB——一个 H100 80 GB 卡只能跑 1 个并发，而且还没算 Q 投影、激活、模型权重的占用。这是不可能商用的。

V4 的解决方案不是”换更高密度的硬件”，而是重新设计 KV cache 的几何结构。每层的 KV 被切成两段：

flowchart LR
  subgraph KVCache["每层 KV cache 几何"]
    direction LR
    Win["滑窗段<br/>window_size=128"]
    Comp["压缩段<br/>max_seq_len // ratio"]
    Win -.物理拼接.-> Comp
  end

  subgraph Window["滑窗段（前 128 槽位）"]
    W1["最近 128 token 的 KV<br/>环形覆盖"]
  end

  subgraph Compressed["压缩段（max_seq_len/ratio 槽位）"]
    C1["每 ratio 个 token 压成 1 组<br/>对长距离信息做语义归约"]
  end

这个双段结构的物理存储就是 Attention.__init__ 里的：

kv_cache_size = args.window_size + (args.max_seq_len // self.compress_ratio if self.compress_ratio else 0)
self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, kv_cache_size, self.head_dim))

kv_cache_size 把两段加在一起。比如 1M context、ratio=4：128 + 1048576/4 = 262272 槽位。每槽 head_dim=512 BF16，每层占 256 MB；61 层占 ~16 GB。

如果 ratio=128 呢？128 + 1048576/128 = 8320 槽位，每层占 8 MB；61 层占 ~488 MB。

V4 通过让每一层独立选择 ratio（4 / 128 / 0），把所有层的 KV 加权平均压到 ~2 GB / 序列——比 dense 64 GB 减少 32 倍。

3.2 Compressor 类全景

V4 的 Compressor 类（inference/model.py）的 __init__ 完整摆出来：

class Compressor(nn.Module):
    def __init__(self, args: ModelArgs, compress_ratio: int = 4, head_dim: int = 512, rotate: bool = False):
        super().__init__()
        self.dim = args.dim
        self.head_dim = head_dim
        self.rope_head_dim = args.rope_head_dim
        self.nope_head_dim = head_dim - args.rope_head_dim
        self.compress_ratio = compress_ratio
        self.overlap = compress_ratio == 4
        self.rotate = rotate
        coff = 1 + self.overlap
        self.ape = nn.Parameter(torch.empty(compress_ratio, coff * self.head_dim, dtype=torch.float32))
        self.wkv = Linear(self.dim, coff * self.head_dim, dtype=torch.float32)
        self.wgate = Linear(self.dim, coff * self.head_dim, dtype=torch.float32)
        self.norm = RMSNorm(self.head_dim, args.norm_eps)
        self.kv_cache: torch.Tensor = None  # 由 Attention 注入

        # decode-phase incremental state
        self.register_buffer("kv_state",
            torch.zeros(args.max_batch_size, coff * compress_ratio, coff * self.head_dim, dtype=torch.float32),
            persistent=False)
        self.register_buffer("score_state",
            torch.full((args.max_batch_size, coff * compress_ratio, coff * self.head_dim), float("-inf"), dtype=torch.float32),
            persistent=False)
        self.freqs_cis: torch.Tensor = None

把它拆成五块：

• 静态参数：wkv、wgate、ape、norm——决定如何把 token 压缩成”压缩 KV”
• 运行时 KV cache：self.kv_cache ——压缩后的 KV 写入这里（注意是被 Attention 注入的，Compressor 自己不分配）
• decode 增量状态：kv_state 与 score_state ——decode 阶段累积尚未”凑齐 ratio 个”的 KV
• rotate 标志：是否用 Hadamard 变换（仅 Indexer 的 Compressor 才用）
• overlap 标志：是否启用”重叠窗口压缩”（仅 ratio=4 才启用）

接下来按 prefill / decode 两条路径分别拆。

3.3 prefill 路径：一次性把整段 token 压完

prefill 阶段（start_pos == 0）的 Compressor.forward 处理整段 KV：

def forward(self, x: torch.Tensor, start_pos: int):
    bsz, seqlen, _ = x.size()
    ratio, overlap, d, rd = self.compress_ratio, self.overlap, self.head_dim, self.rope_head_dim
    dtype = x.dtype
    x = x.float()
    kv = self.wkv(x)                # [B, S, coff*d]
    score = self.wgate(x)            # [B, S, coff*d]

    if start_pos == 0:
        should_compress = seqlen >= ratio
        remainder = seqlen % ratio
        cutoff = seqlen - remainder
        offset = ratio if overlap else 0

        # 处理"凑不齐 ratio"的尾段
        if overlap and cutoff >= ratio:
            self.kv_state[:bsz, :ratio] = kv[:, cutoff-ratio : cutoff]
            self.score_state[:bsz, :ratio] = score[:, cutoff-ratio : cutoff] + self.ape
        if remainder > 0:
            kv, self.kv_state[:bsz, offset : offset+remainder] = kv.split([cutoff, remainder], dim=1)
            self.score_state[:bsz, offset : offset+remainder] = score[:, cutoff:] + self.ape[:remainder]

        score = score[:, :cutoff]
        kv = kv.unflatten(1, (-1, ratio))
        score = score.unflatten(1, (-1, ratio)) + self.ape

        if overlap:
            kv = self.overlap_transform(kv, 0)
            score = self.overlap_transform(score, float("-inf"))

        # 用 softmax(score) 加权求和 → 压缩 KV
        kv = (kv * score.softmax(dim=2)).sum(dim=2)

这一段的代数过程：

1. wkv(x) 把每个 token 投到 coff * head_dim（如果 overlap=True，coff=2，相当于每个 token 准备了”两套”压缩贡献）
2. wgate(x) 算门控分数
3. score + ape——加上 absolute position embedding（每个 ratio-相对位置一组）
4. unflatten(1, (-1, ratio))——把 [B, S] 重排成 [B, S/ratio, ratio]，每 ratio 个 token 聚成一组
5. softmax(score, dim=2)——在组内做 softmax 归一
6. (kv * softmax_score).sum(dim=2)——加权求和，得到每组的”压缩 KV”

最终 kv 形状是 [B, S/ratio, head_dim]——每 ratio 个原始 token 被压成 1 组压缩 KV。

3.3.1 absolute position embedding (ape) 是什么

self.ape: [ratio, coff * head_dim] 是一个可学习的 embedding 表，给 ratio 内每个相对位置一个偏置。

为什么需要？因为 Compressor 在做”组内加权求和”时，ratio 内的 token 顺序信息会丢失——softmax + sum 是对称的。如果不加位置偏置，“3 4 5 6” 和 “6 5 4 3” 会得到相同的压缩 KV，这显然不对。

ape 提供的就是”组内的相对位置感”——给第 i 个位置的 score 加一个偏置 ape[i]，让顺序信息进入 softmax 概率分布。

3.3.2 overlap = True 的特殊处理

ratio=4 的层启用 overlap=True。它的含义是：相邻两组的压缩 KV 共享一半信息——

原始 token: [t0, t1, t2, t3, t4, t5, t6, t7, ...]

不 overlap 的压缩 (ratio=4):
  group 0: 压缩 t0-t3
  group 1: 压缩 t4-t7
  group 2: 压缩 t8-t11
  ...

overlap 的压缩:
  group 0: 压缩 t0-t3 (来自 wkv 输出的"前一半")
  group 0.5 (overlap): 压缩 t2-t5 (来自 wkv 输出的"后一半")
  group 1: 压缩 t4-t7
  ...

coff = 1 + overlap = 2 让 wkv 输出双倍维度，前一半给”正常压缩”，后一半给”overlap 压缩”。overlap_transform 函数把张量重新排列成”前后窗口都对应到正确的 token 区间”。

为什么要 overlap？因为 ratio=4 的层是稀疏注意力的主要工作层——overlap 让”组与组的边界处的信息”不会丢失（边界 token 同时出现在两个相邻组）。这是稀疏 attention 在长上下文上保持精度的一个细节。

3.4 decode 路径：增量压缩与状态累积

decode 阶段（start_pos > 0），每次只送 1 个 token，Compressor 必须增量地累积 KV，每 ratio 步触发一次”压缩落盘”：

else:  # start_pos > 0 (decode)
    should_compress = (start_pos + 1) % self.compress_ratio == 0
    score += self.ape[start_pos % ratio]

    if overlap:
        self.kv_state[:bsz, ratio + start_pos % ratio] = kv.squeeze(1)
        self.score_state[:bsz, ratio + start_pos % ratio] = score.squeeze(1)

    if should_compress:
        # 把累积了 ratio 个 token 的状态做一次 softmax 加权求和
        kv_state = torch.cat([self.kv_state[:bsz, :ratio, :d], self.kv_state[:bsz, ratio:, d:]], dim=1)
        score_state = torch.cat([self.score_state[:bsz, :ratio, :d], self.score_state[:bsz, ratio:, d:]], dim=1)
        kv = (kv_state * score_state.softmax(dim=1)).sum(dim=1, keepdim=True)
        # 把"当前窗口"滚动成"前一窗口"
        self.kv_state[:bsz, :ratio] = self.kv_state[:bsz, ratio:]
        self.score_state[:bsz, :ratio] = self.score_state[:bsz, ratio:]
    else:
        self.kv_state[:bsz, start_pos % ratio] = kv.squeeze(1)
        self.score_state[:bsz, start_pos % ratio] = score.squeeze(1)

这段代码的状态机非常微妙：

• kv_state 维度是 [B, 2*ratio, 2*head_dim]（overlap 时）
• 前 ratio 个槽位是”上一窗口的累积 KV”
• 后 ratio 个槽位是”当前窗口的累积 KV”
• 每 ratio 步，触发 should_compress=True，把累积的 KV 做一次 softmax 加权求和，写入 self.kv_cache[start_pos // ratio]
• 同时，“当前窗口” 滚动成 “上一窗口”，给下一轮腾出空间

这种”前一窗口 + 当前窗口”的双 buffer 设计是 overlap 模式的关键——确保压缩 KV 的”边界平滑”在 decode 阶段也能保持与 prefill 数学一致。

3.4.1 数学等价性

prefill 一次性算完 + decode 逐 token 算完，最终得到的压缩 KV 序列必须完全相同。否则模型在”先 prefill 长 prompt → 后 decode” 与 “纯 decode” 这两种方式下会出现行为漂移。

V4 的工程师通过 kv_state 与 score_state 的精确状态机，让 prefill 和 decode 的输出在数学上 bit-exact 等价。这一点是稀疏 attention 工业化的硬性要求——prefill 和 decode 走不同 codepath 但必须输出相同结果。

3.5 per-layer compress_ratio 的工程哲学

V4 的 compress_ratios 是个长度 62 的整数数组——num_hidden_layers=61（主 attention 层）+ num_nextn_predict_layers=1（MTP 层）= 62。config.json 实测分布：值仅 3 种（0 / 4 / 128），其中 128 出现 31 次、4 出现 30 次、0 仅 1 次（在数组最末——对应 MTP 层）：

[128, 128, 4, 128, 4, 128, 4, ..., 4, 128, 4, 128, 4, 128, 4, 0]
   ↑   ↑                                              ↑    ↑
 第0层 第1层                                       第60层  第61层(MTP)

读这串数字的关键是理解ratio 的语义：

• ratio = 0：本层不压缩，KV 不进入 kv_cache 的压缩段——只走滑窗
• ratio = 4：每 4 个 token 压一组——稀疏注意力的”主力工作层”，启用 Indexer
• ratio = 128：每 128 个 token 压一组——极端高压缩，主要用于”远距离粗略回忆”

V4 团队没有选择”全 ratio=4”或”全 ratio=128”的均匀方案，而是选交替混合：

flowchart LR
  subgraph 模型层级
    direction TB
    L0["第 0-1 层<br/>ratio=128"]
    L1["第 2 层<br/>ratio=4"]
    L2["第 3 层<br/>ratio=128"]
    L3["第 4 层<br/>ratio=4"]
    L4["..."]
    L5["第 60 层 (主模型末层)<br/>ratio=4"]
    L6["第 61 层 (MTP)<br/>ratio=0"]
  end
  L0 -->|粗略远记忆| L1
  L1 -->|精细稀疏选择| L2
  L2 -->|粗略远记忆| L3
  L3 -->|精细稀疏选择| L4
  L4 --> L5
  L5 -->|主模型 dense 输出| Out["主 head 输出"]
  L5 -.MTP 副 head.-> L6
  L6 -->|MTP 完整 KV 视野| MTPOut["MTP head 输出"]

这种”粗 → 细 → 粗 → 细 → 不压缩”的交替结构，体现了一个工程直觉：

1. 粗层（ratio=128）：保留远距离的”主旨信息”——比如 1M context 中前 100K token 在讲什么主题
2. 细层（ratio=4）：在 token 之间做精细的注意力交互——稀疏 attention 在这里发挥作用
3. 不压缩层（ratio=0，仅 MTP 层）：MTP 副 head 用完整 KV 视野——避免压缩损失影响 next-next token 预测的训练监督信号（具体动机是工程推断，源码无注释明示）

这种设计让模型同时具备”远距离粗略记忆”和”近距离精细推理”两种能力，在 1M context 下保持表现。

3.5.1 layer-wise 不均匀的工程动机

如果所有层都 ratio=4，KV cache 在 1M context 下要 16 GB（按 §3.1 的算式）。如果都 ratio=128，KV cache 只要 488 MB，但远距离信息被过度压缩，模型在长 context 任务上会塌陷。

V4 的混合方案给出了一个”成本 / 能力” 的最优解——主模型 61 层中约 31 层 ratio=128 节省显存（远距粗略），约 30 层 ratio=4 保证精度（稀疏选取）。MTP 层（第 62 个元素）的 ratio=0 是给副 head 留的”无损视野”——主模型本身的 61 层并没有 ratio=0 的层。

3.5.2 这串数字会演化吗

V4 Preview 版的 compress_ratios 是一个”工程现状”，不是不可调的常数。GA 版本可能会基于训练 / 推理实测调整这串数字。但数字的整体模式（粗 / 细交替 + 末层 0）几乎一定会保留——这是从 V3.2-Exp 的实验里得到的经验。

3.4·补 状态机的双 buffer 推演（一个具体例子）

decode 阶段的 kv_state / score_state 双 buffer 状态机最难理解。我们用一个具体的例子推演——假设 ratio=4、overlap=True（coff=2），从 start_pos=0 开始 decode 8 个 token：

初始状态（prefill 完成后留下的状态）：

kv_state    = [前一窗口的 4 槽位] + [当前窗口的 4 槽位]
score_state = [前一窗口的 4 槽位] + [当前窗口的 4 槽位]  (初始 -inf)

start_pos=0：第 1 个新 token 进来

should_compress = (0+1) % 4 != 0    → False
kv_state[ratio + 0 % 4] = kv_state[4] = 新 token 的 kv  (overlap 模式下写"当前窗口"的第 0 槽)
score_state[4] = score + ape[0]

start_pos=1：第 2 个 token

should_compress = (1+1) % 4 != 0    → False
kv_state[5] = 新 token 的 kv
score_state[5] = score + ape[1]

start_pos=2：第 3 个 token

should_compress = (2+1) % 4 != 0    → False
kv_state[6] = 新 token 的 kv
score_state[6] = score + ape[2]

start_pos=3：第 4 个 token，触发压缩！

should_compress = (3+1) % 4 == 0    → True
kv_state[7] = 新 token 的 kv  (先存)
score_state[7] = score + ape[3]

# 然后做压缩
# 注意 cat 的细节：取 kv_state 前 4 个的"前 d 维" + 后 4 个的"后 d 维"
kv_pool    = cat([kv_state[:4, :d], kv_state[4:, d:]])   # [8, d]
score_pool = cat([score_state[:4, :d], score_state[4:, d:]])

# softmax 加权求和
new_compressed_kv = (kv_pool * score_pool.softmax(dim=0)).sum(dim=0)

# 写入 kv_cache 的下一个槽位
self.kv_cache[start_pos // ratio = 0] = new_compressed_kv

# 滚动："当前窗口" 变成 "前一窗口"，给下一轮腾空间
kv_state[:4]    = kv_state[4:]
score_state[:4] = score_state[4:]

start_pos=4-7：第 5-8 个 token 重复上述循环，第 8 个 token 时再次触发压缩，写入 kv_cache[1]。

这个状态机看似复杂，但核心思想很简单——每次 decode 都把新 KV 累积到”当前窗口”，每 ratio 步触发一次压缩并把”当前”滚成”前一”。cat 的”前一窗前 d 维 + 当前窗后 d 维”的拼接方式是 overlap 模式的特殊处理——它对应 prefill 阶段 overlap_transform 函数的逻辑。

数学上能保证：prefill 一次性算 8 个 token、和 decode 逐个算 8 个 token，最终 kv_cache[0]、kv_cache[1] 的内容完全一致。这就是 V4 实现”prefill 与 decode 数学等价”的全部秘密。

3.5·补 V4 KV 几何与操作系统虚拟内存的对应关系

V4 的 KV 几何如果用一个熟悉的类比，最贴切的对应是操作系统的虚拟内存分级：

这个类比在多个层面成立：

容量与速度的反比关系：滑窗段最快访问（每次 attention 都查），但容量小（128 槽）；压缩段慢一些（要先经过 Compressor 加权求和）但容量大（百万 token / ratio）。

访问局部性：刚被滑窗段淘汰出去的 token，最有可能被压缩段保留——这是”时间局部性”的对应。每 ratio 个 token 压成一组的设计，相当于”页大小 = ratio”。

层级粒度：L1 / L2 / 主存的粒度逐级变粗（cache line → 页面 → segment），V4 的 KV 几何层级粒度也逐级变粗（token → ratio 组 → 整段压缩）。

TLB 类比：Indexer 在 V4 中扮演的角色类似 TLB（页表查找单元）——它快速给压缩 KV 打分，告诉主 attention “应该去查哪些压缩 KV”。

这个类比对工程师建立直觉特别有用：当你想搞清楚”V4 的某层在做什么内存层级的事”，把它对应到虚拟内存的层级就一目了然。第 4 章会展开 Indexer 这个”TLB”的内部工作机制。

3.6 Compressor 在 Attention 与 Indexer 中的双重身份

V4 的源码里有两套 Compressor：

• 一套在 Attention 里，给主 attention 用：self.compressor = Compressor(args, self.compress_ratio, self.head_dim)，rotate=False
• 一套在 Indexer 里，给稀疏选取的 score net 用：self.compressor = Compressor(args, compress_ratio, self.head_dim, True)，rotate=True

它们有什么差别？只在 rotate 标志：

if self.rotate:
    kv = rotate_activation(kv)
    fp4_act_quant(kv, fp4_block_size, True)
else:
    act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True)

• Attention 的 Compressor：FP8 量化（保留语义精度）
• Indexer 的 Compressor：先做 Hadamard 旋转，再 FP4 量化（更激进的压缩）

为什么 Indexer 的 Compressor 走 FP4？因为 Indexer 的输出只是”打分”，最终 topk 选取之后并不直接参与 attention 输出——它的精度损失对最终结果影响小，可以更激进地压。Hadamard 旋转把信息”打散”到所有维度，让 FP4 量化的精度损失更均匀，进一步提升精度 / 容量比。

这是 V4 的一个微妙工程：两套同源模块用不同精度——结构相同、参数独立训练、量化策略不同。第 4 章会展开 Indexer 的设计；第 12 章会展开 Hadamard + FP4 的精度链路。

3.6·补 Compressor 与 attention sink 的协同

第 2 章讲到 V4 的 attn_sink 是稀疏注意力的”兜底参数”——保证当稀疏选取找不到合适 KV 时数值不崩。Compressor 的设计与 attn_sink 形成了一对互补的”安全网”：

• Compressor 提供”长距离主旨记忆”：即便 Indexer 的 top-1024 选取错了，Compressor 至少把”这段 context 的整体语义”保留在压缩 KV 里
• attn_sink 提供”什么都不选的安全出口”：如果 Indexer 选到的 KV 都不相关，sink 接住这一步 attention 的注意力质量

两者一起构成 V4 稀疏注意力的”双重防御”：宏观信息（Compressor）+ 数值兜底（attn_sink）。这个设计的工程目标是”哪怕 Indexer 训练得不完美，模型也不会崩溃”——这是把稀疏注意力推到 1.6T 规模时必须有的工程冗余。

这里也可以引申一个观察：V4 在每个稀疏化决策处都有数值冗余兜底——

• 滑窗 KV 提供”近距离精确记忆”——保证当稀疏选取出错时，至少近期 token 不会丢
• Compressor 提供”远距离粗略记忆”——保证整体语义不丢
• attn_sink 提供”零选择”出口——保证数值稳定
• MTP 层 ratio=0 的完整 KV 提供”副 head 的 dense 视野”——保证 MTP 训练监督信号质量

四重冗余叠加，让 V4 即便在 Indexer 出现”大面积选取错误”时也能保持可用。这是稀疏路径与 dense 路径”工程纪律”上的最大差别。

3.6·补·补 KV cache 的字节级账本

把 V4 Pro 跑 1M context 的 KV cache 占用按字节级算清楚，理解显存压力的真实分布。

给定参数：

• n_layers = 61
• max_seq_len = 1,048,576
• window_size = 128
• head_dim = 512
• 假设所有层都 BF16 存（每槽 2 bytes × 512 = 1 KB）

单层占用 = (window_size + max_seq_len // ratio) × head_dim × 2 bytes

典型 V4 Pro compress_ratios 配置（约 30 层 ratio=4 + 30 层 ratio=128 + 1 层 ratio=0）：

30 层 × 256 MB = 7680 MB
30 层 × 8 MB   = 240 MB
1 层 × 1024 MB = 1024 MB
─────────────────────
总计 ≈ 8944 MB ≈ 8.9 GB / 序列

注：MTP 层（第 62 个元素，ratio=0）也算上时该层 KV 是 dense（每 token 一组），但 MTP 不在主模型 61 层之内。

对比 V3 的 KV cache（dense MLA, kv_lora_rank=512）：

61 层 × (1,048,576 × 512 × 2 bytes) = 65.5 GB / 序列

V4 是 V3 的约 13%——略低于 README 给的 10%。10% 这个数字应该是基于”全 ratio=128 + 头几层走 ratio=4”的更激进配置算出来的。

当 batch_size=8 时（典型生产负载）：

V3：8 × 65.5 = 524 GB —— 单卡放不下，必须 8+ 张 H100 张量并行 + KV offload V4：8 × 8 = 64 GB —— 单张 H100 能塞下，2 张就能舒服跑

KV 占用从”必须分布式”变成”单卡可塞”，是 V4 在长上下文 + 高并发场景下的革命性突破。第 19 章会讨论这对推理引擎部署架构的影响。

3.6·延展 Compressor 与 Mamba / 状态空间模型的对比

V4 的 Compressor 与 Mamba2、xLSTM 等”状态空间模型”在思想上有共同点——把序列信息压成定长状态——但工程实现完全不同：

V4 的 Compressor 选择非 recurrent：每个压缩组独立计算、可以并行化、可以随机访问。这与 Transformer 的”每个 token 独立 KV”哲学一致——只是把粒度从”每个 token” 放到”每 ratio 个 token”。

代价是：相比 Mamba 的 O(1) per-token 推理，V4 的 O(n/ratio) 推理在超长上下文下成本仍然线性增长。但 V4 通过 ratio=4-128 的非均匀混合，把这个常数压得很低——整体 KV 占用仍然是 ratio = 128 在主导。

这种”非 recurrent 压缩 KV”的选择，让 V4 保留了 Transformer 的全部并行优势，同时获得近线性的 KV cache 增长。这是 SSM 路线与 Transformer 路线之间的一个工程妥协——V4 没有放弃 Transformer，而是给它配了一套”分级 KV 内存”。

3.7 一段动手实验：自己实现一个 mini Compressor

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

class MiniCompressor(nn.Module):
    """简化版 Compressor，不含 overlap，便于理解核心逻辑"""
    def __init__(self, dim=512, head_dim=128, ratio=4):
        super().__init__()
        self.ratio = ratio
        self.head_dim = head_dim
        self.wkv = nn.Linear(dim, head_dim)
        self.wgate = nn.Linear(dim, head_dim)
        self.ape = nn.Parameter(torch.zeros(ratio, head_dim))
        self.norm = nn.LayerNorm(head_dim)

    def forward(self, x):
        # x: [B, S, dim], assume S % ratio == 0
        B, S, _ = x.shape
        kv = self.wkv(x)                                  # [B, S, head_dim]
        score = self.wgate(x)                             # [B, S, head_dim]

        kv = kv.view(B, S // self.ratio, self.ratio, -1)
        score = score.view(B, S // self.ratio, self.ratio, -1) + self.ape

        weights = score.softmax(dim=2)
        compressed = (kv * weights).sum(dim=2)            # [B, S/ratio, head_dim]
        return self.norm(compressed)


# 测试
mc = MiniCompressor()
x = torch.randn(2, 16, 512)
out = mc(x)
print(out.shape)   # 应该是 [2, 4, 128]

跑通这个 mini 版本后，再回看 V4 源码里 overlap 的处理，会觉得”原来不过是给 wkv / wgate 加倍输出维度，前后两半各管一段重叠区”。

3.7·补 Compressor 在 Tensor Parallel 多卡下的状态分布

V4 Pro 在生产部署时通常用 8 卡 TP（Tensor Parallel）。这种部署下，每张卡持有 1/8 的 attention head——n_heads / world_size = 16 个 local heads。Compressor 的状态分布需要回答一个问题：kv_state 与 kv_cache 是按 head 切分的吗？

读源码可以确认：Compressor 的状态不按 head 切分。

# Attention 类中
self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, kv_cache_size, self.head_dim))

self.head_dim 是 512（不是 n_heads × head_dim）——所有 head 共享同一组 KV（V4 是 num_key_value_heads=1 的 MQA-style）。这意味着：

• 每张 TP 卡上的 kv_cache 是完整的、未切分的 ——[B, kv_cache_size, 512]
• 每张卡上的 Compressor 是完整的副本 ——所有 rank 跑相同的压缩计算
• 这浪费了 KV cache 的并行性，但省了通信 ——稀疏 attention 不需要在 TP rank 间同步 KV

这种”KV 全副本、Q/O 张量并行”的混合切分，是 MQA-style 模型在 TP 下的标准做法。它的工程权衡是：

• ✅ KV cache 的访问不需要 all-reduce
• ✅ Compressor / Indexer 的 score 计算不需要跨 rank 同步
• ❌ KV cache 内存被冗余 8 倍——但因为 V4 的 KV 已经被压到 ~8 GB / 序列，8 倍冗余仍可接受

如果换成 GQA（多个 KV head），KV 切分到不同 rank 是合理的，但通信开销会上升。V4 选 MQA + 完整副本，是把”通信简单 + 实现简单”放在了”显存效率”之上。

第 15 章会展开 ColumnParallelLinear / RowParallelLinear 在 V4 中的具体使用，以及 MoE 的 expert parallel 如何与 attention 的 TP 配合。

3.8 延伸阅读

• DeepSeek-V3.2-Exp 模型卡——DSA（DeepSeek Sparse Attention）的最早期描述
• Mamba / Mamba2 论文（arXiv:2312.00752）—— 状态空间模型的”长距离记忆压缩”思路与 V4 Compressor 有平行性
• Linformer / Performer 等”低秩注意力”工作——把 attention 做成 O(n) 的早期尝试
• 本书第 4 章：Indexer 是 Compressor 的下游——把压缩 KV 用作”稀疏 KV 选择”的输入

3.8·补 Compressor 的工程师 FAQ

写到这里，读者最容易问的几个问题集中回答：

Q1: 为什么 Compressor 不是 attention 而是 gated pooling？

attention 是”query 主动查 KV”，gated pooling 是”输入自带门控权重做加权平均”。Compressor 的输入是 hidden state x（不是 query），目标是把 ratio 个 token 压成 1 组——它没有 query 的概念。用 gated pooling（softmax(score) × kv）比用 attention 简单得多，参数也少。

如果用 attention 做压缩，你必须给 Compressor 配一个”压缩 query”——这要么是固定的（学不到 token 自适应）、要么是从输入算出来的（多一组参数）。V4 的 gated pooling 是更经济的选择。

Q2: ape (absolute position embedding) 与主 attention 的 RoPE 是什么关系？

是两套不同的位置编码：

• 主 attention 的 RoPE：作用于完整序列的 1M token，提供”全局位置感”
• Compressor 的 ape：仅作用于”压缩组内部”的 ratio 个相对位置，提供”组内位置感”

两者在源码里是分开的——ape: [ratio, coff*head_dim]，freqs_cis: [max_seq_len, rope_head_dim/2] 复指数。Compressor 在做完 softmax 加权求和后，会再用 apply_rotary_emb 给压缩 KV 套上 RoPE——这时位置是”压缩组对应的 token 位置”（freqs_cis[:cutoff:ratio]）。

Q3: 为什么 ratio=4 启用 overlap，ratio=128 不启用？

overlap 是为了”组与组的边界平滑”。当 ratio=4 时，4 个 token 一组，组与组的边界很近——模型经常需要”跨组”的相关性。overlap 让相邻组共享 2 个 token，减少边界处的信息损失。

当 ratio=128 时，128 个 token 一组——组内本身已经是”粗粒度记忆”，组间边界平滑没那么重要。也不启用 overlap。

工程上还有第二个原因：overlap 让 wkv 输出维度翻倍（coff=2），参数量也翻倍。ratio=128 的层本身存的就少（每层 8 MB），翻倍也才 16 MB；但 ratio=4 的层有 256 MB 量级，翻倍变 512 MB——所以 ratio=4 启用 overlap 在工程上是”显存换精度”的合理交易；ratio=128 不需要这个交易。

Q4: kv_state 与 score_state 必须是 FP32 吗？

源码里这两个 buffer 显式指定 dtype=torch.float32：

self.register_buffer("kv_state", torch.zeros(..., dtype=torch.float32), ...)
self.register_buffer("score_state", torch.full(..., float("-inf"), dtype=torch.float32), ...)

是的，必须 FP32。原因是 decode 阶段每 ratio 步累积一次，跨数百次 decode 累积的状态如果走 BF16 会有精度漂移——softmax 加权求和的结果在数百次累积后偏离 prefill 的对应值。FP32 能保证 prefill / decode 的 bit-exact 等价。

Q5: Compressor 的参数量大不大？

简化算一下 V4 Pro 一层的 Compressor 参数：

• wkv: 7168 × (2 × 512) = 7.34M（overlap=True）
• wgate: 7168 × (2 × 512) = 7.34M
• ape: 4 × (2 × 512) = 4096
• norm: 512
• 总：~14.7M / 层

61 层 × 14.7M ≈ 900M。占 1.6T 总参的 0.06%——非常便宜。再加上 Indexer 的 Compressor，也只有 ~1.8B 总参。

Compressor 用 0.1% 的参数让 V4 的 KV cache 减少 90%——这是一笔极划算的工程交易。

3.8·补·补 Compressor 与 dense MLA 的”性能转折点”

V4 的 Compressor + sparse 路径不是无脑优于 dense MLA。在某些 context 长度下 dense 反而更快。把这条转折点曲线画清楚：

Context ≤ 8K：

dense MLA 的 KV cache 才几百 MB，attention FLOPs 也可承受。Compressor 引入的”压缩 + 反查” 开销反而比直接 dense 慢。这个范围内 V4 比 V3 dense 慢约 5-10%。

Context = 16K - 64K：

dense 的 KV cache 开始变大（GB 量级），FLOPs 也开始凸显。Compressor 开始拿到节省。V4 与 V3 dense 大致持平。

Context = 64K - 256K：

V4 的 Compressor 优势开始拉开——dense 的 O(n²) 在这个范围彻底变得不经济。V4 比 dense 快 2-5x。

Context = 256K - 1M：

dense 几乎不可用——KV cache 爆、FLOPs 爆。V4 是唯一可行选择。速度对比无意义。

Context > 1M（如果未来扩到）：

V4 当前不支持 > 1M——max_position_embeddings=1048576 是硬上限。如果未来扩到 2M / 4M，Compressor 的层级化压缩（多层 ratio）会让 V4 仍然可行。dense 永远不可能。

部署建议：

如果你的产品场景大部分 context < 8K，V4 不是最优选——选 dense 模型（如 Qwen3-32B dense）速度更快。V4 的甜区是 64K+。

理解这条曲线让你做”模型选择”的决策更清醒——不是”V4 在所有场景都最快”，而是”V4 在长上下文场景压倒性优势”。

3.8·延展 Compressor 与 fine-tune 的兼容性

V4 fine-tune 时 Compressor 的几个特殊考虑：

考虑 1：Compressor 参数也要 fine-tune

如果 fine-tune 时冻结 Compressor，新数据分布的”压缩重要性” 与预训练学到的可能不一致——Compressor 输出与下游 attention 不匹配。建议 fine-tune 时 Compressor 也参与梯度更新，但用比 attention 更小的 lr。

考虑 2:compress_ratios 不能改

ratio 是写在 config 的 per-layer 数组——fine-tune 不应该改这串数字。改了等于改 KV cache 几何，需要从头训。

考虑 3：fine-tune 数据的长度分布

如果 fine-tune 数据全是短 prompt（< 4K），Compressor 的远距离压缩能力得不到训练。fine-tune 后模型在短 prompt 上可能有提升、在长 prompt 上反而退化（因为 Compressor 知识被”洗掉”）。建议 fine-tune 数据保留一定比例长 prompt。

考虑 4：与 LoRA 的协同

如果用 LoRA 做 fine-tune（不 full fine-tune），Compressor 是否也加 LoRA 是个选择题。建议不要给 Compressor 加 LoRA——Compressor 已经是低秩结构（gated pooling），再加 LoRA 表达力受限。LoRA 集中在 attention 的 wq_b / wo_b 等位置即可。

这 4 个考虑让 fine-tune V4 时 Compressor 不被破坏——保留预训练学到的”如何压缩长距离信息”能力。

3.9 本章小结

• V4 的 KV cache 是”滑窗 + 压缩段”的双段几何，不是单一的连续 KV
• Compressor 的核心是”按 ratio 把多个 token 压成 1 组” + “ape 给组内位置感”
• prefill 与 decode 的两条 codepath 必须数学等价——kv_state / score_state 双 buffer 是关键
• per-layer 的非均匀 compress_ratios（粗 / 细 / 不压缩交替）是 V4 长上下文成本可控的真正秘诀
• Compressor 在 V4 中有两个身份：一个给 Attention 主路（FP8 + 不旋转），一个给 Indexer score net（FP4 + Hadamard 旋转）

第 4 章我们进入 Indexer——V4 稀疏注意力的”score net”，看它怎么用一个独立的 KV 视角，给压缩 KV 打分并选取 top-1024。