test_example_dequantize_gemm.py 944 Bytes
Newer Older
1
2
3
4
import tilelang.testing

import example_dequant_gemv_fp16xint4
import example_dequant_gemm_fp4_hopper
5
import example_dequant_gemm_bf16_mxfp4_hopper
6
import example_dequant_gemm_bf16_mxfp4_hopper_tma
7
8
9
10
11
12
13
14
15
16
17
18
19


@tilelang.testing.requires_cuda
def test_example_dequant_gemv_fp16xint4():
    example_dequant_gemv_fp16xint4.main()


@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
def test_example_dequant_gemm_fp4_hopper():
    example_dequant_gemm_fp4_hopper.main()


20
21
22
23
24
25
@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
def test_example_dequant_gemm_bf16_mxfp4_hopper():
    example_dequant_gemm_bf16_mxfp4_hopper.main()


26
27
28
29
30
31
@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
def test_example_dequant_gemm_bf16_mxfp4_hopper_tma():
    example_dequant_gemm_bf16_mxfp4_hopper_tma.main()


32
33
if __name__ == "__main__":
    tilelang.testing.main()