examples/flash_attention/example_mha_fwd_bshd.py · 24403aea73bc160a60f20a8f09c1dfb9eb84bb89 · OpenDAS / tilelang

[CI] Add CI test for flash_attention examples (#558) · 24403aea

xs-keju authored Jun 06, 2025



* [CI] Add CI test for flash_attention examples

* Update example_gqa_fwd_bshd.py

* Update example_mha_fwd_bshd_wgmma_pipelined.py

* [CI] Added conditional annotations for tests in flash_attention

* [CI] Added conditional annotations for tests in flash_attention

---------
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>

24403aea

example_mha_fwd_bshd.py 9.47 KB

Replace example_mha_fwd_bshd.py