mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-30 08:03:09 +08:00
[Enhancement] Refactor CUDA vectorized cast generation and remove unsupported FP8 type (#1474)
* 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
co-authored by
LeiWang1999
Zhiwen Mo
LeiWang1999
Zhiwen Mo
parent
2d8bf3e102
commit
174fbe1908
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -3,14 +3,6 @@ import torch
|
||||
import tilelang.testing
|
||||
import tilelang.language as T
|
||||
|
||||
str2dtype = {
|
||||
T.float32: torch.float32,
|
||||
T.float16: torch.float16,
|
||||
T.bfloat16: torch.bfloat16,
|
||||
T.float8_e4m3fn: torch.float8_e4m3fn,
|
||||
T.float8_e5m2: torch.float8_e5m2,
|
||||
}
|
||||
|
||||
|
||||
@tilelang.jit
|
||||
def vectorized_cast_kernel(M: int, dtype_A: str, dtype_B: str):
|
||||
@@ -48,34 +40,39 @@ def parallel_vectorized_cast_kernel(M: int, dtype_A: str, dtype_B: str):
|
||||
return main
|
||||
|
||||
|
||||
def run_vectorized_cast(src_dtype_str: str, dst_dtype_str: str, check_str: str, lanes: int = 2):
|
||||
def run_vectorized_cast(src_dtype: T.dtype, dst_dtype: T.dtype, check_str: str, lanes: int = 2):
|
||||
"""Run the vectorized cast kernel and check the correctness.
|
||||
Args:
|
||||
src_dtype_str: The source data type string.
|
||||
dst_dtype_str: The destination data type string.
|
||||
src_dtype: The source data type.
|
||||
dst_dtype: The destination data type.
|
||||
check_str: Used to ensure vectorized cast is used.
|
||||
lanes: The number of lanes of the source and destination data types.
|
||||
"""
|
||||
|
||||
M = 128 * lanes
|
||||
kernel = vectorized_cast_kernel(M, src_dtype_str, dst_dtype_str)
|
||||
kernel_parallel = parallel_vectorized_cast_kernel(M, src_dtype_str, dst_dtype_str)
|
||||
kernel = vectorized_cast_kernel(M, src_dtype, dst_dtype)
|
||||
kernel_parallel = parallel_vectorized_cast_kernel(M, src_dtype, dst_dtype)
|
||||
|
||||
code = kernel.get_kernel_source()
|
||||
code_parallel = kernel_parallel.get_kernel_source()
|
||||
print(code)
|
||||
assert check_str in code and check_str in code_parallel, f"Cast {src_dtype} to {dst_dtype} with {lanes=} is not vectorized!"
|
||||
|
||||
if src_dtype == T.float4_e2m1fn or dst_dtype == T.float4_e2m1fn:
|
||||
return
|
||||
|
||||
A_float = torch.randn(M, dtype=torch.float32, device="cuda")
|
||||
A = A_float.to(str2dtype[src_dtype_str])
|
||||
B = torch.zeros(M, dtype=str2dtype[dst_dtype_str], device="cuda")
|
||||
C = torch.zeros(M, dtype=str2dtype[dst_dtype_str], device="cuda")
|
||||
A = A_float.to(src_dtype.as_torch())
|
||||
|
||||
A = A_float.to(src_dtype.as_torch())
|
||||
B = torch.zeros(M, dtype=dst_dtype.as_torch(), device="cuda")
|
||||
C = torch.zeros(M, dtype=dst_dtype.as_torch(), device="cuda")
|
||||
|
||||
kernel(A, B)
|
||||
kernel_parallel(A, C)
|
||||
|
||||
torch.testing.assert_close(A.to(str2dtype[dst_dtype_str]), B)
|
||||
torch.testing.assert_close(A.to(str2dtype[dst_dtype_str]), C)
|
||||
|
||||
code = kernel.get_kernel_source()
|
||||
code_parallel = kernel_parallel.get_kernel_source()
|
||||
|
||||
assert check_str in code and check_str in code_parallel, f"Cast {src_dtype_str} to {dst_dtype_str} with {lanes=} is not vectorized!"
|
||||
torch.testing.assert_close(A.to(dst_dtype.as_torch()), B)
|
||||
torch.testing.assert_close(A.to(dst_dtype.as_torch()), C)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -93,15 +90,41 @@ def run_vectorized_cast(src_dtype_str: str, dst_dtype_str: str, check_str: str,
|
||||
(T.float32, T.bfloat16, "__float22bfloat162_rn", 4),
|
||||
(T.bfloat16, T.float32, "__bfloat1622float2", 2),
|
||||
(T.bfloat16, T.float32, "__bfloat1622float2", 4),
|
||||
(T.float8_e4m3fn, T.float32, "__tl_cvt_fp8x2_to_float2", 2),
|
||||
(T.float8_e4m3fn, T.float32, "__tl_cvt_fp8x2_to_float2", 4),
|
||||
(T.float8_e5m2, T.float32, "__tl_cvt_fp8x2_to_float2", 2),
|
||||
(T.float8_e5m2, T.float32, "__tl_cvt_fp8x2_to_float2", 4),
|
||||
],
|
||||
)
|
||||
def test_vectorized_cast(src_dtype, dst_dtype, check_str, lanes):
|
||||
run_vectorized_cast(src_dtype, dst_dtype, check_str, lanes)
|
||||
|
||||
|
||||
@tilelang.testing.requires_cuda
|
||||
@tilelang.testing.requires_cuda_compute_version_ge(8, 9)
|
||||
@pytest.mark.parametrize(
|
||||
"src_dtype, dst_dtype, check_str, lanes",
|
||||
[
|
||||
(T.float8_e4m3fn, T.float32, "__tl_cvt_fp8x2_to_float2", 2),
|
||||
(T.float8_e4m3fn, T.float32, "__tl_cvt_fp8x2_to_float2", 4),
|
||||
(T.float8_e5m2, T.float32, "__tl_cvt_fp8x2_to_float2", 2),
|
||||
(T.float8_e5m2, T.float32, "__tl_cvt_fp8x2_to_float2", 4),
|
||||
],
|
||||
)
|
||||
def test_vectorized_cast_fp8(src_dtype, dst_dtype, check_str, lanes):
|
||||
run_vectorized_cast(src_dtype, dst_dtype, check_str, lanes)
|
||||
|
||||
|
||||
@tilelang.testing.requires_cuda
|
||||
@tilelang.testing.requires_cuda_compute_version_ge(10, 0)
|
||||
@pytest.mark.parametrize(
|
||||
"src_dtype, dst_dtype, check_str, lanes",
|
||||
[
|
||||
(T.float4_e2m1fn, T.float16, "__tl_cvt_fp4x2_to_half2", 2),
|
||||
(T.float16, T.float4_e2m1fn, "__tl_cvt_half2_to_fp4x2", 2),
|
||||
(T.float4_e2m1fn, T.float32, "__tl_cvt_fp4x2_to_float2", 2),
|
||||
(T.float32, T.float4_e2m1fn, "__tl_cvt_float2_to_fp4x2", 2),
|
||||
],
|
||||
)
|
||||
def test_vectorized_cast_fp4(src_dtype, dst_dtype, check_str, lanes):
|
||||
run_vectorized_cast(src_dtype, dst_dtype, check_str, lanes)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tilelang.testing.main()
|
||||
|
||||
Reference in New Issue
Block a user