DeepSelect trades generality for honest 2-to-20x TopK speedups where it matters

TopK selection is one of those operations that looks trivial until it becomes a bottleneck. In sparse attention and token sampling, you are running it constantly - every forward pass, every token, across large batch dimensions - and torch.topk is not built for that pressure. DeepSelect is a custom CUDA kernel that replaces it, and the claimed speedup range of 2 to 20x over vanilla torch.topk is not a single cherry-picked configuration: it varies with batch size and vocabulary size, which is exactly the honest way to report it.

The implementation covers two distinct workloads. The Lightning Indexer scenario targets bfloat16 inputs with arbitrary batch and vocabulary sizes, used inside DeepSeek Sparse Attention as deployed in V3.2 and later models. The Sampling scenario targets float32 inputs with a vocabulary around 128K tokens - the standard token-sampling shape. Both cap topk at 4096, which is a real constraint worth noting: this is not a general-purpose replacement for torch.topk across arbitrary k values.

The performance framing is careful. Because TopK does no floating-point arithmetic, reporting FLOP rates would be meaningless; the benchmark instead measures effective memory bandwidth, which is the right metric for a memory-bound operation. The plots show bandwidth ratios against torch.topk on the same input, not absolute numbers in isolation.

Several design decisions reveal where the authors spent their optimization budget. Setting return_value=False when you only need indices is described as roughly 10% faster - a small but real gain for attention routing where values are discarded. Disabling sorted output is also recommended unless ordering is actually required, since sorting adds overhead. NaN checking is always active and defaults to a hard abort via trap(), which is a deliberate choice: silent NaN propagation in a sparse attention router would be worse than a crash.

The stride alignment requirement on input tensors is the sharpest practical friction point. The row stride must be a multiple of deep_select.get_stride_requirement()[0] bytes, and unaligned inputs require padding. That is not unusual for high-performance CUDA kernels, but it means you cannot drop this in as a zero-friction swap without auditing your tensor layout.

A companion deep-dive document was released alongside the code, covering the algorithm and its implementation in both English and Chinese. That is a meaningful addition - kernel-level optimizations at this specificity are often shipped without explanation, and the documentation suggests the authors expect people to actually read and adapt the code rather than treat it as a black box.

For anyone running inference on DeepSeek sparse models or building custom sampling loops at scale, the bandwidth gains here are real and the scope is honest.