transformer_engine/jax/flax/transformer.py · 5978f1d7544fe0e4b06a036600c80c6a8c8203fe · OpenDAS / TransformerEngine

[JAX] Default to fused attention in JAX DPA (#2363) · 5978f1d7

Kshitij Lakhani authored Nov 07, 2025



* Default to fused attention in JAX DPA
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* Consolidate documentation for DPA in JAX
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com>

* Correctly update the documentation for defaults in JAX DPA
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com>

---------
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

5978f1d7

transformer.py 93.1 KB

Replace transformer.py