Commits · f1512b21c73951a7606375cf4a4ecaa9fc6a2969 · OpenDAS / TransformerEngine

02 Dec, 2025 1 commit

Phuong Nguyen authored Dec 02, 2025



* init triton binding with test case/example

* added Triton as TE-JAX test dependency

* grid with blocksize from autotune
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

f1512b21

27 Sep, 2025 1 commit

[JAX] CollectiveGemm (#2166) · d75bf43f

Phuong Nguyen authored Sep 27, 2025



* init cgemm + unit tests

* UB bootstrap with NCCL, no MPI dependency

* add NVLINK-P2P check + error message

* skip tests if no NVLINK available

* use std::vector to store ncclComm_t

* update misuse of TP warning
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

d75bf43f

08 Jul, 2025 1 commit

[JAX BUILD] Fixes for JAX 0.7.0 (#1936) · 9d031fbd

Phuong Nguyen authored Jul 08, 2025



* Fix jax build
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

9d031fbd

10 Jun, 2025 1 commit

pyproject.toml (#1852) · aedd7e10

Kirthi Shankar Sivamani authored Jun 10, 2025



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

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

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

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

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

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

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

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

* Fix JAX build in baremetal env
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update install inst in readme
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

aedd7e10

09 Jun, 2025 1 commit

Manage dependencies and add missing `einops` req (#1859) · ddcda1ff

Kirthi Shankar Sivamani authored Jun 09, 2025



* Manage deps and add einops
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

ddcda1ff

14 May, 2025 1 commit

Fix JAX build (#1784) · d2973cb8

Kirthi Shankar Sivamani authored May 14, 2025


Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>

d2973cb8

12 May, 2025 1 commit

Reduce framework extensions' binary size and fix debug builds (#1769) · a94e5dde

Kirthi Shankar Sivamani authored May 12, 2025



* Remove default debug info from distutils
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

a94e5dde

28 Apr, 2025 1 commit

Support `nvidia-cu*` wheels for core lib compilation; miscellaneous build improvements (#1717) · 643fb0a0

Kirthi Shankar Sivamani authored Apr 28, 2025



* Add support for nvidia cu* lib wheels
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

* Specify exact package versions
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Add frameworks and nvidia-libs to setup requirements. Add alternates to version finding
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Fix jax wheel install in no toolkit env [wip]
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add missing headers via pip
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Load SOs, revert CMake
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Fix JAX exec without cudatk
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

643fb0a0

18 Apr, 2025 1 commit

[C][Jax] Move cuda kernels from Jax extensions to core (#1697) · 26db7f34

Kirthi Shankar Sivamani authored Apr 18, 2025



* Move jaxx cuda kernels to core
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>

26db7f34

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
14 Aug, 2024 2 commits

[TE/JAX] Add default include path for XLA FFI (#1104) · 8ef3308a

Phuong Nguyen authored Aug 14, 2024



* add default path for ffi include

* add an option to get XLA_HOME from env

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

8ef3308a

[TE/JAX] Prototype for New XLA Custom Calls with FFI (#946) · 4b2b39b4

Phuong Nguyen authored Aug 13, 2024



* implemented custom call with ffi in csrc

* moved headers of misc to misc.h, add ffi.h

* ActLu and DActLu lowering with ffi_lowering

* CastTranspose with ffi_lowering

* enabled cudaGraph

* added 4d input test case to TestActivationLu

* added operand_output_aliases for CastTranspose

* added env var NVTE_JAX_WITH_FFI, default value = 1

* replace casting ActivationEnum by taking its value

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

4b2b39b4

24 Jun, 2024 1 commit

Improve JAX build tool (#942) · 85aeb903

Phuong Nguyen authored Jun 24, 2024



* adding option to select only .cpp files in a dir in the build tool

* change cmake build path

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

85aeb903

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

08 Jun, 2024 1 commit

[JAX] Splitting `csrc/modules.cpp` by category (#883) · 16f3f897

Phuong Nguyen authored Jun 08, 2024



* categorized `csrc/modules.cpp`
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* adapted the build tool
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

16f3f897

06 Jun, 2024 1 commit
- Build system refactor for wheels (#877) · c1b915ae
  Kirthi Shankar Sivamani authored Jun 06, 2024
```
Cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  c1b915ae