apache / apache/tvm

[Bug] [Relax][CUTLASS] CUTLASS tests are silently skipped, masking GEMM codegen and TVM-FFI ABI regressions

Open
#20,031 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

When TVM is built with both `USE_CUDA=ON` and `USE_CUTLASS=ON`, `tests/python/relax/test_codegen_cutlass.py` can still be skipped at module collection time because `tvm.testing.env.build_flag_enabled("USE_CUTLASS")` incorrectly returns `False`.

This false-negative test gate masks at least three regressions in the Relax CUTLASS code generation path:

1. The regular GEMM template references the undefined substitution variable `${A_arg}`, while the template arguments provide `${lhs_arg}`.
2. `tvm/ffi/container/tensor.h`, which provides the required TVM-FFI tensor type conversion support.
3. The generated wrapper exports the old bare function symbol, while the current TVM-FFI library loader expects the `__tvm_ffi_` symbol and the TVM-FFI SafeCall ABI.

### Expected behavior

When TVM is built with `USE_CUTLASS=ON`:

1. `tvm.support.libinfo()["USE_CUTLASS"]` and `tvm.testing.env.build_flag_enabled("USE_CUTLASS")` should report that CUTLASS support is enabled.
2. `tests/python/relax/test_codegen_cutlass.py` should execute instead of being skipped.
3. A minimal Relax `float16` 2-D matmul partitioned for CUTLASS should:
- generate valid CUDA/C++ source;
- compile successfully;
- be loadable by the Relax VM;
- expose the external function using the current TVM-FFI ABI.

### Actual behavior

#### 1. CUTLASS tests are skipped even though CUTLASS is available

The module-level marker in `tests/python/relax/test_codegen_cutlass.py` is equivalent to:
```python
python - <<'PY'
import tvm

libinfo = tvm.support.libinfo()

print("TVM version =", tvm.__version__)
print("libinfo USE_CUTLASS =", libinfo.get("USE_CUTLASS"))
PY
```
output:
```
TVM version = 0.26.dev0
libinfo USE_CUTLASS = OFF
```

Running a CUTLASS test therefore reports it as skipped:
```python
python -m pytest \
'tests/python/relax/test_codegen_cutlass.py::test_matmul_offload[float16-x_shape0-y_shape0-False-none-none]' \
-qrs
```

Output:
```
SKIPPED [...] need cutlass
```
when I modify the `python/tvm/support/libinfo.py` to bypass the problem, the following mistakes were exposed

#### 2. Bypassing the skip exposes invalid regular GEMM source
When a minimal regular GEMM is passed through CUTLASS codegen, the generated source contains an unsubstituted ${A_arg} expression.

The regular GEMM template currently contains code equivalent to:

```
cudaStream_t stream =
${lhs_arg}->device.device_id
```

This replacement should be scoped to the regular GEMM template because A_arg may still be valid in the specialized decode/quantized paths.

#### 3. Generated code cannot cast TVM-FFI arguments to DLTensor*

The generated CUTLASS source includes:

`DLTensor* arg0 = (DLTensor*)(args[0].cast());`

but its generated header list contains only:

```c++
headers = [
"tvm/ffi/function.h",
"tvm/ffi/extra/c_env_api.h",
]
```

Without:

` #include `

the required TVM-FFI tensor type traits/conversion support is unavailable. Compilation then fails because no valid AnyView::cast

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with python/tvm/support/libinfo.py and tests/python/relax/test_codegen_cutlass.py, then run the named pytest case with CUDA and CUTLASS enabled. Trace the build flag and generated CUTLASS source, including the regular GEMM template and generated headers. Done means the test is no longer skipped and the minimal float16 matmul generates, compiles, loads, and exposes the expected TVM-FFI ABI.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers, testing
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.