- 02 Dec, 2025 1 commit
-
-
Phuong Nguyen authored
* 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:
-
- 26 Nov, 2025 1 commit
-
-
Paweł Gadziński authored
* 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:
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:
pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by:
-
- 25 Nov, 2025 1 commit
-
-
Phuong Nguyen authored
* allow dp + fsdp and fixed sr_rng_state partitioning Signed-off-by:
-
- 21 Nov, 2025 2 commits
-
-
Kshitij Lakhani authored
* Remove unnecessary SWA calculation from _segment_ids_pos_to_seqlens_offsets Signed-off-by:
Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by:
-
Kshitij Lakhani authored
* Make BSHD default for Unfused DPA, DPA and MHA in TE JAX Signed-off-by:
-
- 18 Nov, 2025 1 commit
-
-
Paweł Gadziński authored
* fix Signed-off-by:
-
- 14 Nov, 2025 3 commits
-
-
jberchtold-nvidia authored
* Use TE quant if TE fused act is disabled Signed-off-by:
Jeremy Berchtold <jberchtold@nvidia.com> * Keep existing precision Signed-off-by:
-
jberchtold-nvidia authored
* Refactor to avoid storing a global quantization config so direct recipe passing works as intended Signed-off-by:
jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update transformer_engine/jax/quantize/quantizer.py Co-authored-by:
-
Paweł Gadziński authored
* fix Signed-off-by:
-
- 13 Nov, 2025 3 commits
-
-
jberchtold-nvidia authored
* Support for checkpointing quantizations Signed-off-by:
JAX Toolbox <jax@nvidia.com> * Checkpoint quantizations Signed-off-by:
-
Phuong Nguyen authored
* shardy + quantize_layout rework Signed-off-by:
-
Phuong Nguyen authored
* swizzle via nvte Signed-off-by:
-
- 10 Nov, 2025 1 commit
-
-
Teddy Do authored
* Changing default activations in MLP, TransformerLayer, dropout rate after FC1 to 0, and return_layernorm_output to False Signed-off-by:
tdophung <tdophung@nvidia.com> * Fixing the failing tests by hard coding arguments to the previous values instead of relying on newer default values Signed-off-by:
-
- 07 Nov, 2025 3 commits
-
-
Kshitij Lakhani authored
* Default to fused attention in JAX DPA Signed-off-by:
Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> * Correctly update the documentation for defaults in JAX DPA Co-authored-by:
-
Kirthi Shankar Sivamani authored
* Fix cuDNN backend selection for more case. Add CG as a option as well Signed-off-by:
-
Michael Goldfarb authored
-
- 05 Nov, 2025 1 commit
-
-
Paweł Gadziński authored
* fix Signed-off-by:
-
- 31 Oct, 2025 1 commit
-
-
jberchtold-nvidia authored
* Fix mesh resource requirement when no mesh Signed-off-by:
-
- 30 Oct, 2025 1 commit
-
-
Kshitij Lakhani authored
* Fix: Skip determinism tests for bprop for all sm >=100 Signed-off-by:
-
- 28 Oct, 2025 1 commit
-
-
Phuong Nguyen authored
* jax norm + te quant Signed-off-by:
-
- 25 Oct, 2025 1 commit
-
-
Charlene Yang authored
* add max_score for fused/unfused F16 non-CP Signed-off-by:
Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * calculate max per head instead of max over all heads Signed-off-by:
-
- 23 Oct, 2025 1 commit
-
-
jberchtold-nvidia authored
* Make SR rng state always 2D (num_devices, 4) Signed-off-by:
-
- 22 Oct, 2025 2 commits
-
-
jberchtold-nvidia authored
Defer cublas check on fp8 gemms until lowering Signed-off-by: -
jberchtold-nvidia authored
* [JAX] Support recipe flags for disabling SR, RHT, and 2D quantization Signed-off-by:
-
- 18 Oct, 2025 1 commit
-
-
Kirthi Shankar Sivamani authored
* Support wheel build for cuda 13 Signed-off-by:
-
- 14 Oct, 2025 2 commits
-
-
Kirthi Shankar Sivamani authored
* Initial API change Signed-off-by:
Tim Moon <4406448+timmoon10@users.noreply.github.com> * fix nvfp4 test Signed-off-by:
-
Kshitij Lakhani authored
* Add BRCM support when creating a test mask for fused attn Signed-off-by:
-
- 13 Oct, 2025 2 commits
-
-
jberchtold-nvidia authored
assertion check Signed-off-by: -
jberchtold-nvidia authored
* Improve error message for cublas fp8 gemm with incorrect shape Signed-off-by:
-
- 09 Oct, 2025 2 commits
-
-
jberchtold-nvidia authored
Signed-off-by:
-
Kirthi Shankar Sivamani authored
* Update minimum python version to 3.10 and update CI Signed-off-by:
-
- 08 Oct, 2025 1 commit
-
-
Hua Huang authored
* Try async copy of grouped GEMM group_sizes data Signed-off-by:
Hua Huang <huah@nvidia.com> --------- Signed-off-by:
-
- 07 Oct, 2025 1 commit
-
-
Phuong Nguyen authored
* reuse amax for current scaling Signed-off-by:
-
- 06 Oct, 2025 1 commit
-
-
Phuong Nguyen authored
* not fuse bias for output all reduction case + unit tests Signed-off-by:
-
- 03 Oct, 2025 1 commit
-
-
vthumbe1503 authored
Signed-off-by:Varun Thumbe <vthumbe@nvidia.com> *Jax integration for clamped swiglu. This is the continuation of PR which added Clamped Swiglu(used in GPT OSS) support in TE along with Pytorch integration. This PR hooks up the clamped swiglu and dswiglu's nvte APIs to TE Jax.
-
- 02 Oct, 2025 2 commits
-
-
jberchtold-nvidia authored
Fix shard map issue Signed-off-by:
-
jberchtold-nvidia authored
Fix code block in fp8_autocast docstring Signed-off-by:
-
- 01 Oct, 2025 2 commits
-
-
Phuong Nguyen authored
fix rng_state shape Signed-off-by:
-
Phuong Nguyen authored
* rm using_global_amax_of_x Signed-off-by:
-
- 30 Sep, 2025 1 commit
-
-
jberchtold-nvidia authored
Load modules during initialize Signed-off-by:
-
