GitLab

Update copyright to include 2026
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization

1.Fused `moe_permute_with_probs` + `Fp8Padding` and fused `moe_unpermute` + `Fp8Unpadding`,
  that can remove the explicit padding/unpadding of moe expert, improved performance and reduced peak gpu memory usage.
2.Add tests of fused permute/pad and unpermute/unpad.
Signed-off-by: xiaoxi-wangfj <690912414@qq.com>

* [PyTorch/Common] Fuse permute+pad and unpermute+unpad support with_merging_probs
Signed-off-by: xiaoxi-wangfj <690912414@qq.com>

* [PyTorch]format code
Signed-off-by: xiaoxi-wangfj <690912414@qq.com>

* [Common]perf expert_idx loaded once
Signed-off-by: xiaoxi-wangfj <690912414@qq.com>

* fix: pad_offsets can be None
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: xiaoxi-wangfj <690912414@qq.com>

* add padding + merging probs bwd support. Not tested
Signed-off-by: tdophung <tdophung@nvidia.com>

* Fix garbage initialized act grad
Signed-off-by: tdophung <tdophung@nvidia.com>

* all test passing for jax permutation + pad
Signed-off-by: tdophung <tdophung@nvidia.com>

* change tokens_per_experts APIs to num_out_tokens with conservative allocation of worst case padding for output buffer
Signed-off-by: tdophung <tdophung@nvidia.com>

* change test permutation to reduce test time
Signed-off-by: tdophung <tdophung@nvidia.com>

* triggering PR refresh
Signed-off-by: tdophung <tdophung@nvidia.com>

* format code
Signed-off-by: tdophung <tdophung@nvidia.com>

* Remove some tests cases from pytorch side. Add a separate toekn_dispatch test for sanity in case combine accidentally undo an error on dispatch in the roundtrip test. Add distinction between L0 and L2 in test cases in jax
Signed-off-by: tdophung <tdophung@nvidia.com>

* format code
Signed-off-by: tdophung <tdophung@nvidia.com>

* remove chance for inefficiency in moving between CPU and GPU, remove redundant primitive using a new static bool for padding, add assert for align size
Signed-off-by: tdophung <tdophung@nvidia.com>

* fix lint in jax
Signed-off-by: tdophung <tdophung@nvidia.com>

* account for both jax newer and older than version 0.8.2. Adjusted gpu triton binding accordingly
Signed-off-by: tdophung <tdophung@nvidia.com>

* format code
Signed-off-by: tdophung <tdophung@nvidia.com>

* fix typo
Signed-off-by: tdophung <tdophung@nvidia.com>

---------
Signed-off-by: xiaoxi-wangfj <690912414@qq.com>
Signed-off-by: tdophung <tdophung@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: tdophung <tdophung@nvidia.com>

* init
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* lines lenght
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* subtitle --- fix in many files:
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* cross entropy _input -> input rename
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* cross entropy _input -> input rename
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* a lot of small fixes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* torch_version() change
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* add missing module and fix warnings
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* removed training whitespace:
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Update docs/api/pytorch.rst
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>

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

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

* Fix NumPy docstring parameter spacing and indentation

- Standardize parameter documentation to use 'param : type' format (space before and after colon) per NumPy style guide
- Fix inconsistent indentation in cpu_offload.py docstring
- Modified 51 Python files across transformer_engine/pytorch
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
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>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* rename experimental -> custom_recipes
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Decouple python base classes (api)
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* update test_custom_recipe
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Rename experimental -> custom
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Minor
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

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



* Fix import
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Update tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Evgeny Tsykunov <e.tsykunov@gmail.com>

* Update tests/pytorch/test_custom_recipe.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Evgeny Tsykunov <e.tsykunov@gmail.com>

* quantization_base -> quantized_tensor rename
Signed-off-by: Evgeny <etsykunov@nvidia.com>

---------
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny Tsykunov <e.tsykunov@gmail.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>

* optimize permute
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



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

---------
Signed-off-by: Hongxiao Bai <hongxiaob@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: Xin Yao <xiny@nvidia.com>

* add support for new recipe on permute_fusion, rm fp unpermute
Signed-off-by: tongliu <tongliu@nvidia.com>

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

* remove fp8 from index map
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

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



* skip unsupported tests
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: tongliu <tongliu@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: tongliu <tongliu@nvidia.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* add prob permute; fix fp8tensor
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



* revert unnecessary changes in UT
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* remove unnecessary probs dtype convert
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* keep the output nums if probs is not provided
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



* refine the doc string
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* fix lint
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* use fp32 compute type
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* style fix
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* fix empty input return
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* separate prob related functions out
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



---------
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

* add mask-based moe permutation

* change moe_chunk_permute to moe_sort_chunks_by_indices

* fix __all__ in pytorch/permutation.py

* fix func/var names and typos; update tols in UT

---------
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

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

* Upgrade pylint and first round formatting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

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

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

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

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

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

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

* remove dtype from args
* update docs with permutation ops

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

* Add permutation functions

* Add permutation ops

* Remove the dependency on cutlass

* Move permutation.py out of module dir

* Rewrite the unit test and enable skipping if FP8 is unavailable

* Rename exposed C++ API and reorder its parameters + take NVTETensor as inputs

* Use Float8Tensor for FP8 input

* Move dtype to ctx

---------
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
Co-authored-by: Qi Zhang <qizhang@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>