GitLab

* Add LN margin to inference
Signed-off-by: Sangkug Lym <slym@nvidia.com>

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

* Fix symbolic func registration
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* FP8 cuda graphs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>

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

* exclude torch compile from numerics tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* rm fusion from unfused path
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* value_and_grad requires same shape for input and gradients
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use high precision layernorm
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove local_device_ids as it caused unexpected behaviors
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert "Remove local_device_ids as it caused unexpected behaviors"

This reverts commit c54349b2ce1e96ae696cf0d74f5210e55002cf72.
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Remove unused headers
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Unify the fused attn workspace size cpp code
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Reduce the skipped cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename self/cross attention to qkvpacked/kvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update attention mask docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine the attn mask implementations
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* disallow sharding of layernorm learnable parameters; force duplication
Signed-off-by: Keshav <keshavb@nvidia.com>

* fix tests and support tensors for gamma/beta in layernorms
Signed-off-by: Keshav <keshavb@nvidia.com>

* reverting
Signed-off-by: Keshav <keshavb@nvidia.com>

* added tests for rank-1 gamma/beta sharding
Signed-off-by: Keshav <keshavb@nvidia.com>

* fix lint errors
Signed-off-by: Keshav <keshavb@nvidia.com>

---------
Signed-off-by: Keshav <keshavb@nvidia.com>

* Modified MHA and DPA logic to use causal softmax and FA for inference
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Adjusted unfused attention and softmax logic for inference
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Cleaned up the code per pylint
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added test cases to evaluate numerics of incremental decoding
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Apply suggestions from code review
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>

* Apply suggestions from code review [sequence start-end]
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>

* Apply suggestions from code review [inference_params offset update]]
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>

* Fixed bug in KV-cache indices and updated test suite
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added inference_params description and applied suggestions from the code review
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Adjusted absolute tolerances in numerics tests
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Cleaned up the files per pylint
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

---------
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* added support for arbitrary bias shapes for fused_attn
Signed-off-by: Alp Dener <adener@nvidia.com>

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

* Add b1ss/bhss/11ss bias shapes when not requiring dBias
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add bias_b/h to plan cache
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fixed compile errors after PR653 merge
Signed-off-by: Alp Dener <adener@nvidia.com>

* updated JAX unittests for new bias shapes
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed mismatched mask type checking
Signed-off-by: Alp Dener <adener@nvidia.com>

* corrected skip condition
Signed-off-by: Alp Dener <adener@nvidia.com>

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

* corrected skip checks for bias shapes
Signed-off-by: Alp Dener <adener@nvidia.com>

* resolved test issues but neginf with float16 is still problematic with JAX
Signed-off-by: Alp Dener <adener@nvidia.com>

* new bias shapes passing TE JAX CI for seqlen <= 512, seq_q == seq_kv and h_q == h_kv conditions
Signed-off-by: Alp Dener <adener@nvidia.com>

* TE/JAX fused attn tests for new bias shapes passing with neg_inf=-2**27 for Bfloat16 and -2**15 for Float16
Signed-off-by: Alp Dener <adener@nvidia.com>

* code style fixes and test parameter ID cleanup
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed incorrect skip condition for backward fused attn test
Signed-off-by: Alp Dener <adener@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: Alp Dener <adener@nvidia.com>

* Softmax now forces XLA to unshard the hidden dimension with a warning. Unittests updated to check for numerics and warning with bad sharding
Signed-off-by: Alp Dener <adener@nvidia.com>

* correcting cudnn-frontend version
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed mismatched output sharding
Signed-off-by: Alp Dener <adener@nvidia.com>

* combined softmax tests and fixed code style/linting issues
Signed-off-by: Alp Dener <adener@nvidia.com>

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

Support various implementations of RoPE and fix a coordinate representation bug
Signed-off-by: Ming Huang <mingh@nvidia.com>

* added non-reentrant mode support to TE checkpoint
Signed-off-by: Alp Dener <adener@nvidia.com>

* updated get_cuda_rng_tracker kwarg to get_rng_state_tracker to remain consistent with other TE API
Signed-off-by: Alp Dener <adener@nvidia.com>

* docstring cleanup
Signed-off-by: Alp Dener <adener@nvidia.com>

* added mechanism to disable bias_gelu_nvfusion in LayerNormMLP when checkpointing in non-reentrant mode
Signed-off-by: Alp Dener <adener@nvidia.com>

* refactored checkpoint and recompute hook names to match PyTorch implementation
Signed-off-by: Alp Dener <adener@nvidia.com>

* Fixed incorrect reference before assignment
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed argument error in calling native PyTorch checkpoint
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed linting errors for missing docstrings
Signed-off-by: Alp Dener <adener@nvidia.com>

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

* bias GELU fusion consistency between checkpoint test and reference comparison
Signed-off-by: Alp Dener <adener@nvidia.com>

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

* Refine MHA API
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Reuse func from the flax
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* qkv packed draft
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix test_layer with fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add attn_bias_type and enhance a few code flow
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move scale_factor from __call__ to init
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add DPA public API and tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

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

* Add qkv separate fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Apply BSHD_BSHD_BSHD format
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove debug log
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add fused attention layer tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Fine-grained fused attn settings
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove the default value of num_attetnion_head and head_dim
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add teardown for fused attn env
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Unify the Optional notation
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix Pre/Post scale bias comments
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add no_mask tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add checkpoint_name for fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix the fused attn batcher
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@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>

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>

* Added QuickGELUActivation from HuggingFace/Transformers to common and pytorch
Signed-off-by: Alp Dener <adener@nvidia.com>

* Removing 'qgelu' from double-size activations list in LayerNormMLP.
Signed-off-by: Alp Dener <adener@nvidia.com>

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

---------
Signed-off-by: Alp Dener <adener@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>

* Use fused implementation of RoPE in MultiHeadAttention
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add option to avoid updating transpose cache when possible
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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

* Remove unused attr
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>

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

* Added new unfused softmax cuda kernel to support causal attention mask
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added test suite for unfused causal softmax kernel
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Removed test cases with large matrices from the causal softmax test suite
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Cleaned up the code per lint
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added a compute buffer to causal softmax testing suite to store intermediate results without casting
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added more tests cases
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Relaxed absolute tolerance atol
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Relaxed absolute tolerance for BF16
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

---------
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* test alibi between fa and fu
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move alibi slopes and bias to global to avoid repeating calculation
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix alibi slopes/bias generation
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix _is_flash_attention_supported to allow alibi type
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* disable padding mask when alibi is used for fused attn arbi backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add support for custom [n_heads] alibi_slopes in flash, fused, unfused attention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove alibi_type=none tests as they are unnecessary
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* change bias/dbias shape to allow b,1/1,h/b,h in arbi backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* tweak tests for arbi post_scale_bias [1,h,s,s] or alibi_slopes [n_heads]
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change bias/dbias shape in max512 backend - incomplete
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove max512 changes from last commit and disable max512 (and arbi temporarily) for [b, h, s, s]; pending cuDNN backend support
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up and tweak backend selection logic
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace || with () in docstring
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix bias shape for max512 backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* combine slopes/bias generation to one function get_alibi() and fix alibi tests
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 PR557 bugs
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>

* encapsulate global alibi tensors into a dict cache
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* reduce alibi slopes test size
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* use dBias shape to define bias_b/bias_h because jax materializes dBias rather than Bias in bwd abstract
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>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Add zero_centered_gamma option to RMSNorm
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* More improvements to tests
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Tweaking the tolerances
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix LayerNormMLP test
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Update transformer_engine/common/rmsnorm/rmsnorm_api.cpp
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update transformer_engine/common/rmsnorm/rmsnorm_api.cpp
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Tweak tolerances with bfloat16
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

---------
Signed-off-by: Przemek Tredak <ptredak@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>

* Adding support of sequence parallelism
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding RoPE
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix wrong batch_logical_axes
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Rnaming FSDP outer env var
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Poring RoPE to Praxis layers.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Porting GeLU + [FP8 Cast].
Signed-off-by: Ming Huang <mingh@nvidia.com>

* WAR to make XLA successfully match FP8 GEMM on FFN1 with GeLU.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Allowing arbitrary dimension of NVShape for the workspace allocation
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding checkpoint_name to fused functions of mlp.py to get better perf with nn.scan.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Modify with review feedback.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix bugs
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix typo.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fixed for lint
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Follow review feedback to modify code.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix typo.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Port SP to Praxis
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Fix an issue when enabling both GQA and RoPE.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Update docs
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

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

* Remove paddle.fluid usage from tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Removed cudaMalloc/WorkspaceManager in JAX csrc. JAX custom ops now request buffers from XLA for their workspace tensors.
Signed-off-by: Alp Dener <adener@nvidia.com>

* removed unused GEMM C++ API in TE-JAX
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed typo in layernorm_geglu_fp8_mlp and removed unnecessary shape reductions in primitives
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed import order for linting
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed custom op errors due to incorrect static arg nums in JAX jit
Signed-off-by: Alp Dener <adener@nvidia.com>

* shifted cudnnSetStream further down the kernel to avoid error when executing dummy kernel call with nullptr stream
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed linting errors for blank lines
Signed-off-by: Alp Dener <adener@nvidia.com>

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

* fused apply rope
Signed-off-by: Xin Yao <xiny@nvidia.com>

* Apply suggestions from code review
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Xin Yao <yaox12@outlook.com>

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

* make rotary_percent optional
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

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

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

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

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

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

* sync apex: fuse sin/cos
Signed-off-by: Xin Yao <xiny@nvidia.com>

* sync apex: fused rope for thd format
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

* Fix license headers
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* add support for bshd format
Signed-off-by: Xin Yao <xiny@nvidia.com>

* support different seq length
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

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

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

* Make outputs contiguous in SBHD case
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <yaox12@outlook.com>
Signed-off-by: Przemek Tredak <ptredak@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: Przemyslaw Tredak <ptredak@nvidia.com>

[PyTorch] Workaround for incorrect output from torch.cuda.is_bf16_compatible() on V100s and TU102s (#626)

* replaced torch.cuda.is_bf16_compatible() with explicit sm_80 check via torch.cuda.get_device_capability()
Signed-off-by: Alp Dener <adener@nvidia.com>

* implement te.utils.is_bf16_compatible() to replace torch.cuda counterpart
Signed-off-by: Alp Dener <adener@nvidia.com>

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

* added missing parameter materialization on real device for LayerNorm and RMSNorm
Signed-off-by: Alp Dener <adener@nvidia.com>

* added new unittest for deferred initialization and modified parameter materialization to support standalone execution outside of FSDP
Signed-off-by: Alp Dener <adener@nvidia.com>

* restored tensor parallel attributes that were being wiped out by the parameter reset
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed incorrect order of fp8 metadata initialization
Signed-off-by: Alp Dener <adener@nvidia.com>

* added deferred init unittest to the QA script
Signed-off-by: Alp Dener <adener@nvidia.com>

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

Activation offloading to CPU's for the Linear, Layernorm Linear and the Layernorm MLP modules (#571)

* Added support activation offloading to CPU's
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Moving CPU offloading library to TE
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Restructured code, added switch to choose between weight/activation offloading
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Removed arg during constructor
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Fix nit-pick errors
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

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

* Fix to the code block in docs
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Added offloading unit test
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Fixed formatting
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* wgrad fusion fix, minor errors and lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Fixed stray PyT tensors in LayernormMLP getting offloaded
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Fixed typi
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Fix offloading for rmsnorm, rm test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

fix failing tests due to PR #557
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

Signed-off-by: hugo-syn <hugo.vincent@synacktiv.com>

* make TransformerLayer accept a `bshd` or `sbhd` tensor format
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Fixes from feedback
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* more feedback fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* remove incorrect info from docstring
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix from feedback
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Support num_gqa_groups arguments
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add GQA support on the JAX bridge code
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix the kv stride of the arbitrary backend
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Complete rewrite fused attention tests and add GQA coverage
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Support unfused GQA
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Calculate seqlen before the primitive for the better perf
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add GQA layer tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Apply code style checks for te_jax
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Apply code style checks for tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add num_gqa_groups doc
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine the qkv_type
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Correct the variable naming
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Handle Max512 CAUSAL
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add WAR for the latest jax image
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@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>

* Adding Cast custom call
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Applying cast to the kernel of layernorm_fp8_dot
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Applying native cast to the kernel of fp8_dot.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Apply Cast and native cast to layernorm_geglu_fp8_dot
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the bug to enable layernorm_geglu_fp8_dot in LayernormMlp
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Modifiied code with the review feedback.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding 2xACC control to FP8 GEMMs.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Set precision as an static arg
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@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>

* try to use cuDNN fused attention for context parallelism
Signed-off-by: xren <xren@nvidia.com>

* assert CP is only supported with NVTE_F16_arbitrary_seqlen
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* port fused attn api to context parallelism
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add one more assert
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* assert CP does not support padded tokens
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add qkv_format into CP implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove qkv_format from CP function
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix qkv_for,at
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix bwd error with FA v2
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make cp implementation support non-causal masking
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* bug fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove redundant asserts for CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor assert information change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* assert core attn bias has not been supported with CP yet
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make CP work with window_sizes of [-1, -1] and [-1, 0]
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add draft code for fa test with cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* move fused attn test to a specific folder
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add assert_close to flash attn cp test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add more tests for CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add optional arguments for FA v2.4+
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add skip condition for CP test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* class and function naming fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* docstring fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* do not use fused attn if backend does not work with CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* create a separate folder for CP test as it needs multi-GPUs
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add attn_mask_type check in attn_forwrad_func_with_cp
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code format fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

---------
Signed-off-by: xren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

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

* Reduce size of PyT sanity tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add test cases with 126M model and weird dimensions
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add missing arg in GPT 126M test
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>

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

* add sliding window to FA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* change bert test to causal as unfused does not support padding
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix FlashAttention for v2-2.3 versions
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* verify FA swa works
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix mask related restrictions and duplicate code after merge
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add docstring for get_swa func
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move repeated code into a function
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add determinism filter and fix FA warning message
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add message for determinism filter
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* Disable dynamo for Fused Attention
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>