DeepInfer 源码拆解（十一）：注意力族与压缩 KV

容器侧：四个类，各自只管权重形状

mla_attn.py 的模块 docstring 把四个类列得很清楚，逐个看：

ColumnParallelLora 是 MLA 低秩投影的通用容器，描述的是 x -> A -> (norm) -> B 这条链：

1
class ColumnParallelLora(nn.Module):
2
    def __init__(self, in_features, out_features, rank, *,
3
                 num_lora_rope_head=0, rope_head_dim=64,
4
                 norm_module=None, use_fp8_gemm=False):
5
        self.A = ReplicatedLinear(in_features, rank + num_lora_rope_head * rope_head_dim, ...)
6
        self.B = ReplicatedLinear(rank, out_features, ...) if out_features > 0 else None
7
        self.norm = norm_module

两个细节值得停顿：Wq 用 num_lora_rope_head=0（query 侧没有 MQA rope 头），而 Wkv 用 num_lora_rope_head=1——A 的输出维度是 rank + 1 * rope_head_dim，也就是低秩压缩 latent 与 MQA rope 头拼在同一个投影里，这正是 MLA 与「纯低秩 + 单独 rope」写法的区别；另外 Wkv 的 out_features=0，于是 B 为 None——KV 侧没有上投影，因为 V4.1 推理只缓存压缩后的 latent，上投影在注意力核心里按需恢复。

SoftmaxAPEGatedColumnParallelLora 是压缩 KV 的投影容器，docstring 一句话点破 C1 与 C2 的差别：

1
"""KV projection of the compressor: `A` produces the values, `Gate` produces the logits.
2

3
The single-token softmax of C1 is always 1, so only A is needed; C2 uses both A and Gate.
4
"""

压缩率是 1 的路径（单 token 压缩）softmax 恒为 1，所以只需要 A；压缩率是 2 的路径既要值也要门控 logits，A 与 Gate 都要建。注意两者的注释：use_fp8_gemm=False——Gate 与 A 走 bf16，不参与 fp8 量化，因为它们是压缩权重（数值敏感度与主投影不同）。

MQAIndexer 是稀疏索引器的参数容器，构造时用一串断言把自己钉死在一条路径上：

1
assert self.use_rope
2
assert not self.use_q_norm
3
assert self.rope_dim == 64
4
assert self.sparse_stage == 2
5
assert not index_cfg.get("mlp_kproj", False), "DeepInfer KProject only supports a flat k_proj"

也就是说索引器必须带 rope、rope 维度固定 64、只支持 stage-2 稀疏、且 k_proj 必须是扁平线性层（不支持 MLP 形态的 k 投影）。它持有三个投影：q_B（q_lora_rank -> n_head * head_dim，fp8）、weights_proj（hidden -> n_head，bf16，每头一个 logits 权重）、k_proj（kv_lora_rank + qk_rope_head_dim -> head_dim，bf16）。

容器侧的「结构性缺失」：哪些层留空

LoraRopeCompressFlashMHAAutoRecomp.__init__ 有两个参数：is_kv_source_layer 与 is_index_module_source_layer。它们决定容器里留空还是建实例：

1
self.PackedWkv = (
2
    SoftmaxAPEGatedColumnParallelLora(...) if is_kv_source_layer else None
3
)
4
self.index_attn = None
5
if self.sparse_stage == 2 and cfg.get("attn_index", None) is not None and is_index_module_source_layer:
6
    self.index_attn = MQAIndexer(cfg)

注释写得很直白：只有 kv-source 层真的持有 PackedWkv；其它层的 PackedWkv 是 None，前向时从源层取。index_attn 同理。这意味着容器在设计上就承认了「层与层不一样」——不是每层都有完整的注意力参数，而是按拓扑分成源层与消费者层两组。哪些层是源层，正是由 cross_layer.py::SharedLayerTopology 从配置推导：layer_map[l] == l 的层是源层，其余层复用源层的模块。

这个设计还有一个隐含收益：显存和加载成本按源层数量计。如果 61 层里只有若干层持有压缩权重与索引器，参数加载与显存占用就只落在这些层上，其余层在容器侧是空的。

运行时侧：把容器适配成每步要用的东西

attn.py 里 V4AttentionBase.__init__ 做的是「适配」，主要动作有四类：

第一类，把 HF 权重整理成推理布局。 例如 q_a 与 kv_swa 的投影权重可以合并成一次 GEMM（merge_linear_weights），q_B 与索引器的 q_B 也可以合并；但注释立刻给出例外——开启 yoco_swa_share_hidden 时不能融合，因为 Q 与 KV 的输入不再是同一份 hidden。Wv 与 out_proj 则要做量化布局变换（transform_sf_into_required_layout 与 permute_wv_proj），把 fp8 scale 排成融合内核要求的形状。

第二类，把「层类型」翻译成运行时开关。

1
self.layer_cfg = getattr(attn_module, "cfg", AttrDict())
2
self.is_sparse = bool(self.layer_cfg.get("index_topk"))
3
self.has_packed_kv = getattr(attn_module, "n_packed_heads", 0) > 0
4
self.use_swa = self.layer_cfg.get("use_swa", True)
5
assert self.use_swa or self.has_packed_kv, "use_swa=False requires packed KV attention"
6
if self.is_sparse:
7
    assert self.use_swa

三行开关覆盖了注意力族的三个正交维度：有没有压缩 KV、要不要 SWA、有没有本地索引器。断言则固定住合法组合——稀疏必须以 SWA 在场为前提，因为稀疏索引选出的压缩条目要和滑窗 KV 拼在同一次注意力里。

第三类，构造调度键。 这是「运行时」存在的最直接证据：

1
self.sched_meta_key = f"{self.kv_compress_ratio}_{self.has_packed_kv}_{self.use_swa}_{self.layer_cfg.get('index_topk', 0)}"

FlashMLA 的 tile scheduler metadata 按这个键缓存（input_metadata.flash_mla_sched_meta_dict[self.sched_meta_key]）——键里的四个量全部来自运行时适配的结果，容器侧一个都不涉及。

第四类，决定是否启用融合内核。 use_fused_norm_rope_attn_rope_cast 的取值把「稠密 packed-KV 解码」排除在外（该场景禁用融合），并要求头数落在 (64, 128)、平台是 GPU。融合内核把 norm、rope、cast 与前向的 rope 收尾合成一步，是 FP8 路径上的常态。

把四类动作合起来看，问题一的答案就完整了：容器是加载期的事实（权重形状、层间差异、留空与否），运行时是每步的事实（调度键、缓存别名、量化布局、融合开关）。两者的边界之所以画在这里，是因为容器要在权重加载时确定、且必须能被 HF 权重映射稳定引用；而运行时的一切都要在每步前向里随批次形状与调度状态变化。

压缩本身：C1 与 C2 两条路径

Compressor 的构造断言把支持面收窄到 C1/C2：

1
self.ratio_key = str(self.kv_compress_ratio)
2
if self.kv_compress_ratio not in (1, 2):
3
    raise ValueError("Only C1/C2 compression is supported")

forward_pre_norm 按压缩率分支：

• C1（kv_compress_ratio == 1）：断言投影输出是 bf16 且 state_cache is None——单 token 压缩没有需要攒的窗口，不需要状态缓存；若在 CP 下则先按 mirror_cp_reverse_index 做一次 gather。
• C2：断言投影输出是 fp32 且状态缓存在场，然后把投影输出按最后一维 chunk(2) 切成 score 与 latent——这正是 SoftmaxAPEGatedColumnParallelLora 里 A 与 Gate 两路的输出。

压缩的内核分 prefill 与 decode 两个入口，参数列表本身就是一份「压缩需要知道什么」的清单：

1
kv_compressed = compress_and_update_state_prefill(
2
    score, latent, state_cache,
3
    input_metadata.state_block_idx,
4
    input_metadata.compressed_batch_mapping[self.ratio_key],
5
    input_metadata.cu_seqlens_q, input_metadata.cu_seqlens_k,
6
    input_metadata.cu_compressed_seqlens_q[self.ratio_key],
7
    self.kv_compress_ratio,
8
    input_metadata.mirror_cp_reverse_index,
9
)

decode 侧的签名短得多，但多了一个关键参数：

1
kv_compressed = compress_and_update_state_decode(
2
    score, latent, state_cache,
3
    input_metadata.state_block_idx, input_metadata.context_lens,
4
    max_seqlen_q=envs.DEEPINFER_USE_NEXTN_DECODE + 1,
5
    compress_ratio=self.kv_compress_ratio,
6
)

max_seqlen_q = NEXTN + 1 这一行要和第三篇（KV 内存与 SWA 有界重放）对上：投机解码一次 decode step 可能推进多个草稿 token，被拒绝的部分要回滚，而压缩的 ring 缓存里同样留有回退历史。这里的 max_seqlen_q 就是压缩内核侧为这件事预留的声明。

存储路径：一次融合内核完成 norm、rope、cast、store

forward_store_kv_cache 是压缩 KV 落 cache 的唯一路径，注释把融合粒度写得很清楚：一次调用完成 norm、rope、Hadamard 变换、量化与落盘，目标 layout 由 get_kv_cache_format(kv_cache) 推出来（V4.1 上是 MXFP8/FP4 系列布局）。落盘位置是 slot_mapping_store_ca[self.ratio_key]——压缩条目按压缩率分桶存放，C1 与 C2 各自一张槽位表。

这段代码里最值得学的是一个「按需拷贝」的细节：

1
pre_rope_kv = torch.empty_like(kv_compressed) if return_pre_rope_kv else None

注释交代了理由：只有 KProject 索引器需要这份 pre-RoPE latent，其它路径不拷贝。索引器的 K 由主压缩 KV 的 pre-RoPE 版本投影而来（见问题三），所以当这一层没有索引需求时，这份拷贝连同它的显存分配一起被省掉。这是典型的「把特化机会做到单字段」的写法。

索引器在算什么

index_attn.py 的 MQA 类 docstring 一句话给出算法定义：算 relu(q·k)·weights 的 logits，并为每个 query 选出 top-k 个压缩条目。数据格式上它有一组自己的约定（都在构造的注释里）：

• 索引器的 Q 与 K 全程 FP4（E2M1 + UE8M0 scale），weights 与 psum 走 bf16——. # Indexer Q/K always use FP4 (E2M1 + UE8M0 scale); weights and psum use BF16.；
• max_kv_len = runtime_max_model_len // kv_compress_ratio，索引长度按压缩率折算；
• psum 的行宽由 logits 的实际布局决定：稠密路径是 max_kv_len * 2 字节，稀疏候选路径是候选 token 数乘 2，并断言 1024 字节对齐。

还有一个注释解释了这 22 KB 代码里最「玄学」的参数——分块宽度：

1
# NOTE The chunk size is min'ed with 8192 because mqa downclocks when the chunk size is large; the chunk has to be
2
# split up and have some topk operations stacked on it to achieve the best performance.
3
self.compute_chunk_size = min(8192, envs.DEEPINFER_INDEX_ATTN_MEMORY_LIMIT // psum_row_bytes // 8 * 8)

compute_chunk_size 由显存配额 DEEPINFER_INDEX_ATTN_MEMORY_LIMIT 与 psum 行宽推出来，但必须再取一次 8192 的上限：chunk 太大时 MQA 内核会降频，把 chunk 拆小、再把 topk 叠上去才是最优解。这条限制的工程含义是：索引器的显存配额不只是「够不够」的问题，也是「性能对不对」的问题。

两级 topk：候选块与候选内选点

索引器的 topk 不是一次完成的，而是「块级 → token 级」两级：

1. 候选源层先算块分数（candidate_block_max，即块内 max-reduce），用 topk 选出 candidate_topk_blocks 个候选块，写入 cross_layer_runtime.block_candidate_idxs；
2. 同组的后续层只在候选块集合内算 token 级 topk（fp8_fp4_sparse_mqa_logits / fp8_fp4_paged_sparse_mqa_logits），算完再用 restore_candidate_indices_inplace 把「候选内下标」还原成全局下标。

这个两级结构是精度与吞吐的折中：块级筛选用更便宜的内核砍掉绝大部分 KV 页，token 级再用稀疏内核在候选集合里精确选点。构造侧的断言把候选配置钉死在一小组合法值上：candidate_block_size in (8, 16)、topk_block % 8 == 0、score_reduce in ("max",)；消费者层的头部约束是 n_head <= 32 and head_dim == 128。

候选块机制的归属也在构造时确定：

1
self.is_candidate_topk_source = shared_topology is not None and shared_topology.is_candidate_source_layer(self.layer_idx)
2
self.use_candidate_topk = candidate_source is not None and not self.is_candidate_topk_source

候选源产生候选、消费者只消费候选——而「谁是源、谁是消费者」来自共享拓扑，这就自然引到下一个问题。

共享的机制：layer_index_k_map 与一次投影

索引器的 K（下称 index-K）是这样来的：主压缩 KV 的 pre-RoPE latent 经 k_proj 投影、再过 RMSNorm（KProject 的 docstring 原话），然后按 FP4 layout 落进索引器自己的 KV cache（ica_cache）。它的构造分支是共享的核心：

1
# HF V4.1 index-K is projected from the main compressed KV.
2
# Shared index-K consumers reuse the source cache without constructing a K module.
3
if self.store_index_k:
4
    self.index_k_module = KProject(index_attn_module.k_proj, index_attn_module.k_norm)

store_index_k=False 的消费者层干脆不建 K 模块：decode 时直接读源层的 ica_cache，prefill 时复用源层 load 出来的那份张量（cross_layer_runtime.loaded_prefill_index_kv），自己只重算 MQA logits。

为什么值得共享？因为这份 K 的来路决定了它天然是「一份就够」的：所有消费者层的 index-K 都来自同一份主压缩 KV latent，投影矩阵也是同一组权重（源层与消费者的 kv_compress_ratio、index_topk 由拓扑断言强制相等）。每层各算一遍，是把一次线性投影加一次 RMSNorm 重复 N 遍；共享则把这个代价摊到源层一次。

拓扑为什么必须 fail-fast

cross_layer.py 的 SharedLayerTopology 把配置里的三张稀疏地图归一化成覆盖全部层的列表：

• layer_index_module_map：索引器（topk 产出）的共享；
• layer_index_k_map：index-K 的共享，默认跟随 index-module 地图；
• layer_kv_map：主 KV（ca_cache）的共享。

归一化的规则是「自引用 = 源层」，消费者只能指向它前方最近的源层（分段共享，_validate_segmented_map）。在此之上还有一条粒度链校验：

1
self._validate_map_refines("index-topk", self.layer_index_module_map, "index-k", self.layer_index_k_map)
2
self._validate_map_refines("index-k", self.layer_index_k_map, "kv", self.layer_kv_map)

注释解释了为什么必须成立：消费者读的是源层的 ica_cache，而这份 K 又是从源层的压缩 KV 投影出来的，所以每一跳都不能跨过更粗粒度段的边界。换句话说，共享关系必须构成 index-topk ⊆ index-K ⊆ KV 的细化链；否则「复用谁的缓存」和「KV 是谁写的」会对不上。

这条链还有一个方向性的约束——share-kv 不能越过 YoCo 截断边界：

1
assert n_prefill_layer == cfg.num_hidden_layers or self.layer_kv_map[n_prefill_layer] == n_prefill_layer

因为截断模式下，边界前后的 KV 归属不同（第十篇讲过的跨层共享与 YoCo 边界协同），消费者若跨越边界去复用，拿到的 KV 来源语义就变了。

最值得记录的是失败模式本身，源码把它写成了断言的理由：

1
# The cache engine decides whether a layer participates in sharing by "does this layer carry
2
# layer_info_shared_cfg" (_get_shared_mapped_layer_id); if a consumer layer lacks it, the cache
3
# engine treats it as a source and allocates an independent cache, but the topology side never
4
# writes it, so it silently reads uninitialized memory, hence this must be fail-fast.

缓存引擎与拓扑是两套判定：缓存引擎看「这一层有没有携带 layer_info_shared_cfg」来决定是否别名别人的缓存；拓扑只看地图。两边一旦不一致（比如消费者层漏配了那段配置），缓存引擎会把它当源层、给它分独立缓存，而拓扑侧永远不会往那块缓存里写——读到的就是未初始化的显存，且不会报错。所以这一条必须 fail-fast：宁可在模型构建时崩溃，也不能让错误的数据无声地流进注意力。

CrossLayerRuntime 则是每步前向的载体，docstring 列出了六类跨层值：shared_hidden（第一个 YoCo 消费者层的 bf16 hidden）、pre_rope_kv（最近 kv 源层的 pre-RoPE 压缩 KV，给索引器的 KProject 用）、topk_indices（最近索引源层的 topk，share-index）、block_candidate_idxs（候选块）、loaded_prefill_kv（最近的 kv 源层已经 load 好的 prefill KV，消费者直接复用同一张量、只覆盖其中的 SWA 区）、loaded_prefill_index_kv（索引源层已 load 的 index-K）。它只有一个槽位、不做来源匹配——因为拓扑是分段共享的，而层按递增的 layer_id 执行，消费者需要的永远是「前方最近的源层」写下的那份。

三个类

文件里只有三个类：V4AttentionBase（适配层：norm、rope、权重合并、量化布局，及前述四类动作）、V4Attention（主模型注意力：压缩 KV、稀疏索引、SWA、YoCo、跨层共享五条线都汇在这里）、DFlashAttention（草稿注意力：独立的 draft KV）。

V4Attention.__init__：分支在构造期就分层

除了问题一讲的运行时开关，构造期还有一组「YoCo 相关」的判定：

1
self.is_yoco_hidden_consumer_layer = is_yoco_consumer_layer(cfg, layer_idx)
2
self.is_first_yoco_consumer_layer = self.is_yoco_hidden_consumer_layer and self.layer_idx == cfg.n_prefill_layer
3
self.yoco_swa_share_hidden = self.use_swa and self.is_yoco_hidden_consumer_layer and cfg.get("yoco_swa_share_hidden", True)

yoco_swa_share_hidden 的含义是：YoCo 消费者层的 SWA 投影不消费本层 hidden，而消费跨层共享的那份 hidden（第一个消费者层例外——它直接用本层的 attn-norm 结果，见 _resolve_shared_hidden）。这解释了前面那个「Q 与 KV 不能融合」的例外来自哪里。

forward：一次前向的编排

主路径的 forward（约 120 行）是 46 KB 的缩影，按顺序可以压成九步：

1. 元数据扩展。 input_metadata.compute_attn_info(self.max_model_len, self.window_size, block_size)——注意 block_size 是从 kv_cache.ca_cache 反推的（shape[1] * kv_compress_ratio），压缩条目的块尺寸乘回压缩率才是 token 意义上的块尺寸。

2. 解析共享 hidden、投影当前 hidden。 _project_current_hidden_inputs 只让非 YoCo 的 KV 源层在这里投影主 KV——YoCo 共享 KV 走单独的投影（shared_hidden_kv_proj），因为它的输入不是本层 hidden。

3. YoCo compact 分支。 若 consumer_full_metadata 在场（compact prefill），先调 _store_yoco_full_kv：用全量 hidden 更新主/索引 KV，供 compact query 的前向使用；随后主路径跳过 KV 更新，只算 compact query 的 topk。注释强调这条分支的必要性：query 被裁剪了，但 KV 仍要覆盖完整上下文。

4. Q 与 SWA 投影。 _project_q_and_swa 一次拿到 q_a 与 kv_swa（能融合时是一路 GEMM）；在 CP 下，压缩器输入的 all_gather 先发起（_make_cp_gather_getter），与下面的 Q 投影重叠。

5. SWA 张量准备。 _prepare_swa_kv 用融合内核完成 norm、rope、e4m3 量化（round_sf=True, sf_format="ue8m0"）与落 cache；prefill 侧还负责把本 rank 的 KV 在 CP all_gather 后接回来。

6. Q 侧投影与拆包。 q_b_proj 一次算出主注意力 Q 与索引器 Q（merge_linear_weights 合并过），再 split 回两段；随后 _q_b_norm_rope 按融合内核开关决定是否省略。

7. _compute_kv_ca_index：把 KV 侧三件事串成一条。 顺序是「主 KV store（返回 pre-RoPE latent）→ index-K store/topk → 把 topk 写进 cross_layer_runtime」。共享判定在这里体现：kv 源层 _store_main_kv 并把 pre_rope_kv 写进运行时；is_index_k_source 的层从运行时取它；is_sparse 但不是索引源的层直接复用 cross_layer_runtime.topk_indices。若这一层自带 index_attn，则调用索引器（store_kv=not yoco_full_kv_stored——compact 分支下 KV 已经由全量路径写过，跳过重复 store）。收尾时 cross_layer_runtime.topk_indices = kv_ca_index，供后续消费者层复用。

8. prefill：拼装 KV 再跑注意力。

1
# Assemble the input kv for prefill attention: [swa cache | current swa kv | compressed cache]
2
total_kv = total_k_swa_cached + total_k_swa_cur + total_k_compressed

三段拼接：已缓存的 SWA、本次新增的 SWA、压缩条目，分别由 load_paged_kv_cache 与 kv_swa.copy_ 分段填充；kv_index_cat 把 SWA 与压缩两段索引拼成一份，送给稀疏 prefill 内核。这里藏着一条顺序约束，注释写得很明确：

1
# The SWA store slots may overlap the ranges loaded above, so prefill must store only after the cache load.

SWA 的 store 槽位可能与上面的 load 区间重叠——同一块环形缓冲既被读也被写，所以必须先把旧内容读完，再写新内容，顺序颠倒就会读到自己的新值。这种约束不会出现在签名里，只能靠注释与代码顺序传递。

9. decode：SWA 与压缩 KV 各走各的 cache。

1
attn_kv_args = dict(
2
    key_cache=kv_cache.swa_cache,
3
    indices_in_kvcache=kv_swa_index,
4
    key_cache_extra=kv_cache.ca_cache,
5
    indices_in_kvcache_extra=kv_ca_index,
6
    index_length_extra=dense_index_len,
7
)

SWA 在场时，注意力内核的主 cache 是 SWA、额外 cache 是压缩条目；非 SWA 层则反过来只有 ca_cache 一路。keys_cache_extra 这个形参在 FlashMLA 侧的对应物就是第三篇提到的「extra k cache」。decode 还会带上 batch_invariant=envs.DEEPINFER_BATCH_INVARIANT——批量不变推理的开关直接贯穿到注意力内核，这是总览篇讲过的确定性要求在这条路径上的落点。

把九步里的分支条件抽出来，就是问题四的答案——attn.py 的代码量主要不是算法，而是分支矩阵：

DFlashAttention：草稿路径的注意力变体

DFlashAttention 服务于 DFlash/MTP 草稿，前向入口是 forward_drafts。与主路径的两个关键差别：

草稿 KV 是「一次 decode step 的临时物」。注释原话：

1
# Draft keys live for one decode step and are overwritten by the next step.

草稿 KV 写进 swa_cache 的一小段专用槽位（dflash_newx_slot_mapping_store），下一步直接覆盖——不进主 KV cache，也不参与前缀复用。这与第六篇（投机解码篇）讲的草稿生命周期是同一件事，在这里能落到具体的槽位映射上。它的调度元数据用固定键 "DFlash" 缓存，独立于主路径的 sched_meta_key。

另外构造参数 prefill_only 会只保留 kv_norm 就返回——P/D 分离部署下，纯 prefill 角色不需要草稿侧的 Q 投影。