DeepSelect：DeepSeek 稀疏注意力 TopK 高性能算子

从一次尴尬的模型推理卡顿说起

当你用 DeepSeek-V4 处理一篇万字长文时，有没有注意过一个"隐形瓶颈"——TopK 算子。

Sparse Attention（稀疏注意力）机制中，模型每生成一个 token，都需要从数万甚至数十万候选中选出 top-2048 个最相关的上下文 token。"选出最大的 K 个"这个操作听起来再简单不过，但当 vocab_size（词表大小）达到 128K、单次 batch 处理数十条序列时，torch.topk() 成了端到端延迟的主要贡献者。开发者们通常把注意力放在 Flash Attention、KV Cache 压缩上，而 TopK 优化往往被忽视——因为它"太简单了"。

DeepSelect 就是来解决这个"隐形瓶颈"的。

背景：DeepSeek 为什么需要自定义 TopK

DeepSeek 团队在 V3.2 及后续版本的稀疏注意力（DeepSeek Sparse Attention, DSA）架构中，引入了 Lightning Indexer 模块。这个模块的核心是一个 TopK 算子——从所有历史 token 中选出最相关的 2048 个，让模型只"关注"这些 token。

标准 torch.topk() 的问题在于：它为通用场景设计，没有针对以下特定约束优化：

• bfloat16 输入：推理阶段的标准精度
• 任意 batch_size 和 vocab_size：从单条短序列到大批量长文本
• topk ≤ 4096：DSA 的 topk 固定在 512-2048 范围
• 随机访问模式：每次只需要 top-k 个，不需要完整排序

通用算子对所有场景一视同仁，而 DeepSelect 只为上述精确场景优化，因此能榨出 2~20 倍的加速。

核心算法：随机块扫描 + Radix-Select

DeepSelect 的算法设计非常巧妙，核心思想是"早停"（early termination）——不需要遍历完整个输入序列才能确定 top-k。

算法维护一个动态阈值 T（初值为 -∞），然后：

1. 随机扫描块：按随机顺序处理大小为 B 的数据块，而不是顺序遍历
2. 阈值筛选：只将大于当前阈值 T 的元素加入候选缓冲区
3. 周期压缩：当候选缓冲区超过 k+B₂ 时，在 shared memory 中执行 Radix-Select 重新选出 top-k，并更新阈值 T
4. 随着 T 越来越大，后续元素通过筛选的概率越来越低，算法自然"早停"

数学上证明了每个元素平均只被读取约 4 次（最优情况），远低于朴素的多次扫描方案。由于随机块顺序与输入数据独立，算法在最坏情况下也有性能保证。

两大应用场景

Lightning Indexer 场景（bfloat16）

用于 DeepSeek 稀疏注意力中，为每个新 token 选择最相关的历史 token。输入为 bfloat16 精度，支持任意 batch_size 和 vocab_size，主要优化 topk=512（即 DeepSeek V4 的配置）。

建议关闭 sorted_index（除非必须按 index 排序），以及设置 return_value=False 当只需要索引不需要值——分别可获得额外性能收益。

Sampling 场景（float32）

用于采样阶段，从约 128K 的词表中选出候选 token。输入为 float32，vocab_size 约为 128K。Ascend 版本目前仅支持 bfloat16。

性能：2~20 倍加速不是虚标

DeepSelect 的性能数据由仓库内置 benchmark 直接测量 report，提供的是相对于 torch.topk() 的有效内存带宽比值，而非绝对数字——因为 TopK 不做浮点运算，FLOPs 不是合适指标。

在 Lightning Indexer 场景（bffloat16, topk=512）下，随着 vocab_size 增大，优势愈发明显：小 batch 时加速约 2-4 倍，大 batch 大词表时可超过 10 倍。在 Sampling 场景（float32, vocab_size=129280）下，加速比稳定在 4-20 倍区间。

技术栈与架构

DeepSelect 采用 Python + CUDA C++ 的混合架构：

• Python 封装层（deep_select/interface.py）：自动探测硬件平台（NVIDIA CUDA 或华为 Ascend），动态加载对应后端，暴露统一的 topk() 接口
• CUDA kernel（csrc/cuda_kernels/v3/）：预生成大量模板实例化代码，涵盖不同 batch_size、topk、thread 配置的组合，由 setup.py 调用 scripts/generate_instantiations.py 自动生成
• Ascend NPU kernel（csrc/ascend_kernels/）：华为昇腾平台专属实现

代码库包含完整的算法深度解析文档（docs/DeepSelect-deep-dive.zh.md），包含严格的数学证明和实现细节，在同类开源项目中极为罕见。

局限性：专注带来性能，放弃通用性

DeepSelect 有明确的适用边界：

• topk ≤ 4096：超过 4096 的场景不被支持，也不计划支持
• 仅支持两个 dtype：bfloat16（Lightning Indexer）和 float32（Sampling）
• Ascend float32 缺失：华为昇腾版本仅支持 bfloat16
• 无排序选项（Ascend）：Ascend 版本不支持 sorted=True

这些约束是刻意的——为了最大化性能，不兼容的输入类型和场景直接拒绝编译。

行业意义：推理优化进入"最后一公里"

DeepSelect 的出现代表了一个趋势：LLM 推理优化正在从宏观架构层（Attention 变体、KV Cache 压缩）下沉到微观算子层。当 Flash Attention、Speculative Decoding 等优化已经广泛落地，各家模型纷纷集成 cuDNN 的 DSA kernel 时，DeepSeek 选择从 TopK 这个"被忽视的角落"切入，恰恰填补了一个真实存在的性能缺口。

它也间接证明了中国 AI 基础设施团队不仅能跟进最新架构，还能针对特定场景做深度的、经过数学证明的定制优化，而非泛泛的"性能调优"。vLLM 已经宣布集成 DeepGEMM 和 Lightning Indexer CUDA kernel，期待未来更多推理引擎跟进 DeepSelect。