GitLab

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

* Initial refactor; linker error
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix linking issue and make mpi conditional
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Use max SMs at the last RS chunk in pipelined overlap
Co-authored-by: Sangkug Lym <slym@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Make userbuffers support opt-in

Decouple userbuffers from MPI. Refactor MPI handling in build system. Standardize names to "userbuffers".
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

* Add FP8 support for Ada
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

* better message for no fp8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* same thing for onnx test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Add tensorflow build

Improve build instructions

Fix pybind enum usage

Fix Python_EXECUTABLE cmake var

Move scale_inv calculations to FW
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Apply clang-format
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Format python files
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add TF build CI
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Lint checks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Another round of lint checks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Fix TF image tag
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Use the existing recipe file
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add license claim blocks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Fix a bug about bias dtype conversion
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add mnist example and cleanup old examples
Signed-off-by: kaixih <kaixih@nvidia.com>

* Autopep8 the tests
Signed-off-by: kaixih <kaixih@nvidia.com>

* Autopep8 the examples
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add example in Readme
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add unit tests and linting for TensorFlow
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add causal mask for non-fused case
Signed-off-by: kaixih <kaixih@nvidia.com>

* Fix the mismatched TF vs TE masks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Addressing CI tests
Signed-off-by: kaixih <kaixih@nvidia.com>

* Run lint test
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add missing import
Signed-off-by: kaixih <kaixih@nvidia.com>

* Skip fp8 tests for pre-Hopper GPUs
Signed-off-by: kaixih <kaixih@nvidia.com>

* Remove non-pytest tests
Signed-off-by: kaixih <kaixih@nvidia.com>

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

---------
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: kaixih <kaixih@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.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>