apache / apache/tvm

[Bug][Relax][CUDA] topk on large input launches kernel with invalid gridDim.y and fails with CUDA_ERROR_INVALID_VALUE

Open
#19,549 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

### Summary

A minimal `torch.topk` model compiled through TVM Relax for a CUDA target fails at runtime with `CUDA_ERROR_INVALID_VALUE`.

The model is very small:

```python
values, _ = torch.topk(x, k=10, dim=-1)
return values
```
For input shape `(60, 1000, 1000)`, `tvm.compile(...)` succeeds, but the Relax VM fails when launching the generated CUDA kernel topk_kernel_2.

The failing launch configuration is:
```
grid=(1,240000,1), block=(256,1,1)
// func_name=topk_kernel_2
```
Here `gridDim.y = 240000`, which exceeds the usual CUDA grid y-dimension limit. It looks like the CUDA TopK lowering/scheduling path generates an invalid launch configuration for this large input.

### Minimal reproduction
```python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import platform
import traceback
import numpy as np
import torch
import tvm
from tvm import relax

class TopKModel(torch.nn.Module):
def forward(self, x):
values, _ = torch.topk(x, k=10, dim=-1)
return values

def make_tvm_array(arr, dev):
if not isinstance(arr, np.ndarray):
arr = np.array(arr)
if not arr.flags["C_CONTIGUOUS"]:
arr = np.ascontiguousarray(arr)

# Some TVM builds expose tvm.nd.array, while newer/custom FFI builds may not.
if hasattr(tvm, "nd") and hasattr(tvm.nd, "array"):
try:
return tvm.nd.array(arr, device=dev)
except TypeError:
return tvm.nd.array(arr, dev)

try:
from tvm.runtime import ndarray as _nd
if hasattr(_nd, "array"):
try:
return _nd.array(arr, dev)
except TypeError:
return _nd.array(arr, device=dev)
except Exception:
pass

if hasattr(tvm, "runtime") and hasattr(tvm.runtime, "tensor"):
try:
return tvm.runtime.tensor(arr, device=dev)
except TypeError:
return tvm.runtime.tensor(arr, dev)

raise RuntimeError("Cannot construct TVM NDArray/Tensor in this TVM build")

def main():
print("=" * 80)
print("Environment")
print("=" * 80)
print("python:", sys.version.replace("\n", " "))
print("platform:", platform.platform())
print("torch:", torch.__version__)
print("tvm:", getattr(tvm, "__version__", ""))
print("tvm path:", getattr(tvm, "__file__", ""))
print("CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES", ""))
print("torch.cuda.is_available:", torch.cuda.is_available())
print("tvm.cuda(0).exist:", tvm.cuda(0).exist)

target = tvm.target.Target(
"cuda -keys=cuda,gpu -arch=sm_86 -max_num_threads=1024 -thread_warp_size=32"
)
dev = tvm.cuda(0)

print("target:", target)

model = TopKModel().eval()

# Export on CPU. The compiled TVM module will run on CUDA.
x = torch.rand(60, 1000, 1000, dtype=torch.float32)

with torch.no_grad():
eager_out = model(x)

print("input shape:", tuple(x.shape))
print("eager output shape:", tuple(eager_out.shape))

ep = torch.export.export(model, (x,))

from tvm.relax.frontend.torch import from_exported_program

ir_mod = from_exported_program(ep)

print("=" * 80)
print("tvm.compile")
print("=" * 80)

ex = tvm.compile(
ir_mod,
target=target,
relax_pipeline="default",
tir_pipeline="default",
)

print("compile: OK")

print("=" * 80)
print("Relax VM run")
print("=" * 80)

vm = relax.VirtualMachine(ex, dev)
fn = vm["main"]

x_tvm = make_tvm_array(x.numpy(), dev)

try:
out = fn(x_tvm)

# Force async CUDA launch errors to surface.
try:
dev.sync()
except Exception:
pass

print("run: OK")
if hasattr(out, "numpy"):
out_np = out.numpy()
print("output shape:", out_np.shape)
print("output dtype:", out_np.dtype)
else:
print("output type:", type(out))

except Exception:
print("run: FAILED")
traceback.print_exc()

if __name__ == "__main__":
main()
```

### Actual behavior

The module compiles successfully:
```
tvm.compile
compile: OK
```
But the Relax VM fails during CUDA kernel launch:
```
RuntimeError: CUDALaunch CUDA_ERROR_INVALID_VALUE
grid=(1,240000,1), block=(256,1,1)
// func_name=topk_kernel_2
```
The generated CUDA source is also printed in the exception. The relevant part is that the failing kernel is topk_kernel_2, and it is launched with:
```
grid=(1,240000,1), block=(256,1,1)
```

### Environment

TVM: 0.23.0
LLVM: 17.0.6
Python: 3.10.16 (from stack paths)
NumPy: 2.2.6

### Triage

* needs-triage
* bug

cc @junrushao

Contributor guide

No contributing guide indexed for this repository

Research direction

Run the provided torch.topk reproduction through tvm.compile and the Relax VM, then inspect the generated topk_kernel_2 launch and the CUDA TopK lowering or scheduling path. The fix is done when the (60, 1000, 1000) case launches within valid CUDA grid limits and the VM produces the expected top-k output without CUDA_ERROR_INVALID_VALUE.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.