GitLab

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

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

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

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

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

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

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

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

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

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

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

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

* add default path for ffi include

* add an option to get XLA_HOME from env

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

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

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

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

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

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