transformer_engine/jax/attention.py · 3b1f5a11f28521c834b42eb38092e8360eb63afe · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Add fast path for causal masking with segment IDs. (#1601) · 3b1f5a11
Michael Goldfarb authored Mar 31, 2025
```
Add fast path for causal masking with segment IDs.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
```
3b1f5a11

attention.py 42.9 KB