Commits · 838345eba4fdd2a169dd9e087d39c30a360e684a · OpenDAS / TransformerEngine

17 Dec, 2024 1 commit

[JAX] Fused attention unit tests fixes and refinements (#1352) · 7f5c784e

Reese Wang authored Dec 17, 2024



* Add util functions to attn_mask_type
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add util functions to qkv_layout
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix THD cross reference code
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove explicit segment_pad, encoding it to segment_ids
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add jax.jit, replace _token with segment_ids, rename bias shape enum
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add comment for make_mask
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add doc strings for the added functions
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove cache for fa deterministic which causes UT failed
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename fixture to avoid conflict
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

7f5c784e

12 Nov, 2024 1 commit

[TE/JAX] XLA FFI calls for Softmax and FusedAttnBackward (#1319) · 237b4930

Hua Huang authored Nov 12, 2024



* FFI for all softmax functions
Signed-off-by: Hua Huang <huah@nvidia.com>

* FFI for FusedAttnBackward and Dequantize

FusedAttnBackward passed all testes in test_fused_attn.py.
Dequantize is not used currently; finish it for completeness.
Signed-off-by: Hua Huang <huah@nvidia.com>

* Fix FusedAttnBackward FFI pybind & simplify
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Revert changes to tests/jax/test_fused_attn.py
Signed-off-by: Hua Huang <huah@nvidia.com>

---------
Signed-off-by: Hua Huang <huah@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

237b4930

11 Nov, 2024 1 commit

[JAX] Support Ring Attention (Context Parallelism) (#1059) · bfddb483

Ming-Xu Huang authored Nov 11, 2024



* Implement ring attention primative for Jax.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Michael Goldfarb <mgoldfarb@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

bfddb483

07 Nov, 2024 1 commit

[JAX] Added prepare phase for the FusedAttnForwardFFI (#1313) · e5ffaa76

Phuong Nguyen authored Nov 06, 2024



* added prepare phase for the FusedAttnForwardFFI

* enabled FusedAttnForwardFFI by default

* moved prepare phase into pybind

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

e5ffaa76

01 Nov, 2024 1 commit
- [JAX] Fix for Disable FusedAttn with FFI by default (#1304) · 4b8ffef4
  Phuong Nguyen authored Nov 01, 2024
```
rm default value for NVTE_JAX_FUSED_ATTN_WITH_FFI
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
```
  4b8ffef4
31 Oct, 2024 2 commits

[TE/JAX] Disable FusedAttn with FFI by default (#1298) · 23caab3f
Phuong Nguyen authored Oct 31, 2024
```
* disable fused attn with ffi

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
```
23caab3f

[TE/JAX] Custom call with FFI - lowering all attributes with bind all (#1289) · 9dddb36d

Phuong Nguyen authored Oct 31, 2024



* lowering a dict of attrs

* improve err message with line and func info

* implement a product() for ffi dimensions

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

9dddb36d

25 Oct, 2024 1 commit

[C/PyTorch] Add max_t support for THD (#1244) · 7fb22c37

Charlene Yang authored Oct 25, 2024



* WIP: add max_t support for THD
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* WIP: save tensors for debug and point to new FE
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* fix stats in bwd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* fix stats in fwd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* add docstring for DPA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add docstring
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: first try on adding max_b and max_t
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Revert "[pre-commit.ci] auto fixes from pre-commit.com hooks"

This reverts commit c3d522e9f5aef3c8ddfec5bf6ff24c3db97bb059.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Revert "WIP: first try on adding max_b and max_t"

This reverts commit 3bc01ebaf2aa846fd16634e2d33b0d0f5803a076.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update docstring and fix max_seqlen logic for thd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* revert two lines of change in docstring
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: add get_max_b/t
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* fix max_seqlen code and docstring
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* sucess: add max_b/max_t
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove debug code
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change max_b/max_t buckets
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* fix b vs orig_b
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix b vs orig_b with 0 fill
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update FE for T3HD/TH3D
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add max_b to conversion kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* fix lint
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* fix changes after last merge
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* add Jax support for max_t
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* update FE to 1.8.0-rc
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update FE to 1.8.0
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* code review/formating fixes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* fix Stats shape for <9.6
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* return nullptr for offset_stats when cudnn < 9.6
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* add more version control
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

7fb22c37

24 Oct, 2024 2 commits

[JAX] XLA Custom Calls with FFI for FusedAttnFwd, Quantize, Transpose,... · 18c2234c

Hua Huang authored Oct 24, 2024


[JAX] XLA Custom Calls with FFI for FusedAttnFwd, Quantize, Transpose, ActLuFP8, LayerNormForwardFP8FFI, and LayerNormBackwardFFI (#1263)

* Add TransposeFFI, test passed
Signed-off-by: Hua Huang <huah@nvidia.com>

* Add ActLuFP8FFI; fix TransposeFFI
Signed-off-by: Hua Huang <huah@nvidia.com>

* Add QuantizeFFI
Signed-off-by: Hua Huang <huah@nvidia.com>

* Add FusedAttnForwardFFI and some unit tests
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Minor fix
Signed-off-by: Hua Huang <huah@nvidia.com>

* Add LayerNormForwardFP8FFI & LayerNormBackwardFFI
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Revise FusedAttnForwardFFI()
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Add FFI_CudaGraph_Traits

All tests passed, ready for merge
Signed-off-by: Hua Huang <huah@nvidia.com>

* Bug fix for FFI data type mismatch

Also add a safeguard on the entrance to FFI function
Signed-off-by: Hua Huang <huah@nvidia.com>

---------
Signed-off-by: Hua Huang <huah@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

18c2234c

[JAX] Fix correctness of JAX fused attention with CP and improve numerics... · 20c75295

Michael Goldfarb authored Oct 24, 2024


[JAX] Fix correctness of JAX fused attention with CP and improve numerics check in unit tests (#1282)

Fix correctness of JAX fused attention with CP.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>

20c75295

16 Oct, 2024 1 commit

fix assertion bug for SWA API in TE-JAX (#1242) · 43b9e1ee

Md Fahim Faysal Khan authored Oct 15, 2024



fixed assertion bug for SWA
Signed-off-by: Md Fahim Faysal Khan <mdfahimfaysa@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

43b9e1ee

15 Oct, 2024 1 commit
- Check for backend support in Jax context parallel fused attention test (#1227) · 20c55e46
  Michael Goldfarb authored Oct 15, 2024
```
Update test to check support for context parallel attention.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
```
  20c55e46
10 Oct, 2024 1 commit

[JAX] Expose sliding window attn to TE-JAX API (#1205) · 85e60e64

Hua Huang authored Oct 10, 2024



* Expose JAX sliding window attn API
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* No SWA in context parallel; fix RNG seed in test
Signed-off-by: Hua Huang <huah@nvidia.com>

* Handle SAW API discrepancy in cuDNN and Python
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Add SAW API for flax, all tests passed

Will update tests/jax/test_praxis_layers.py next
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Update test_praxis_layers.py for SWA, test passed
Signed-off-by: Hua Huang <huah@nvidia.com>

* Use tuple window_size; update for PR #1212
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Add and adjust some pytest.skip
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Revised following Reese Wang's comments

Still need further debugging:
FAILED test_fused_attn.py::TestFusedAttn::test_backward[NO_SWA-DROP_0.0-4-128-256-16-16-64-BF16-CROSS-KV_PACKED-NO_MASK-NO_BIAS] - AssertionError:
FAILED test_fused_attn.py::TestFusedAttn::test_backward[NO_SWA-DROP_0.0-4-128-256-16-16-64-BF16-CROSS-KV_PACKED-NO_MASK-POST_SCALE_BIAS-1HSS] - AssertionError:
FAILED test_fused_attn.py::TestFusedAttn::test_backward[NO_SWA-DROP_0.0-4-128-256-16-16-64-BF16-CROSS-SEPARATE-NO_MASK-NO_BIAS] - AssertionError:
FAILED test_fused_attn.py::TestFusedAttn::test_backward[NO_SWA-DROP_0.0-4-128-256-16-16-64-BF16-CROSS-SEPARATE-NO_MASK-POST_SCALE_BIAS-1HSS] - AssertionError:

These errors does not exist in the previous commit
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Fix no-SWA test case errors in previous commit
Signed-off-by: Hua Huang <huah@nvidia.com>

* Add Padding mask w/ sliding windows sanity tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use float32 for the reference code softmax calculation
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Hua Huang <huah@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Reese Wang <rewang@nvidia.com>

85e60e64

17 Sep, 2024 1 commit

[JAX] Context Parallel Attention with All-Gather (#1106) · 9101a78f

Michael Goldfarb authored Sep 17, 2024



Implementation of context parallel fused attention using all-gather.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>

9101a78f

19 Aug, 2024 1 commit

Stop using global mesh for custom_partitioning. (#1112) · ee541e83

Frédéric Bastien authored Aug 19, 2024


Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

ee541e83

08 Aug, 2024 1 commit

[JAX] Support non-deterministic algo for cuDNN FA (#1056) · 86f27e12

Reese Wang authored Aug 08, 2024



* Support non-deterministic algo
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine the helper function name
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move fixture to conftest.py
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

86f27e12

06 Aug, 2024 1 commit

[JAX] Reduce lowering time after cuDNN 90300 (#1032) · d74e65f5

Reese Wang authored Aug 07, 2024



* Support actlen = 0 after cuDNN 9.3.0
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add runtime_segment < max_segment tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

d74e65f5

03 Jul, 2024 1 commit

[JAX] Add experimental internal used THD(packed) fused attn API (#964) · 687697a7

Reese Wang authored Jul 03, 2024



* Integrate experimental ragged offset
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use per sequence based offsets
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Remove v/o_seq_offsets
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add FP16 sanity tests and remove forward tests from the automatically run tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance input checks
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Separate fused attn to 2 differnt APIs and add the docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add experimental to the docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add runtime segments check
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove finished TODO
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

687697a7

14 Jun, 2024 1 commit

Apply formatting (#929) · 9416519d

Kirthi Shankar Sivamani authored Jun 13, 2024



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

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

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

9416519d

13 Jun, 2024 1 commit

[JAX] Splitting cpp_extensions.py (#899) · 5986342a

Phuong Nguyen authored Jun 13, 2024



* Splitted cpp_extensions.py, renamed mlp.py and fused_attn.py
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* fixed import in tests
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

5986342a