mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
Yuqi Dong
13bdcd605e
* [Refactor]: Change the params in pytest to avoid oom error during ci * format * fix * Update test_example_cast.py * Update parameters in test_example_cast * Update test_example_flash_attention.py * update * format * fix * fix * format
20 lines
492 B
Python
20 lines
492 B
Python
import tilelang.testing
|
|
|
|
import example_gqa_decode
|
|
import example_mha_inference
|
|
|
|
|
|
# TODO(lei): fix the correctness of gqa decode on sm90
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_le(8, 9)
|
|
def test_example_example_gqa_decode():
|
|
example_gqa_decode.main()
|
|
|
|
|
|
def test_example_example_mha_inference():
|
|
example_mha_inference.main(BATCH=1, H=32, Q_CTX=128, KV_CTX=2048, D_HEAD=128, causal=False)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|