Accelerate PyTorch TopK Operations with DeepSelect

What Is DeepSelect?

DeepSelect is a high-performance implementation of the TopK kernel used in DeepSeek Sparse Attention and sampling. According to the project README, it can deliver a 2–20x speedup over vanilla torch.topk for the workloads supported by the library.

Rather than attempting to optimize every possible TopK configuration, DeepSelect concentrates on two scenarios: a Lightning Indexer workload using torch.bfloat16 inputs and a sampling workload using torch.float32 inputs.

Supported Workloads

Lightning Indexer scenario

• Input type: torch.bfloat16
• Batch size: from 1 upward, with optimizations for both small and large batches
• Vocabulary size: from 1 upward, with optimizations for both small and large vocabularies
• TopK size: small and no greater than 4096

For better performance, leave sorted_index disabled unless the results must be ordered. If values are unnecessary, set return_value=False to skip their output.

Sampling scenario

• Input type: torch.float32
• Batch size: from 1 upward
• Vocabulary size: approximately 128K
• TopK size: small and no greater than 4096

The repository measures performance using effective memory bandwidth rather than floating-point throughput because TopK does not perform floating-point arithmetic. Its benchmark reports performance relative to torch.topk on the same input.

Key Features

• Specialized TopK kernels for DeepSeek Sparse Attention and sampling workloads
• Support for torch.bfloat16 and torch.float32 inputs in their respective target scenarios
• Configurable torch.int32 or torch.int64 index output
• Optional value output for lower overhead when only indices are needed
• Variable-length rows through a per-row exclusive upper bound
• User-owned output buffers for indices
• Built-in NaN checking

Installation

Clone the repository, initialize its submodules, and install the package from the project directory:

git clone https://github.com/deepseek-ai/DeepSelect.git
cd DeepSelect
git submodule update --init --recursive
pip install -v .

The examples use CUDA tensors, so create the inputs on device="cuda".

Basic TopK Usage

Import PyTorch and DeepSelect, create a two-dimensional input tensor, and call deep_select.topk:

import torch
import deep_select

batch_size = 4
vocab_size = 204800
topk = 1024

x = torch.randn(
    batch_size,
    vocab_size,
    dtype=torch.bfloat16,
    device="cuda",
)

values, indices = deep_select.topk(
    x,
    topk,
    sorted_index=True,
    indices_type=torch.int32,
    return_value=True,
)

print(values.shape)   # (4, 1024)
print(indices.shape)  # (4, 1024)

The input must have shape (batch_size, vocab_size). Its final dimension must be contiguous, and its row stride must be a multiple of deep_select.get_stride_requirement()[0] bytes.

The returned values have the same data type as x. The indices use the type selected through indices_type, which can be torch.int32 or torch.int64. With sorted_index=True, each row’s returned indices are arranged in ascending order.

Returning only indices

If downstream code does not need the selected values, disable value output:

values, indices = deep_select.topk(
    x,
    topk,
    sorted_index=False,
    indices_type=torch.int32,
    return_value=False,
)

The README notes that skipping value output is approximately 10% faster. Avoid requesting sorted output unless ordering is required, because sorting adds overhead.

Handling Stride Alignment

DeepSelect imposes alignment requirements on input rows and output rows. Inspect them at runtime:

input_alignment, output_alignment = deep_select.get_stride_requirement()

print("Input row-stride alignment:", input_alignment, "bytes")
print("Output row-stride alignment:", output_alignment, "bytes")

If an input’s row stride is not a multiple of the required input alignment, pad the tensor before passing it to DeepSelect. The last dimension must remain contiguous.

DeepSelect allocates both outputs by default. Their strides are aligned to deep_select.get_stride_requirement()[1] bytes, which means the returned tensors may be non-contiguous. Code consuming the outputs should not assume standard contiguous strides.

Using an existing index buffer

The API accepts output_idx= when indices should be written into a buffer managed by the caller. That buffer must satisfy the same output stride requirement. Consult deep_select/interface.py for the complete function signature before configuring a custom buffer.

Working with Variable-Length Rows

Use end to assign each row an exclusive upper bound. This is useful when rows in the same batch contain different valid lengths:

import torch
import deep_select

batch_size = 2
vocab_size = 129280

x = torch.randn(
    batch_size,
    vocab_size,
    dtype=torch.float32,
    device="cuda",
)

end = torch.tensor(
    [129280, 100000],
    dtype=torch.int32,
    device="cuda",
)

values, indices = deep_select.topk(
    x,
    1000,
    end=end,
    sorted=True,
    indices_type=torch.int64,
)

Here, the first row can use all 129280 elements, while the second row is limited to the first 100000 elements. If a row is shorter than the requested TopK size, DeepSelect pads its outputs with value_oob_fill_value and idx_oob_fill_value.

NaN Handling

NaN checking is always enabled. By default, abort_when_nan_found=True, so the kernel invokes trap() and aborts when it detects a NaN. Rows whose valid length is less than or equal to the requested TopK size are not NaN-checked.

Validate or sanitize input data before launching the kernel if NaNs are possible, especially when retaining the default abort behavior.

Performance and Advanced Tips

1. Stay within the supported TopK range. DeepSelect supports small TopK values up to and including 4096. Larger values are unsupported.
2. Match the intended workload. Use torch.bfloat16 for the Lightning Indexer scenario and torch.float32 for the sampling scenario described by the project.
3. Skip sorting when possible. Ordering by index or value costs performance, so request it only when downstream processing requires it.
4. Skip values when possible. Set return_value=False when only indices are consumed.
5. Check alignment early. Plan tensor shapes and padding around get_stride_requirement() instead of fixing incompatible tensors in a performance-critical loop.
6. Do not assume contiguous outputs. DeepSelect aligns output row strides, and the resulting tensors may be non-contiguous.
7. Benchmark the real shape. TopK performance depends on the input type, batch size, vocabulary size, and TopK size. Test the exact configuration used by the application.

Running the supplied benchmark

The repository includes a performance mode in its test script:

python3 tests/test.py --perf-only

This benchmark compares DeepSelect with torch.topk on the same inputs and reports the ratio using effective memory bandwidth.

Conclusion

DeepSelect provides specialized TopK kernels for compatible DeepSeek Sparse Attention and sampling workloads. Start with an aligned two-dimensional CUDA tensor, keep topk at or below 4096, and disable sorting or value output when unnecessary. For production integration, verify stride requirements, account for potentially non-contiguous outputs, and benchmark the exact tensor shapes used by your application.