[PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization#1921

Merged

tdophung merged 21 commits intoNVIDIA:mainfrom

021ai:fused_perm_pad

Dec 27, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization#1921

[PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization#1921
tdophung merged 21 commits intoNVIDIA:mainfrom
021ai:fused_perm_pad

Commits on Dec 10, 2025

[PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization

Commits on Dec 11, 2025

[PyTorch/Common] Fuse permute+pad and unpermute+unpad support with_merging_probs

[PyTorch]format code

Commits on Dec 12, 2025

[Common]perf expert_idx loaded once

Merge branch 'main' into fused_perm_pad

Commits on Dec 17, 2025

fix: pad_offsets can be None

Merge branch 'main' into fused_perm_pad

Commits on Dec 19, 2025

add padding + merging probs bwd support. Not tested

Fix garbage initialized act grad

all test passing for jax permutation + pad

change tokens_per_experts APIs to num_out_tokens with conservative allocation of worst case padding for output buffer

change test permutation to reduce test time

triggering PR refresh

Commits on Dec 20, 2025

format code

Remove some tests cases from pytorch side. Add a separate toekn_dispatch test for sanity in case combine accidentally undo an error on dispatch in the roundtrip test. Add distinction between L0 and…

format code

remove chance for inefficiency in moving between CPU and GPU, remove redundant primitive using a new static bool for padding, add assert for align size

Commits on Dec 22, 2025

fix lint in jax

account for both jax newer and older than version 0.8.2. Adjusted gpu triton binding accordingly

format code

Commits on Dec 23, 2025

fix typo

[PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization#1921tdophung merged 21 commits intoNVIDIA:mainNVIDIA/TransformerEngine:mainfrom 021ai:fused_perm_pad021ai/TransformerEngine:fused_perm_padCopy head branch name to clipboard

Commits

Commits on Dec 10, 2025

Commits on Dec 11, 2025

Commits on Dec 12, 2025

Commits on Dec 17, 2025

Commits on Dec 19, 2025

Commits on Dec 20, 2025

Commits on Dec 22, 2025

Commits on Dec 23, 2025

[PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization#1921
tdophung merged 21 commits intoNVIDIA:mainfrom
021ai:fused_perm_pad