tests/pytorch/attention/test_attention.py · 490a5f41ada5788bc6dd94ba54ab024e465e0ec6 · OpenDAS / TransformerEngine

[PyTorch] Fix attention backend and tests for `sm120` (#2320) · 490a5f41

Kirthi Shankar Sivamani authored Oct 30, 2025



* Fix attention backend and tests for sm120
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Disable MLA only for backward
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

490a5f41

test_attention.py 103 KB

Replace test_attention.py