DeepSeek开源DeepSelect：比torch.topk快二十倍！
内容提要
DeepSeek 开源高性能 TopK 算子库 DeepSelect，针对稀疏注意力中的闪电索引器和采样器两个场景优化，速度可达 torch.topk 的 2 至 20 倍。该库不绑定推理框架，可通过 pip 直接替换，结果稳定，并支持关闭排序和返回值以进一步提速，助力 V4.1 Flash 推理大幅加速。
延伸解读
torch.topk 的性能痛点与版本波动
文章指出 torch.topk 在 GPU 上效率极低，某些尺寸下有效内存带宽仅 12.7 GB/s，远低于显卡理论带宽。更严重的是，PyTorch 从 2.6 升级到 2.7 时，torch.topk 在 Nvidia L20 上延迟从 11.30 毫秒涨到 17.32 毫秒，在 H20 上提取十亿浮点数前一千个从 36 毫秒暴涨到 1.6 秒。这种性能倒退和波动让依赖它的推理管线面临风险，也解释了为何需要专用算子库。
DeepSelect 的差异化设计：专注场景与按需裁剪
DeepSelect 不追求通用，只服务闪电索引器（bfloat16，topk≤4096）和采样器（float32，词表约128K，topk≤4096）两个场景，针对输入约束重新设计内存搬运和异步流水线。它提供 sorted_index 和 return_value 参数，允许关闭排序或仅返回索引，后者可再快约10%。这种按需裁剪避免了 torch.topk 总是排序和返回值的多余开销，同时保持结果稳定，且不绑定任何推理框架。
独立替换能力与竞品对比
文章对比了 vLLM 的 flashinfer 多 CTA topk（约40微秒）和 SGLang 单 block 实现（65微秒以上），指出它们速度虽可但绑定各自框架且不保证结果稳定。Nvidia TensorRT-LLM 的 Guess-Verify-Refine Top-K 针对 Blackwell 优化，但必须替换整个推理栈。DeepSelect 通过 pip 安装即可替换 torch.topk，不绑定框架，结果稳定，速度是 torch.topk 的2到20倍，这种独立替换能力是其关键差异。
性能优势的边界与未知风险
文章提到 DeepSeek V4.1 Flash 输出速度达每秒280到500 token，远超上代 V4 Flash 的120到140 token，DeepSelect 在底层贡献了力量。但作者也指出，目前公开测试都是固定输入尺寸和显卡型号下的单次基准，在大规模并发、batch size 动态变化、显存碎片化等真实场景中，DeepSelect 能否保持2到20倍优势及结果稳定性仍是黑箱，技术报告只给出了理想条件下的数字。
Q&A
DeepSelect是什么？它主要用来做什么？
DeepSelect是DeepSeek开源的一个高性能TopK算子库，专门用于加速稀疏注意力机制（DSA）和采样器中的关键计算步骤，即从大量token中快速筛选出最相关的TopK个token。
DeepSelect相比torch.topk有哪些优势？
DeepSelect相比torch.topk有2到20倍的速度提升，结果稳定，不绑定任何推理框架，可通过pip直接替换，并且支持关闭排序和返回值以进一步提速。
DeepSelect针对哪两个场景进行了优化？
DeepSelect针对闪电索引器场景和采样器场景进行了优化。闪电索引器场景用于稀疏注意力模块，输入数据类型为torch.bfloat16，topk值≤4096；采样器场景用于token采样，输入数据类型为torch.float32，词表大小约128K，topk值≤4096。
为什么torch.topk在GPU上速度慢？
torch.topk速度慢主要是因为内存搬运效率低和进行了多余的排序操作。例如，在处理某些尺寸数据时，有效内存带宽只有12.7 GB/s，远低于显卡理论带宽1800 GB/s。此外，torch.topk总是返回排序后的结果和数值，即使下游不需要。
DeepSelect如何实现比torch.topk更快的速度？
DeepSelect通过针对特定场景重新设计内存搬运方式，利用数据分块和异步流水线重叠计算与读写，并且允许关闭排序（sorted_index=False）和返回值（return_value=False）来减少不必要的操作，从而大幅提升速度。
DeepSelect的安装和使用方式是什么？
DeepSelect可以通过git clone下载，更新子模块后使用pip install -v .安装。使用时，只需将torch.topk替换为deep_select.topk，API类似：deep_select.topk(x, topk, sorted_index=True, indices_type=torch.int32, return_value=True)。
DeepSelect的结果稳定性为什么重要？
在模型推理中，结果不稳定意味着同样的输入可能产生不同的TopK结果，导致同一问题得到不同回答，这在生产环境中是致命的。DeepSelect保证了结果稳定，而vLLM的flashinfer和SGLang的单block实现都不保证结果稳定。
DeepSelect与TensorRT-LLM的Top-K优化有何不同？
TensorRT-LLM的Guess-Verify-Refine Top-K针对Blackwell架构优化，但绑定在TensorRT-LLM推理管线中，使用它需要更换整个推理栈。DeepSelect不绑定任何推理框架，可以独立替换torch.topk，且结果稳定。
