[MoE][PyTorch] Add prob permutation to mask-based MoE permutation; Fi…

…x FP8 related codes (#1468)

* add prob permute; fix fp8tensor

Signed-off-by: Hongxiao Bai <[email protected]>

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

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

* revert unnecessary changes in UT

Signed-off-by: Hongxiao Bai <[email protected]>

* remove unnecessary probs dtype convert

Signed-off-by: Hongxiao Bai <[email protected]>

* keep the output nums if probs is not provided

Signed-off-by: Hongxiao Bai <[email protected]>

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

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

* refine the doc string

Signed-off-by: Hongxiao Bai <[email protected]>

* fix lint

Signed-off-by: Hongxiao Bai <[email protected]>

* use fp32 compute type

Signed-off-by: Hongxiao Bai <[email protected]>

* style fix

Signed-off-by: Hongxiao Bai <[email protected]>

* fix empty input return

Signed-off-by: Hongxiao Bai <[email protected]>

* separate prob related functions out

Signed-off-by: Hongxiao Bai <[email protected]>

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

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

---------

Signed-off-by: Hongxiao Bai <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Xin Yao <[email protected]>
Co-authored-by: Phuong Nguyen <[email protected]>

docs/api/pytorch.rst

@@ -48,10 +48,14 @@ pyTorch
 
 .. autoapifunction:: transformer_engine.pytorch.moe_permute
 
+.. autoapifunction:: transformer_engine.pytorch.moe_permute_with_probs  
+
 .. autoapifunction:: transformer_engine.pytorch.moe_unpermute
 
 .. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_by_index
 
+.. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_by_index_with_probs
+
 .. autoapifunction:: transformer_engine.pytorch.initialize_ub
 
 .. autoapifunction:: transformer_engine.pytorch.destroy_ub