- 04 Dec, 2025 1 commit
-
-
Tim Moon authored
* Initialize empty tensors with shape=[0] instead of shape=[]. Signed-off-by:
Tim Moon <tmoon@nvidia.com> * Fix runtime crash in LayerNorm Still seeing correctness issues. Signed-off-by:
greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by:
Tim Moon <4406448+timmoon10@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/common/swizzle/swizzle.cu Co-authored-by:
pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by:
-
- 02 Dec, 2025 3 commits
-
-
Kunlun Li authored
* Add primary weighs fp8 support for mxfp8 Signed-off-by:
kunlunl <kunlunl@nvidia.com> * Fix unit test and add better error log to unit test Signed-off-by:
-
Phuong Nguyen authored
* add grouped_tensor classes and helpers Signed-off-by:
Phuong Nguyen <phuonguyen@nvidia.com> * rm non-contiguous option and dptrs Signed-off-by:
-
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:
-
- 26 Nov, 2025 2 commits
-
-
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:
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:
-
Tim Moon authored
Do not initialize recipe state in base op class Op attrs may not be set. Move recipe state initialization to linear op constructor. Signed-off-by:
-
- 25 Nov, 2025 7 commits
-
-
Paweł Gadziński authored
* main Signed-off-by:
-
Pingtian Li authored
* fix backward_dw cuda graph order Signed-off-by:
Pingtian Li <pingtianl@nvidia.com> * add validation for num_layers_per_chunk Signed-off-by:
-
vthumbe1503 authored
* fix ci issue Signed-off-by:
Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert back testing changes Signed-off-by:
-
Zhongbo Zhu authored
* minor fix of torch view dtype Signed-off-by:
Zhongbo Zhu <zhongboz@nvidia.com> * multi-tensor RHT amax, compiles Signed-off-by:
-
vthumbe1503 authored
remove linear redundant check Signed-off-by: -
Phuong Nguyen authored
* allow dp + fsdp and fixed sr_rng_state partitioning Signed-off-by:
-
Teddy Do authored
* Change order of arguments to make jax works Signed-off-by:
tdophung <tdophung@nvidia.com> * make num_experts a tl.constepxr again Signed-off-by:
-
- 21 Nov, 2025 8 commits
-
-
Sudhakar Singh authored
* Add support for THD+CP+SWA through A2A comms Signed-off-by:
Sudhakar Singh <sudhakars@nvidia.com> * unblock the `padding`+`THD`+`CP(A2A)` with SWA case in A2A forward Signed-off-by:
-
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:
-
Jan Bernlöhr authored
Signed-off-by:
janbernloehr <jan@bernloehrs.de> Co-authored-by:
-
oliver könig authored
* ci: Build and attach bdist wheels to release page Signed-off-by:
oliver könig <okoenig@nvidia.com> * free up space Signed-off-by:
-
Jack authored
Signed-off-by:Jack <lityangweiguang@163.com>
-
Paweł Gadziński authored
* fix Signed-off-by:
-
Paweł Gadziński authored
* fix Signed-off-by:
-
- 19 Nov, 2025 2 commits
-
-
Kirthi Shankar Sivamani authored
* Minor CPU overhead changes Signed-off-by:
-
Jianbing authored
* refactor mxfp8_cast_only kernel Signed-off-by:
Jianbing Dong <jianbingd@nvidia.com> * fix ptx.cuh after format Signed-off-by:
Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
-
- 18 Nov, 2025 4 commits
-
-
Jaime authored
[PyTorch] Implement Selective Activation Checkpointing for LayerNormMLP with checkpoint flag (#2311) * custom tests for selective activation checkpointing for layernorm mlp Signed-off-by:
Jaime Cardenas <jaime@evolutionaryscale.ai> * add selective layernorm mlp to te.pytorch Signed-off-by:
Jaime <102792198+jaimec00@users.noreply.github.com> Signed-off-by:
-
Paweł Gadziński authored
* fix Signed-off-by:
-
Kirthi Shankar Sivamani authored
* Cache device tensors properly Signed-off-by:
-
Paweł Gadziński authored
* fix Signed-off-by:
-
- 17 Nov, 2025 3 commits
-
-
Charlene Yang authored
* [Common] Deleted unused header (#2324) Deleted unused header Signed-off-by:
Oleg Goncharov <ogoncharov@nvidia.com> Signed-off-by:
Peter Dykas <wdykas@nvidia.com> * [JAX] L1_jax_distributed_test suit with individual executions (#2321) * L1 rework Signed-off-by:
Jeremy Berchtold <jberchtold@nvidia.com> * Set HF_HUB_OFFLINE to disable any network calls to HF when the pre-downloaded dataset is available Signed-off-by:
Przemek Tredak <ptredak@nvidia.com> Signed-off-by:
Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> * Correctly update the documentation for defaults in JAX DPA Co-authored-by:
Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Co-authored-by:
root <root@gpu-h100-0496.cm.cluster> Co-authored-by:
jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Co-authored-by:
Kunlun Li <94586211+kunlunl@users.noreply.github.com> Co-authored-by:
Michael Goldfarb <mgoldfarb@nvidia.com> Co-authored-by:
wdykas <73254672+wdykas@users.noreply.github.com>
-
Evgeny Tsykunov authored
* Enable reference current scaling recipe Signed-off-by:
Evgeny <etsykunov@nvidia.com> * minor Signed-off-by:
-
Kirthi Shankar Sivamani authored
Initial changes to remove pytorch overheads Signed-off-by:
-
- 14 Nov, 2025 5 commits
-
-
jberchtold-nvidia authored
* Use TE quant if TE fused act is disabled Signed-off-by:
-
Paweł Gadziński authored
* init Signed-off-by:
root <root@ptyche0312.ptyche.clusters.nvidia.com> Co-authored-by:
-
jberchtold-nvidia authored
* Refactor to avoid storing a global quantization config so direct recipe passing works as intended Signed-off-by:
-
Paweł Gadziński authored
* fix Signed-off-by:
-
Robin Zhang authored
* reset cudagraph Signed-off-by:
Robin Zhang <robinz@nvidia.com> * use closure instead of mutable default values Signed-off-by:
-
- 13 Nov, 2025 5 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:
-
Evgeny Tsykunov authored
Fix amax computation using output_t data in normalization Signed-off-by: -
Lifu Zhang authored
Signed-off-by:
Lifu Zhang <lifuz@login-lyris01.lyris.clusters.nvidia.com> Co-authored-by:
-
