mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-26 07:22:54 +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
17 lines
445 B
Python
17 lines
445 B
Python
import tilelang.testing
|
|
import example_group_per_split_token_cast_to_fp8
|
|
import example_per_token_cast_to_fp8
|
|
|
|
|
|
def test_example_group_per_split_token_cast_to_fp8():
|
|
example_group_per_split_token_cast_to_fp8.main(
|
|
M=1024, N=1024, BG=2, blk_m=4, batch_sizes=[128, 896])
|
|
|
|
|
|
def test_example_per_token_cast_to_fp8():
|
|
example_per_token_cast_to_fp8.main(M=2048, N=512, blk_m=8)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|