import tilelang.testing

import example_mha_sink_fwd_bhsd
import example_mha_sink_bwd_bhsd
import example_gqa_sink_bwd_bhsd
import example_gqa_sink_fwd_varlen
import example_gqa_sink_bwd_varlen


@tilelang.testing.requires_cuda
def test_example_mha_sink_fwd_bhsd_full_attn():
    example_mha_sink_fwd_bhsd.main()


@tilelang.testing.requires_cuda
def test_example_mha_sink_fwd_bhsd_sliding_window():
    example_mha_sink_fwd_bhsd.main(window_size=128)


@tilelang.testing.requires_cuda
def test_example_mha_sink_bwd_bhsd():
    example_mha_sink_bwd_bhsd.main()


@tilelang.testing.requires_cuda
def test_example_mha_sink_bwd_bhsd_sliding_window():
    example_mha_sink_bwd_bhsd.main(window_size=128)


@tilelang.testing.requires_cuda
def test_example_gqa_sink_bwd_bhsd():
    example_gqa_sink_bwd_bhsd.main()


@tilelang.testing.requires_cuda
def test_example_gqa_sink_bwd_bhsd_sliding_window():
    example_gqa_sink_bwd_bhsd.main(window_size=128)


@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
def test_example_gqa_sink_varlen():
    example_gqa_sink_fwd_varlen.main()  # non-causal
    example_gqa_sink_bwd_varlen.main()  # causal


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