DeepInfer 源码拆解（三十四）：候选块 topk 管线——块级筛选、紧凑下标与全局还原

从第十一篇的「两级 topk」接着往下

第十一篇（DeepInfer 源码拆解（十一）：注意力族与压缩 KV）在讲稀疏索引器时给出过一个骨架：索引器的 topk 不是一次选完，而是先按块选候选、再在候选集合内选 token，两级结构是精度与吞吐的折中。那篇的重点在「谁共享谁、为什么共享」，候选块只占了一小节。

这一篇把候选块这条线单独拆开。它的代码量不大——三个 TileLang 内核加起来不到四百行，真正难的是四套下标坐标系如何在预填充、解码、跨层复用和分页 KV 之间来回换算。先把结论摆出来：

• 候选块的本质，是把「每层都要扫全长」的索引器计算，改成「候选组内一个源层扫全长、其余层只扫一个定长候选集」。 在 DeepSeek-V4.1-Flash 的配置里，候选集是 2048 个块 × 每块 8 个压缩 token，共 16384 个位置，而索引器最终只取 512 个。这两个数字决定了整条管线的形状。
• 块级筛选与 token 级筛选用的是两套打分语义：块级取块内最大 logits，并且把「最后一个有效块」抬成 $+\infty$，确保它一定进入候选；token 级才在候选集合里做真正的 top-k。前者是粗筛，允许漏掉块内非最大的位置；后者负责精度。
• 下标要在四种坐标系之间换算：序列内压缩 token 下标、候选集合内的紧凑下标、分页 KV 的物理槽位、以及预填充里带 ks 偏移的全局下标。三个内核分别负责其中一段，任何一段的边界条件写错，都不会报错，只会让注意力读到错误的位置。

（图源：DeepInfer 官方仓库 doc/images/architecture_overview.png。本文讨论的候选块内核位于 model_executor/layers/tilelang_ops 算子层，由其上的稀疏索引器运行时调用。）

这条管线到底省了什么

先把问题摆清楚。稀疏注意力（DeepSeek Sparse Attention，DSA）在 DeepSeek-V3.2 中被正式提出，核心是一个闪电索引器（lightning indexer）加一个细粒度 token 选择（fine-grained token selection）机制：索引器为每个 query 计算它与所有历史 token 的相关性分数，再选出 top-k 个 token 交给主注意力。索引器本身被压得很便宜——头数少、能用 FP8/FP4 实现，但它的计算量随上下文长度线性增长，而且每一层都要算一遍。

DeepInfer 面对的模型是 DeepSeek-V4.1-Flash。从模型配置（config.json 的 text_config）能读出几个关键数字：

把这张表放在一起，故事就清楚了：

• 第 0、1 层是纯滑窗（compress_ratios 为 0），没有压缩 KV，也没有索引器；第 2–19 层压缩率为 2，第 20–39 层压缩率为 1。
• 索引器模块按「段共享」分布在 8 个源层上，但候选块机制只从第 20 层开始（candidate_source_layer_id = 20）。第 2、8、14 层这三个索引器要自己扫全长（它们的「消费层」是 2–7、8–13、14–19 段里的各层）；第 20 层之后的索引器源是 24、28、32、36 四层，它们作为消费层复用第 20 层产出的候选块，第 20 层自己则是候选源、按全长扫描。
• 候选集大小是 $2048 \times 8 = 16384$ 个压缩 token，而最终 topk 只有 512 个——候选集比结果宽 32 倍，留下足够的余量让「先粗筛、再精筛」两步走。

（图源：DeepSeek-V3.2 论文 arXiv:2512.02556 的 Figure 2 v32_arch.svg。绿色部分即索引器打分与 top-k 选择，候选块机制把「打分」这一步的计算域从全长缩到候选集。）

省了什么？把每层索引器的 logits 计算量写成一个简单模型。设序列有 $L$ 个压缩 token，每行 query 的 logits 开销正比于它扫描的位置数。稠密路径下索引器要扫整个 $L$；候选块路径下它只扫候选集 $C = 2048 \times 8 = 16384$ 个位置（超过 $L$ 时自然退化为全长）。于是，对第 20–39 层这些消费层：

（自制示意图：按 $C=16384$ 的候选集画出消费层索引器的相对开销，纵轴按 $L=128$ 归一化。横轴为均匀取样的长度点，不是严格对数刻度。）

这个图有两个容易误读的地方，需要点破：

1. 横轴和纵轴都是「相对量」，不是真实耗时。 真实耗时还取决于 logits 行的 stride 对齐、内核的 wave 量化、以及每个 query 是否同长。这张图只说明「扫描域」这一层开销的上限。
2. 候选块的收益是分段有效的：上下文长度小于候选集时（也就是 $L \le 16384$ 的压缩 token，对应原始长度视压缩率而定），候选块路径基本没有收益，但也不会变慢——此时内核走的就是全长扫描，候选集合退化为「所有块」。这解释了为什么 candidate_block_max 在有效块数 $< topk\_blocks$ 时不需要额外防护。

这里还有一个更隐蔽的收益：索引器 logits 的中间结果不落回全局内存。 稠密路径下，每层要把一整行 $L$ 个 bf16 logits 写进 psum 缓冲再做 topk；候选路径下，psum 行宽固定为候选 token 数，与 $L$ 无关。对长上下文来说，这部分中间结果的显存带宽省得比计算还多。

四套下标与一次数据流

候选块管线的难点不在算法，在下标。整条管线至少涉及四套下标：

1. 序列内压缩 token 下标（seq-local compressed index）：一个 query 眼中的「第几个压缩条目」。它是从 0 开始、只对有效压缩条目计数的逻辑下标，也是模型配置里 index_topk 的口径。
2. 候选集合内的紧凑下标（compact index）：把「候选块集合」看成一个虚拟数组 $[0, 2048 \times 8)$，下标 $\mathit{cid}$ 表示候选集合里的第 $\mathit{cid}$ 个位置。它和序列下标的换算是：

其中 $B$ 是 candidate_block_size（V4.1 为 8），$b_j$ 是候选集合里第 $j$ 个块的序列内块号。这个换算是内核三的核心。 3. 分页 KV 的物理槽位（physical slot）：分页 KV 的存储单位是页，逻辑 token 下标要经块表换算成物理地址，形式为 $\text{block\_table}[t / B_{\text{page}}] \cdot B_{\text{page}} + (t \bmod B_{\text{page}})$，其中 $B_{\text{page}}$ 是一页包含的压缩 token 数。它是内核外的 convert_kv_index 干的。 4. 预填充的全局下标（gathered global index）：预填充时多个序列的 KV 被 gather_kv 拼成一段连续张量，每个序列有自己的起点 ks，所以序列内下标要加上 ks 才是这段拼接张量里的全局位置。

一次前向里，数据从 query 出发，依次经过这四套坐标：

1
query -> q_fp8or4（索引器 Q 量化）
2
      -> 索引器 K：KProject 投影 + 量化，写进 ica_cache（复用第十一篇的 fused_transform 存储路径）
3
      -> p_sum = relu(q·k)·w 的 logits（DeepGEMM：稠密或稀疏）
4
      -> [仅候选源层] candidate_block_max -> 块分数
5
      -> [仅候选源层] DeepSelect topk -> 候选块号 block_candidate_idxs
6
      -> [消费层] 稀疏 MQA 只算候选块内的 logits
7
      -> topk（DeepSelect）在紧凑坐标里选出 topk_tokens 个
8
      -> [消费层] restore_candidate_indices_inplace -> 序列内下标（或加 ks 全局下标）
9
      -> [解码] convert_kv_index -> 分页物理槽位
10
      -> 主注意力的 index 参数（FlashMLA 稀疏路径）

注意这条链上的共享：候选块的产生只需要在源层做一次，后续消费层直接读 cross_layer_runtime.block_candidate_idxs。这与第十一篇讲的索引 K 跨层共享同构——计算量最大的那一步只做一次，其余层复用结果。

还要注意链尾的 -1：四套坐标里任何一步「这个位置不存在」都用 -1 表示。它在 topk 的阶段由 idx_oob_fill_value 写入，在 restore 和 compact_lens 里被当作无效块跳过，在 convert_kv_index 里被原样保留，最后交给 FlashMLA 当作「跳过」标记。整条链上 -1 承担同一个语义，这是后面「工程边界与陷阱」里反复出现的主题。

内核一：candidate_block_max

这个内核做一件简单的事：把一行 $L$ 个 bf16 logits 按 block_size 切成若干块，取每块内的最大值，得到一个行宽为 $\lceil L / B \rceil$ 的块分数数组。但它的接口描述里有几个不那么简单的约定，先看文档字符串：

1
def candidate_block_max(logits, valid_lens, block_size):
2
    """Compute block scores and valid block counts for candidate selection.
3

4
    Each score is the maximum valid token score in the block, capped at the
5
    largest finite bf16 value. The last valid block gets +inf so
6
    that candidate selection always includes it. Scores outside the valid block
7
    range are unspecified.
8
    ...

三句话对应三个设计决策：

为什么用块内最大值（max-reduce）而不是平均值或求和？ 索引器最终要选的是高分 token，而主注意力的质量取决于「高分行是否被选中」。块内最大值保证了一个块只要含有一个高分位置，它的块分数就高，不会被块内其它低分位置稀释。平均值的风险是：一个块里既有非常高的分、也有大量低分位置，平均之后排名掉出候选集，那个高分位置就丢了——而它本来就是最该被保留的。这是「粗筛只能漏掉低分行」的保守策略，代价是分数分布均匀的块会更容易被选中，候选集覆盖面变宽。

为什么最后一个有效块要抬成 $+\infty$？ 这是跨层共享的必要条件之一。候选集是层间共享的：源层按配置选出 candidate_topk_blocks 个块，消费层沿用同一组块。如果源层按 max 排出的候选集恰好不含「最后一个有效块」，消费层就可能连最新写入的 KV 都看不到。虽然注意力本身有滑窗兜底（V4.1 的 sliding_window 是 128），但索引器的产出决定了主注意力的稀疏路径是否覆盖最近的位置，而解码时最新 token 恰恰是最重要的。抬成 $+\infty$ 是用一个数值约定来保证「最新块永远在候选集里」，不需要任何额外的下标检查。

为什么超出有效块范围的值是 unspecified？ 文档里写得很明确：Only columns before valid_blocks[row] are guaranteed initialized; selection must use end=valid_blocks. 内核只在有效范围内写，有效范围之外的显存是脏的。topk 内核必须传 end，否则排序会把脏数据当成候选。这一点在下面 DeepSelect 的接口里反复出现。

内核的返回类型也值得注意：

1
    block_scores: [num_rows, aligned_num_blocks], bf16.
2
        num_blocks = ceil(max_model_len / block_size); rows are padded to a
3
        multiple of 1024 bytes for deep_select. Only columns before
4
        valid_blocks[row] are guaranteed initialized; selection must use
5
        end=valid_blocks.
6
    valid_blocks: [num_rows], int32 number of valid blocks per row.

返回的是完整存储 score_storage（行宽对齐），而不是裁到 num_blocks 的切片。原因在注释里：1024 字节对齐是 deep_select 的要求。bf16 元素是 2 字节，1024 字节就是 512 个元素。如果一个行的块数不是 512 的倍数，就把行宽向上补齐——DeepSelect 的行级 stride 假设由此满足。这也是为什么 _topk_candidate_block 里可以直接把 block_scores 传给 topk 而不必再复制一份。

调用方 MQA._topk_candidate_block 把这几件事串起来：

1
    def _topk_candidate_block(self, p_sum, k_len, offset, cross_layer_runtime):
2
        topk_blocks = self.candidate_topk_cfg.topk_block
3
        block_size = self.candidate_topk_cfg.block_size
4
        output = cross_layer_runtime.block_candidate_idxs[offset : offset + p_sum.size(0)]
5
        block_scores, valid_blocks = candidate_block_max(p_sum, k_len, block_size)
6
        topk(
7
            block_scores,
8
            topk_blocks,
9
            sorted=False,
10
            end=valid_blocks,
11
            indices_type=torch.int32,
12
            sorted_index=True,
13
            output_idx=output,
14
            idx_oob_fill_value=-1,
15
            return_value=False,
16
        )

逐参数解释：

• p_sum 是这一层的全部 query 行（预填充时一个 chunk 的全部 token，解码时整个 batch）。offset 是预填充切片的分段起点，k_len 是对应行的有效键长。
• topk_blocks 是候选块数量（2048），sorted=False 表示不要求按值排序——只要按块号排序即可，sorted_index=True 负责后半句。
• end=valid_blocks 是关键：告诉 DeepSelect「每行只有前 valid_blocks[row] 个块有效」，避免脏数据参与排序。
• output_idx=output：直接写进跨层运行时状态里预分配的 block_candidate_idxs，省去一次拷贝。
• idx_oob_fill_value=-1：不足 topk 个候选块时补 -1，后续内核一律把负数当作「无效块」。
• return_value=False：不需要分数，只取下标。DeepSelect README 明确建议关掉值输出以提速。

深挖：块分数内核的调度与对齐

candidate_blocks.py 的顶层只是一个调度层，真正的 TileLang 实现位于 tilelang_ops/cuda/candidate_blocks.py。这个内核的设计有四处值得单独讲。

1
    @T.prim_func
2
    def _candidate_block_max(logits, valid_lens, block_scores, valid_blocks,
3
                             task_ids, num_tasks) -> None:
4
        # 第一段：prep，按组估计每个请求的任务数，生成任务编号
5
        with T.Kernel(T.ceildiv(num_request_groups, prep_threads), threads=prep_threads) as prep_id:
6
            ...
7
        # 第二段：计算，按任务编号把 logits 分块取入，做 reduce_max
8
        with T.Kernel(num_ctas, threads=threads) as cta_id:
9
            for iteration in T.serial(cta_id, total_tasks, num_ctas):
10
                task = task_ids[iteration]
11
                ...

1
                for local_request, local_block in T.Parallel(requests_per_cta, blocks_per_request):
2
                    row_id = request_group_id * requests_per_cta + local_request
3
                    block_id = block_group_id * blocks_per_request + local_block
4
                    if block_id == valid_blocks_frag[local_request] - 1:
5
                        scores_frag[local_request, local_block] = T.infinity(T.bfloat16)

1
    last_blocks = (valid_blocks - 1).clamp_min(0).unsqueeze(1).to(torch.int64)
2
    last_scores = block_scores.gather(1, last_blocks)
3
    last_scores = torch.where(valid_blocks.unsqueeze(1) > 0, torch.full_like(last_scores, torch.inf), last_scores)
4
    block_scores.scatter_(1, last_blocks, last_scores)

1
    assert torch.equal(valid_blocks, reference_valid_blocks)
2
    num_blocks = reference_scores.shape[1]
3
    valid_mask = torch.arange(num_blocks, device="cuda").unsqueeze(0) < reference_valid_blocks.unsqueeze(1)
4
    assert torch.equal(scores[:, :num_blocks][valid_mask].view(torch.int16), reference_scores[valid_mask].view(torch.int16))

选择：DeepSelect 的 topk 怎么接进来

候选块的「选出哪几块」和「每行选哪几个 token」都调用同一个 topk：DeepSelect。它是 DeepSeek 为 DSA 和采样器定制的高性能 TopK 算子，于 2026 年 9 月开源（v1.0.0），官方说明其对 torch.topk 有 2–20 倍加速。索引器场景的参数（bf16 输入、topk=512、batch 与长度不受限）与 V4.1 完全吻合，所以 DeepInfer 直接用它。

这里不重复 DeepSelect 的算法推导，但要把它的三个接口约定与候选块管线的耦合讲清楚，因为踩坑点都在这里：

约定一：end 必须传，表示每行的有效范围。 上游 psum 和块分数都不保证越界位置是干净的。_topk_candidate_block 传的是 valid_blocks，_topk_prefill / _topk_decode 传的是 topk_lens（候选时是 candidate_mqa_lens，稠密时是 klen）。代码里还有一行注释：

1
            topk_index = topk(
2
                p_sum,
3
                self.topk_tokens,
4
                sorted=False,
5
                end=topk_lens,  # We must pass `end` to the topk kernel since OOB p_sum is not cleaned.

约定二：output_idx 与 idx_oob_fill_value。 预填充里 topk_index 是一个预分配的张量，output_idx=topk_index[chunk] 让内核直接写入；idx_oob_fill_value=-1 把「有效候选不足 topk 个」的位置填成 -1。主注意力的稀疏路径把 -1 当成「跳过」，所以这个约定必须贯穿到最后的 convert_kv_index（它也会把超出 context_lens 的位置写成 -1）。

约定三：output_idx_offset 只用于稠密路径。 稠密时序列内下标要加 ks 变成全局下标，topk 内核可以一次性加上：

1
                    output_idx=topk_index[chunk],
2
                    output_idx_offset=ks[chunk] if candidate_blocks is None else None,

候选路径不能在 topk 里加偏移，因为此时下标还是「紧凑下标」，加上 ks 没有意义。正确顺序是：先在紧凑坐标里 topk，再用 restore_candidate_indices_inplace 还原成序列内下标，并在同一内核里加 ks（通过 output_idx_offset 参数）。这条注释 if candidate_blocks is None 是整条链上最容易读漏的分支。

解码路径不在这里加偏移：解码时下标最终要换成物理槽位，偏移在 convert_kv_index 里随块表一起完成（见「一次前向的时序」一节的解码尾部）。ks 只属于预填充的「拼接 KV」语境，解码的 KV 是分页的，没有 ks 这一说。

（图源：DeepSelect 官方仓库 assets/perf_bf16_cuda.png，bf16 索引器场景、topk=512、CUDA 平台，纵轴为相对 torch.topk 的加速比。）

这张图按 batch size 分子图，展示索引器场景下相对 torch.topk 的加速比。它对应的正是候选块管线里的两种形状：预填充是「行多、每行一个 query chunk」，解码是「行多、每行只有一两个 query」。DeepSelect 的文档还特别提到一种退化情形——batch 很小时最后一个 wave 里活跃 SM 很少（wave quantization），它会用线程块簇（Threadblock Cluster）把一条序列拆到多个 SM 上，各 CTA 先算局部 top-k，再由 CTA0 归并。候选块的源层在长上下文、小 batch 时正好落进这一类。

内核二：get_candidate_compact_lens

候选块选出来以后，消费层的稀疏索引器不是「按块」算 logits，而是把候选块拼成一个连续的行来算。原因在 DeepGEMM 的 fp8_fp4_sparse_mqa_logits 接口：它接收 candidate_blocks 的形状 [num_rows, topk_blocks] 和 candidate_block_size，为每行输出宽度为 topk_blocks * block_size 的 logits——也就是紧凑行。块与块之间的空洞（没被选中的块）不占行宽。

紧凑行带来一个新问题：每行的有效宽度不是固定的 16384，而是「候选块里实际有效的 token 数」。最后一块可能只填了一半；候选块号可能指向序列末尾之外；topk_blocks 里可能有 -1 占位。topk 内核需要一个每行的 end，这个 end 就是紧凑有效长度。

get_candidate_compact_lens 用 32 行代码解决这件事：

1
def get_candidate_compact_lens(candidate_blocks, valid_lens, block_size):
2
    """Return compact token counts for the selected candidate blocks.
3

4
    candidate_blocks: [num_rows, topk_blocks], int32 sequence-local block
5
        indices in arbitrary order; negative values are ignored.
6
    valid_lens: [num_rows], int32 sequence lengths.
7
    block_size: Number of key tokens in one candidate block, either 8 or 16.
8
    """

参考实现就是一行向量化公式：

1
    block_starts = candidate_blocks * block_size
2
    tokens_per_block = (valid_lens.unsqueeze(1) - block_starts).clamp_(min=0, max=block_size)
3
    tokens_per_block.masked_fill_(candidate_blocks < 0, 0)
4
    return tokens_per_block.sum(dim=1, dtype=torch.int32)

把它拆开：

其中 $K$ 是 topk_blocks，$L_r$ 是第 $r$ 行的有效压缩长度（压缩 token 数），$b_{r,j}$ 是该行第 $j$ 个候选块的块号。这个式子的含义非常直观：第 $j$ 个候选块从位置 $b_{r,j}B$ 开始，块内有效 token 数是「从块起点到序列末尾还剩多少」，再夹到 $[0, B]$。

三个边界情况都在这个式子里自然处理：

• 最后一块不满：$L_r - bB < B$，取剩余数。
• 候选块完全越界：$L_r - bB \le 0$，取 0。
• -1 占位：masked 掉，贡献 0。

内核实现与参考实现的差别只在调度：

1
    threads = 32
2
    B = 128
3
    with T.Kernel(num_rows, threads=threads) as row_id:
4
        for chunk_id in T.serial(T.ceildiv(topk_blocks, B)):
5
            T.copy(candidate_blocks[row_id, chunk_id * B : (chunk_id + 1) * B], candidate_frag)
6
            for i in T.Parallel(B):
7
                candidate_slot = chunk_id * B + i
8
                if candidate_slot < topk_blocks:
9
                    block_id = candidate_frag[i]
10
                    if block_id >= 0:
11
                        remaining = valid_len - block_id * block_size
12
                        token_sum[0] += T.clamp(remaining, 0, block_size)

一行一个 CTA、32 线程（一个 warp）。这是典型的「行短、行多」内核：2048 个块按 128 一块分 16 次循环，每次一个 warp 处理 128 个元素，归约用 reducer 做。为什么不用多行一个 CTA？因为每行只需要一个 warp 的算力，多行会浪费；为什么不用 128 线程？因为 128 个元素正好一个 warp 一条指令处理完，多了反而要跨 warp 归约。topk_blocks 最多 4096（构造期断言），所以循环最多 32 次，开销可以忽略。

1
        if cross_layer_runtime.candidate_mqa_lens is None:
2
            cross_layer_runtime.candidate_mqa_lens = get_candidate_compact_lens(
3
                candidate_blocks, valid_lens, candidate_block_size)

1
            cross_layer_runtime.candidate_mqa_schedule_metadata.clear()
2
            cross_layer_runtime.candidate_mqa_lens = None

内核三：restore_candidate_indices_inplace

topk 在紧凑坐标里跑完，输出的是 [num_rows, topk_tokens] 的紧凑下标。但下游的 FlashMLA 稀疏注意力要的是序列内的压缩 token 下标（预填充还要加 ks 变成全局下标），解码还要再转成物理槽位。所以必须把紧凑下标还原。

还原公式就是前面下标一节的换算，套在每个位置上：

其中 $c_i$ 是该行第 $i$ 个紧凑下标，$b_{r,j}$ 是第 $j$ 个候选块号，$o_r$ 是可选的行偏移（预填充的 ks）。参考实现：

1
def restore_candidate_indices_inplace_ref(compact_indices, candidate_blocks, block_size, output_idx_offset=None):
2
    compact_capacity = candidate_blocks.shape[1] * block_size
3
    valid = (compact_indices >= 0) & (compact_indices < compact_capacity)
4
    safe_indices = compact_indices.clamp(0, compact_capacity - 1)
5
    candidate_slots = torch.div(safe_indices, block_size, rounding_mode="floor")
6
    block_offsets = safe_indices.remainder(block_size)
7
    block_ids = candidate_blocks.gather(1, candidate_slots.to(torch.int64))
8
    restored = block_ids * block_size + block_offsets
9
    if output_idx_offset is not None:
10
        restored += output_idx_offset.unsqueeze(1)
11
    compact_indices.copy_(torch.where(valid & (block_ids >= 0), restored, -1))

几个细节需要展开：

「in-place」不是优化，是语义。 函数名为 inplace，返回值是 None。调用处直接：

1
                if candidate_blocks is not None:
2
                    restore_candidate_indices_inplace(topk_index[chunk], candidate_blocks[chunk],
3
                                                      candidate_block_size, ks[chunk])

因为紧凑下标和序列内下标都是 int32、宽度相同，原地改写省去了一次额外分配。但原地改写的代价是输入必须只读一次——内核里每行先把 candidate_blocks[row] 整行搬进共享内存，避免在改写 compact_indices 时和块号读取冲突。

先 clamp 再取块。 参考实现在取 candidate_slots 之前先把下标夹到 [0, capacity)，是为了让无效值（-1 或越界）的 candidate_slots 和 block_offsets 仍是合法下标、不触发 gather 越界；有效性由 valid 掩码单独判断。内核里则直接用条件判断：

1
                    compact_idx = compact_frag[i]
2
                    restored_indices[i] = -1
3
                    if compact_idx >= 0 and compact_idx < topk_blocks * block_size:
4
                        candidate_slot = compact_idx // block_size
5
                        block_id = candidate_shared[candidate_slot]
6
                        if block_id >= 0:
7
                            restored_indices[i] = block_id * block_size + compact_idx % block_size
8
                            if has_offset:
9
                                restored_indices[i] += offset

block_id >= 0 这一层判断不可省。 紧凑下标可能合法（落在 [0, capacity)），但它指向的第 j 个候选块是 -1（topk 的 idx_oob_fill_value）。此时还原结果必须写 -1，而不是 -1 * B + offset。参考实现里 valid & (block_ids >= 0) 两个条件缺一不可。

output_idx_offset 与 topk 里的偏移不是一回事。 topk 的 output_idx_offset 是在稠密路径下加的；这里的偏移是在候选路径下加的，而且要加在还原之后。两者互斥：调用方根据 candidate_blocks is None 二选一（前一节已经写过）。

1
    threads = 128
2
    B = 512
3
    with T.Kernel(T.ceildiv(topk_tokens, B), num_rows, threads=threads) as (chunk_id, row_id):

1
                if self.is_candidate_topk_source:
2
                    self._topk_candidate_block(p_sum, klen[chunk], start_q, cross_layer_runtime)

调度元数据：候选块如何进入 DeepGEMM

前面说「消费层只扫候选块」，这件事真的落在 DeepGEMM 的稀疏 MQA 接口上：fp8_fp4_sparse_mqa_logits（预填充）和 fp8_fp4_paged_sparse_mqa_logits（解码）。这两个接口不负责「找候选块」，它们只接收 candidate_blocks 和 candidate_block_size，按候选集合算 logits。真正麻烦的是传进去的调度元数据（schedule_metadata）：DeepGEMM 要据此把 query 分配给 SM、把候选块映射到 KV 页、决定每个 wgmma 消费哪段数据。

1
    def _get_sparse_prefill_metadata(self, q_dtype, num_kv_tokens, input_metadata, chunk,
2
                                     candidate_blocks, candidate_block_size, cross_layer_runtime):
3
        schedule_key = ("prefill", self.ratio_key, chunk.start, chunk.stop,
4
                        candidate_block_size, candidate_blocks.shape[1], q_dtype)
5
        sparse_cache = cross_layer_runtime.candidate_mqa_schedule_metadata
6
        if schedule_key not in sparse_cache:
7
            ks = input_metadata.ks[self.ratio_key][chunk]
8
            ke = input_metadata.ke[self.ratio_key][chunk]
9
            candidate_base = torch.div(ks, candidate_block_size, rounding_mode="floor")
10
            absolute_candidate_blocks = candidate_blocks[chunk] + candidate_base.unsqueeze(1)
11
            sparse_cache[schedule_key] = deep_gemm.get_sparse_mqa_logits_metadata(
12
                ks, ke, num_kv_tokens, absolute_candidate_blocks, q_dtype,
13
                candidate_block_size,
14
                **({} if PLATFORM.is_npu else {"use_unaligned_ks": True}),
15
            )
16
        return sparse_cache[schedule_key]

1
            # Unaligned ks: block i is anchored at floor(ks / block_size) + i, and the
2
            # kernel resolves the residual ks % block_size against the raw (unpadded) KV.

1
        assert compressed_block_size % candidate_block_size == 0, "candidate blocks must not cross KV cache pages"

1
        schedule_key = ("decode", self.ratio_key, compressed_block_size,
2
                        candidate_block_size, candidate_blocks.shape[1], q_dtype)

1
        if candidate_blocks is None:
2
            dense_cache = input_metadata.mlp_relu_schedule_metadata
3
            if self.ratio_key not in dense_cache:
4
                dense_cache[self.ratio_key] = deep_gemm.get_paged_mqa_logits_metadata(...)

1
        if self.candidate_topk_cfg is not None:
2
            candidate_width = self.candidate_topk_cfg.topk_block
3
            assert candidate_width % 8 == 0
4
            cross_layer_runtime.candidate_mqa_schedule_metadata.clear()
5
            cross_layer_runtime.candidate_mqa_lens = None
6
            cross_layer_runtime.block_candidate_idxs = torch.full(
7
                (q_fp8or4.size(0), candidate_width), -1, device=q_fp8or4.device, dtype=torch.int32
8
            )
9
            cross_layer_runtime.candidate_topk_block = candidate_width
10
            cross_layer_runtime.candidate_block_size_before_compress = (
11
                self.candidate_topk_cfg.block_size * self.kv_compress_ratio)

1
            if local_index:
2
                blocks = shared.index_candidate_blocks[layer]
3
                if blocks > 0:
4
                    block_size = shared.index_candidate_block_size[layer]
5
                    keys = np.minimum((keys + block_size - 1) // block_size, blocks) * block_size
6
                attn += (1 if m.indexer_use_fp4 else 2) * idx_h * idx_d * int(keys.sum())

拓扑与配置：哪些层能当候选源

候选块机制的正确性有一半依赖拓扑校验。SharedLayerTopology._validate_candidate_config 逐条检查：

1
        previous_source = 0
2
        for layer_id in sorted(self.layer_candidate_source_map):
3
            source = self.layer_candidate_source_map[layer_id]
4
            assert source >= previous_source
5
            previous_source = source

候选源按层序单调不减。 候选块是一份跨层共享的运行时状态（block_candidate_idxs），每个源层写出时都会覆盖同一份缓冲。如果层序上出现「源回退」——后面的层指向更早的源——那么在一次前向里就会发生「源 A 写 → A 的消费层读 → 源 B 写 → A 的消费层再读」，状态被中途覆盖。单调不减保证候选组按层序切成连续段：每个源只服务它之后、下一个源之前的那一段层，源先写、消费层后读，不会交错。

1
        for source, candidate_cfg in self.candidate_source_cfg.items():
2
            assert self.layer_index_module_map[source] == source
3
            assert self.layer_candidate_source_map[source] == source
4
            assert layer_attn_cfg[source].get("index_topk", 0) > 0

候选源必须是索引模块源，且必须是它自己的候选源。 前者保证「产出候选块的那一层真的在算索引」，后者保证「候选源自己也按块分数路径走」，而不是被归到别的源下面当消费层。这两条把候选块机制钉死在索引器之上，不会出现「一个不持有索引器的层充当候选源」这种矛盾配置。

1
            assert topk_block > 0 and topk_block % 8 == 0
2
            assert 0 < block_size <= 16 and score_reduce in ("max",) and source_compress_ratio > 0

参数约束。 topk_block % 8 == 0 与两个内核共同依赖的行/列对齐契约一致（前面讲过的行 stride 与 restore 的 8 元素 stride）；score_reduce 目前只接受 max——这是为将来留的口子（TODO 注释也说了「further tighten」）；block_size 在 (0, 16]，内核里再收紧到 (8, 16) 两个值。

消费层的约束：

1
            assert compress_ratio == source_compress_ratio
2
            block_size_before_compress = self.candidate_source_cfg[source].candidate_block_size * source_compress_ratio
3
            consumer_block_size = block_size_before_compress // compress_ratio
4
            assert consumer_block_size <= 16
5
            candidate_capacity = self.candidate_source_cfg[source].candidate_topk_blocks * consumer_block_size
6
            assert candidate_capacity >= layer_cfg.index_topk

最后一条最值得看：候选容量必须不小于消费层的 topk。 候选集有 2048 × 8 = 16384 个位置，topk 只要 512 个，余量是 32 倍。如果某个层的 index_topk 配得比候选容量还大，topk 就必然填不满（只能返回 -1），稀疏注意力收到的 index 会有大量无效位——虽然不会算错，但稀疏度完全失效，效果退化到滑窗。断言把它变成启动期错误。

配置从哪来？hugging_face/config.py 的 _shared_layer_cfg 直接读模型的 config.json：

1
    if candidate_source is not None:
2
        shared.layer_candidate_source_map = {layer: candidate_source for layer in range(candidate_source, n_layer)}
3
        shared.candidate_source_cfg = {
4
            candidate_source: AttrDict(
5
                candidate_topk_blocks=text["candidate_topk_blocks"],
6
                candidate_block_size=text["candidate_block_size"],
7
                topk_block_score_reduce="max",
8
            )
9
        }

注意 layer_candidate_source_map 的构造是稠密字典：从 candidate_source_layer_id 到最后一层，每层都映射到同一个源。这与前面「候选源单调不减」的稀疏字典设计并置来看，说明归一化后消费层全部指向同一个源——V4.1 只有一个候选源（第 20 层）。多源情况下 SharedLayerTopology 也支持（分段），但当前模型没用到。

另外两个映射来自模型的 index_source_layer_ids（[2, 8, 14, 20, 24, 28, 32, 36]），通过 _source_map 把每层映射到「不晚于它的最近源层」。所以：

• 第 2–7 层共享第 2 层的索引器；
• 第 8–13 层共享第 8 层；
• ……
• 候选组内（20–39）的所有层在配置上都映射到第 20 层作为候选源，但真正持有索引器、会消费候选块的是 24、28、32、36 四个源层。

这形成了一个嵌套关系：候选组比索引模块段更粗。第 20 层的候选组覆盖 20–39 层，而这 20 层里包含 5 个索引模块段（20–23、24–27、28–31、32–35、36–39）。所以候选组可以跨多个索引模块段，但有三条约束必须同时满足：候选源自身是一个索引模块段的源；每个消费层的压缩率与候选源一致；候选源在层序上不晚于任何消费层。这三条恰好就是 _validate_candidate_config 的三组断言。

一次前向的时序

预填充里，一个候选组的一次前向是这样的（以第 20 层为候选源；配置上 21–39 层都映射到它，但真正执行索引器、消费候选块的是 24、28、32、36 四层）：

1
layer 20（候选源 + 索引模块源，use_candidate_topk = False）
2
  ├─ 量化 Q；写索引 K（必要时）
3
  ├─ 清空候选状态 / 预分配 block_candidate_idxs；清 candidate_* 字段
4
  ├─ DeepGEMM 稠密 MQA logits（全长）→ p_sum
5
  ├─ candidate_block_max(p_sum) → block_scores                         [内核一]
6
  ├─ DeepSelect topk(block_scores) → block_candidate_idxs（跨层共享）
7
  ├─ 记录 candidate_topk_block / candidate_block_size_before_compress
8
  ├─ DeepSelect topk(p_sum, end=klen, +ks) → 序列全局下标（稠密偏移）
9
  └─ >> 把 topk 下标交给主注意力（index 参数）
10

11
layer 21–39（同组消费层，use_candidate_topk = True）
12
  ├─ 读 block_candidate_idxs / candidate_topk_block / size_before_compress
13
  ├─ get_candidate_compact_lens（首个消费层算一次，其余复用）             [内核二]
14
  ├─ get_sparse_mqa_logits_metadata（按 chunk 惰性缓存）
15
  ├─ DeepGEMM 稀疏 MQA logits（仅候选块）→ p_sum'
16
  ├─ DeepSelect topk(p_sum', end=compact_len) → 紧凑下标
17
  ├─ restore_candidate_indices_inplace(+ks) → 序列全局下标               [内核三]
18
  └─ >> 把 topk 下标交给主注意力（index 参数）

两条路径的差别集中在三件事上：logits 的扫描域（全长 vs 候选集）、topk 的坐标（序列内 vs 紧凑）、以及候选的产出位置（源层产出、消费层复用）。is_candidate_topk_source 与 use_candidate_topk 在构造期互斥：

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

这条互斥保证了两条路径不会在同一层上叠加：源层按稠密口径算，消费层按候选口径算，中间只有 block_candidate_idxs 一个传递物。

解码路径与预填充的差别只在「q 很小、没有 chunk 切分」；源层仍按稠密口径算，消费层仍按候选口径算。候选块只算一次、整步复用：解码时每步的 query 很少（普通解码 1 个，投机验证时最多 NEXTN + 1 个），候选块由源层在当步算出，消费层在同一步内直接读 block_candidate_idxs。解码的调度元数据按 (ratio, compressed_block_size, block_size, blocks, q_dtype) 缓存，不需要按 chunk 切分。

解码尾部还有一个转换：

1
        convert_kv_index(
2
            topk_index.unsqueeze(1),
3
            compressed_block_size,
4
            topk_lens.unsqueeze(1),
5
            input_metadata.block_tables,
6
        )

convert_kv_index 把序列内下标换成物理槽位，并把 i >= context_lens 的位置写 -1。参考实现：

1
def convert_kv_index_ref(topk_idxs, block_size, context_lens, block_tables):
2
    logical_block = torch.div(topk_idxs, block_size, rounding_mode="floor")
3
    physical_block = block_tables.gather(1, logical_block.reshape(b, s * k).long()).reshape_as(topk_idxs)
4
    physical_idx = physical_block * block_size + topk_idxs % block_size
5
    valid = torch.arange(k, device=topk_idxs.device).view(1, 1, k) < context_lens.unsqueeze(-1)
6
    return torch.where(valid, physical_idx, -1).to(torch.int32)

这里 block_size 是分页 KV 的页大小，不是候选块大小。两个「块」概念在这里交汇：候选块（索引器的打分单位）与页块（KV cache 的物理单位）。前者的 8 个压缩 token 落在后者的同一页内——正是那条 compressed_block_size % candidate_block_size == 0 断言的物理含义。

工程边界与陷阱

整理这一篇里值得写进「踩坑清单」的点：

1. -1 是整条管线共用的「无效」编码。 topk 的 idx_oob_fill_value、候选块的无效占位、convert_kv_index 的越界位置，都用 -1 表示「这个位置不存在」。好处是所有内核不必区分空值来源；坏处是任何把 -1 当普通下标用的代码都会读到 [-1]（Python/Torch 里是最后一个元素），在内核里会越界。restore 和 compact_lens 都对 -1 单独判断，改这两段时必须保持。

2. end 不能省。 DeepSelect 的 topk 排序会读到有效范围外的脏数据；候选路径的 end 是 candidate_mqa_lens（按候选块对齐的紧凑长度），稠密路径是 klen。两者不能互换：紧凑长度只对候选集合有意义。

3. 块分数内核的 +inf 在 bf16 里是 0x7F80。 如果后面有代码对块分数做 exp 或除以它，会得到 NaN。块分数只进 topk，不进任何算术，这是刻意的：如果想让「最后一块」不占强权重，应该在 reduce 之前处理，而不是在分数上加常数。

4. 共享内存量：候选块整行 8 KB。 restore 内核每行把 [topk_blocks] 整行搬进共享内存。topk_blocks = 4096 时是 16 KB，还在共享内存的合理范围内；但如果将来把候选块数调到 16384，共享内存就会接近极限，需要改成「按 x 轴 CTA 共享、或者从全局内存直接读」。

5. 源层清空与消费层复用的顺序契约。 源层在 forward 开头清空 candidate_mqa_schedule_metadata 与 candidate_mqa_lens，这两个字段记录的是「当前候选组在本步算出的东西」；消费层靠惰性填充复用它们。前提是「源层先执行、消费层后执行」，这由拓扑的单调源保证。如果顺序被破坏（例如手工改配置让某个消费层排在源层之前），消费层读到的 block_candidate_idxs 还是 None，会在取 .shape 时直接抛异常——这里至少是快速失败，而不是静默算错。

6. candidate_block_size_before_compress 的命名陷阱。 它看起来像「压缩前的块大小」，但实际存的是 block_size × compress_ratio，也就是块覆盖的原始 token 数。消费层再除以自己的压缩率得到压缩条目数。读这段代码时如果按字面理解成「V4.1 的 8 乘以……」，会算错一倍。

7. 拓扑断言和缓存引擎是两套独立判定。 第十一篇已经讲过：缓存引擎判断「这一层是不是共享」看的是配置里有没有 layer_info_shared_cfg；拓扑校验只是静态读配置，两者不一致时，拓扑不会报错，缓存引擎会静默读到一个从未写入的缓存。候选块同理：消费层拿到的 block_candidate_idxs 是运行时状态，与静态拓扑无关，所以拓扑必须保证「源层先于消费层执行」。

8. use_unaligned_ks 只在 GPU 上开。 昇腾分支不传这个参数；跨平台调试时不要假设两边的调度元数据构造完全一致。

小结与下一篇

回到开头的问题：候选块管线在省什么、怎么省、代价是什么。

省什么。 索引器每层都要扫全长，这是稀疏注意力里唯一没有被 token 选择「摊薄」的部分。候选块把「全部层扫全长」改成「一个源层扫全长（还要加块分数），其余层扫定长候选集」。在 V4.1 的配置下，候选组覆盖第 20–39 层，其中真正持有索引器的层是 20、24、28、32、36 五个；源层（第 20 层）按全长付出代价，其余四个索引器层作为消费层，扫描域被压到 16384 个压缩 token。就这个候选组而言，全长扫描 从 5 次降到 1 次，另加一次块分数归约和四次定宽扫描；第 2、8、14 层不受候选块影响，仍各自扫全长。

怎么省。 三个内核串成一条链：candidate_block_max 把 logits 归约成块分数（max-reduce + 最后块 $+\infty$），DeepSelect 在块分数上选候选块，get_candidate_compact_lens 算出紧凑行宽，restore_candidate_indices_inplace 把紧凑下标还原。调度元数据按 chunk 和候选组缓存，源层清、消费层填。

代价。 三处：一是候选块是源层「一次性」决定的，后续层没有机会补充——如果源层的块分数漏掉了某个真正的高分行，后面所有层都会漏掉它；二是拓扑约束变得复杂（单调源、压缩率一致、候选容量 ≥ topk），配错时只在断言处暴露；三是三段下标换算增加了实现复杂度，而它们中最容易出错的边界（end、-1、unaligned ks）都是运行时才显现的。

到此，注意力主线里的「索引器 → 候选块 → topk → 稀疏注意力」已经闭环。下一篇按路线图继续深入模型层的另一块：mlp.py 的 616 行——pre_rope_kv 到底在哪个函数里产出、为什么共享 index-K 的消费层连 KProject 都不建，细粒度专家的分组 GEMM（_grouped_gemm 如何落到 DeepGEMM 的连续布局接口）、共享专家为什么在每一层都用 replicated 线性层而不是切分，以及 swiglu_limit（V4.1 配置里是 10.0，源码里以 swiglu_clamp_value 的名字出现在每一路 GEMM 调用里）如何夹住门控输出。

参考资料

1. DeepInfer 官方仓库（DeepSeek V4.1-Flash 推理框架）
2. candidate_blocks.py（调度层与参考实现）
3. tilelang_ops/cuda/candidate_blocks.py（三个 TileLang 内核的 CUDA 实现）
4. index_attn.py（MQA 索引器运行时：候选块、紧凑下标与调度元数据）
5. cross_layer.py（共享拓扑的候选源校验与 CrossLayerRuntime）
6. hugging_face/config.py（从模型 config.json 生成共享映射与候选配置）
7. flash_mla_utils.py（convert_kv_index：序列内下标到物理槽位）
8. DeepSeek-V4.1-Flash-1017 模型配置（config.json：索引源层、候选块参数与压缩率）
9. DeepSelect：DSA 的 TopK 算子（算法推导与性能数据）
10. DeepSelect 算法与实现分析
11. DeepGEMM（稀疏 MQA logits 内核与调度元数据接口）
12. TileLang（本文三个内核的编程框架）
13. DeepSeek-V3.2 论文：DSA 的闪电索引器与细粒度 token 选择（arXiv:2512.02556）
14. DeepInfer 部署问答（官方 E/P/D 架构与稀疏注意力上线配置）
15. FlashMLA（稀疏 MLA 注意力内核）