Commits · 1398fa5f36bfb780acfb7348dcd297a84fd3e705 · OpenDAS / TransformerEngine

14 Aug, 2025 1 commit

[Common] Reduce CUDA driver calls (#2067) · c582f6be

Xin Yao authored Aug 14, 2025



* reduce driver calls
Signed-off-by: Xin Yao <xiny@nvidia.com>

* reduce driver calls
Signed-off-by: Xin Yao <xiny@nvidia.com>

* adjust tests to capture this
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>

c582f6be

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

22 Jul, 2025 1 commit

[Common] Improved performance of mxfp8 cast kernels (#1628) · cb504cda

Oleg Goncharov authored Jul 22, 2025



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

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

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



* Minor code refactoring to avoid unnecessary checks
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



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

* Fixed dBias accumulation error due to initialization. Minor code refactoring
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Test case to reproduce the init error
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Fixed rowwise dbias error
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Changed ptx API
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Added a struct for two packed FP8 values
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Rolled back to scalar code for columnwise scaling due to its better performance
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

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



* Rebased on main
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Fixes per code review
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Removed constexpr in C++ test suite to build faster
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Computed activations are now numerically truncated to InputType before scaling. Improved test suite.
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



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

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

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

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



* Modified mismatches checks of MXFP8 to address FP8 numerics
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Implemented Jeremy's fixes to JAX test suite with an intermediate downcast
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Reduced the dims of the test tensors to improve CI runtime
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed memory alignment issue. Compute dbias without downcast.
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Fixed misaligned memory issue also in gated kernels. Reduced size of MXFP8 gated tests
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



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

cb504cda

13 Jun, 2025 1 commit

[common] Added support of FP4 data type (#1779) · 7b94bd99

Oleg Goncharov authored Jun 13, 2025



* Added support of FP4 data type
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Refactoring to BitsNum in progress
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed compilation errors. All C++ tests passed
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed a typo
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Added FP4 guard to TMA tensor descriptor data type
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Fixed errors in JAX C++ extensions
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Removed dummy NVFP4 C++ test file
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

* Refactored the code per the review notes. Fixed JAX build error.
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Removed unnecessary static casts
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Typo fix
Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>

* Pass correct num bits to create_2D_tensor_map; fixes CI
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

---------
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

7b94bd99

29 May, 2025 1 commit

Avoid memory allocations and deallocations when creating NVTETensor (#1813) · 4292653c

Przemyslaw Tredak authored May 29, 2025



* Changed the Tensor allocation strategy
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Disable debug flag
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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



* Fix the double free error
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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



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

* Fixed pyTorch recipe extension
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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



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

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

* Hide TensorAllocator and fix the usage in LayerNorm
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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



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

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

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



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

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

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

4292653c

22 May, 2025 1 commit

Build support for cuda 13 (#1809) · 00328ac7

Kirthi Shankar Sivamani authored May 22, 2025



* Build support for cuda 13
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix build for cudnn 8.9*; cuda 12.1
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

00328ac7

05 May, 2025 1 commit

[C][PyTorch] Move multi tensors kernels from PyTorch extensions to core (#1744) · c972f5a7

Kirthi Shankar Sivamani authored May 05, 2025



* Move multi tensors kernels from PyTorch extensions to core
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add int16 type to core (for storing fp32 param remainders)
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Fix perf, memory, vars
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Re-add device guard for multi-device
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix junk output dtype for non-per tensor
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fixes for test and upgrade mcore version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

c972f5a7

24 Apr, 2025 1 commit

Introduce nvte_memset to provide a fill kernel that is faster than... · 62d1b2bd

jberchtold-nvidia authored Apr 24, 2025


Introduce nvte_memset to provide a fill kernel that is faster than cudaMemsetAsync for small sizes (#1716)

* nvte_memset fills single float value
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Support larger sizes than a single value and add tests
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

---------
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

62d1b2bd

26 Feb, 2025 1 commit

Added memory alignment check to cast_fp8_1D (#1507) · 5d85857a

Oleg Goncharov authored Feb 26, 2025



* Added TMA alignment check to cast_fp8_1D
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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



* Use tensor const-ref instead of tensor const-ptr
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

5d85857a

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
21 Aug, 2024 1 commit

Update FP8 scale-inverse in kernels with FP8 output (#1083) · 8e3561bf

Tim Moon authored Aug 20, 2024



* Perform scale-inv update in cast-transpose kernels
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Perform scale-inv update in cast and activation kernels
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Perform sclae-inv update in LayerNorm and RMSNorm kernels
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Perform scale-inv update after FP8 GEMMs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fuse casts and scale-inv updates in linear module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fuse casts and scale-inv updates in layernorm-linear module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Simplify kernel to update FP8 scale-inv
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Debug amax update in layernorm kernels
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



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

* Debug ONNX export

Use quantization scaling factor in ONNX quantize op.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Review suggestion from @ptrendx
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Debug mismatched dtypes
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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

8e3561bf