開発
DeepSelect：DeepSeek Sparse Attentionとサンプリング向け高性能TopKカーネル
原題: DeepSelect: High-Performance TopK Kernels for DeepSeek Sparse Attention and Sampling
DeepSeekモデルのTopK処理を高速化
DeepSeekは2026年9月10日、DeepSelect v1.0.0と実装解説を公式GitHubで公開しました。
DeepSelectはDeepSeek Sparse Attentionとサンプラーで使うTopK処理のCUDA実装です。
対象となるSparse AttentionはDeepSeek V3.2、V4、V4.1で採用されています。
DeepSeekは標準のtorch.topkと比べて2〜20倍の高速化を確認したとしています。
この数値はDeepSeek自身の性能測定で、独立評価ではありません。
Lightning Indexer向けbfloat16と、サンプリング向けfloat32の処理を対象にしています。
topkは4096以下という制限があり、入力の行ストライドにもアラインメント条件があります。
Pythonから値とインデックスを取得でき、可変長行や出力バッファの指定にも対応します。
リポジトリを取得し、pipでソースからインストールします。
コードはMITライセンスで公開されています。