tests/jax/test_fused_attn.py · c2937c5abacb85326f093e74bb282fb491b30b3d · OpenDAS / TransformerEngine

[JAX] Add THD + SWA unit tests (#1390) · b898cbe1

Reese Wang authored Jan 08, 2025



* Fix SWA mask for THD and forcing seqlen_kv >= seqlen_q for SWA
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Generalize sliding window mask
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix pylint
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>

b898cbe1

test_fused_attn.py 29.8 KB

Replace test_fused_attn.py