DeepSelect

DeepSelect 是一个针对稀疏注意力与采样场景优化的高性能 TopK 算子 CUDA 实现，专为 DeepSeek V3.2/V4 系列模型中的 Sparse Attention（DSA）和 logits 采样设计。它支持 bfloat16（Lightning Indexer 场景）和 float32（Sampling 场景）输入，在小 TopK（≤4096）、大词汇表（最高 128K+）和可变批大小下，相比 PyTorch 原生 topk 实现获得 2–20 倍加速。核心特点包括内存带宽优化、可选跳过值输出、索引排序开关及严格对齐的内存布局要求。适用于大语言模型推理中的稀疏注意力计算与高效 token 采样等低延迟、高吞吐关键路径。