Skip to content

Release 2.15#656

Open
VeeraRajasekhar wants to merge 23 commits into
release_v2.15_rocmfrom
release-sync-v2.15-260630
Open

Release 2.15#656
VeeraRajasekhar wants to merge 23 commits into
release_v2.15_rocmfrom
release-sync-v2.15-260630

Conversation

@VeeraRajasekhar

Copy link
Copy Markdown
Contributor

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

ptrendx and others added 22 commits April 20, 2026 11:45
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
…(#2905)

Zero out padded region when swizzling via group quantize

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* initial implementation for mxfp8

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

* semi-working FP8; broken F16

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

* clean up last commit

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

* comment out F16 pass

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

* pull in grouped_quantize for MXFP8

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

* grouped tensor - pytorch

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

* quantize mxfp8

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

* fix shapes/strides

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

* fix unfused; clean up

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

* split d to d_qk/d_v; attempt at bwd

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

* fix last merge

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

* update FE

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

* attempt at SWA/MLA

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

* remove prints

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

* remove leftover prints

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

* Revert "update FE"

This reverts commit d9ff5662aa4b4b6267c77baf614aada6602fa133.

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

* update FE

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

* fix MLA O strides; add bottom_right_diagonal

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

* attempt at bwd

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

* fix get_quantizers; attempt at bwd

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

* fix fprop; add o_format

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

* attempt at bwd with o_format/d_out_format/dqkv_layout

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

* fix dtype/o_format/etc in bwd calls

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

* fix generateMatrixStridesWithFormats and _v1; fix padding for mxfp8

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

* fix upon last commit for paddedsizes

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

* add mxfp8 env var

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

* disable FA for mxfp8

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

* add mha test

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

* attempt at bwd; force determinism; fix shapes

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

* remove prints

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

* update FE

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

* update FE from pre-merge branch to post-merge develop

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

* allow MXFP8 linear + f16 attn

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

* test cp a2a

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

* remove prints temporarily

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

* test cp p2p

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

* minor fixes for mla

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

* open up a2a for mla

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

* test ag

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

* tweaks for last commit

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

* enable mla ag

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

* fix merge

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

* fix merge

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

* revert to main grouped tensor impl

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

* minor tweaks to return to main

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

* remove prints

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

* fix combine_and_quantize for f16

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

* minor tweaks

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

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

* fix ds descale_o

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

* Revert "fix ds descale_o"

This reverts commit cd0bd82e239ff01210338b4e34cb8784109d22ec.

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

* minor fixes for p2p and ag

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

* tweak cp test skips

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

* update FE

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

* fix bwd KV tensors

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

* tweak recipe control and backend selection

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

* tweak quantizer logic

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

* minor fixes after last two commits

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

* improve generate strides

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

* minor fixes for previous commit

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

* fix bwd for current/delayed

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

* tweak test configs

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

* fix dO/dO_f16 strides

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

* fix tests: SWA logic/test configs

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

* fix ag

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

* add fp8 sink attn

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

* fix a2a comm for F16

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

* remove nan/inf print in test

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

* fix fa a2a

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

* fix fa a2a+p2p f16

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

* update FE to include new fixes

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

* fix thd for bwd

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

* refactor a2a for fu/fa

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

* update FE to fix d64

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

* refactor ag

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

* refactor p2p/a2a+p2p; mostly regarding shapes

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

* add shadow f16 fwd

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

* update FE to fix SWA/BRCM

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

* switch to GH FE temporarily

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

* switch back to GL FE

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

* update FE to latest commit

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

* update group tensor usage after merge main

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

* env vars for qdq(q,k), o_f16 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

* allow other recipes than mxfp8

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

* fix grouped tensor for MLA

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

* change cp test configs

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 shadow f16 bwd

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

* fix a2a+p2p for sbhd

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

* fix last commit and causal flag for fa

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

* enable fp8 sink and disable fp8_mha

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

* minor cleanup for cp/non-cp

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

* update FE for FP8 sink

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

* fix TE for FP8 sink

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

* temporary: random sink/print sink

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

* Revert "temporary: random sink/print sink"

This reverts commit 706095f802e04cbdd5d88ee53849cc5ec938203f.

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

* replace d_out_format with do_format

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

* fix compare_and_assert for None cases

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

* remove logic for b and simplify logic for dqkv types

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

* minor fix for ndim_q/kv

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

* add explanation of fp8_output/grad in MHA

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

* tidy up FP8 checks for bhsd/learnable

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

* remove leading underscores in nvte_convert_qkv_format

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

* simplify logic in generateMatrixStridesWithLayout

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

* clean up strides/ifelse-recipe logic

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

* tweak checks in utils.py

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

* tweak UnfusedDPA

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

* enable testing for ag+swa and disable fp8_mha

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

* tweak FusedAttn, fp8/f16 tensor naming/docstring

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

* replace d_out_format with do_format

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

* fix lint

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

* clean up a2a

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

* clean up ag

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

* clean up p2p/a2a+p2p

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

* tweak test configs

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

* qdq dO in bwd shadow f16 path

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

* tweak qdq dO logic

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

* remove prints in shadow paths

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

* update FE to allow non-determinism

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

* fuse qkv transposes; first pass

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

* remap parallelism to grid(bh, splits, 3) block(s/splits x d); use nvec = 128 bits

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

* allocate contiguous block for qkv

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

* fix grouped tensor row/col scale_inv offsets

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

* use fused permute kernels

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

* quantize row/col as needed in fwd/bwd, non-cp/cp

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

* Revert "quantize row/col as needed in fwd/bwd, non-cp/cp"

This reverts commit ca5376956e8b8f662c7fa88661695b3e9eda4f8f.

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

* Reapply "quantize row/col as needed in fwd/bwd, non-cp/cp"

This reverts commit f19e852be3463210f2b3be5839ae8931e5ad92d0.

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

* fix v_col format when row is quantized

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

* add back necessary bwd quants for shadow paths/cp a2a

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

* remove ZInv for all layouts except T3HD

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

* fix cp p2p with zinv

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

* temporarily switch to GH FE main

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

* switch back to GL FE

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

* fix ag after merge main

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

* add condition for qdq(do) to not affect other tests

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

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

* fix amax dqkv

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

* fix fp8_recipe in DPA utils

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

* remove use of amax for mxfp8

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 o_format/do_format/dqkv_layout to cache indicators for fp8 and f16

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

* enable sink attn + FP8 in CP

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

* update FE to GH v1.22.0

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

* fix for inconsistent kwarg name in permute to grouped tensor

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

* add TMA permute

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

* Revert "add TMA permute"

This reverts commit 2532a50e829144bee290fc94acb8f3f154a62ea9.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* TMA load for bhsd transposes

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

* fix some lint

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

* temp: quant+perm+swizzle, rope, perm_fused

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

* remove mla_rope for now; clean up quant+permute+pad_swizzle; create multi_tensor_swizzle

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

* fix last commit

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

* implement narrow-m for col swizzle; reorder to pad+perm+swizzle

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

* fused pad into perm; remove at::zeros as zeros done in perm kernels

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

* remove shadow code

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

* minor fix for permute shapes

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

* check smem size before entering narrow-k/m kernels

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

* expand permute to multi_tensor_

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

* refactor qkv/do quant; create a fast_path call

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

* fix lint

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

* cleanup grouped tensor fix

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

* remove _with_amax for create_unquantized_tensor

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

* fix last commit

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

* reimplement inplace_multi_tensor_swizzle

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

* fix last commit; set swizzled flag in python

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

* remove permute_to_grouped_tensor_bwd; clean up fwd

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

* add doxygen for multi_tensor_swizzle

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

* clean up nvte_convert_qkv_format

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

* fixes based on code review

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

* group layouts/formats in APIs

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

* rename nvte_convert_qkv_format to nvte_convert_qkv_shape

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

* remove MXFP8 create_unquantized_tensor

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

* rename permute_to_grouped_tensor to transpose_to_bhsd

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

* add multi_tensor_swizzle_xx_unchecked and split the calls/paths

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

* straighten up indexing for multi_tensor_pad

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

* batch up kernel calls per-16-tensors for pad and permute

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

* remove nvec128; rename nvec64 back to nvec

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

* add Macros/arch specifics for compilation

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

* fix lint

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

* attempt 1: MLA RoPE

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

* Revert "attempt 1: MLA RoPE"

This reverts commit 79229248718d26a0ae7029206adc26c687bb42a7.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix kv_cache tests for Fused, is_page=True

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

* attempt 2: MLA RoPE

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

* use DIVUP/_TO_MULTIPLE for pad_s_d_for_mxfp8

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

* remove CUDNN_VERSION 8900 macros

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

* add narrow-k/m swizzle tests

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

* compile flash_attn.cu with special archs

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

* Revert "attempt 2: MLA RoPE"

This reverts commit 3b854b29a3677de2005fecff821d801ccd9bf5d4.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* make contiguous instead of check is_contiguous

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

* remove unused s_q/s_kv

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

* remove unused issue_tma_store_strided

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

* add version gate for mxfp8 for CPP users

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

* replace nvte_get_qkv_shape with AttentionShape

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

* populate nvte_ changes to Jax

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

* update FE to 1.22.1

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

* fix minor merge issue

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

* revert to FE 1.21 since it's what mxfp8 needs

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

* udpate jax attention shapes

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

* Revert "revert to FE 1.21 since it's what mxfp8 needs"

This reverts commit f09961a03bd7f5a316474b5d77b8292c7a49c1a6.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* pick FE 1.22 to support mxfp8 and avoid rng issue in 1.22.1

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

* fix CP AG test on Hopper

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>
Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
fix fp8 and is_bwd_fp8 relationship

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Fix FA4 selection when FA3 is unavailable.

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>
…ed quantization kernel (#2921)

Fix the race in the dbias computation

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* remove ctype to eliminate memory usage from the cudnn kernel

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

* Remove c_dtype from fusible ops test

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

---------

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* [Common, PyTorch] Add triton mHC kernels & pytorch operators

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* fix

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* nit

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* make linter happy

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* nit

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* ah OK

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* new configs to improve perf

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* add APIs to docs

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* fix typos, check deterministic, refactor

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* fix

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

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

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

* reset rng for all tests

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* add docstring

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* fix api doc

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* whoops

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* grad_x doesn't have to zero

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* nit

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* nit

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* force pytorch to not use bf16 for reduction

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* use TE's general_gemm instead

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

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

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

* Looks like this is how to make TE use fp32 acc

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

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

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

---------

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…True (#2936)

* fix

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

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

* cleanup

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

* zero_out should also be tested

Signed-off-by: Varun Thumbe <vthumbe@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: root <root@gb-nvl-059-compute03.nvidia.com>
… (#2924)

* Fix contiguous path for k=2880

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

* format

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

* Review suggestion from @Oleg-Goncharov

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

* Add test for swizzle + padding fusion

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

* Address review comments

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

---------

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

* Avoid removing usages from quantized weight in linear op

Quantized weight tensor may be used across steps, so removing a usage is not safe.

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

* Tweak test to catch bug when alternating train and infer steps

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

* Avoid removing usages from quantized weights in grouped linear op

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

* Restore pre-forward quantizer config in ops

Turns out we still need this in case the quantizer is used before the forward, e.g. in previous ops or CPU offloading.

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

* Blindly preserve quantizer usages in quantized weight params.

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>
…n expert (#2947)

Add workaround for cuteDSL stride requirement for zero token expert

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…pace for THD sequences (#2522)

* Get seqlens and offsets in O(N) space instead of O(N*N) space

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

* Re enable fast causal path

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

* Fix: seqoffsets calculation for THD

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

* Clean up code. Add new comments. Fix unecessary pasing of seg pos to the seqoffsets calculation API

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

* Optimize and fix the slow O(T*T) path for seqlens and seqoffsets calculation for THD non-cp and Cp p2p ring
    - Newer path is O(T*max_segments) per seq
    - Newer path works well with CP p2p ring

    Fix BRCM cross attn by routing to new slow path rather than fast causal path

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

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

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

* Fix lint failure

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

---------

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kshitij  Janardan Lakhani <klakhani@login-ptyche02.ptyche.clusters.nvidia.com>
Co-authored-by: JAX Toolbox <jax@nvidia.com>
… (#2948)

* Switch to cuDNN-FE min version 1.23.0 to enable fused grouped MLP

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

* Fix tests

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

---------

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

* accumulate bias in fp32 instead of bf16 in ref impl dbias to avoid accumulated numerical error

Signed-off-by: tdophung <tdophung@nvidia.com>
…(#2955)

* Better documentation for single param and envvar guard

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

* fix doc

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

* Fix test envvar

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

---------

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: ksivamani <ksivamani@nvidia.com>
- Resolve 9 merge conflicts (1 submodule + 8 content), preserving ROCm
  guards (CK backend selection, IS_HIP_EXTENSION gating) and adopting
  upstream refactors (swizzle check_scale_inv_shapes / unchecked variants,
  fused_attn FP8 validation removal, assert_close test helpers).
- Take upstream cudnn-frontend submodule pin (strictly ahead, unmodified).
- Port fused_attn_rocm/fused_attn.cpp nvte_fused_attn_fwd/bwd to upstream's
  new format-aware C API (o_format/do_format/dqkv_layout/scale_inv formats).
  ROCm CK/AOTriton backends derive layouts from qkv_layout and reject
  FP8/MXFP8, so new params are accepted for parity but unused.
- flash_attn.cu: guard upstream's new TMA (Tensor Memory Accelerator) transpose
  path (CUtensorMap / cp.async.bulk / ptx::mbarrier_*) behind
  #ifndef __HIP_PLATFORM_AMD__; ROCm uses the vectorized/shared-memory fallback
  paths. Force the TMA dispatch off on ROCm (sm_arch returns >=100 there).
- swizzle.cu: move get_max_dynamic_smem() out of the CUDA-only guard so the
  ROCm-compiled swizzle kernels that call it can link (hipify maps the smem
  attribute to hipDeviceAttributeSharedMemPerBlockOptin).
@VeeraRajasekhar VeeraRajasekhar self-assigned this Jul 1, 2026
@VeeraRajasekhar VeeraRajasekhar marked this pull request as ready for review July 2, 2026 02:06
@VeeraRajasekhar VeeraRajasekhar requested a review from Copilot July 2, 2026 02:06
@VeeraRajasekhar VeeraRajasekhar added the ci-level 3 CI test level 3 label Jul 2, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Release cut for TransformerEngine 2.15.0, including updates across PyTorch/JAX attention paths, grouped MLP/linear FP8+MXFP8 behavior, and scale swizzle utilities, plus expanded CI/QA coverage.

Changes:

  • Add/extend MXFP8-related attention support (new QKV formats/layouts, softmax offset grad plumbing, CP updates) across Python and C++/CUDA backends.
  • Refine grouped linear/MLP behavior (env-gated single-param mode, cuDNN frontend gating, updated quantizer usage, wgrad path updates).
  • Add multi-tensor scale swizzle helpers + tests; adjust build/QA scripts for new coverage and release versioning.

Reviewed changes

Copilot reviewed 61 out of 63 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
transformer_engine/pytorch/utils.py Add env-gated resolver for grouped-linear single-param flags (with warnings).
transformer_engine/pytorch/triton/init.py Export mhc from the Triton wrapper package.
transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py Adjust scale shape/offset computations for quantized grouped tensors (incl. block scaling).
transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py Update fused grouped-MLP forward gating and cuDNN kwargs handling.
transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py Update fused grouped-MLP backward gating, wgrad kernel selection, and tensor packing edge cases.
transformer_engine/pytorch/ops/basic/grouped_linear.py Gate single grouped weight/bias via env var; refine weight quantizer usage + recipe-state reset behavior.
transformer_engine/pytorch/ops/basic/basic_linear.py Refine weight quantizer usage + recipe-state reset behavior for BasicLinear.
transformer_engine/pytorch/ops/_common.py Consolidate cuDNN frontend version gate for fused grouped-MLP support.
transformer_engine/pytorch/newton_schulz.py Change coefficient API to structured tuples and flatten only at C++ call boundary.
transformer_engine/pytorch/module/grouped_linear.py Gate single grouped weight/bias via env var in module wrapper.
transformer_engine/pytorch/csrc/util.h Add unchecked multi-tensor swizzle API declaration.
transformer_engine/pytorch/csrc/extensions/swizzle.cpp Implement unchecked swizzle path and add pybind-facing multi-tensor swizzle helpers; improve grouped swizzle padding handling.
transformer_engine/pytorch/csrc/extensions/pybind.cpp Bind multi-tensor swizzle + fused-attn tensor utilities into the PyTorch extension.
transformer_engine/pytorch/csrc/extensions.h Extend fused attention signatures and declare multi-tensor transpose/pad APIs.
transformer_engine/pytorch/cpp_extensions/fused_attn.py Extend fused-attn Python wrapper API with output/scale formats and richer backward outputs.
transformer_engine/pytorch/attention/multi_head_attention.py Adjust FP8 output decisions to account for MXFP8 scaling constraints.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py Extend DPA recipe routing for MXFP8; update layout detection and CP call wiring.
transformer_engine/jax/csrc/extensions/attention.cpp Fix ragged vs non-ragged fused-attn tensor shapes and pass new format args through.
transformer_engine/jax/attention.py Replace O(T²) mask path with O(T·S) per-segment aggregation path for THD seqlen/offset metadata.
transformer_engine/common/util/pybind_helper.h Add new QKV format/layout enum values to Python bindings.
transformer_engine/common/transformer_engine.cpp Make scale shape checking optionally skippable in input validation.
transformer_engine/common/include/transformer_engine/swizzle.h Add unchecked multi-tensor swizzle API for cases with intentionally mismatched scale shapes.
transformer_engine/common/include/transformer_engine/fused_attn.h Extend fused-attn API with new formats/layouts and add a C++ AttentionShape helper.
transformer_engine/common/fused_attn/utils.h Add MXFP8 padding helpers, stride generation by format, and extend fused-attn descriptor keys.
transformer_engine/common/fused_attn/utils.cu Add stride generation support for BHSD layout.
transformer_engine/common/fused_attn/fused_attn.cpp Plumb new fused-attn format/layout args; extend backend eligibility for BHSD and newer cuDNN versions.
transformer_engine/common/fused_attn/fused_attn_fp8.h Update FP8 fused-attn forward/backward signatures to match new format/layout plumbing.
transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h Remove cuDNN version compile-guard wrappers from header.
transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu Remove cuDNN version compile-guard wrappers from compilation unit.
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h Extend f16 arbitrary-seqlen signatures with output/grad formats and dqkv layout.
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Plumb new format/layout args through arbitrary-seqlen kernels and caching keys.
transformer_engine/common/fused_attn_rocm/fused_attn.cpp Accept (but ignore) new fused-attn format/layout args on ROCm for API parity.
transformer_engine/common/common.h Add integer vectorized-load macro and update CheckInputTensor signature default arg.
transformer_engine/common/CMakeLists.txt Move flash_attn.cu into arch-specific CUDA source list.
transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh Add synchronization before dbias reuse path.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh Fix scale indexing for swizzled scales + zero OOB padding; add sync before dbias reuse path.
tests/pytorch/utils.py Handle None outputs in compare helper.
tests/pytorch/distributed/run_newton_schulz.py Update reference to new coefficient tuple format.
tests/pytorch/attention/test_attention.py Expand attention test matrix (new dims, new scaling mode, softmax offset grad, new formats).
tests/pytorch/attention/test_attention_with_cp.py Expand CP attention test coverage (new configs, determinism plumbing, MXFP8 scaling mode).
tests/pytorch/attention/run_attention_with_cp.py Add MXFP8 quantizers + determinism env plumbing; improve NaN/Inf asserts.
tests/jax/test_custom_call_compute.py Make grouped-dense bias grad reference accumulate in fp32 to match primitive.
tests/cpp/operator/test_swizzle.cu Add compact grouped-scale swizzle test and broaden narrow-path coverage.
tests/cpp/operator/test_multi_swizzle.cu Add new C++ tests for multi-tensor swizzle/unswizzle/roundtrip.
tests/cpp/operator/CMakeLists.txt Include new multi-swizzle test; disable it on ROCm for now.
qa/L0_pytorch_unittest/test.sh Enable grouped-linear single-param mode for relevant unit tests; add MHC unittest invocation.
qa/L0_pytorch_lint/test.sh Add a lint QA job for cpplint + pylint.
qa/L0_pytorch_debug_unittest/test.sh Enable grouped-linear single-param mode for debug unittests.
docs/api/pytorch.rst Add API docs entries for Triton MHC fused functions.
ci/pytorch.sh Enable grouped-linear single-param mode for fusible-ops and sanity tests in CI.
build_tools/VERSION.txt Bump version from 2.15.0.dev0 to 2.15.0.
Comments suppressed due to low confidence (1)

transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py:317

  • grouped_gemm_wgrad_kernel() is annotated/used as an optional kernel selector, but it can currently raise ImportError when grouped_gemm_wgrad_wrapper_sm100 is unavailable. Since callers pass cudnn_wgrad_kernel_fn=self.grouped_gemm_wgrad_kernel() unconditionally, this can crash at runtime even though a safe fallback path exists (i.e., returning None so the non-cuDNN wgrad path is used). Catch the import failure and return None instead.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

# See LICENSE for license information.

"""PyTorch wrappers for Triton kernels."""
from transformer_engine.pytorch.triton import mhc
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci-level 3 CI test level 3

Projects

None yet

Development

Successfully merging this pull request may close these issues.