transformer_engine/jax/attention.py · 2643ba1df43397cc84c9da5fe719a66d87ad9a0a · OpenDAS / TransformerEngine

[JAX] Expose context parallel params to jax DPA api (#1292) · d7256866

Md Fahim Faysal Khan authored Nov 04, 2024



Exposed context parallel params to DPA api
Signed-off-by: Md Fahim Faysal Khan <mdfahimfaysa@nvidia.com>
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>

---------
Signed-off-by: Md Fahim Faysal Khan <mdfahimfaysa@nvidia.com>
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
Co-authored-by: Michael Goldfarb <mgoldfarb@nvidia.com>

d7256866

attention.py 23.3 KB

Replace attention.py