Commits · dfe5b7dfc2288afc5d2f247709b1e0328af331e4 · OpenDAS / TransformerEngine

03 Oct, 2025 1 commit

[Common][Pytorch] Add support for the FP8 Block Scaling (ie. Deepseek) recipe on Blackwell (#2157) · dfe5b7df

Jan Bielak authored Oct 03, 2025



* Update to_string(NVTEScalingMode) to include block scaling
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Add `nvte_swizzle_block_scaling_to_mxfp8_scaling_factors`
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Convert FP8 block scaling tensors to MXFP8 tensors on Blackwell and newer in GEMM
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Allow Blackwell and newer in Deepseek recipe compatbility check
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Allow data_rows % 4 != 0 in 1d kernel
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Load scaling factors in unswizzled order in 1d kernel
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Enforce use of power of two scaling
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Skip the FP8 block scaling exact GEMM test on Blackwell
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Skip further tests with pow_2_scales=False
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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



* Initial implementation of tensor conversion for grouped gemm
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Skip non power of two scaling cpp unit tests
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Fix handling of all gather
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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



* Apply suggestions from code review
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Use compute capability 10.0 for logic with Blackwell
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

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

---------
Signed-off-by: Jan Bielak <jbielak@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>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

dfe5b7df

06 Aug, 2025 1 commit

[PyTorch] Multi-tensor swizzle scaling factors for MXFP8 and fuse padding zeros (#2019) · c0d2f1a5

Xin Yao authored Aug 07, 2025



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

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

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

* pad zeros in swizzle kernels
Signed-off-by: Xin Yao <xiny@nvidia.com>

* unify single- and multi-tensor swizzle
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



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

* fix bug for col swizzle
Signed-off-by: Xin Yao <xiny@nvidia.com>

* check context & fix signifiers
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>

c0d2f1a5

21 Apr, 2025 1 commit

rtx5090 arch fix support (#1659) · c7702309

Sudhakar Singh authored Apr 21, 2025



* rtx5090 arch fix support
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* apprend `nvte` to the function name so that its visible in framework specific dirs
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

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

* add filter for nvte_is_supported_nontn_fp8_gemm
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* properly expose the api
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

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

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



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

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

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



* move the function to apt header/c files
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

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

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



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

---------
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

c7702309

17 Apr, 2025 1 commit

[PyTorch] Move swizzle scaling factor to cpp (#1683) · 4e036c8c

Xin Yao authored Apr 18, 2025



* move swizzle scaling factor to cpp
Signed-off-by: Xin Yao <xiny@nvidia.com>

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



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

4e036c8c

07 Feb, 2025 1 commit
- Update main branch with TE 2.0 code, update version to 2.1.0.dev0 · 544dd14b
  Przemek Tredak authored Feb 07, 2025
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
  544dd14b
02 Jan, 2025 1 commit
- Update copyright to include 2025 (#1388) · c9ea6be9
  Kirthi Shankar Sivamani authored Jan 02, 2025
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  c9ea6be9
16 Oct, 2024 1 commit

[PyTorch] Build custom ORT ops before running ONNX export tests (#1252) · f6b766bd

Tim Moon authored Oct 15, 2024



* Build custom ORT ops before running ONNX tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove ONNX from context parallelism tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Export ONNX ops that do compute in FP32

Matches internal impl of TE kernels.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add build script for custom ORT ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

f6b766bd

18 Jun, 2024 1 commit

Remove leftover implementations for optional userbuffers support (#932) · 16c82574

Kirthi Shankar Sivamani authored Jun 17, 2024



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

16c82574

14 Jun, 2024 1 commit

Apply formatting (#929) · 9416519d

Kirthi Shankar Sivamani authored Jun 13, 2024



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

9416519d

13 Jun, 2024 1 commit

[C/PyTorch] Removed MPI dependence in Userbuffers (#901) · e706e5fa

Alp Dener authored Jun 13, 2024



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

e706e5fa

03 Jan, 2024 1 commit
- Change the copyright to include 2024 (#583) · cd798c97
  Przemyslaw Tredak authored Jan 02, 2024
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
  cd798c97
20 Oct, 2023 1 commit

Better way of checking cuDNN version (#485) · d097883e

Przemyslaw Tredak authored Oct 20, 2023



* Ability to check cuDNN version from Python
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Modify the fused attention test to not use the CUDNN_VERSION env
variable which is specific to NGC containers
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

d097883e

27 Jul, 2023 1 commit

Exposing RMSNorm in pyTorch (#306) · fdc09f42

Przemyslaw Tredak authored Jul 27, 2023



* Exposing RMSNorm in pyTorch extensions
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* First pass at the Python API
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Added numerics tests and fixed issues
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

* Added ONNX export and tests for RMSNorm
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

* Added normalization option to the TransformerLayer
Added tests
Fixed test failures
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix documentation
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Fix IMA and invalid type error
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Increase RMSNorm threshold for bf16 case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix ONNX tests
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>

fdc09f42