GitLab

* Create pytorch/dot_product_attention module and pytorch/d_p_a/utils.py
Move attention logging into a separate class in pytorch/d_p_a/utils.py
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

* Create FlashAttentionUtils class in pytorch/d_p_a/utils/py for versioning info
Move versioning info out of pytorch/attention.py
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

* Move AttentionParams and get_attention_backend from attention.py to d_p_a/utils.py
Fix tests and imports for the above refactor change
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

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



* Move get_qkv_layout(), get_full_mask(), get_alibi(), get_attention_quantizers() to d_p_a/utils.py
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

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



* Move tensor packing and unpacking helper functions from pyt/attention.py to d_p_a/utils.py
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

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



* Move cumulative seqlens and indices methods from pyt/attention.py to d_p_a/utils.py
Rename cumulative functions from using _cu_ to using _cumul_ to differentiate from CUDA cu calls protocol
Rename tensor packaging methods with leading underscore to make them as internal to file
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

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



* Remove unnecessary imports in pytorch/attention.py and d_p_a/utils.py
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

* Create d_p_a/inference.py and move InferenceParams from pyt/attention.py to it
Modify tests and other files to import InferenceParams correctly
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

Modify docs api for InferenceParams
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

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



* Create d_p_a/rope.py and move RoPE methods from  pytorch/attention.py to it
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

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



* Code cleanup
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

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



* Fix qa testing induced bug
Code clean up
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

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



* Fix incorrect pack_tensor arg type
Code clean up
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

* nit: Resolve lint errors
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

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



* Remove typedef FAUtils for FlashAttentionUtils
Use attn_log instead of att_log
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

Fix lint error
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

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



* nit: Fix the function name from get_cumul to the earlier get_cu
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

* nit: Fix typos, explicit imports and remove extra comments
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

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

* fix recompilation of out and lse correction in p2p+bshd/sbhd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix recompilation of get_seq_chunk_ids_for_reordering
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* fix recomplilation of reorder_seq_chunks_for_a2a
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* recover a change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor change to softmax_lse correction
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* cache cu_seqlens for BSHD/SBHD format
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* do not need to allocate out buffer for BSHD/SBHD
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code refactoring
Signed-off-by: Xiaowei Ren <xren@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: Xiaowei Ren <xren@nvidia.com>

* refactor init out correction
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix a docstring
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code refactoring
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix init out correct dtype
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add pad_between_seqs to DPA API
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add pad_between_seqs to the API of MHA and transformer layer
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add pad_between_seqs to the API of MHA and transformer layer
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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

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

* change API for hierarchical CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* move fp8 code before qkv reshape
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* try to insert A2A for hierarchical CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make fwd work
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove a redundant sync
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make bwd of hierarchical CP work
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix dout a2a in bwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix q_f16 with fp8
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* assert hierarchical CP implementation does not support THD format
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* bug fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* assert hierarchical CP does not support attn bias
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add unit test for hierarchical CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix cp_comm_type in unit test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* bug fix and code cleaning
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* an assert info change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* dout shape fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* move function definitions to the front of the first call
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix tensor view comments
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* refine CP unit test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* save cp_size_a2a and rank_a2a in fwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add more explainations of cp_group in doc_string
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* CPU perf optimization in linear autograd function

Avoid enable_grad context when possible in cast function. Cache distributed group properties.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* CPU perf optimization in prepare_forward function

Avoid torch.nn.Module impl of __setattr__.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Avoid module import in TE module forwards
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use fast getter for params
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Reuse tensor dims in linear autograd func
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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



* Apply optimizations to grouped linear
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug test failures
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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



* Debug test failures
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Avoid deepcopy in tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move _fast_setattr logic to __setattr__ method
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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



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

* Docs fixes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* docs fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* docs fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* clean code for CP function args
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add a placeholder for Ulysses implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* commit code change to CP+A2A
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* finish the draft fwd implementation of Ulysses
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add draft bwd implementation of Ulysses
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make swa work with ulysses
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* commit FP8 code for Ulysses
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix qkv type in the bwd of FP8+CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix qkv_dtype of FP8+CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code refactoring
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor code change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* config cp correction dtype of FP8+CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* code style change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* save chunk_ids
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* try to make Ulysses A2A async
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make more a2a async
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix a2a_outputs
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix chunk_ids generation for A2A
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* avoid code duplication of a2a before attn
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove code duplication of a2a after attn
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add cp_stream in A2A implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* bug fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix qkv of fp8_fwd + bf16_bwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix kernel order in cp a2a communication
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code cleaning for CP a2a
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix merging with main
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix a2a communication order
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* adjust sequence chunk reordering for a2a
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add docstring for A2A implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* change an assert info
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add unit tests of A2A implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add more A2A unit test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix CP unit tests
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add more cp unit tests
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix window size of no_mask
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fused attn does not support swa+no_mask
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* change num_gqa_groups to 2 for A2A implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* function and variable renaming
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code cleaning for CP all-gather implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* some function renaming
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove redundant code
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* commit code change for kv all-gather implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix all-gather implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add a window size check
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code cleaning
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add unit test of all_gather+no_mask
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix all-gather cp implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code cleaning
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* code format fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code format fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix FP8 with A2A implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add paper references to CP implementations with all-gather and all-to-all
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* change pdf to abs
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* elaborate cp_comm_type
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix CP docstring
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* add support for padding in UnfusedDPA
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 support for padding_causal/_bottom_right
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix padding_causal/_bottom_right
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



* need to test max512 backend
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 last commit
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 mask logic in unfused
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



* use actual_seqlen for alibi/causal_bottom_right padding
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



* minor fixes and convert causal to causal_bottom_right for inference
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



* use causal in kv cache inference test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* simplify get_alibi logic
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



* simplify the non-padding path for get_alibi
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



* avoid batch_size loop in generating padding_causal/_bottom_right masks
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



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

feat(pytorch): Allow TransformerLayer and MultiheadAttention to accept sequence length parameters (#1066)

* Added ability for seqlen for transformer and mha layer
Signed-off-by: Lukasz Pierscieniewski <lukaszp@nvidia.com>

* Documentation for new parameters
Signed-off-by: Lukasz Pierscieniewski <lukaszp@nvidia.com>

* Add tests for THD layout, assert for THD layout with KV-Cache
Signed-off-by: Lukasz Pierscieniewski <lukaszp@nvidia.com>

* Fixed tests
Signed-off-by: Lukasz Pierscieniewski <lukaszp@nvidia.com>

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

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



* Move THD logic in shape calculation, add missing optional in params
Signed-off-by: Lukasz Pierscieniewski <lukaszp@nvidia.com>

* Skip the THD test on GPUs older than Ampere
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

---------
Signed-off-by: Lukasz Pierscieniewski <lukaszp@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Przemek Tredak <ptredak@nvidia.com>

* add window_size to AttnFuncWithCP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add seq_offsets_qkvo for cudnn thd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add seq_offsets_qkvo to AttnFuncWithCP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix seq_offsets calculation of cudnn thd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove a thd assert
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix bias for thd test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add thd test for cudnn FA with CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* skip GQA/MQA test for cuDNN THD
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make sure seq_offsets are computed with qkv_group of hd_hd_hd while CP>1
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix seq_offsets inputs
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove two comments
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix attn mask type for cudnn thd with cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix attn_mask_type check
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix attn_mask_type for cudnn fa with thd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix a typo
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix out dout in bwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* assert cudnn+thd does not support attn bias
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* check if attn_mask_type has padding
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* change cp test batch size to 2
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix code format
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix two assert info
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix assert comment
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix assert comments
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix assert comments
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* assert swa+CP cannot work with thd format
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add a new CP function for swa
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add a missing dgrads
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add draft fwd function for swa+cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* enable flash attention for swa+cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove an assert of swa+cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* call SWAFuncWithCP for swa+cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* use 2hd layout
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* change qkv_format check
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add a code comment
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* tensor shape bug fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* tensor shape fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add function to compute cu_seqlens of a cp rank
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add cu_seqlens and cu_seqlens_padded to context parallelism
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix FlashAttention output sequence length
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix cu_seqlens_kv_per_step calculation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* zero dQKV for ending padded tokens
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* zero dQKV tensors of FlashAttention
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix softmax_lse correction
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove padded tokens of KV to save comounication
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* do not need to zero dkv for FlashAttention any mroe
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* zero out tensors
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove redundant code
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix CP unit test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix kv shape of cp test with thd format
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* update cp unit test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* add simple code framework
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* try not to have a separate CP function for SWA
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* backup some code change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* back up code
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* clean up fwd implementation of SWAFuncWithCP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove redundant code
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code cleaning
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix assert info
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* reduce kv chunk concat overheads
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make AttnFuncWithCP and SWAFuncWithCP have same API
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add a docstring
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* preliminary implementation of SWAFuncWithCP forward seems working
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix output shape of SWAFuncWithCP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code refactoring for FlashAttention and add a code placeholder for bwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* use gather_along_first_dim
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* finish the preliminary implementation of bwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove redundant code
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix assert condition
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add draft implementation of SWA+CP with FusedAttention
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix attention mask type of swa+cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code cleaning
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add qkv_layout
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add missing window_size argument
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix kv shape of swa+cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* bug and typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix dout shape
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add multi stream in fwd of swa+cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* save chunk_ids_to_kv_ag in fwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add multi stream in bwd of swa+cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor fix to cp stream sync
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* rename AttnFuncWithCP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* check if window size is None
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix docstring of AttnFuncWithCP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add env var for users to choose KV ag or KV p2p
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* update cp tests
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix window size in cp unit test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix pytest skip messages
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add cp_comm_type into API
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* code cleaning
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* assert sequence length divisible requirements
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* add support table of context parallelism
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* typo and code format fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* do not print multiple disabling messages
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* bug fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* fix device in torch.arange and adjust code for the PR of MLA
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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



* fix typos and clean asserts
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

---------
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Co-authored-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>
Co-authored-by: Xiaowei Ren <xren@cs-cw-dfw-login-01.cm.cluster>

Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* add set_deterministic and minor fixes for swa
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* change window_size to int64_t; fix swa/determinism tests; cache _attention_backends
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add window_size to get_backend; fix jax and paddle
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fixes; add set_deter to bwd_impl
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix FP8 tests due to determinism
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add support matrix for SWA and bias
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



* minor fixes and 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



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

* add wording on window_size special cases
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor tweak on wording
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* call bwd with deterministic=true for jax/paddle
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add determinism words in documentation
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>

* update to FE 1.5.1 and add bottom right causal
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



* adjust logic for backend selection
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.5.2
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add get_attention_backend function
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 get_attention_backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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



* tweak get_attention_backend and fix unit tests
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



* minor fixes for unfused, get_backend, etc
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 transformer_engine/pytorch/attention.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* explicitly skip FP32 and padding tests because there is no support
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fix for window size check
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update check_set_window_size and add enc_dec_attn_mask_type/enc_dec_window_size
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



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



---------
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>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Remove optional UB build leftovers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* rm unused import
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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



---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.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>

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

* Fixed Llama tutorial. Changed batch size and added fused=True.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Tutorial updated but not complete yet.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Tutorial notebook reseted - removed fuse=true
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Removed fused=true
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Batch size back to 8
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Typo and commented out line
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* fixed whitespace
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* fixed whitespace
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Added comment to attention line. Fixed potential bug with loading weights - now loading works correctly, confirmed by the generation code.
Signed-off-by: root <root@ipp2-1661.nvidia.com>

* Comments
Signed-off-by: root <root@ipp2-1661.nvidia.com>

* Models cast added again
Signed-off-by: root <root@ipp2-1661.nvidia.com>

* Weight download info
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Moved parameter gate_proj_size to config
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* gate_proj_size removed and put immediate_size instead
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Llama 3 added to tutorial
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Typos fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Typos fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Fixed model loading
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Loading fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Different dim for attention
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Reversed other commit
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Changed name to kv_channels
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Fixed typo
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Back to kv_channels in transformer layer
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Back to kv_channels in transformer layer
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Small bug fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Small bug fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Test fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* changed file modes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* lint fix and resolved conflict
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* lint fix and resolved conflict
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Lint fix, hopefully last
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>
Signed-off-by: root <root@ipp2-1661.nvidia.com>
Co-authored-by: root <root@ipp2-2373.nvidia.com>
Co-authored-by: root <root@ipp2-1588.nvidia.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: root <root@ipp2-0037.nvidia.com>
Co-authored-by: root <root@ipp2-1661.nvidia.com>
Co-authored-by: root <root@ipp2-2371.nvidia.com>
Co-authored-by: root <root@ipp2-1589.nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-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>

* added squared relu in te-torch
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

* FP8 cuda graphs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>

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

* exclude torch compile from numerics tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* More numerics fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* rm fusion from unfused path
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>

* Enable DGRAD RS overlap
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>

* fix lint; apply suggestions
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* TP-RS overlap with send/recv

Atomic GEMM based TP-RS overlap with send/recv
Signed-off-by: Sangkug Lym <slym@nvidia.com>

Specify userbuffer overlap method of each overlap instance
Signed-off-by: Sangkug Lym <slym@nvidia.com>

P2P TP-RS overlap with fp8 GEMM outputs
Signed-off-by: Sangkug Lym <slym@nvidia.com>

Fix TP-RS overlap with send/recv
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* cleanup
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* cleanup
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* linting
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* fix typo
Signed-off-by: Sangkug Lym <slym@nvidia.com>

---------
Signed-off-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Added QuickGELUActivation from HuggingFace/Transformers to common and pytorch
Signed-off-by: Alp Dener <adener@nvidia.com>

* Removing 'qgelu' from double-size activations list in LayerNormMLP.
Signed-off-by: Alp Dener <adener@nvidia.com>

* indent fix
Signed-off-by: Alp Dener <adener@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>

* Pass knobs for TP comm overlap instead of env vars
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>

* Comment out debugging print
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>

* Remove docstring
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>

* Remove debugging output
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>

---------
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>
Co-authored-by: Jaemin Choi <jaeminc@nvidia.com>

* test alibi between fa and fu
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move alibi slopes and bias to global to avoid repeating calculation
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix alibi slopes/bias generation
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix _is_flash_attention_supported to allow alibi type
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* disable padding mask when alibi is used for fused attn arbi backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add support for custom [n_heads] alibi_slopes in flash, fused, unfused attention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove alibi_type=none tests as they are unnecessary
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* change bias/dbias shape to allow b,1/1,h/b,h in arbi backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* tweak tests for arbi post_scale_bias [1,h,s,s] or alibi_slopes [n_heads]
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change bias/dbias shape in max512 backend - incomplete
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove max512 changes from last commit and disable max512 (and arbi temporarily) for [b, h, s, s]; pending cuDNN backend support
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up and tweak backend selection logic
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace || with () in docstring
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* combine slopes/bias generation to one function get_alibi() and fix alibi tests
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 PR557 bugs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Update transformer_engine/pytorch/attention.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

* encapsulate global alibi tensors into a dict cache
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* reduce alibi slopes test size
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* use dBias shape to define bias_b/bias_h because jax materializes dBias rather than Bias in bwd abstract
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>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* make TransformerLayer accept a `bshd` or `sbhd` tensor format
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Fixes from feedback
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* more feedback fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* remove incorrect info from docstring
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix from feedback
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

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

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

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

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

* change bert test to causal as unfused does not support padding
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix FlashAttention for v2-2.3 versions
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* verify FA swa works
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix mask related restrictions and duplicate code after merge
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* move repeated code into a function
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add determinism filter and fix FA warning message
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

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

Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

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

* [PyTorch] TransformerLayer: add parallel_attention_mlp to support Falcon models
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* [PyTorch] add test for parallel_attention_mlp to test_numerics
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* [PyTorch] TorchGPT: fix dropout for parallel_attention_mlp

Now uses nn.functional.dropout because depending on the path there are one or two dropouts.
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* Apply suggestions from code review
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* [PyTorch] test_gpt_accuracy: fix spelling in construction of TorchGPT
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

---------
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

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

* rename set_context_parallel_running to set_context_parallel_group
Signed-off-by: xren <xren@nvidia.com>

* bug fix
Signed-off-by: xren <xren@nvidia.com>

---------
Signed-off-by: xren <xren@nvidia.com>

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

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

* Initial commit
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Repro for RS output mismatch with Single GEMM + Split pipelined RS
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* minor changes for AG->GEMM pipelined overlap
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Add Atomic Gemm cublasApi attributes and initial implementation of AG->Atomic GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* AtomicGemm+RS functional with workaround
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* add amax update to layernorm_linear for FP8 unit test accuracy
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Enable reducescatter2_userbuff_strided variants
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Bug fix
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* AG+AtomicGemm overlap functional but gemm doesnt overlap with comm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Add userbuffers_sendrecv kernel variants
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* TransformerLayer API changes to enable AtomicGemm+RS overlap
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Code cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Code cleanup2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* [UB] AllGather Atomic GEMM overlap using userbuffer_sendrecv kernels
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Code cleanup + bug fix for multiatomic sendrecv kernel
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Bug fixes
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* [UB] Add shuffling for better AG AtomicGEMM overlap
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Bug fix for AG AtomicGemm overlap
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Bug fix for multiAtomicAG and singleAtomicAG
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Use chunk_i+1 as recv_chunk for multiatomic_AG with shuffling
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Launch AtomicGEMM after first-chunk AG
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Rebase to main
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Add FP8 ReduceScatter kernels, AtomicGEMM+FP8 RS not functional
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Revert "Add FP8 ReduceScatter kernels, AtomicGEMM+FP8 RS not functional"

This reverts commit 80a47a76355440cd5fb4314c96fe9fda632d87f9.
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Add support for NVLS-MC and FP8 Reduce Scatter
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Bug fix
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Atomic and Multiatomic FP8 RS functional
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Remove debug print
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* UB comm initialization hang fix
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Code cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Create new GEMM API for Atomic GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

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

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

* Bug fix
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Revert NVLS-MC
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Check cu* versions for running atomic gemms
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Add experimental warning
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Add warning to c api
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* initial changes [wip]
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add padding mask support for FA
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Address review comments
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* rm causal mask from tests and add padding
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix some conflicts
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Add unpadding mask
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix padding mask
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* [wip] fix API
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add packing and unpacking
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* fix atomic_add bf16 torch.compile
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Generate non all True masks
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Fix core attention export and FusedAttn filter
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix all ONNX tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Optimizations and caching fixes in torch.dynamo
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Fixes and reviews
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

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

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

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

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

* more changes; WIP
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* more changes; WIP
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fixes and wrong results
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* sb3hd/bs3hd working on top of 3xsbhd/bshd/thd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* remove qkvso_strides on torch side; cover it in generateQKVStrides
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* all 15 layouts pass
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* minor fixes and test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* removed most debug info/clean up
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add note to deprecate some qkv layouts
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix code for unit tests in test_fused_attn.py
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove a couple more comments
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* fix onnx for core attn; not fixed
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove nvtx and add env var for workspace opt
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove testing for env var
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace zeros/zeros_like with empty/empty_like
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix nvtx marker name for _q_k_v API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove sm80 when compiling for h100
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add mapping from qkv layout to layout group and qkv format
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up enums mapping and remove trailing spaces
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* simplify workspace opt control logic; only need env var
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix fp8 test, and minor modifications for other tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* avoid overwriting model configs in unit test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* random fixes/improvements: get_qkv_format/etc, default values, docstrings, comments
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* change workspace opt logic back to FORCE_WORKSPACE_OPT
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix FP8 tests and generateStrides function
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix get_backend logic for max512/arbitrary
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix unit tests; need cleanup
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up unit tests for layouts, and fix minor lint issue
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor tweaks for CI testing: onnx string issue and test fused attn first
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove one unsupported layout from max512 and add a check to qkvpacked API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix te layer test; reduce test time
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* revert compiler option changes; add back sm80 for even h100
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove some unit tests or make them optional to reduce CI time
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove more unit tests temporarily
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove _q_k_v in naming and add NVTE_ERROR for FP8 Aux_CTX_Tensors size checks
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove temp tests from last commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace with te::getenv
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove redundant contiguous()
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove thd->bs3hd user warning to avoid GPU sync
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* adjust fused attn bs in tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* temporary fix for onnx issue; more fixes in PR 437
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove unused variables
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: Charlene Yang
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>