GitLab

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

* use separate qkv
Signed-off-by: jaywan <jaywan@nvidia.com>

add support for GQA
Signed-off-by: jaywan <jaywan@nvidia.com>

minor changes
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

change rtol
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

fix reshape issue
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

add rmsnorm and rotary position embedding
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

update rmsnorm
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

refactor layernorm and rmsnorm
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

support swiglu
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

add fused rope
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

minor changes
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

add rope api to __init__
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

minor changes
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

fix fp8 dtype issue
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* simplify ut cases
Signed-off-by: jaywan <jaywan@nvidia.com>

* Update transformer_engine/paddle/layer/attention.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Shijie <505749828@qq.com>

* fix name issue
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

---------
Signed-off-by: Shijie Wang <jaywan@nvidia.com>
Signed-off-by: jaywan <jaywan@nvidia.com>
Signed-off-by: Shijie <505749828@qq.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Implement fused kernel for FP8 scale update
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add fused kernel for amax and scale update

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

* Replace paddle.fluid imports with paddle.base
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move fused kernel to core library
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Use FP8 update kernel in Paddle
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug FP8 scale update in Paddle
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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

* Make update kernel in-place for PyTorch
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Revert cudnn-frontend commit
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* use separate qkv
Signed-off-by: jaywan <jaywan@nvidia.com>

* add support for GQA
Signed-off-by: jaywan <jaywan@nvidia.com>

* minor changes
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* change rtol
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* fix reshape issue
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

---------
Signed-off-by: jaywan <jaywan@nvidia.com>
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* Actively free tensor in bwd
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* - Add inplace support for fp8 casting
- Allow skipping weight update in fp8 meta update
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Support weight caching for Linear
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add weight caching for LayernormLinear
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add weight caching for LayerNormMLP
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add weight caching for Transformer layer
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add PP unittests
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix CI
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

---------
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add SP for linear
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add SP for LayerNormLinear
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add SP for LayerNormMLP
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add SP API for transformer layer
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add sequence_parallel attr
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add SP unittests for Transformer and Attention
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix compatibility with PaddleNLP
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

---------
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* fix cudnn FA softmax shape
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* set inplace rng_state
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

---------
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* Deprecate QKV_INTERLEAVED use in JAX
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Deprecate QKV_INTERLEAVED use in Paddle
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance qkv enum mappings
Signed-off-by: rewang <rewang@nvidia.com>

* Fix LD_LIBRARY_PATH issue
Signed-off-by: rewang <rewang@nvidia.com>

* Arbitrary seqlen kernels only support self attention currently
Signed-off-by: rewang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: rewang <rewang@nvidia.com>

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

* fix mask conversion and rng_state
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* refactor fused attn
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* use CUB to do prefix sum
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* fuse dropout add
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* minor changes
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* optimize kernel
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

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

* Change scaling factor from E8M0 to E8M23
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Eliminate amax_and_scale_update bubbles
Signed-off-by: rewang <rewang@nvidia.com>

* Add CUDA check
Signed-off-by: rewang <rewang@nvidia.com>

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

* Add FP8 support

- Add FP8 recipe
- Add FP8 path for nn layers
- Add MNIST FP8 example
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Update README
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix LayerNormMLP FP8 backward
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix FP8 training in float32 accumulation
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix FP8 checkpointing for non forward execution cases (same as #323)
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Refactors and improvements for better code stype, readability and organization
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Remove unnecassary pylint override
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

---------
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* add flash attn tests
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* update flash attn
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* fix random seed
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

---------
Signed-off-by: Shijie Wang <jaywan@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* add more ops
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* add skipif
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* fix bug
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* minor change
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* minor change on coding style
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

---------
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* Add cast_transpose

Add gelu, gelu_fp8

Add cast_transpose_bgrad_dgelu

Add layernorm_fwd and layernorm_fwd_fp8

Add layernorm_bwd
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix missing header
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

---------
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* First step of PaddlePaddle integration
- Add build option for paddle
- Add basic test framework
- Add 3 basic operators: cast_from_fp8, cast_to_fp8, gemm
Signed-off-by: Tian Zheng <tizheng@nvidia.com>
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix review comments
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Support paddle build
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add paddle build support for new building framework
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix review comments
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Clean up build process for Paddle stub file
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Minor fixes
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix pylint "wrong-import-order" warning
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix review comments
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Skip BF16 GEMM tests for unsupported arch
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

---------
Signed-off-by: Tian Zheng <tizheng@nvidia.com>
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>