GitLab

* 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>

* 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>

Add THD + GQA supports for cuDNN >= 9.6
Signed-off-by: Reese Wang <rewang@nvidia.com>

* 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>

Modify unit tests to work around cuDNN 9.4 regression.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>

Add the missing 1HSS tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* 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>

* 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>

* 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>

* 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>

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

* fix inconsistency for attn mask; now True means participating in attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix sliding window window_size for decoder+padding combination
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* revert paddle changes regarding mask
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* revert softmax to 1-mask;0-keep
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* enforce 1-mask out; 0-keep rule for jax masks
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* revert pytorch mask changes; some kept in tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* revert to jax fused attn on main
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* inverse mask logic for get_cu_seqlens/_and_indices in PyTorch implementation and mask generation in unit tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* temporarily disable update_weight_scale_inv
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* enforce window_size for decoder
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add docstring for mask definition 1-mask out;0-keep
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* tweak make_decoder_mask and make_mask in jax tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* skip dBias for shapes other than 1HSS; otherwise dq/dk/dv NaNs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* expand attn_biases from list to variables in save_for_backward
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix use of variable before assignment in jax dact_lu
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove window size definition for decoder
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add change notes in README for padding mask in PyTorch
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* tweak padding mask notes in README
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* expand list to tensors for save_for_backwards
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* Add layernorm_fp8_dot unit test
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update the softmax primitives support conditions
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add tests for the softmax primitives
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Round1 refactor of test_layer
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Split dropout arguments of ref code and add hidden/intermediate dropout elementwise comparison
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add dropout_braodcast_dim, self_attn_mask tests and clean a few code
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Abstract test layer and fix a rope reference code diff
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add epsilon and float32 tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add relpos_bias and attention dropout tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Loose the atol
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add doc string for test_layer
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add doc string for test_layer
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Avoid to left bias parameters in graph when use_bias=False
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

* Unify the fused attn workspace size cpp code
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Reduce the skipped cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename self/cross attention to qkvpacked/kvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update attention mask docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine the attn mask implementations
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* added support for arbitrary bias shapes for fused_attn
Signed-off-by: Alp Dener <adener@nvidia.com>

* Fix linting
Signed-off-by: Alp Dener <adener@nvidia.com>

* Add b1ss/bhss/11ss bias shapes when not requiring dBias
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add bias_b/h to plan cache
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fixed compile errors after PR653 merge
Signed-off-by: Alp Dener <adener@nvidia.com>

* updated JAX unittests for new bias shapes
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed mismatched mask type checking
Signed-off-by: Alp Dener <adener@nvidia.com>

* corrected skip condition
Signed-off-by: Alp Dener <adener@nvidia.com>

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

* corrected skip checks for bias shapes
Signed-off-by: Alp Dener <adener@nvidia.com>

* resolved test issues but neginf with float16 is still problematic with JAX
Signed-off-by: Alp Dener <adener@nvidia.com>

* new bias shapes passing TE JAX CI for seqlen <= 512, seq_q == seq_kv and h_q == h_kv conditions
Signed-off-by: Alp Dener <adener@nvidia.com>

* TE/JAX fused attn tests for new bias shapes passing with neg_inf=-2**27 for Bfloat16 and -2**15 for Float16
Signed-off-by: Alp Dener <adener@nvidia.com>

* code style fixes and test parameter ID cleanup
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed incorrect skip condition for backward fused attn test
Signed-off-by: Alp Dener <adener@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: Alp Dener <adener@nvidia.com>

* Refine MHA API
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Reuse func from the flax
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* qkv packed draft
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix test_layer with fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add attn_bias_type and enhance a few code flow
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move scale_factor from __call__ to init
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add DPA public API and tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

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

* Add qkv separate fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Apply BSHD_BSHD_BSHD format
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add fused attention layer tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Fine-grained fused attn settings
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove the default value of num_attetnion_head and head_dim
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add teardown for fused attn env
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Unify the Optional notation
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix Pre/Post scale bias comments
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add checkpoint_name for fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix the fused attn batcher
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

Signed-off-by: hugo-syn <hugo.vincent@synacktiv.com>

* Support num_gqa_groups arguments
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add GQA support on the JAX bridge code
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix the kv stride of the arbitrary backend
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Complete rewrite fused attention tests and add GQA coverage
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Support unfused GQA
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Calculate seqlen before the primitive for the better perf
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Apply code style checks for te_jax
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Apply code style checks for tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add num_gqa_groups doc
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Correct the variable naming
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Handle Max512 CAUSAL
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add WAR for the latest jax image
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Integrate cuDNN frontend v1 to fused attention and miscellaneous fixes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix jax/paddle for unit tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* simplify stride generation
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix and/or logic in get_backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove v.contiguous() since get_qkv_layout covers it
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* skip fp8 tests for sm89
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* revert mask type to comma-separated list
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* integrate v1/pre-release-5
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* cleanup prerelease5 integration and fix FA2.1 commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* force dropout to 0 if not training
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* testing bias/alibi and padding+causal; add alibi to unfused DPA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* set flag_arb to false when non determinism is not allowed
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* followup on prev commit; remove redundant python env var setting
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: minor tweaks for tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix determinism logic for fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

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

* remove use of non-deter var for backend selection
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix workspace size in bwd and uncomment bias test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* disable arbitrary backend + post scale bias in Jax; waiting on PR 525
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up bhsd order
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* swap bias/rng_state order in aux_ctx_tensor and add bias to aux_ctx_tensor in _qkvpacked/_kvpacked API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove support for padding_causal + cross for max512
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change alibi bias to float32 for bias_1_4/5 tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further clean up tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix thd fwd output shape for FlashAttention and add backend info for DPA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix definition of workspace limit when dbias is present
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further tweak DP_WORKSPACE_LIMIT definition
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* disallow alibi+no_mask for sdpa flash and update alibi tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update jax/paddle after PR525 and fix DP_WORKSPACE_LIMIT for dbias Jax tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* disable dbias for non-hopper archs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* remove build dir in setup.py
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change selection logic to prefer fused attn on sm90
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix h and s order in header
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update to cudnn fe v1 branch
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove manual setting of workopt path due to dbias after v1 update
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add post_scale_bias and alibi to sdpa flash support matrix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix support matrix in header files
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move headers back to .cu and change seed/offset to int64
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update Megatron commit in L1 test and remove all prints in fused attn test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix fp8 arg in L1 Megatron script
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* print only when debug flag is on
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove checkpointing loading to avoid loading other tests results
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* Add rng_state output for cross fused attention
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add rng_state and output for the flash attention backward
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add bias for the jax cross attn API
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add bias in the backward for the arbitrary fused attn backend
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

[C/JAX] Support more mask types for the arbitrary seqlen kernels and minor changes of JAX bias (#469)

* Move bias to float32
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Increase neg infinity abs values
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enable varlen tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

* Support variable sequence length after cuDNN 8.9.6
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use unique_ptr instead of shared_ptr
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add a new mask type: PADDING_CAUSAL_MASK
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Support flash padding mask after 8.9.6
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the Max512 handling for causal masking and add the related tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update the fused attn support lists
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove padding_aware from the caching
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Reduce the pad ratio tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix a bug with cuDNN 8.9.5
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Release backend resource after the module level unit test
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Clean the jax live arrays before running the unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

* Deprecate QKV_INTERLEAVED use in JAX
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Deprecate QKV_INTERLEAVED use in Paddle
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance qkv enum mappings
Signed-off-by: rewang <rewang@nvidia.com>

* Fix LD_LIBRARY_PATH issue
Signed-off-by: rewang <rewang@nvidia.com>

* Arbitrary seqlen kernels only support self attention currently
Signed-off-by: rewang <rewang@nvidia.com>

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

* Debug PyTorch and Paddle tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Only run Paddle layer tests with cuDNN fMHA on supported archs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug PyTorch fMHA tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Reduce JAX FP8 GEMM sizes

Avoid split-k kernels on Ada.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Disable JAX fused self-attention test on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Run supported fused attention tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Run supported fused attention JAX tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Enable Paddle fused attention on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update reference scale calculation in TensorFlow test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Restore backend support to reference FP8 attention impl in PyT test

Review suggestion from @cyanguwa
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix merge conflicts
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug Paddle tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Loosen tolerances for Paddle attention tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Assume causal mask implies equal seqlens in Paddle attention tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fused attention kernel only supports sm80 and sm90
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update transformer_engine/jax/csrc/modules.cpp
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* arbitary fused kernel supports sm86/sm89 after 8.9.3
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Forward is_fused_attn_kernel_available to cpp backend
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove cpp is_fused_attn_available API
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Fix flash attention dropout probability with inference
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add output as the fused attention ctx tensor
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add rng_state as the fused attention ctx tensors
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add flash attention supported lengths to the fused attention
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refactor attention primitive to reuse abstract shaped array
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Detect backend type to allocate appropriate ctx size
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Skip dropout correctness instead of return success
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use cudaMemsetAsync and enhance the error handling
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add flash attention kernel elts_per_thread update
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove redundant max 512 suffix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Keep only DType and remove NVTEDType from python
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Re-calculate workspace size for self attention
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance bias/dbias shape guard
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the seed/rng_state checker
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use jax.core.ShapedArray as jax.abstract_arrays is deprecated
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the unittest docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Enable fused attention dropout
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Cast the uint32 key/counter to int64
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update dropout support in fused attention docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revise devPtrCuSeqlen* to align the naming
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Support different Jax PRNG impls
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert CastAsync since it is not used
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Implement is_training for 16-bit fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add fused attn with dropout sanity unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the comments readability and rng_state checker
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change the attention dropout shape to align other frameworks
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Make encoder tests deterministic
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change the default seed for the jax encoder tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Maintain offset in TE
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the resource safety
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert rng_state type to allow only i64
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Handle the corner case for elts_per_threads calculation
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Populate rng state by kernels
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename rng_state as seed in cpp_extensions
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update the attention dropout comment
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Unfused scale+softmax if bias is present
Signed-off-by: Reese Wang <rewang@nvidia.com>

* WAR a causal masking + no_bias bug and add the unittests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix the optional args (bias) sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Disable fused attn in JAX by default, enable it with NVTE_USE_FUSED_ATTN
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add thread local for the plan cache
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename dbeta to dbias for the readability
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add scaled softmax with dropout test cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Updated NVTE_FUSED_ATTN variable name
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add fused attention unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use NVTE_* enums
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use NVTE_Mask_Type and remove FMHADescriptor
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move common functions to utils
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change namespace to fused_attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move fused_attn_max_512_fwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Move fused_attn_max_512_bwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove redundant blank line
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix a potential bug for cu_seqlen converter
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Refine the unfused attention warning message
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Remove the deprecated header
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Rename to fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add attention related mask
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add attn_mask_type and attn_bias_type
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refactor jax primitive API
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove dsoftmax from the customcall
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add None guard for bias and dropout_rng
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add version guard
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add is_fused_attn_kernel_available() to correctly dispatch the attention impl
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

* Add the missing blank lines
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change the order of FADescriptor members
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the readability of fused_attn_max_512.cu
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Generalize the input dimension unpacking
Signed-off-by: Reese Wang <rewang@nvidia.com>

* 16 bits fused attention requires 8.9.1
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update fused attention support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Handle None type when sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change to the padding ratio
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Performance optimization for non-bias cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert the cudnn-frontend PRIVATE keyword which was used for debugging
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert "Update fused attention support matrix"

This reverts commit 4effe67d0f08f733919a329ce5ab421958740f4a.
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Treat b * s as total_seqs to align ragged cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add FP16/BF16 max_seqlen <= 512 fused attention to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine test_fused_attn.py

* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Unify the cuDNN compile version
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add dropout to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Slightly adjust the headers
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Typo fix: remove redundant either
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Consolidating fused attention requirements
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error line report
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add more annotations to the custom calls
Signed-off-by: Reese Wang <rewang@nvidia.com>

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