examples/flash_attention/example_mha_bwd_wgmma_pipelined.py · 481cae424724f82da40e7056b4c4d2c44f047dc2 · OpenDAS / tilelang

[Example] Revert the atomic/split&sum templates in MHA backward examples (#943) · 481cae42

Tong WU authored Oct 06, 2025



* revert split+sum template for MHA backward

* lint

* Update example_mha_bwd.py

* Update example_mha_bwd_wgmma_pipelined.py

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

481cae42

example_mha_bwd_wgmma_pipelined.py 15.2 KB

Replace example_mha_bwd_wgmma_pipelined.py