[MoE][PyTorch] Add mask-based MoE permutation (#1373)

* add mask-based moe permutation * change moe_chunk_permute to moe_sort_chunks_by_indices * fix __all__ in pytorch/permutation.py * fix func/var names and typos; update tols in UT --------- Signed-off-by: Hongxiao Bai <[email protected]> Co-authored-by: Phuong Nguyen <[email protected]> Co-authored-by: Tim Moon <[email protected]>
NVIDIA · Jan 27, 2025 · 2fce82b · 2fce82b
1 parent c2c3d54
commit 2fce82b
Show file tree

Hide file tree

Showing 6 changed files with 1,625 additions and 50 deletions.
diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst
@@ -52,6 +52,8 @@ pyTorch
 
 .. autoapifunction:: transformer_engine.pytorch.moe_unpermute
 
+.. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_by_index
+
 .. autoapifunction:: transformer_engine.pytorch.initialize_ub
 
 .. autoapifunction:: transformer_engine.pytorch.destroy_ub