GitLab

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

* make d_out contiguous for bwd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove cudnnDestroy to let torch handle it
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>

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

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

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

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 workspace optimization for arbitrary_seqlen fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add use_workspace_opt to cudnn plan cache and fix workspace estimate
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* modify workspace opt logic; move zero fill to FP8 API only; other minor 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 try/catch
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix std string error when input is nullptr
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* Add = for required vs allowed workspace comparison
Signed-off-by: cyanguwa <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 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>

* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix docs for fused-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>

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