GitLab

* Add basic infrastructure for Sequential module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add linear op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add FP8 support in linear op

Runs, but need to validate. Runtime errors with non-FP8 params and FP8 compute, or FP8 params and non-FP8 compute.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add reshape op and unit test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add bias op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add unfused linear op

Test does not pass with FP8.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug unfused linear op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add test for linear+bias op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add separate abstract classes for unfused and fused ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Consolidate unfused ops in submodule
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add linear-bias fused op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use fused cast-transpose in linear ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Disable GEMM+bias fusion with FP32 activations

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

* Add parallel unit test for unfused linear op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Refactor parallel tests to reduce job launches
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add all-reduce, all-gather, and reduce-scatter ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove unused file
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Add support for FP8 scale updates

Still need to implement amax reductions.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Fuse GEMM+bias in row TP

Add documentation for unfused ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename pipeline to fuser

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

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

* Preserve cached FP8 transpose between ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add option for fused wgrad accumulation
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Directly output FP8 from linear if needed
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix cuDNN front-end commit
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use updated FP8 tensor API for transpose caching
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use updated API for FP8 scale updates
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add tests for non-default FP8 recipes
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename UnfusedOperation to BasicOperation
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add unit test to check amax reduction with fusable op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Operator autograd state no longer needs to be initialized
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Initial functional implementation of linear op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug fused linear+bias op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove autograd context from functional linear impl
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use functional linear impl in fused linear+bias op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename subdirectory from "fuser" to "ops"

Avoid confusion with kernel fusers and graph compilers.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update with Float8Tensor changes in #820
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove unnecessary CPU overheads
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Correctly pass FP8 metadata from next op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Add convenience functions to manipulate Sequential class
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update name of PyTorch extensions module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Clear saved tensor data in linear op after bprop
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Update name of PyTorch extensions module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix test name in QA script
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update name of PyTorch extensions module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Run distributed tests even when only 1 GPU is available
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Only run distributed tests with 2 GPUs if there are >=2 GPUs
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



* Review suggestions from @sudhakarsingh27 and @ksivaman

Fix spelling of "fusible". Avoid "input" name in internal APIs.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update transformer_engine/pytorch/ops/__init__.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

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

Build for python < 3.8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove implicit padding and unpadding
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

Fix size mismatch error in fp8 transpose.
Signed-off-by: Dennis Liu <denliu@nvidia.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>

* removed libcuda.so link at compile time for TE/PyTorch extension
Signed-off-by: Alp Dener <adener@nvidia.com>

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

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



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

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

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



* updated get_symbol() in TE/common/cuda_utils.h to new impl based on cudaGetDriverEntryPoint
Signed-off-by: Alp Dener <adener@nvidia.com>

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

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



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

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

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



---------
Signed-off-by: Alp Dener <adener@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>

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

* enable unfused attn for cross attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* omit cudnn 9.1.1 and 9.2.1 due to bugs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* set cu_seqlens_padded to cu_seqlens by default
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* Revert "enable unfused attn for cross attn"

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

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

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

* Revert "remove mha_fill for FP8"

This reverts commit 83ffc44114dc6eb3d426d742b6c5a4d34805ec04.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* lower cudnn version to >=9.2.1
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>

cache was added in Python 3.9.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* fix compilation
Signed-off-by: Edenzzzz <wtan45@wisc.edu>

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

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

Signed-off-by: Edenzzzz <wtan45@wisc.edu>

* Fix CUDA 12.1 compilation bug
Signed-off-by: Wenxuan(Eden) Tan <wtan45@wisc.edu>
Signed-off-by: Edenzzzz <wtan45@wisc.edu>

---------
Signed-off-by: Edenzzzz <wtan45@wisc.edu>
Signed-off-by: Wenxuan(Eden) Tan <wtan45@wisc.edu>
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* GroupedGEMM via multi-stream cublas

* fix A/B is nullptr while D is not nullptr

* add fp8 grouped gemm

* register with TorchScript

* add the GroupedLinear layer

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Jiang Shao <jiangs@nvidia.com>
Co-authored-by: Qi Zhang <qizhang@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

Release GIL in PyTorch pybind11 functions
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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

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

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

* replace with_padding with padded
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 after merge
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 fix for fused attn fwd/bwd calls
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 Jax
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 spacing in docstring
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix pytorch tests; fix paddle api
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



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

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

* fix jax
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 attn with CP
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

fix tp_initialized error
Signed-off-by: Charlene Yang <8636796+cyanguwa@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>

* Add the option to use SM for P2P comm in TP overlap
Signed-off-by: Sangkug Lym <slym@nvidia.com>

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

* Python formatting with black
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Format C++ with clang-format
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



* Update transformer_engine/pytorch/csrc/comm_gemm_overlap.h
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

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

* subclass DPA with BaseModule and test with test_gpt_checkpointing
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

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

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

* add hook in case core_attention._extra_state is missing
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* check named buffers in BaseModule; remove FP8 scratchpad override function; test FP8 for sm90+
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: test size, interval in recipe, named_buffer loop
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move BaseModule from FusedAttention to DPA
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>

* A hot fix to disable CE deadlock check
Signed-off-by: Pavel Shamis (Pasha) <pasharesearch@gmail.com>

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

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



---------
Signed-off-by: Pavel Shamis (Pasha) <pasharesearch@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-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>

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

* remove linters, fix clang-format
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* revert accidental formatting changes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

* add fp8 attn include in the correct file
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* use subprocess instaed of os.sys
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add example script and update notebook
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* minor tweaks
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 related comments
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* move get_cudnn_version from common to pytorch utils
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* added DL framework callbacks for bootstrapping userbuffers without MPI
Signed-off-by: Alp Dener <adener@nvidia.com>

* removed userbuffers availability check in TE modules since userbuffers is now always compiled
Signed-off-by: Alp Dener <adener@nvidia.com>

* added comm+GEMM overlap example with LayerNormMLP
Signed-off-by: Alp Dener <adener@nvidia.com>

* lintin and review fixes
Signed-off-by: Alp Dener <adener@nvidia.com>

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

* added header guards
Signed-off-by: Alp Dener <adener@nvidia.com>

* removed defunct userbuffers checks in build_utils and setup.py
Signed-off-by: Alp Dener <adener@nvidia.com>

* added exposed API in modules/base.py to __all__
Signed-off-by: Alp Dener <adener@nvidia.com>

* removed transformer_engine/CMakeLists.txt and shifted all TE/common compile into transformer_engine/common/CmakeLists.txt
Signed-off-by: Alp Dener <adener@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Add norm_factor arg into DotProductAttention
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

* Change kwarg name from `norm_factor` to `softmax_scale`
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

* Change all norm_factor representation into softmax_scale
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

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

* Update attention.py changing typo
Signed-off-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>

---------
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

reverting autocast API back to PyTorch v2.3.1 and below
Signed-off-by: Alp Dener <adener@nvidia.com>

* expose multi_tensor_* kernels
Signed-off-by: Xin Yao <xiny@nvidia.com>

* fix lint
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>

skip switching to nvfuser for torch >= 2.2
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

added @torch._disable_dynamo fixed deprecation warnings with torch autocast API for TE checkpoint
Signed-off-by: Alp Dener <adener@nvidia.com>

restricted fsdp asserts on primary fp8 weights to TE modules
Signed-off-by: Alp Dener <adener@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>

---------
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Avoid select operation in cast-transpose extension
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Avoid select operation in cast-transpose-dbias extensions
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Avoid select op in LayerNorm and RMSNorm
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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

* New TE wrapper for PyTorch FullyShardedDataParallel to make TE modules distribute their activations after the forward pass and gather them before the backward pass
Signed-off-by: Alp Dener <adener@nvidia.com>

* simplified TE module setup for FSDP comms
Signed-off-by: Alp Dener <adener@nvidia.com>

* FSDP scatter/gather for tensors saved into autograd ctx now working for base TE modules
Signed-off-by: Alp Dener <adener@nvidia.com>

* make sure activation recompute disables FSDP scatter/gather
Signed-off-by: Alp Dener <adener@nvidia.com>

* make sure Fp8 weight buffers are sharded at the end of the backward pass and gathered before forward
Signed-off-by: Alp Dener <adener@nvidia.com>

* Fixed typo in attribute name
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed bug in finding FSDP-wrapped TE modules
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed typo in fp8 weight tensor name
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed incorrect # of gradients
Signed-off-by: Alp Dener <adener@nvidia.com>

* Added fp8 amax gradient hook tensor to the parameter reset
Signed-off-by: Alp Dener <adener@nvidia.com>

* get rid of erroneous dummy tensor leftover from incorrect rebase
Signed-off-by: Alp Dener <adener@nvidia.com>

* Linting fixes
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixing git snafu and removing debug statements
Signed-off-by: Alp Dener <adener@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Remove interval arg from recipe
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Remove usage of interval and use explicit kwarg for testing recipes
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>

* Modify CUDA graph tests to use grad accumulation steps
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Initialize grad buffers before capturing CUDA graph in CUDA graph tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Only use BS=2 in CUDA graph tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update tests/pytorch/test_cuda_graphs.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

Replace int8_t in PyTorch extensions with int64_t
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* add seq_offsets_o and use new offset calculation
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* addition to previous commit; fix unit test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* WIP: test padding between sequences
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: fix tests for padding between sequences
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix tests for sbhd/bshd layouts; clean up
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update cudnn-frontend and add tests for max_seqlen_q=1 and d=256 for inference
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* test sbhd/bshd layouts for sq1, d256 inference case
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* replace wording from accumulative to cumulative
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add offset tensors to custom fp8 mha tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add sm>=90 constraint for thd support
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix cuDNN support for sq=1, d=256
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* modify cudnn version and restrict MQA/GQA support for THD
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add dummy tensor to pass jax build
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add dummy tensor to pass paddle build
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>

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

* add multi-tensor kernels
Signed-off-by: Xin Yao <xiny@nvidia.com>

* add FusedAdam
Signed-off-by: Xin Yao <xiny@nvidia.com>

* add test to qa
Signed-off-by: Xin Yao <xiny@nvidia.com>

* add FusedSGD
Signed-off-by: Xin Yao <xiny@nvidia.com>

* fix lint
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Initial refactor of FP8 workspaces in Linear module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove extra kernel launch
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Minor perf optimizations

Tensor base class functions in Float8Tensor have significant overhead.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Refactor FP8 workspaces in LayerNormLinear and LayerNormMLP
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Document FP8 workspace function
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Revert changes to FP8 recipe tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add support for lazy FP8 transpose caching

Previous caching behavior (always fill cache) incorrectly filled cache during CUDA graph warmup steps.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Debug ONNX export

ONNX FP8 cast ops assumed that FP8 scales were created during model export (i.e. not initialized during training).
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Make sure Float8Tensor.transpose_2d is backward compatible
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Revert changes to ONNX export operations

Work around ONNX test failures by filling FP8 scale tensors instead of copying into them.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug scale factor update in Float8Tensor transpose_2d
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

Make sure RoPE frequencies are in FP32
Signed-off-by: Tim Moon <tmoon@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>

TE checkpoint now preserves the torch autocast context from the forward pass during the recompute phase
Signed-off-by: Alp Dener <adener@nvidia.com>