GitLab

* Do not include logging macros in installed C headers
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Debug C++ tests

Use Google style for header includes.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update CUDA driver macros

Incorporating changes from #389.
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use core error checking macros in PyTorch extensions

Hack to get around macro redefinition warning.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix missing arg when getting CUDA driver error string
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Reuse logging header in frameworks
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@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>

* 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 zero_center_gamma/functional pass
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

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

* Refactored code style for improved readability and consistency
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Docs enhancement for zero_centered_gamma
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add escape for line break and remove some bad if conditions
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

* add building workflow for jax modules
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* replace bit_cast with reinterpret_cast
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add nvtx to cmake check list
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor layernorm fwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor rmsnorm fwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor layernorm_bwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* set pytorch as default in setup.py
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* rename extension from *.cc to *.cpp

cpplint cannot recognize *.cc file, so rename the extension
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor style, to align TE/PyTorch
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add pybinding, unittest and qa
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix license
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* disable c-extension-no-member and no-name-in-module
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add dataclass avoid pylint error
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update transformer_engine/__init__.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* Update tests/jax/test_custom_call_shape.py

fix typo
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* Update tests/jax/test_custom_call_shape.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* add building workflow for jax modules
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* replace bit_cast with reinterpret_cast
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add nvtx to cmake check list
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor layernorm fwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor rmsnorm fwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor layernorm_bwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* set pytorch as default in setup.py
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* rename extension from *.cc to *.cpp

cpplint cannot recognize *.cc file, so rename the extension
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor style, to align TE/PyTorch
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add pybinding, unittest and qa
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix license
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* disable c-extension-no-member and no-name-in-module
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add dataclass avoid pylint error
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update transformer_engine/__init__.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* Update tests/jax/test_custom_call_shape.py

fix typo
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* Update tests/jax/test_custom_call_shape.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* fix conflict due to PR62
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix c-extension-no-member and no-name-in-module

1. add transformer_engine_jax into extension-pkg-whitelist
2. convert pylintrc from CRLF to LF format
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update setup.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* remove pylint:disable and refactor import order
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

---------
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>