deepseek-ai/DeepSelect

DeepSelect is a high performance implementation of the TopK kernel used in DeepSeek Sparse Attention (DSA) (which is used in DeepSeek V3.2, DeepSeek V4, and DeepSeek V4.1 models) and the sampler. It supports NVIDIA CUDA and Huawei Ascend platforms, and achieves 2 20x speedup compared to vanilla torch.topk. DeepSelect 是为 DeepSeek 稀疏注意力（DeepSeek Sparse Attention, DSA）和采样器（Sampler）定制的高性能 TopK 算子，支持 NVIDIA CUDA 平台与华为昇腾（Ascend）平台。与原生 torch.topk 相比，它可实现 2～20 倍的加速。 TopK workloads vary widely, and the fastest algorithm & implementation highly depends on the input dtype, batchsize, vocabsize, and topk. This repository only focuses on the following cases: This scenario covers: Recommendations: This scenario covers: Currently, only the CUDA implementation supports torch.float32 as the input dtype. The Ascend implementation supports torch.bfloat16 only. Measured with the benchmark in tests/test.py (python3 tests/test.py --perf-only), which reports the ratio against torch.topk on the same input. The metric is effective memory bandwidth: TopK does no floating-point math, so a FLOP rate…