test_example_flash_attention.py 2.43 KB
Newer Older
1
2
3
import tilelang.testing

import example_gqa_bwd
4
import example_gqa_bwd_wgmma_pipelined
5
import example_mha_bwd
6
import example_mha_bwd_bhsd
7
8
9
10
11
12
13
14
import example_mha_fwd_bhsd_wgmma_pipelined
import example_gqa_fwd_bshd
import example_mha_fwd_bshd
import example_gqa_fwd_bshd_wgmma_pipelined
import example_mha_fwd_bshd_wgmma_pipelined
import example_mha_fwd_varlen
import example_mha_bwd_wgmma_pipelined
import example_mha_fwd_bhsd
15
16
17
18
19
20
import example_gqa_bwd_tma_reduce_varlen


@tilelang.testing.requires_cuda
def test_example_gqa_bwd_tma_reduce_varlen():
    example_gqa_bwd_tma_reduce_varlen.main()
21
22
23
24
25
26
27


@tilelang.testing.requires_cuda
def test_example_gqa_bwd():
    example_gqa_bwd.main()


28
29
30
31
32
33
@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
def test_example_gqa_bwd_wgmma_pipelined():
    example_gqa_bwd_wgmma_pipelined.main()


34
35
@tilelang.testing.requires_cuda
def test_example_mha_bwd():
36
    example_mha_bwd.main(BATCH=1)
37
38


39
40
@tilelang.testing.requires_cuda
def test_example_mha_bwd_bhsd():
41
    example_mha_bwd_bhsd.main(BATCH=1)
42
43


44
45
46
@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
def test_example_mha_bwd_wgmma_pipelined():
47
    example_mha_bwd_wgmma_pipelined.main(BATCH=1)
48
49
50


@tilelang.testing.requires_cuda
51
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
52
def test_example_gqa_fwd_bshd_wgmma_pipelined():
53
54
    example_gqa_fwd_bshd_wgmma_pipelined.main(
        batch=1, heads=16, seq_len=1024, dim=128, is_causal=False, groups=16, tune=False)
55
56
57
58


@tilelang.testing.requires_cuda
def test_example_gqa_fwd_bshd():
59
60
    example_gqa_fwd_bshd.main(
        batch=1, heads=16, seq_len=1024, dim=128, is_causal=False, groups=16, tune=False)
61
62
63


@tilelang.testing.requires_cuda
64
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
65
66
67
68
69
70
71
72
73
74
def test_example_mha_fwd_bhsd_wgmma_pipelined():
    example_mha_fwd_bhsd_wgmma_pipelined.main()


@tilelang.testing.requires_cuda
def test_example_mha_fwd_bhsd():
    example_mha_fwd_bhsd.main()


@tilelang.testing.requires_cuda
75
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
76
def test_example_mha_fwd_bshd_wgmma_pipelined():
77
    example_mha_fwd_bshd_wgmma_pipelined.main(batch=1, heads=32, seq_len=256)
78
79
80
81


@tilelang.testing.requires_cuda
def test_example_mha_fwd_bshd():
82
    example_mha_fwd_bshd.main(batch=1, seq_len=256)
83
84
85
86
87
88
89
90
91


@tilelang.testing.requires_cuda
def test_example_mha_fwd_varlen():
    example_mha_fwd_varlen.main()


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