DeepSelect: DeepSeek's High-Performance TopK Kernel for DSA

On September 10, 2026, the DeepSeek organization published the repository deepseek-ai/DeepSelect on GitHub and, the same day, released both v1.0.0 and a bilingual deep-dive document on the algorithm. As of release day the repo had 152 stars, its primary language is CUDA, and both creation and first push happened on September 10. This is not another "we trained a model" story; it is something lower-level and easier to overlook: a high-performance TopK operator with its companion sampler.

The same-day echo matters. DeepSeek also open-sourced V4.1 Flash that day, and DSA is precisely the sparse-attention mechanism underneath it, while this library is the operator base DSA depends on. The "model open-sourcing" is only the part above water; DeepSelect is the operator implementation below it. We also published a V4.1 Flash hotspot explainer in this batch (read it here).

What DeepSelect Is: The TopK Operator Base Under DSA

DeepSelect is not a general-purpose TopK library; it is a high-performance implementation of the TopK kernel used in DeepSeek Sparse Attention (DSA) and the sampler. DSA runs in DeepSeek V3.2, V4, and V4.1. Inside DSA, TopK selects the most relevant positions from many context tokens; during sampling, TopK selects candidate tokens from a large vocabulary. DeepSelect serves these two "pick the largest k" needs far faster than native torch.topk, with an official 2 to 20 times speedup (README: "2 ~ 20x speedup").

Untangle the chain: the user calls a DeepSeek model; the model runs DSA; DSA repeatedly calls TopK; DeepSelect swaps that layer for a hand-written CUDA kernel. So its value first shows as cheaper, faster DeepSeek-series inference, not a casual TopK drop-in. Its first nature is to serve DeepSeek-series models, a point we revisit below.

Why TopK Becomes the Attention Bottleneck

Many engineers intuitively focus on the "matrix multiply" when thinking about attention, assuming the bottleneck is always GEMM. But in sparse attention, TopK can become a non-negligible slice of end-to-end latency. The reason is direct: DSA must pick the most relevant positions from a long list of context tokens, and that "picking" is itself a TopK; candidate-token filtering during sampling is another TopK over a large vocabulary (around 128K). Both can recur at every layer and every decode step.

The native torch.topk is a general implementation that leaves headroom for every dtype, shape, and topk size, so its constant factors cannot be pushed to the extreme. When the input is a row hundreds of thousands long and you need only a very small k (say 512), the general "sort the whole row then take top k" approach does much wasted work. DeepSelect's starting point is exactly this: since the input distributions, dtypes, and topk sizes of the two workloads are highly predictable, you can tailor the algorithm, replacing the uneconomic full sort with a "single pass plus threshold convergence" idea, and thus capture that 2-to-20x speedup headroom.

This also explains why the official docs stress one point: TopK workloads vary widely, and the fastest algorithm and implementation depend heavily on the input dtype, batch_size, vocab_size, and topk. DeepSelect explicitly focuses on only two scenarios rather than claiming it "covers everything", and this restraint is precisely what lets it be fast.

The RadixSelect Algorithm: One Pass, Random Blocks, Threshold Convergence

DeepSelect's algorithm, DeepSelectTopk, keeps a top-k threshold T initialized to -inf and loops through three steps:

1. Scan: Process the input one block of size B at a time, in a random block order.
2. Filter: Keep only elements greater than T, appending them to the candidate buffer topk_candidate.
3. Compact: When the buffer grows large (or at the end), run a radix-select-based TopK in shared memory to reduce it to k elements, and set T to the smallest among them.

The pseudocode (verbatim from the source) is:

DeepSelectTopk(x[0:N), k, B, B2):
    require 1 <= k <= N and B >= 1 and B2 >= 1

    topk_candidate = []                     # shared memory
    topk_threshold = -inf                   # current top-k threshold
    p = random_permutation(ceil_div(N, B))  # independent of x

    for block in p:
        lo, hi = block * B, min((block + 1) * B, N)

        # Keep only elements that may still enter the top-k.
        for j in range(lo, hi):
            if x[j] > topk_threshold:
                topk_candidate.append((x[j], j))

        # Periodically compact the candidate buffer.
        if len(topk_candidate) >= k + B2:
            topk_candidate = RadixSelectTopK(topk_candidate, k)
            topk_threshold = min(v for v, _ in topk_candidate)

    return RadixSelectTopK(topk_candidate, k)

Initially the threshold is -inf, so all elements are accepted before the first compaction; afterward topk_threshold is the current k-th largest value and never decreases, so later elements are less likely to pass and the buffer grows more slowly. A reasonable config for k = 512 is B = B2 = 1024.

A few engineering-critical properties: at any time len(topk_candidate) <= k + B2 + B; every element of x is read exactly once in contiguous blocks of size B (friendly to memory bandwidth); excluding the random permutation, the algorithm needs only O(k + B + B2) extra space, small enough for shared memory. To avoid performance degradation on unfavorable inputs, the block order must be random; the next section gives this a rigorous guarantee via an expected upper bound.

What the "Expected Upper Bound on Total Elements Processed" Really Means

With a fixed input, the only randomness is a uniform random block permutation independent of the input; we ask how many elements all the RadixSelectTopK calls process in total under the worst distribution.

Define m = ceil(N / B), L = k + B + B2, and H_m = sum(1/i) for i from 1 to m (the m-th harmonic number). The document proves: conditioned on the set U_i of the first i processed blocks, the expected number of elements the i-th block appends satisfies E[A_i | U_i] <= L / i; by linearity of expectation the total appended count E[A] <= L * H_m. Let R be the number of in-loop TopK calls; each removes at least B2 elements, so R * B2 <= A; each appended element contributes once to an input while each in-loop call retains k elements carried into a later call, hence the total processed count W = A + kR. The final result:

E[W] <= (1 + k / B2) * L * H_m
     = (1 + k / B2) * L * (ln m + O(1))

For an engineer this means: under a random block order, the expected amount of "redundant elements repeatedly processed" is controlled by a logarithmic term ln(N/B) in the input size, not by linear growth with N. Even with hundreds of thousands or millions of elements, only O((1 + k/B2)(k + B + B2) * log(N/B)) truly enters radix-select for repeated comparison.

The performance summary adds: when B, B2 = Theta(k), the TopK stage's extra computation is only O(k * log(N/k)), far smaller than N. This is the essential advantage of "one pass plus threshold convergence" over "full-row sort": you pay a logarithmic cost only for the small set that can enter the top-k. After careful hardware-specific optimization, a substantial speedup over torch.topk in DSA workloads is no surprise.

Why Measure by Effective Memory Bandwidth, Not FLOPs

The README states the methodology frankly: run the benchmark in tests/test.py (python3 tests/test.py --perf-only), which reports the speedup ratio against torch.topk on the same input; the core metric is "effective memory bandwidth". The reason: TopK does no floating-point math, so a FLOP rate is meaningless here. A TopK kernel should be judged by how much HBM bandwidth it uses (the official charts draw Lightning Indexer on a shared 0-to-7 TB/s axis), not by FLOPs. This also explains why DeepSelect's focus lands on fewer global-memory reads, contiguous blocks, the candidate set in shared memory, and constant factors squeezed with bit manipulation and PTX - all of that serves bandwidth.

Concrete benchmark settings: Lightning Indexer uses bfloat16, topk = 512, one subplot per batch size; Sampling uses float32, vocab_size = 129280, topk = 512.

Supported Scenarios and Boundaries: What It Does Not Do Matters More

For a very early operator library, the most valuable information is often not "what it can do" but its hard constraints. DeepSelect explicitly supports only two scenarios, each with rigid limits.

Other boundaries to keep in mind:

• Hard topk cap of 4096: requests larger than 4096 are out of scope. This is a deliberate trade-off, not an oversight.
• Strict dtype split: Lightning Indexer takes only bfloat16, Sampling takes only float32; do not mix them.
• Optimized for specific shapes only: the repo tunes only for the two workloads above and is not guaranteed fast for arbitrary shapes.
• Row-stride alignment requirement: the input tensor x must have its row stride aligned to deep_select.get_stride_requirement()[0] bytes, last dimension contiguous; unaligned inputs need padding. Both outputs are allocated by the call, with strides aligned to get_stride_requirement()[1] bytes (possibly non-contiguous); pass output_idx= to write into your own buffer, which must meet the same stride requirement.
• Variable-length rows: end sets a per-row exclusive upper bound; rows shorter than topk are padded with value_oob_fill_value / idx_oob_fill_value.
• NaN handling is always on: with the default abort_when_nan_found=True the kernel invokes trap() and aborts; rows whose length is <= topk are never NaN-checked.

Two README tuning tips save real performance: disable sorted_index unless the output must be ordered by index or by value (either ordering costs performance); and set return_value=False when values are not needed, skipping the value output and running faster (about 10%).

On long-context cost, our companion comparison in this batch (read it here) compares the bills of different long-context approaches; a low-level operator like DeepSelect is the key link turning "long-context sparsification" into deployable inference. The integration SOP (read it here) is worth reading alongside.

A Cold Look: 152 Stars and Its Real Value

Let us pour a little cold water to set correct expectations.

First, 152 stars means it is very early. That count on release day reflects the attention of "an official DeepSeek release" more than community-wide production validation of its stability. The citation lists authors Yi Qian, Shengyu Liu, and Yichen Li (2026); the project is very new, and its interface, boundaries, and performance may change in later versions.

Second, its direct usefulness for non-DeepSeek models is limited. DeepSelect is tailored for DSA (V3.2 / V4 / V4.1) and the corresponding sampling workloads; the input dtype, topk cap, and shape assumptions are all bound to DeepSeek's inference path. Dropping in an unrelated task with arbitrary dtype and shape will be neither in its optimization scope nor guaranteed supported. Treating it as a "general TopK acceleration library" misunderstands its positioning.

Third, its value leans more toward the ecosystem than the general. For DeepSeek-series users it opens the DSA operator base so you can audit, rewrite, and squeeze more performance for your own hardware; for the community it is a high-quality reference implementation of hand-writing a CUDA TopK for a specific workload, and the PTX and bit-manipulation tricks inside are teaching material. Its significance is not "make everyone's TopK faster" but "make DeepSeek's sparse attention run, and run cheaply", while handing operator optimizers a learnable template.

FAQ

Q1: How is DeepSelect different from the ordinary torch.topk? A1: Two differences. Positioning: torch.topk is general, while DeepSelect focuses only on the DSA Lightning Indexer and sampling workloads with its "one pass plus threshold convergence" RadixSelect algorithm. Performance: the official measurement is 2 to 20 times faster, using effective memory bandwidth as the metric rather than FLOPs. If your input is one of its two supported scenarios, the gain is clear; otherwise it is neither supported nor guaranteed faster.

Q2: Why must DeepSelect's topk be less than or equal to 4096? A2: A deliberate trade-off, not an omission. In the DSA Lightning Indexer and sampling scenarios the needed k (such as 512 in the benchmark) is far below 4096, and the candidate buffer, compaction trigger, and shared-memory layout are all designed around "small topk". Requests larger than 4096 are out of scope, with no correctness or performance guarantee.

Q3: Which input data types does DeepSelect support? A3: Only two classes, strictly split: Lightning Indexer takes torch.bfloat16, and Sampling takes torch.float32. The input tensor's row stride must be aligned to deep_select.get_stride_requirement()[0] bytes and its last dimension must be contiguous; unaligned inputs need padding. Both outputs are allocated by the call, with strides aligned to get_stride_requirement()[1] bytes.

Q4: Can I use DeepSelect directly in my own non-DeepSeek model? A4: You can call it, but first confirm your workload is one of its two supported scenarios (bfloat16 Lightning Indexer, or float32 Sampling with vocabulary around 128K), with topk <= 4096 and inputs meeting the stride and contiguity requirements. If your model is not DeepSeek-series and its distribution does not match these two scenarios, DeepSelect is out of scope and may be neither faster nor supported. Its first nature is to serve DSA (V3.2 / V4 / V4.1).

Q5: How do I verify the speedup of DeepSelect on my machine? A5: Use the built-in benchmark: python3 tests/test.py --perf-only, reporting the speedup ratio against torch.topk on the same input, with effective memory bandwidth as the metric. Lightning Indexer uses bfloat16, topk=512; Sampling uses float32, vocab_size=129280, topk=512. For more performance, follow the README: disable sorted_index unless ordering is required, and set return_value=False when values are not needed.

References

• DeepSeek DeepSelect GitHub repository (152 stars, CUDA, DeepSeek official org, created and pushed on 2026-09-10): https://github.com/deepseek-ai/DeepSelect
• Official README (positioning, supported scenarios, performance metric, installation and usage): https://github.com/deepseek-ai/DeepSelect/blob/main/README.md
• Official Chinese algorithm deep-dive docs/DeepSelect-deep-dive.zh.md (algorithm background / algorithm / expected upper bound / performance summary / implementation)
• Official English algorithm deep-dive docs/DeepSelect-deep-dive.md
• GitHub API measured data (2026-09-10): 152 stars, tag v1.0.0, bilingual algorithm docs released the same day