Commits · de6afe24b2b33d13a86770d761c165804ab0400e · OpenDAS / TransformerEngine

11 Aug, 2025 1 commit

[PyTorch] Fix high-precision dtype for MXFP8 AG (#2058) · de6afe24

Kirthi Shankar Sivamani authored Aug 11, 2025



* Fix high-precision dtype for MXFP8 AG
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

de6afe24

09 Aug, 2025 1 commit

Use userbuffers for MXFP8 wgrad all-gather overlap (#1982) · 077e26c3

Daniel Stokes authored Aug 09, 2025



* fix: Add stream synchronization before destroying MPI communicator (#1979)
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* feat: Implement column-wise userbuffer overlap for comm+GEMM operations

Add support for overlapping column-wise allgather communication with GEMM
operations to improve training performance:

* **Core infrastructure changes:**
  - Update bulk_overlap_columnwise_ag() to accept explicit stream parameter
  - Modify userbuffers send/recv loops to use rank-ordered iteration
  - Add userbuffers_send_all/recv_all function declarations

* **Python integration:**
  - Add bulk_overlap_ag_with_external_gemm() C++ extension function
  - Expose new overlap function via pybind11 bindings
  - Update overlap method configurations to include more ring_exchange ops

* **LayerNorm MLP optimization:**
  - Enable column-wise quantization for FC2 gradient output
  - Implement overlap of allgather communication with FC2 DGRAD GEMM
  - Use fill_userbuffers_buffer_for_all_gather for efficient buffering

This optimization allows overlapping communication and computation phases
more effectively, reducing training wall-clock time by hiding allgather
latency behind GEMM execution.
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* fix: Working userbuffer overlapping API
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* fix: Fix overwriting bulk overlap UB object for layernormLinear
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* fix: Update external overlap to use tp size instead of nvsize to determine number of copies
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* fix: Fix linter error
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* fix: Explanatory comments of overlap logic
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* fix: Fix the UB fused ops tests
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* fix: Fix linter errors
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

---------
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

077e26c3

08 Aug, 2025 1 commit

[PyTorch debug] Improve precision debug tools performance (#1909) · b6b3abce

Paweł Gadziński authored Aug 08, 2025



* turn on userbuffers for layers without debug
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* code drop
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* working change
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
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: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
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: Pawel Gadzinski <pgadzinski@nvidia.com>

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



* tests and fixes
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: Pawel Gadzinski <pgadzinski@nvidia.com>

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



* fixes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* update nvinspect version
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: Pawel Gadzinski <pgadzinski@nvidia.com>

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



* add
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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



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



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



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



* fixes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fixes
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: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
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: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix ci
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

b6b3abce

07 Aug, 2025 1 commit

[PyTorch] Fix numeric overflow caused by int-type parameters and return value... · dd083bdf

ldl authored Aug 07, 2025


[PyTorch] Fix numeric overflow caused by int-type parameters and return value in the roundup function (#2034)
Signed-off-by: lvdunlin <lvdunlin@xiaomi.com>
Co-authored-by: lvdunlin <lvdunlin@xiaomi.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

dd083bdf

06 Aug, 2025 3 commits

[PyTorch] fix input_quantizer usage for save_original_input; fix blockwise FP8... · de69ca0e

hx authored Aug 07, 2025


[PyTorch] fix input_quantizer usage for save_original_input; fix blockwise FP8 convert_and_update_tensor (#1978)

* fix input_quantizer in save_original_input bwd
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* fix get shape of blockwise tensor with only compact colwise data
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* fix blockwise FP8 convert_and_update_tensor
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



---------
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>
Co-authored-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>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

de69ca0e

[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

[PyTorch] Fix zero initialization in permute kernel for padded slots (#2026) · 7101f4be
xiaoxi-wangfj authored Aug 06, 2025
```
Signed-off-by: xiaoxi-wangfj <690912414@qq.com>
```
7101f4be

04 Aug, 2025 2 commits

[PyTorch] Fix cudagraph static_input and static_grad_input reuse (#2018) · 06947e87

buptzyb authored Aug 05, 2025



* fix graph static grad input reuse
Signed-off-by: Robin Zhang <robinz@nvidia.com>

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

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



* update
Signed-off-by: Robin Zhang <robinz@nvidia.com>

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

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



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

06947e87

[Core] Kernel that swaps first two tensor dimensions (#1998) · c3f8a9f5

Tim Moon authored Aug 04, 2025



* Add basic kernel for swapping first two tensor dims
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add NVRTC kernel for swapping first dims
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add PyTorch extension for swap first dims kernel
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Tweak variable names
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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



* Make sure writes are contiguous
Signed-off-by: Tim Moon <tmoon@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>

c3f8a9f5

01 Aug, 2025 1 commit

Manually launch wgrad accumulation and reduce in backward_dw() instead of backward() (#1976) · 1258bbe0

Hongbin Liu authored Aug 01, 2025



* disable wgrad accumulation and reduce in backward() And manually launch it in backward_dw()
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* format
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* refactor
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* refactor
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* set skip_backward_post_hook to True only if delay_wgrad_compute is True
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* format
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

---------
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>
Co-authored-by: Hongbin Liu <hongbinl@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

1258bbe0

31 Jul, 2025 1 commit

Fuse amax computation into normalization kernel for current scaling (#2013) · 51eb6362

Jan Bielak authored Jul 31, 2025



* Compute amax in normalization kernels as long as the pointer is provided, even if using non quantized output
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Fuse amax computation into normalization forward
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Use TE lahyernorm kernel instead of raising error about unsupported cuDNN feature
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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



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

51eb6362

30 Jul, 2025 2 commits

Refactor normalization.cpp to use quantizer logic introduced in #1952 (#2006) · 11ac24cf

Jan Bielak authored Jul 29, 2025



Refactor normalization.cpp to use quantizer logic introduced in #1952 instead of manual quantization
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

11ac24cf

[PyTorch] Fix bug with clearing op outputs during backward (#2008) · 020428f0
Tim Moon authored Jul 29, 2025
```
Fix merge conflict bug with clearing op outputs
Signed-off-by: Tim Moon <tmoon@nvidia.com>
```
020428f0

29 Jul, 2025 4 commits

Fuse amax computation into activation kernel (#2004) · feda5b55

Jan Bielak authored Jul 29, 2025



* Compute amax in activation kernels when the output pointer is provided, even for non-fp8 outputs
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit 9f13fe2fefc58cae93bc467d87d01ecf792a0381)

* Initialize metatensor values
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Fuse computation of amax into the activation kernel for fp8 current scaling
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit 2b54327ac9c931a5340983a79e99de5caa0399dd)
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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



* Zero out amax in `create_hp_tensor_with_amax` instead of relying on `Float8CurrentScalingQuantizer.__init__` to zero-initialize it
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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



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

feda5b55

Rename `do_not_clear` to `_do_not_clear` (#1977) · f858dc35

Jan Bielak authored Jul 29, 2025


Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

f858dc35

[PyTorch] Refactor C++ quantizer infrastructure (#1952) · cb5013bd

Tim Moon authored Jul 29, 2025



* remove reciprocal op
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Refactor Quantizer::create_tensor function
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



* Fix bug when constructing FP8 tensor
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



* Add quantize function to C++ quantizers
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Prototype function to coerce Python quantized tensors to match quantizer
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use quantizer class in tex.quantize
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add FP8 current scaling support for activation backward
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Disable quantized GEMM output with FP8 current scaling
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add coerce_tensor functions for MXFP8 and DSv3
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



* Avoid quantizing empty tensors
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use consistent shapes for FP8 transposes
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* In attention impl, construct FP8 tensors with pre-initialized scale-invs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Initialize MXFP8 scales to zero
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



* Store copy of quantizer when creating quantized tensors
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Make sure quantized tensors have private quantizer

Avoid problems with in-place ops after quantizer usages are changed externally.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename "coerce_tensor" to "convert_and_update_tensor"
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



* Make sure CUDA context is available when launching NVRTC kernel
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Expose CUDA context creation function externally
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

cb5013bd

Fix the use-after-free bug in unfused normalization (#2002) · 5a495a39
Przemyslaw Tredak authored Jul 28, 2025
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
5a495a39

25 Jul, 2025 5 commits

[PyTorch] Add ops for dropout and constant scale (#1995) · c6c1f50e

Tim Moon authored Jul 25, 2025



* Add ops for dropout and constant scale
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



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

c6c1f50e

Fixed double buffering issue for assymetric layers (#1984) · 38c26dd8

Selvaraj Anandaraj authored Jul 25, 2025



* Fixed double buffering issue for assymetric layers
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>

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

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



---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

38c26dd8

[C][PyTorch] Remove deprecated `device_id` arg for multi tensor API (#1994) · 1470116e

Kirthi Shankar Sivamani authored Jul 25, 2025



* Remove deprecated device arg
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

1470116e

[PyTorch] Enable generic QK norm support (+ RMSNorm/LayerNorm) (#1966) · 374849e3

Evgeny Tsykunov authored Jul 25, 2025



* Support RMSNorm for QK
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

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



* rms -> RMSNorm, l2 -> L2Normalization (align with current pattern)
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Support LayerNorm + init refactor
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Before/after RoPE
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

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



* Fix pylint
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

374849e3

[PyTorch] Optimize cudagraph static_grad_outputs reuse (#1992) · e950ceb0

buptzyb authored Jul 25, 2025



* optimize static grad outputs
Signed-off-by: Robin Zhang <robinz@nvidia.com>

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

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



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

e950ceb0

24 Jul, 2025 3 commits

Fix cudnn versioning support in PyTorch DPA and Fused attn (#1991) · 71b2dd48

Kshitij Lakhani authored Jul 24, 2025



Fix cudnn versioning in support in PyTorch DPA and Fused attn
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>

71b2dd48

Add `in_place` kwarg to extra tensor ops (#1983) · ee841084

Jan Bielak authored Jul 23, 2025



* Mark output tensors as not deletable in backward
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Add `in_place` kwarg to `MakeExtraOutput`
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Rename `AddInPlace` to `AddExtraInput` and add an `in_place` kwarg
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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



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

ee841084

[PyTorch] Improve L2Normalization basic op (#1964) · dab931a7

Evgeny Tsykunov authored Jul 24, 2025



* Increase intermediate precision and reuse tensors from fwd
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

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



* JIT warmup only when required
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

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



* Recompute only rsqrt_norm
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

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



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

dab931a7

23 Jul, 2025 2 commits

Fix the device for cuDNN/cuBLAS handles (#1974) · 4296b7d0

Charlene Yang authored Jul 23, 2025



* fix current device for cuDNN/cuBLAS handles
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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



* add unit test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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



* use weight device and improve tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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



---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

4296b7d0

[PyTorch] Reset recipe state in fusible operations when FP8 amax history length changes (#1985) · fdb87afc

Tim Moon authored Jul 23, 2025



* Fix bug where TE ops were not updating fp8_meta dicts
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename reset_recipe_state function
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



* Update error message when initializing meta device quantized weight without recipe
Signed-off-by: Tim Moon <tmoon@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>

fdb87afc

22 Jul, 2025 2 commits

Refactor `te.ops` (#1951) · e0204fbb

Jan Bielak authored Jul 22, 2025



* Refactor _OperationFuserAutogradFunction.forward to use less parameters
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit f8f59b1bb184e89468058521df4cfff029ad909c)

* Rename `BackwardBiasActivation` to `BackwardActivationBias`
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit 397c58fc296f801fe4ad600aadc2daff3b78be45)

* Use forward operation order in backward fused operations
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit 2d37a9385069b066e6cdeff3eb9173c2079cb791)

* Rename `prev_op_grad_input_quantizer` to `prev_op_grad_output_quantizer`
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit d7ab5dfb23e216866f7f4fc4d7a99f625d329f1e)

* Make OperationFuser persistent
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit 77984d9715d31e87519dc6ea1e02c483a81355a7)

* Distribute extra inputs to and collect extra outputs from multiple module groups in Sequential
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
(cherry picked from commit 0716aaad542e59f2c1ac4620167965a0334bbf71)

* Take requires_grad into account when fusing operations
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Change get_quantizer to return None if no quantization recipe is used
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Refactor pre_first_forward
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Fix for failing `test_make_graphed_callables[fp8_recipe0-*-True-*-linear_op]`
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Fix linting errors
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Apply suggestions from code review
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Fix fp8 meta tensors in CUDA Graph capture
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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



* Fix failing distributed userbuffers tests
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-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>

e0204fbb

[PyTorch] Debug linear layer when saving original input and using debug quantizer (#1963) · 315b47db

Tim Moon authored Jul 21, 2025



* Debug linear layer when saving original input and using debug quantizer
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Workaround bugs with quantizing with only column-wise usage
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



* Remove unused imports
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Avoid unnecessary row-wise data
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Workaround bugs with quantizing with only column-wise usage

FP8 does not support transpose-only cast.
Signed-off-by: Tim Moon <tmoon@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>

315b47db

21 Jul, 2025 2 commits

Fix the condition error when checking fp8 attn in `get_attention_backend` (#1965) · ab5cc407

yuzhongw-nvidia authored Jul 22, 2025



Update utils.py

Fix the condition error of the FP8 attention in `get_attention_backend`
Signed-off-by: yuzhongw-nvidia <yuzhongw@nvidia.com>
Co-authored-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com>

ab5cc407

[PyTorch] Reset FP8 weight workspace if usages are invalid (#1972) · 78a38212
Tim Moon authored Jul 21, 2025
```
Reset FP8 weight workspace if usages are invalid
Signed-off-by: Tim Moon <tmoon@nvidia.com>
```
78a38212

17 Jul, 2025 3 commits

[Common] Optimize KV cache related kernels (#1914) · f8933bbf

Charlene Yang authored Jul 17, 2025



* optimize kv_cache reindex and copy kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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



* avoid reindexing from python side
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* rename variable from previous commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

f8933bbf

[PyTorch] Add save_original_input in Linear/GroupedLinear to save memory (#1865) · 07afda98

hx authored Jul 17, 2025



* save original input
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



* fix input_quantizer usage in Linear bwd
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* minor fix
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* refine the docstring
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* Merge remote-tracking branch 'origin/main' into save_bf16_in_fp8_gemm
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



* decouple linear bwd with save_original_input; clean up UTs
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



---------
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>
Signed-off-by: hx <hongxiaob@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

07afda98

mxfp8 (for all gemm layouts) is not supported on 120+ arch yet (#1939) · fa91ed72

Sudhakar Singh authored Jul 16, 2025



* mxfp8 is not supported on 120+ arch yet
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* change the default recipe for arch 120
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

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

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



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

fa91ed72

16 Jul, 2025 2 commits

[Pytorch] Dynamo ONNX export support (#1497) · 0a1499fa

Paweł Gadziński authored Jul 16, 2025



* some initial code
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* onnx support
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* mxfp8 support
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fixed returning layernorm etc
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* formatting
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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



* license fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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



* tests passing
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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



* refactor
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* lint
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fixes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* added pip install to test.sh
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Update transformer_engine/pytorch/export.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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



* float8currentscaling quantizer exception
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



* added to wheels
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* onnx versions
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* installations in tests
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: root <root@prenyx0221.a51.clusters.nvidia.com>

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

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



* lint fix
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: root <pgadzinski@nvidia.com>

* fixes
Signed-off-by: root <pgadzinski@nvidia.com>

* fixes
Signed-off-by: root <pgadzinski@nvidia.com>

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

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



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



* fixes
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: Pawel Gadzinski <pgadzinski@nvidia.com>

* Update setup.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>

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

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



* onnxscript version chnage
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



* fixes
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 CI
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@gmail.com>

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

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

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Signed-off-by: root <root@prenyx0221.a51.clusters.nvidia.com>
Signed-off-by: root <pgadzinski@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@gmail.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: root <root@prenyx0221.a51.clusters.nvidia.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@gmail.com>

0a1499fa

Bump up FA to 2.8.1 (#1949) · 6c526794
vcherepanov-nv authored Jul 15, 2025
```
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
```
6c526794

14 Jul, 2025 2 commits

[PyTorch] Fix bugs in router fusion (#1944) · 397c4be6

Autumn1998 authored Jul 14, 2025



* fix underterminsic problem in CI
Signed-off-by: tongliu <tongliu@nvidia.com>

* fix bug on mbs>1
Signed-off-by: tongliu <tongliu@nvidia.com>

* fix bug on sm dispatcher
Signed-off-by: tongliu <tongliu@nvidia.com>

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

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



* fix CI initial values
Signed-off-by: tongliu <tongliu@nvidia.com>

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

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



---------
Signed-off-by: tongliu <tongliu@nvidia.com>
Co-authored-by: tongliu <tongliu@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>

397c4be6

[PyTorch] Optimize the performance of permute fusion kernels (#1927) · dc97cc9e

hx authored Jul 14, 2025



* optimize permute
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



* fix lint
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: Hongxiao Bai <hongxiaob@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>
Co-authored-by: Xin Yao <xiny@nvidia.com>

dc97cc9e

12 Jul, 2025 1 commit

Add backward fusions of dbias+quantize and dbias+dactivation+quantize to `te.Sequential` (#1942) · 37da2d3b

Jan Bielak authored Jul 11, 2025



* Fix clearing tensor data in backward removing is_first_op
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Misc fixes
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Use Linear weight dtype and device for compute consistently
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Add backward dbias + quantize fusion
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Pass recipe to OperationFuser to allow recipe-dependent fusions
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Remove redundant view from activations
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Add bias activation backward fusion
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Apply suggestions from code review
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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



---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-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>

37da2d3b

11 Jul, 2025 1 commit
- [PyTorch] Make `MXFP8Tensor` unpickling function backward compatible (#1943) · 31fc29ab
  Tim Moon authored Jul 10, 2025
```
Make MXFP8Tensor unpickling function backward compatible
Signed-off-by: Tim Moon <tmoon@nvidia.com>
```
  31fc29ab