#1351·tilelang

[BUG] illegal memory access when using topk_selector

Author: yweng0828Created Nov 27, 2025Updated Sep 10, 2026

Labelsbug

Required prerequisites

I have read the documentation https://tilelang.com.
I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

What version of TileLang are you using?

0.1.6.post2+cuda.gitf0c721a4

System information

3.12.3 (main, Aug 14 2025, 17:47:21) [GCC 13.3.0] linux 0.1.6.post2+cuda.gitf0c721a4 2.9.1+cu128

Problem description

When I use tl_topk(), if the input data values are close (for example, only the last 10 bits are different), an illegal memory access will occur.

Reproducible example code

Please using examples/deepseek_v32/topk_selector.py. And add the following code.

The Python snippets:

python


# Returns float values that differ in the last 10 bits
def get_10bit_data(bs: int, seq_len: int) -> torch.Tensor:
    torch.manual_seed(42)

    top_22_bits_mask = 0xFFFFFC00
    last_10_bits_mask = 0x000003FF
    fixed_top_22_bits = 0x3F900000

    # Generate random bits for the last 10 bits
    random_bottom_bits = torch.randint(
        0, 2**10, (bs, seq_len), dtype=torch.int32, device="cuda"
    )

    # Combine: fixed top 22 bits with random last 10 bits
    score_bits = (fixed_top_22_bits & top_22_bits_mask) | (
        random_bottom_bits & last_10_bits_mask
    )

  # Convert back to float
  score = score_bits.view(torch.float32)
  return score


batch = 2
seq_len = 32 * 1024
topk = 2048
input = get_10bit_data(batch, seq_len).cuda()
starts = torch.zeros(batch, dtype=torch.int32).cuda()
ends = torch.ones(batch, dtype=torch.int32).cuda() * seq_len

indexes_ref = torch.topk(input, topk, dim=-1)[1] 
print(f"indexes_ref shape: {indexes_ref.shape}, indexes_ref: {indexes_ref}") # good

indexes = tl_topk(input, starts, ends, topk) 
print(f"indexes shape: {indexes.shape}, indexes: {indexes}") # illegal memory access

Traceback

pytb

Traceback (most recent call last):
  File "/workspace_new/tilelang/examples/deepseek_v32/topk_selector.py", line 284, in <module>
    test_topk_selector()
  File "/workspace_new/tilelang/examples/deepseek_v32/topk_selector.py", line 235, in test_topk_selector
    print(f"indexes shape: {indexes.shape}, indexes: {indexes}")
                                                     ^^^^^^^^^
  File "/workspace_new/tilelang/venv/lib/python3.12/site-packages/torch/_tensor.py", line 1109, in __format__
    return object.__format__(self, format_spec)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace_new/tilelang/venv/lib/python3.12/site-packages/torch/_tensor.py", line 568, in __repr__
    return torch._tensor_str._str(self, tensor_contents=tensor_contents)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace_new/tilelang/venv/lib/python3.12/site-packages/torch/_tensor_str.py", line 722, in _str
    return _str_intern(self, tensor_contents=tensor_contents)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace_new/tilelang/venv/lib/python3.12/site-packages/torch/_tensor_str.py", line 643, in _str_intern
    tensor_str = _tensor_str(self, indent)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace_new/tilelang/venv/lib/python3.12/site-packages/torch/_tensor_str.py", line 375, in _tensor_str
    formatter = _Formatter(get_summarized_data(self) if summarize else self)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace_new/tilelang/venv/lib/python3.12/site-packages/torch/_tensor_str.py", line 142, in __init__
    value_str = f"{value}"
                  ^^^^^^^
  File "/workspace_new/tilelang/venv/lib/python3.12/site-packages/torch/_tensor.py", line 1108, in __format__
    return self.detach().item().__format__(format_spec)
           ^^^^^^^^^^^^^^^^^^^^
torch.AcceleratorError: CUDA error: an illegal memory access was encountered
Search for `cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Expected behavior

No response

Additional context

Please also pay attention to accuracy issues after the repair. Thank you.

Source: tile-ai/tilelang

View original on GitHub View discussion on GitHub