# DeepEP topk_idx_t Custom Dtype Requirements: Implementation Guide

> Learn the requirements for DeepEP topk_idx_t custom dtype. Implement efficiently with torch.int32 or torch.int64 based on TOPK_IDX_BITS. Get the guide now.

- Repository: [DeepSeek/DeepEP](https://github.com/deepseek-ai/DeepEP)
- Tags: how-to-guide
- Published: 2026-04-25

---

**DeepEP requires `topk_idx` tensors to use the custom integer dtype `deep_ep.topk_idx_t`, which compiles to either `torch.int32` or `torch.int64` based on the `TOPK_IDX_BITS` environment variable, defaulting to 64-bit integers.**

In the `deepseek-ai/DeepEP` repository, `topk_idx_t` serves as the specialized integer type for storing expert indices during Mixture-of-Experts (MoE) routing. Understanding these DeepEP topk_idx_t custom dtype requirements ensures compatibility with the library's dispatch kernels and buffer management systems.

## What Is topk_idx_t in DeepEP?

The `topk_idx_t` type defines how DeepEP stores indices for the top-k experts selected per token. Located in `csrc/kernels/configs.cuh` (lines 71–77), this typedef selects between `int32_t` and `int64_t` based on the `TOPK_IDX_BITS` preprocessor macro. By default, the system uses 64-bit integers to accommodate large expert counts, but you can reduce memory footprint by switching to 32-bit indices at compile time.

## Compile-Time Size Configuration

You control the bit-width of `topk_idx_t` before building the extension, not at runtime.

The configuration flows through two critical files:

- **`csrc/kernels/configs.cuh`** defines the C++ typedef based on the `TOPK_IDX_BITS` macro
- **[`setup.py`](https://github.com/deepseek-ai/DeepEP/blob/main/setup.py)** (lines 84–88) captures the environment variable and passes it as a compiler flag

To build DeepEP with 32-bit indices instead of the default 64-bit:

```bash
export TOPK_IDX_BITS=32
python setup.py install

```

To revert to 64-bit:

```bash
export TOPK_IDX_BITS=64
python setup.py clean --all
python setup.py install

```

## Python-Level Type Alias

After compilation, the dtype appears in Python as `deep_ep.topk_idx_t`. The [`deep_ep/__init__.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/__init__.py) file (line 7) maps this alias to the correct PyTorch dtype:

- When `TOPK_IDX_BITS=64`: resolves to `torch.int64`
- When `TOPK_IDX_BITS=32`: resolves to `torch.int32`

Always reference this alias rather than hardcoding `torch.int64` or `torch.int32` to ensure your code works across different DeepEP builds.

```python
import deep_ep
import torch

print(deep_ep.topk_idx_t)  # torch.int64 (default) or torch.int32

```

## Tensor Requirements and Constraints

Any tensor representing `topk_idx` must satisfy strict constraints enforced in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py) (lines 301–306). The dispatch kernels require:

1. **Exact dtype match**: Must equal `deep_ep.topk_idx_t`
2. **2-dimensional shape**: Dimensions must be `[num_tokens, num_topk]`
3. **Memory contiguity**: Tensor must be contiguous in memory

These constraints prevent expensive device-side conversions and ensure alignment with CUDA kernel expectations.

## Runtime Validation

Before launching GPU kernels, DeepEP validates tensor properties using `EP_HOST_ASSERT` statements within dispatch functions like `intranode_dispatch` and `internode_dispatch`. As implemented in the buffer module, these assertions verify:

- Dtype matches the compiled `topk_idx_t`
- Dimensionality equals 2
- Memory layout is contiguous
- Top-k size matches between arguments

A validation failure triggers an immediate assertion error with a descriptive message, halting execution before any GPU work begins.

## Practical Implementation Examples

### Building with 32-Bit Support

```bash

# Clean previous artifacts

python setup.py clean --all

# Configure for 32-bit indices

export TOPK_IDX_BITS=32
python setup.py install

# Verify in Python

python -c "import deep_ep; print(deep_ep.topk_idx_t)"  # torch.int32

```

### Creating a Valid topk_idx Tensor

```python
import torch
import deep_ep

num_tokens = 8192
topk = 8

# Create indices using the correct dtype

topk_idx = torch.randint(
    0, 256, 
    (num_tokens, topk), 
    dtype=deep_ep.topk_idx_t
)

# Ensure contiguous memory layout

topk_idx = topk_idx.contiguous()

# Verify properties

assert topk_idx.dtype == deep_ep.topk_idx_t
assert topk_idx.dim() == 2
assert topk_idx.is_contiguous()

```

### Validation Helper Function

```python
def validate_topk_tensor(tensor: torch.Tensor, expected_topk: int):
    """Validate tensor against DeepEP topk_idx_t requirements."""
    assert tensor.dtype == deep_ep.topk_idx_t, \
        f"Expected {deep_ep.topk_idx_t}, got {tensor.dtype}"
    assert tensor.ndim == 2, \
        f"Expected 2D tensor [num_tokens, topk], got {tensor.ndim}D"
    assert tensor.is_contiguous(), \
        "Tensor must be contiguous in memory"
    assert tensor.shape[1] == expected_topk, \
        f"Expected topk={expected_topk}, got {tensor.shape[1]}"
    
    return True

```

## Summary

- **Compile-time only**: Set `TOPK_IDX_BITS` to 32 or 64 before running [`setup.py`](https://github.com/deepseek-ai/DeepEP/blob/main/setup.py) to control the dtype size in `csrc/kernels/configs.cuh`
- **Python abstraction**: Always use `deep_ep.topk_idx_t` instead of literal torch dtypes to ensure compatibility across builds
- **Strict constraints**: Tensors must be 2D `[num_tokens, num_topk]`, contiguous, and use the exact compiled dtype
- **Early validation**: Runtime checks in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py) prevent dtype mismatches from reaching the GPU

## Frequently Asked Questions

### Can I change the topk_idx_t bit-width after installing DeepEP?

No. The `topk_idx_t` dtype is determined at compile time when [`setup.py`](https://github.com/deepseek-ai/DeepEP/blob/main/setup.py) processes the `TOPK_IDX_BITS` environment variable and generates the compiler flags. To change from 64-bit to 32-bit (or vice versa), you must uninstall the current build, set the environment variable, and reinstall. The [`deep_ep/__init__.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/__init__.py) module exposes the compiled dtype constant for runtime reference, but you cannot alter the underlying C++ type without rebuilding.

### Why does DeepEP enforce contiguity for topk_idx tensors?

The dispatch kernels in DeepEP use direct memory pointers and stride assumptions optimized for contiguous data layouts. Non-contiguous tensors would require additional indexing calculations or memory copies inside the CUDA kernels, significantly impacting the performance-critical paths of intranode and internode dispatch operations. The `EP_HOST_ASSERT` checks in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py) enforce this requirement before GPU launch to prevent silent errors or performance degradation.

### What happens if I pass a torch.int64 tensor to a 32-bit DeepEP build?

The operation will fail with an assertion error during the dtype check phase. DeepEP's runtime validation compares `tensor.dtype` against `deep_ep.topk_idx_t` using `EP_HOST_ASSERT` statements. If you compiled with `TOPK_IDX_BITS=32`, the Python alias resolves to `torch.int32`, and passing a `torch.int64` tensor triggers a mismatch error before any GPU computation begins. Always verify compatibility using `tensor.dtype == deep_ep.topk_idx_t`.

### Is there a performance advantage to using 32-bit indices?

Yes, particularly for large-scale deployments. Using `TOPK_IDX_BITS=32` reduces memory bandwidth pressure during expert routing, cutting the size of index tensors in half compared to 64-bit indices. This optimization proves especially valuable in high-throughput scenarios with many tokens and large top-k values, though you must ensure your expert count remains below the 2^31 limit imposed by 32-bit signed integers.