Skip to content

Upgrade/v2.17 gated 20260810 to main - #102

Open
zhaoyinglia wants to merge 527 commits into
flagos-ai:mainfrom
zhaoyinglia:upgrade/v2.17-gated-20260810
Open

Upgrade/v2.17 gated 20260810 to main#102
zhaoyinglia wants to merge 527 commits into
flagos-ai:mainfrom
zhaoyinglia:upgrade/v2.17-gated-20260810

Conversation

@zhaoyinglia

Copy link
Copy Markdown
Collaborator

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

jomitchellnv and others added 30 commits March 13, 2026 12:29
…kScaling and Float8BlockScaling quantized model init. (NVIDIA#2753)

* Updates FusedAdam with FSDP2 and MXFP8

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>

* removes xfailing unit test for MXFPr MXFP8

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>

* addresses comments related to reset parameters and guard against self.capturable

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>

* adds e2e unit test

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>

* adds test to non meta device init

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>

* attempts to add float8block scaling fsdp hooks

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>

* adds e2e test for Float8BlockScaling

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>

* addresses review comments and code cleanup

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>

* more review comments addressed

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>

* removes unused block_len param

Signed-off-by: Jonathan Mitchell <jomitchell@umb-b300-dp-147.ipp4a1.colossus.nvidia.com>

* fixes failing unit test because we still need to xfail nvfp4 dcp

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>

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

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

* lint - replacing todo with note

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1429.ipp1a1.colossus.nvidia.com>

---------

Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>
Signed-off-by: Jonathan Mitchell <jomitchell@umb-b300-dp-147.ipp4a1.colossus.nvidia.com>
Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1429.ipp1a1.colossus.nvidia.com>
Co-authored-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com>
Co-authored-by: Jonathan Mitchell <jomitchell@umb-b300-dp-147.ipp4a1.colossus.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: Jonathan Mitchell <jomitchell@ipp1-1429.ipp1a1.colossus.nvidia.com>
Signed-off-by: Peter St. John <pstjohn@nvidia.com>
* fix for async dcp checkpointing

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

* Apply suggestions from code review

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Peter St. John <peterc.stjohn@gmail.com>

* Update transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
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

* Address Greptile review feedback: defensive guards for edge cases

- Add _quantizer None guard in new_empty dispatch
- Replace self.is_cpu with explicit _data/_transpose checks in __reduce_ex__
- Make get_metadata() safe for cleared tensors (both _data and _transpose None)

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

---------

Signed-off-by: Peter St. John <pstjohn@nvidia.com>
Signed-off-by: Peter St. John <peterc.stjohn@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…ped Tensor Swizzling (NVIDIA#2669)

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

* remove changes not needed for bf16

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

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

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

* keep only pytorch binding for now

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

* linting error

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* add fast accumulator support

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

* MXFP8 grouped GEMM + tensor-scaled FP8 fixes

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

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

* Change version to 13.3

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* fix the test

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

* Random padding condition shouldnt be done for mxfp8

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Remove incorrect comment

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* CUBLAS > 13.2 is enough

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* CUBLAS version needed for MXFP8 indeed seems to be 13.3

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* all changes for grouped gemm

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

* Accidental line removal added back. Plus need changes ci t trigger

Add documentation for scaling factors in common.h

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Update cuBLAS version requirement for MXFP8 support

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* grouped gemm: address code review comments

- Replace nvte_set/get_grouped_tensor_swizzled_scales with nvte_set_grouped_tensor_param
- Add host-side validation: A and B must use same scaling mode (both MXFP8 or both tensor scaling)
- Add host-side validation: A and B must both be FP8 or both non-FP8; restrict inputs to FP8/BF16
- Restrict output (C/D) to BF16/FP32; remove FP16 from supported types
- Refactor workspace allocation: replace manual offset arithmetic with moving pointer pattern
- Use void* + NVTEScalingMode in setup kernel instead of separate float*/char* scale params
- Extract use_columnwise(swap_dims) helper to eliminate duplicated MXFP8 columnwise blocks
- Split set_fp8_scale_pointers into set_fp8_scale_pointers / set_mxfp8_scale_pointers
- Remove scale_inv_ptrs from GroupedOperandSelection; pass workspace pointers directly
- Move swizzled-scales validation into validate_grouped_gemm_inputs for fail-fast behavior
- Add use_split_accumulator to GroupedMatmulConfig (Hopper only, default false)
- Add FP8 test case with per-tensor scales; add BF16/MXFP8 shape-varying test cases

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

* address reviee comments

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

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

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

* missed merged conflict handling

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

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

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

* minor change

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* forgot adding a or

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

* resolve merge conflicts

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* address review comments

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

* address minor review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* remove unecessary code

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

* one line that broke everything :(

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

* address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* address review comments

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

* Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* address review comments

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

* unecessary

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* revert caching changes

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* fix minor bug from greptile

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* revert for now

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

* address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* address review comments

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

* address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* address review commentsgp

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

---------

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-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: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
…VIDIA#2757)

* [Common] Fix linker error for to_string(DType) in distributed tests

Make transformer_engine::to_string(DType) inline in common.h so that
translation units outside libtransformer_engine.so can resolve it
without requiring the symbol to be exported.

Regression introduced by 61f9594 which added to_string(DType) calls
into TRANSFORMER_ENGINE_TYPE_SWITCH_* macros, causing test object files
to reference the symbol that the linker version script hides.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

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

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

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

---------

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…-Cast Fusion Kernel (NVIDIA#2555)

* first draft

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* pass numerical unit test

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* format

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* add benchmark script

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* lint and format

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* compile guard

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* warning fix

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* resolve greptile comment

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* minor style fixes

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix namespace

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* resolve some comments

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix comment

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* attempt to fix compile CI with guard

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

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

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

* better naming for tests

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix deprecate messsage

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* more compile guard

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* new API name

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix format all in one

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* try to fix compile CI again

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* AI code review comments

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* to pass oldest compile CI with cuda 12.1

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* add more guards to nvfp4

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* make multiply inverse default numerics

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* update numerics of nvfp4 partial cast as well

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* resolve comments

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* add NVTE_BUILD_NUM_PHILOX_ROUNDS after rebase

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* simplify compile guard messsages

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

---------

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…Linear` (NVIDIA#2761)

* Load multi-param checkpoint from single-param config in GroupedLinear

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

* Multi-param to single param case

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

* Multi-param to single param case

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

* Better varnames

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

---------

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…A#2765)

* Fix GMM cuBLAS version and SM arch checks

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

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

---------

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
…DIA#2773)

* Pin python 3.13 in vermin check

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

* Update vermin version for python 3.14 support

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

* Use sha instead of tag

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

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

* work finished

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>

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

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

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

* removed warning.warn

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

* [PyTorch] Remove dead None-check for num_out_tokens in moe_permute_mask_map_forward

num_out_tokens is typed as int in the custom_op signature and can never
be None; the check was incorrectly carried over from the class-based
upstream version during merge conflict resolution.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <pgadzinski@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
…are detached (NVIDIA#2772)

[PyTorch] Change the restore tensor API to ensure tensors are detached from ctx

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…ges it (NVIDIA#2781)

Install pytest in onnx L1 test as Pyt container no longer packages it

Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>
…_descriptors (NVIDIA#2782)

* Fix zero-sized groups in update_tma_descriptors

Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>

* Update test_cast_mxfp8_grouped.cu

Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>

* Apply suggestions from code review

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@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: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
NVIDIA#2693)

* Enable sm120 support for fused attn if cuDNN is 9.18.1+

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

* Force intermediate tensors such as S, Sum_Exp, and Max to be BHS1 shape instead of TH1 for sm120

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

* Add support for sm120 correct batch, seq dims

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* Add support for sm120 BHS1 style max logit even QKV are THD to avoid incorrect max logit calculation (includes padded tokens in max calculation)

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* Disable fused and flash attn for sm120 filter:kv cache

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* For CP P2P attn, set softmax_lse_in_packed_format to False if sm120+

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* Assert in TE if T3HD/TH3D layout is used on sm120 before cuDNN F16 sdpa arbitrary kernel call

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* Modify is_ragged_q && cudnn_runtime_version >= 90600 check to also include a check for sm120

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

* nit: Code clean up

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* Disable fused attn for T3HD and TH3D

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* nit: Add missed sm120 guard

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* Modify sm120 condition to be very specific to sm120 and not generalized to sm120+

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

* nit: Fix missing sm120 check in fwd

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* Move the check for sm120 T3HD/TH3D to nvte_get_fused_attn_backend() instead of higher layers in TE stack

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* nit: Check for matching sm120 and not sm120+

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 <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* code drop

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

* code drop

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

* docs

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

* nvfp4 internals support

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

* lint fixes

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

* Update transformer_engine/debug/features/dump_tensors.py

Co-authored-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

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

* Update transformer_engine/debug/features/dump_tensors.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>

* Update transformer_engine/debug/features/dump_tensors.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.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

* Update tests/pytorch/debug/test_log.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>

* Update transformer_engine/debug/features/dump_tensors.py

Co-authored-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

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

* fix

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

* Remove dump_quantized_internals support from DumpTensors

Drop the dump_quantized_internals config option, the _get_quantized_internals
method, and all helper functions for extracting scales/raw data from
Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, and NVFP4Tensor.

Remove corresponding tests: test_dump_tensors_nvfp4_unpacked_codes and
NVFP4_DUMP_TENSORS_CONFIG, and scale/data assertions from test_dump_tensors_sanity.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
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

* Address Greptile review comments

- Add dot ('.') to _sanitize_name to handle common PyTorch dotted layer
  names like 'encoder.layer.0.attention'
- Add docstring note about pickle dependency for the 'quantized' key
- Add comment explaining weights_only=False in test
- Remove redundant local RecipeState import in test_nvfp4_numeric

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Remove portability suggestion from quantized key docstring

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Compute rank lazily in _expected_root_dir

Avoids relying on stale self.rank when ensure_initialized is called
before initialize() has set the rank. Consistent with how nvdlfw_inspect
logger resolves rank.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* detach tensors before saving; verify dump filename in test

Detach both high_precision and quantized tensors before saving to avoid
serializing the autograd graph. For QuantizedTensor this is a zero-copy
view (make_like), so no extra GPU allocation.

Add filename format assertion to test_dump_tensors_sanity to catch
regressions in _sanitize_name or the naming convention.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
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 empty dump_dict log; assert QuantizedTensor type in test

Log a message when no tensors are available to dump so the user
has an explicit signal that no file was written.

Assert that the quantized key round-trips as a QuantizedTensor
to catch regressions in detach() or serialisation path.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
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

* Update transformer_engine/debug/features/dump_tensors.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>

* Address review: iter subdirs, remove dead rank field, add allclose test and MSE example

- Organize dumps into per-iteration subdirectories (iter_000000/) to keep
  file count manageable per directory.
- Remove unused self.rank attribute from TensorLogger.
- Add torch.allclose assertion in test to verify serialization correctness.
- Add docstring example showing how to load dumps and compute MSE.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

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

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

* fix: use detach().clone() to avoid shared storage in DumpTensors

Using tensor.detach() creates a view sharing the same underlying
storage. If any in-place operation modifies the tensor after the
dump, the saved data would be silently corrupted. Use .clone()
to ensure the dump captures an independent copy of the data.

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

* test: use torch.equal instead of torch.allclose for serialisation round-trip

The saved tensor is an exact bit-for-bit copy (detach().clone()), so
torch.equal is the correct check. torch.allclose with its default
tolerances could mask a genuine dtype conversion or precision loss
introduced by a future change to the serialisation path.

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

* fix: add tp_size to DumpTensors.inspect_tensor and fix KeyError in call_feature backward compat pop

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

---------

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <pgadzinski@nvidia.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* change distributed tests infra for fsdp2

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* verbose flag for reporting

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* add back coments

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

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* another minor fix

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* not needed for this PR

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* address review comments

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

* unecessary comments
* add cudnn dln+add

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

* try fixing cudnn build issue

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

* guard against cudnn version

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

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

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

* change itype to wtype for add in rmsnorm_bwd

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

* remove dead code

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

* remove dangling todo

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

---------

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* add blackwell support filter for 9.7<=cudnn<9.18.1

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

* simplify conditionals

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix conditionals again

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix conditionals again

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* update the error log

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* remove the python filter and correct the cpp filter

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…quired (NVIDIA#2798)

* Disable fused attention for sm120 if determinism is required

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* nit: disable fused attn for sm120 determinism, if training

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

---------

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
… always and `Max` when `return_max_logit=True` (NVIDIA#2677)

* cudnn now returns Stats always and Max only with `return_max_logit=true`

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix a typo that caused a bug

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* update doc strings

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

* fix more docs

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fixes from the feedback

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* update cudnn-frontend to v1.19.1

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* update the cudnn frontend

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix a wrong omission

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>
Move cuDNN FE to v1.21.0

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

* Add warning if using BSHD and max_segments_per_seq > 1

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

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

* Update transformer_engine/jax/attention.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>

* Update transformer_engine/jax/attention.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>

* Apply suggestions from code review

Co-authored-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com>
Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>

* Remove warning test

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

---------

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com>
)

* Refactor to group_sizes per tensor

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Support first_dims and last_dims instead of a single group_sizes per
tensor

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Refactor GMM FFIs to store static attrs as structs

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

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

* Cleanup C++ v2 FFI

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Fix int64 workspace usage

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Address greptile comments

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Refactor wgrad-specific checks to be generic for GMM in gemm.py

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Refactor XLA FFI struct setup

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Fix edge case in TE v1 GMM

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

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

* Fix issues on Hopper

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Refactor

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

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

* Address comments

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

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

* Lint

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Fixes for Hopper

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Address review comments

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

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

* Grouped quantization test fixes

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

---------

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Pass input_output_alias to TritonAutotunedKernelCall

Signed-off-by: JAX Toolbox <jax@nvidia.com>

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

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

* Add jax version guard for the input_output_aliasing fix

Signed-off-by: tdophung <tdophung@nvidia.com>

---------

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

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* one review comment form greptile

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* instead part of the comment not needed

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

* address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* Update transformer_engine/pytorch/tensor/float8_blockwise_tensor.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* No need to set it to None

Remove unnecessary columnwise data and scale inv assignments.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

---------

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
* cudnn now returns Stats always and Max only with `return_max_logit=true`

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix a typo that caused a bug

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* update doc strings

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

* fix more docs

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fixes from the feedback

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* update cudnn-frontend to v1.19.1

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* update the cudnn frontend

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix a wrong omission

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

* bugfix: mask out padding tokens when THD

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

* fixes from greptile feedback

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

* minor nit

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fixes from feedback

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Enabled persistency with WorkID Query feature

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added a struct with tunable parameters

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added persistency with static scheduling

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed test cases

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Ready for benchmarking

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed out-of-boundary error

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Tuned kernel parameters

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Refactoring

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Refactoring 2

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Refactoring 3

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Removed the dynamic (WorkID Query) persistency

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Ready for PR

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Fixes per the review

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Ready for benchmark

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Ready for benchmark - Regular kernel

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added the source code to the profiler

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added constructors to Job and Block descriptors

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Removed the prefetch overlapping between jobs

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Cache tensor ID

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* ShapeRepresentation is not a template parameter

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Removed redundant fence_proxy

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Refactoring

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Used mixed precision FMA

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added Quantize parameters

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added the fast math branch

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added the fast math to cpp test suite

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Align tests

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Use STS instead of generic ST

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Add zero-tensor cases

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Used LDS instead of generic LD in colwise path

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Used LDS instead of generic LD in rowwise

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Ready for merge

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Uncommented test cases

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added FP16 Fast math path to rowwise processing

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Refactoring

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Fixed lint

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Fix

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixes

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fix

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed test suite

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed test suite

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixes per the review

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Modifications per the review

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Assert the buffer size

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Added fast math RCP for bf16

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fast math for BF16 is now default

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Fixed compilation error when compiling on previous archs

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Boundary condition fix

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed compilation error

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Refactoring. Moved helpers to core-common

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Refactoring

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Refactoring per the review

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Addressed the PR review comments

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Fixed the compilation error when PTX was compiled for CUDA 13.0

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Fixed pytorch extensions

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

---------

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
timmoon10 and others added 26 commits June 11, 2026 13:18
…A#3120)

Revert "[PyTorch] Add op-level activation offload opt-out API (NVIDIA#3108)"

This reverts commit 9b06f26.

Signed-off-by: Tim Moon <tmoon@nvidia.com>
…VIDIA#3117)

* Refactor grouped MLP into joint fused op

Consolidate the experimental grouped MLP forward and backward CuTe DSL fusions into a single joint fused operation. Move grouped-MLP-specific helper logic out of ops/_common.py and update tests to assert the joint forward/backward fusion object.

Co-authored-by: Codex <codex@openai.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Review suggestions from @greptile-apps

Also fix linter warnings.

Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------

Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Codex <codex@openai.com>
…le names (NVIDIA#3086)

Signed-off-by: Teddy Do <tdophung@nvidia.com>
Signed-off-by: tdophung <tdophung@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Enable colwise only 2d nvfp4

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

* Enable colwise only for the main 2d kernel

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

* Resolve comments

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

* Fix test comment

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

* Enable rectangular shapes in tests

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>
* Copy grouped MLP tests from TE ops tests

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Reduce TE ops test cases

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Reduce overparametrized grouped MLP tests

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Clean up grouped MLP test leftovers

Remove unused imports and helpers left after splitting grouped MLP tests out of the fusible ops suite.

Co-authored-by: Codex <codex@openai.com>
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

* Review suggestion from @greptile-apps

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

---------

Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Codex <codex@openai.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Expert Parallelism: common C API + NCCL EP v0.1 backend

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

* Expert Parallelism: persistent ncclEpHandle cache with allow_handle_mem_reloc gating

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

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

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

* Build: NCCL_HOME discovery supports Debian/Ubuntu multiarch lib paths

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

* bump NCCL

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

* Expert Parallelism: require token_dtype in NVTEEpGroupConfig and enforce at dispatch

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

* Expert Parallelism: document ep_comm lifetime, v0.1 single-GPU scope, static layer registration

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

* Expert Parallelism: drop version label from initialize scope note

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

* Expert Parallelism: pointer-keyed LRU handle cache; drop register_layer + NVTEEpHandle struct (NVTE_EP_HANDLE_CACHE_SIZE=-1 disables eviction)

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

* bump nccl to latest v0.1

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

* tests/cpp_distributed: drop unused NCCL EP header include path

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

* common/ep: fold nvte_ep_* stubs into ep_api.cpp under #if NVTE_WITH_NCCL_EP

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

* common/ep: dlopen libnccl_ep.so so libtransformer_engine.so loads without it

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

* common/ep: add BUILD_RPATH=NCCL_EP_LIB_DIR for in-tree dev builds

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

* common/ep: polish ep.h docstrings; drop unused NVTE_CHECK_NCCL from logging.h

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

* common/ep: expose zero_copy in NVTEEpGroupConfig; map to NCCL_EP_ZERO_COPY_{ON,OFF}

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

* tests/cpp_distributed: exercise zero_copy=ON in EPZeroCopyTest.IdentityAllSymm

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

* tests/cpp_distributed: tighten EPZeroCopyTest comments

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

* common/CMakeLists: correct NCCL resolution comment (not bundled with CUDA Toolkit)

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

* common/CMakeLists: shorten NCCL/GIN headers comments

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

* setup,common: bundle libnccl_ep.so.0 next to libtransformer_engine.so for wheel install

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

* .gitmodules: drop nccl branch pin and align indentation with other submodules

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

* setup: gate NCCL EP build on arch >= 90 or native; drop sm_90 fallback

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

* common,setup,tests: discover nccl.h via find_path/NCCL_INCLUDE_DIR; drop submodule header mirror

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

* common/ep: simplify make_nccl_ep_tensor to take NVTETensor and optional CommWindow

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

* common/ep: move te_dtype_to_nccl_dtype out of EPBackend into anon namespace

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

* common/ep: reword multicast check; drop NVLS framing

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

* common,tests: replace unicode em-dash and box-drawing chars with ASCII in EP files

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

* bump nccl to latest v0.1

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

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

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

* nccl commit to 2.31.0a4-1

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

* common/CMakeLists: point NCCL_EP_INCLUDE_DIR at build/include staged headers

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

* common/CMakeLists: clarify NCCL EP missing-header instructions

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

* common/ep: use int64_t instead of long for handle-cache size env (cpplint runtime/int)

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

* common/ep: fix dangling sizes pointer in make_nccl_ep_tensor (NVTEShape lifetime)

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

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

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

* common/ep: static-link libnccl_ep.a into libtransformer_engine.so; drop dlopen loader and runtime bundling

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

* common/ep: lock handle_mem_size, guard tensor ndim, forward-declare ncclWindow in public header

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

* qa,tests/cpp_distributed: emit per-rank JUnit XML and isolate sub-test failures in L1 CI

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

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

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

* common/ep: harden NCCL EP init/probe paths and CI build short-circuit

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

* common/ep: drop Sphinx backticks and consolidate EP docstring

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

* setup: keep _discover_nccl_home local; revert utils.py move

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

* qa/L1_cpp_distributed: build and run each suite independently

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

* qa/L1_cpp_distributed: pass --allow-run-as-root --oversubscribe to mpirun in run_test_ep.sh

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

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

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

* common/ep: check initialized_ under mutex_ in per-step EPBackend ops

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

* setup/ep: rename NVTE_BUILD_WITH_NCCL_EP to NVTE_WITH_NCCL_EP, harden NCCL EP build

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

* tests/cpp_distributed: anchor NCCL discovery on libnccl, derive include dir from its prefix

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

* tests/cpp_distributed: template EP unit tests on token dtype; switch routing to global counter

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

* ep: drop misleading rebuild advice in NCCL EP errors and submodule header version log

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

* common/ep: drop multicast support check; NCCL EP does not require multicast

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

* setup/ep: error out when NVTE_CUDA_ARCHS yields no SM>=90 arch for NCCL EP

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

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

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

* make core to be RTLD_LAZY

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

---------

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

Revert "Expert Parallelism: common C API + NCCL EP backend (NVIDIA#3034)"

This reverts commit c3396ee.

Signed-off-by: Tim Moon <tmoon@nvidia.com>
* update cublaslt grouped gemm filter

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

* update nvfp4 filter and tests

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

* test correctness

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* better test

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

---------

Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
…mel > INT_MAX (NVIDIA#3136)

Widen multi_tensor_apply tensor sizes to int64

TensorListMetadataBase::sizes was declared int32 but populated from
Tensor::numel(), so a tensor with numel > INT_MAX truncated to a negative
size and the consumer kernels then computed out-of-bounds offsets, hitting
an illegal memory access at the next sync.

Store sizes as int64_t and read them into an int64_t (or the existing
index_t on the already-templated Adam kernels) before the
n -= chunk_idx * chunk_size subtraction. Widen the chunk_size kernel
argument in the non-Adam consumers to int64_t as well so the
chunk_idx * chunk_size element offset is computed in 64-bit.

Fixes NVIDIA#2918

Signed-off-by: Javier de Jesus <javier.dejesusj9@gmail.com>
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Expert Parallelism: JAX primitives + VJPs

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…VIDIA#3151)

Revert "Add wheel support for Newton-Schulz method via cuSolverMp (NVIDIA#3004)"

This reverts commit 20e185c.

Signed-off-by: ksivamani <ksivamani@nvidia.com>
…m-mem zero-copy (NVIDIA#3035)

* Expert Parallelism: PyTorch wrapper + autograd ops with symm-mem zero-copy

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

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…NS_PER_RANK (NVIDIA#3150)

* nccl with relax num_dispatch_tokens%64!=0

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

* Skip EP tests/examples on nodes without NVLink

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

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…` with `total_recv_tokens_per_rank` placeholder (NVIDIA#3154)

* versioning EP C configs

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

* Rename EP prepare token_counts to recv_tokens_per_expert

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

* Add total_recv_tokens_per_rank placeholder to nvte_ep_prepare

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

* Adapt PyTorch EP binding to versioned nvte_ep C config API

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

* Rename EP group config max_num_sms to num_comm_sms

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

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…L2 Jax dist (NVIDIA#3159)

* Keep the routing map format alive

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

* Fix incorrectly launched multi process EP tests in L2 Jax instead of L2 jax dist

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

---------

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* disable 9.23.0/.1 for mxfp8 attention

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>
# Conflicts:
#	.github/actions/build-pytorch-wheel/Dockerfile
#	.github/workflows/build.yml
#	.github/workflows/lint.yml
#	.gitignore
#	README.rst
#	benchmarks/linear/benchmark_grouped_linear.py
#	build_tools/VERSION.txt
#	build_tools/wheel_utils/Dockerfile.aarch
#	build_tools/wheel_utils/Dockerfile.x86
#	build_tools/wheel_utils/build_wheels.sh
#	docs/api/pytorch.rst
#	docs/conf.py
#	docs/debug/3_api_features.rst
#	docs/envvars.rst
#	docs/examples/op_fuser/op_fuser.rst
#	docs/features/low_precision_training/index.rst
#	docs/index.rst
#	examples/jax/collective_gemm/common.py
#	examples/jax/collective_gemm/run_test_cgemm.sh
#	examples/jax/collective_gemm/test_gemm.py
#	examples/pytorch/quantized_model_init/fully_shard.py
#	qa/L0_jax_unittest/test.sh
#	qa/L0_pytorch_debug_unittest/test.sh
#	qa/L0_pytorch_unittest/test.sh
#	qa/L1_jax_distributed_unittest/test.sh
#	qa/L1_pytorch_distributed_unittest/test.sh
#	qa/L1_pytorch_thunder_integration/test.sh
#	qa/L2_jax_distributed_unittest/test.sh
#	qa/L3_pytorch_FA_versions_test/test.sh
#	setup.py
#	tests/cpp/operator/test_cast_mxfp8_grouped.cu
#	tests/cpp/operator/test_cast_nvfp4_transpose.cu
#	tests/cpp/operator/test_grouped_gemm.cu
#	tests/cpp/operator/test_swizzle.cu
#	tests/cpp/test_common.cu
#	tests/cpp/test_common.h
#	tests/jax/test_custom_call_compute.py
#	tests/jax/test_fused_attn.py
#	tests/jax/test_fused_router.py
#	tests/pytorch/attention/run_attention_with_cp.py
#	tests/pytorch/attention/test_attention.py
#	tests/pytorch/attention/test_attention_with_cp.py
#	tests/pytorch/debug/test_api_features.py
#	tests/pytorch/debug/test_log.py
#	tests/pytorch/debug/test_perf.py
#	tests/pytorch/debug/test_sanity.py
#	tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py
#	tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py
#	tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py
#	tests/pytorch/distributed/run_numerics_exact.py
#	tests/pytorch/distributed/test_cast_master_weights_to_fp8.py
#	tests/pytorch/distributed/test_comm_gemm_overlap.py
#	tests/pytorch/distributed/test_torch_fsdp2.py
#	tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
#	tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py
#	tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py
#	tests/pytorch/nvfp4/test_nvfp4_group_quantize.py
#	tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py
#	tests/pytorch/nvfp4/test_nvfp4_module_exact.py
#	tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py
#	tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py
#	tests/pytorch/test_cpu_offloading.py
#	tests/pytorch/test_custom_recipe.py
#	tests/pytorch/test_float8_current_scaling_exact.py
#	tests/pytorch/test_float8blockwisetensor.py
#	tests/pytorch/test_fused_router.py
#	tests/pytorch/test_fusible_ops.py
#	tests/pytorch/test_grouped_tensor.py
#	tests/pytorch/test_numerics.py
#	tests/pytorch/test_onnx_export.py
#	tests/pytorch/test_permutation.py
#	tests/pytorch/test_quantized_tensor.py
#	tests/pytorch/test_recipe.py
#	tests/pytorch/test_sanity.py
#	tests/pytorch/utils.py
#	transformer_engine/__init__.py
#	transformer_engine/common/CMakeLists.txt
#	transformer_engine/common/__init__.py
#	transformer_engine/common/cast/cast.cu
#	transformer_engine/common/cast/core/common.cuh
#	transformer_engine/common/cast/dispatch/dequantize.cuh
#	transformer_engine/common/cast/dispatch/quantize.cuh
#	transformer_engine/common/cast/fp8/gated_fp8.cuh
#	transformer_engine/common/cast/fp8/quantize_fp8.cuh
#	transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh
#	transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
#	transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh
#	transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
#	transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
#	transformer_engine/common/cast/mxfp8/swizzle.cuh
#	transformer_engine/common/cast/nvfp4/core_nvfp4.cuh
#	transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh
#	transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh
#	transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh
#	transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh
#	transformer_engine/common/comm_gemm/comm_gemm.cpp
#	transformer_engine/common/common.cu
#	transformer_engine/common/common.h
#	transformer_engine/common/fused_attn/fused_attn.cpp
#	transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu
#	transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h
#	transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu
#	transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h
#	transformer_engine/common/fused_attn/fused_attn_fp8.cu
#	transformer_engine/common/fused_attn/fused_attn_fp8.h
#	transformer_engine/common/fused_attn/utils.h
#	transformer_engine/common/fused_router/fused_moe_aux_loss.cu
#	transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu
#	transformer_engine/common/fused_router/fused_topk_with_score_function.cu
#	transformer_engine/common/fused_router/utils.h
#	transformer_engine/common/gemm/cublaslt_grouped_gemm.cu
#	transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu
#	transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu
#	transformer_engine/common/hadamard_transform/group_hadamard_transform.cu
#	transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu
#	transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu
#	transformer_engine/common/hadamard_transform/hadamard_transform.cu
#	transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu
#	transformer_engine/common/include/transformer_engine/fused_attn.h
#	transformer_engine/common/include/transformer_engine/gemm.h
#	transformer_engine/common/include/transformer_engine/swizzle.h
#	transformer_engine/common/include/transformer_engine/transformer_engine.h
#	transformer_engine/common/include/transformer_engine/utils.h
#	transformer_engine/common/multi_tensor/compute_scale.cu
#	transformer_engine/common/multi_tensor/scale.cu
#	transformer_engine/common/normalization/common.cpp
#	transformer_engine/common/recipe/__init__.py
#	transformer_engine/common/recipe/nvfp4.cu
#	transformer_engine/common/swizzle/swizzle.cu
#	transformer_engine/common/transformer_engine.cpp
#	transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
#	transformer_engine/common/triton/permutation.py
#	transformer_engine/common/util/utils.cu
#	transformer_engine/debug/features/api.py
#	transformer_engine/debug/features/fake_quant.py
#	transformer_engine/debug/features/log_fp8_tensor_stats.py
#	transformer_engine/debug/features/per_tensor_scaling.py
#	transformer_engine/jax/attention.py
#	transformer_engine/jax/cpp_extensions/__init__.py
#	transformer_engine/jax/cpp_extensions/gemm.py
#	transformer_engine/jax/cpp_extensions/quantization.py
#	transformer_engine/jax/cpp_extensions/router.py
#	transformer_engine/jax/csrc/extensions.h
#	transformer_engine/jax/csrc/extensions/attention.cpp
#	transformer_engine/jax/csrc/extensions/gemm.cpp
#	transformer_engine/jax/csrc/extensions/inspect.cpp
#	transformer_engine/jax/csrc/extensions/misc.h
#	transformer_engine/jax/csrc/extensions/pybind.cpp
#	transformer_engine/jax/csrc/extensions/quantization.cpp
#	transformer_engine/jax/csrc/extensions/router.cpp
#	transformer_engine/jax/debug/experimental/inspect.py
#	transformer_engine/jax/dense.py
#	transformer_engine/jax/flax/__init__.py
#	transformer_engine/jax/flax/module.py
#	transformer_engine/jax/flax/transformer.py
#	transformer_engine/jax/permutation.py
#	transformer_engine/jax/quantize/tensor.py
#	transformer_engine/jax/router.py
#	transformer_engine/jax/triton_extensions/permutation.py
#	transformer_engine/jax/triton_extensions/utils.py
#	transformer_engine/jax/version_utils.py
#	transformer_engine/pytorch/__init__.py
#	transformer_engine/pytorch/attention/dot_product_attention/backends.py
#	transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
#	transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
#	transformer_engine/pytorch/attention/dot_product_attention/utils.py
#	transformer_engine/pytorch/attention/multi_head_attention.py
#	transformer_engine/pytorch/constants.py
#	transformer_engine/pytorch/cpp_extensions/fused_attn.py
#	transformer_engine/pytorch/cpp_extensions/gemm.py
#	transformer_engine/pytorch/cpu_offload.py
#	transformer_engine/pytorch/cpu_offload_v1.py
#	transformer_engine/pytorch/csrc/common.h
#	transformer_engine/pytorch/csrc/extensions.h
#	transformer_engine/pytorch/csrc/extensions/attention.cpp
#	transformer_engine/pytorch/csrc/extensions/cast.cpp
#	transformer_engine/pytorch/csrc/extensions/gemm.cpp
#	transformer_engine/pytorch/csrc/extensions/misc.cpp
#	transformer_engine/pytorch/csrc/extensions/pybind.cpp
#	transformer_engine/pytorch/csrc/extensions/router.cpp
#	transformer_engine/pytorch/csrc/extensions/swizzle.cpp
#	transformer_engine/pytorch/csrc/quantizer.cpp
#	transformer_engine/pytorch/csrc/type_converters.cpp
#	transformer_engine/pytorch/csrc/util.h
#	transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py
#	transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py
#	transformer_engine/pytorch/distributed.py
#	transformer_engine/pytorch/experimental/quantization_nvfp4.py
#	transformer_engine/pytorch/jit.py
#	transformer_engine/pytorch/module/base.py
#	transformer_engine/pytorch/module/grouped_linear.py
#	transformer_engine/pytorch/module/layernorm_linear.py
#	transformer_engine/pytorch/module/layernorm_mlp.py
#	transformer_engine/pytorch/module/linear.py
#	transformer_engine/pytorch/ops/__init__.py
#	transformer_engine/pytorch/ops/_common.py
#	transformer_engine/pytorch/ops/basic/__init__.py
#	transformer_engine/pytorch/ops/basic/activation.py
#	transformer_engine/pytorch/ops/basic/basic_linear.py
#	transformer_engine/pytorch/ops/basic/bias.py
#	transformer_engine/pytorch/ops/basic/grouped_linear.py
#	transformer_engine/pytorch/ops/basic/swiglu.py
#	transformer_engine/pytorch/ops/fused/__init__.py
#	transformer_engine/pytorch/ops/fused/backward_activation_bias.py
#	transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py
#	transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py
#	transformer_engine/pytorch/ops/fuser.py
#	transformer_engine/pytorch/optimizers/fused_adam.py
#	transformer_engine/pytorch/permutation.py
#	transformer_engine/pytorch/quantization.py
#	transformer_engine/pytorch/quantized_tensor.py
#	transformer_engine/pytorch/router.py
#	transformer_engine/pytorch/tensor/_quantization_helpers.py
#	transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
#	transformer_engine/pytorch/tensor/float8_tensor.py
#	transformer_engine/pytorch/tensor/grouped_tensor.py
#	transformer_engine/pytorch/tensor/mxfp8_tensor.py
#	transformer_engine/pytorch/tensor/nvfp4_tensor.py
#	transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py
#	transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
#	transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py
#	transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py
#	transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py
#	transformer_engine/pytorch/tensor/utils.py
#	transformer_engine/pytorch/transformer.py
#	transformer_engine/pytorch/triton/__init__.py
#	transformer_engine/pytorch/triton/permutation.py
#	transformer_engine/pytorch/utils.py
@CLAassistant

CLAassistant commented Aug 10, 2026

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 23 committers have signed the CLA.

✅ zhaoyinglia
❌ sraman-rgb
❌ allenphilipj
❌ HollowMan6
❌ cael-ling
❌ Wohox
❌ harryzhou2000
❌ francesco-bertolotti
❌ lhb8125
❌ timmoon10
❌ Oleg-Goncharov
❌ KshitijLakhani
❌ LeSingh1
❌ ksivaman
❌ zianglih
❌ tdophung
❌ fallintoplace
❌ vcherepanov-nv
❌ phu0ngng
❌ yaox12
❌ negvet
❌ javierdejesusda
❌ cyanguwa
You have signed the CLA already but the status is still pending? Let us recheck it.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.