trending/tilelang
1
0
Fork 0
mirror of https://github.com/tile-ai/tilelang.git synced 2026-10-25 07:14:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Enhancement] Refactor CUDA vectorized cast generation and remove unsupported FP8 type (#1474)

Browse Source 
* Refactor CUDA vectorized cast generation and remove unsupported FP8 type

* test fix

* lint fix

* Refactor CUDA vectorized cast function naming for clarity

* Add support for float4_e2m1fn type conversions in CUDA vectorized casts

- Implemented conversions between float4_e2m1fn and float32, half2, and float2 in utils.cc and cuda_fp4.h.
- Updated test_tilelang_language_vectorized_cast.py to validate new conversions and ensure correctness.
- Enhanced dtype conversion in dtypes.py to handle float4_e2m1fn appropriately, logging a warning for unsupported types in PyTorch.

* Enhance vectorized cast tests for new data types

- Added tests for vectorized casting of float8 and float4 data types, ensuring compatibility with CUDA compute versions.
- Refactored existing test functions to improve clarity and organization, separating tests for different data types.
- Updated parameterization to include additional test cases for new conversions.

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
This commit is contained in:
Xiangwen Wang
2025-12-22 22:44:03 +08:00
committed by GitHub
co-authored by LeiWang1999 Zhiwen Mo
parent 2d8bf3e102
commit 174fbe1908
8 changed files with 272 additions and 290 deletions
Download Patch File Download Diff File Expand all files Collapse all files
testing/python/debug/test_tilelang_debug_print.py
+10 -4
View File
@@ -12,13 +12,12 @@ def debug_print_buffer(M=16, N=16, dtype=T.float16):
shared_buf = T.alloc_shared([M, N], dtype)
T.print(shared_buf)
jit_kernel = tilelang.compile(program, target="cuda", execution_backend="tvm_ffi")
jit_kernel = tilelang.compile(program)
profiler = jit_kernel.get_profiler()
profiler.run_once()
def test_debug_print_buffer():
debug_print_buffer(dtype=T.bool)
debug_print_buffer(dtype=T.int8)
debug_print_buffer(dtype=T.int16)
debug_print_buffer(dtype=T.int32)
@@ -31,10 +30,17 @@ def test_debug_print_buffer():
debug_print_buffer(dtype=T.float32)
debug_print_buffer(dtype=T.float64)
debug_print_buffer(dtype=T.bfloat16)
@tilelang.testing.requires_cuda
def test_debug_print_buffer_cuda_fp8():
debug_print_buffer(dtype=T.float8_e4m3fn)
debug_print_buffer(dtype=T.float8_e4m3fn)
debug_print_buffer(dtype=T.float8_e4m3fnuz)
debug_print_buffer(dtype=T.float8_e5m2)
@tilelang.testing.requires_rocm
def test_debug_print_buffer_rocm_fp8():
debug_print_buffer(dtype=T.float8_e4m3fnuz)
debug_print_buffer(dtype=T.float8_e5m2fnuz)