standardize moe permute unpermute (#73365)

pesionzhao · web-flow · commit 9955a07af516 · 2025-06-17T17:40:19.000+08:00
* standardize api interface

* restrict cuda version

* standardize unit test

* standardize variable name

* reduce topk
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -5977,7 +5977,6 @@ void MoePermuteInferMeta(const MetaTensor& X,
                          const MetaTensor& XScale,
                          const MetaTensor& expert_routemap_topk,
                          const MetaTensor& expert_prob_topk,
-                         const int topk,
                          const int num_experts,
                          const std::vector<int>& tokens_per_expert,
                          const int padding_multiplex,
diff --git a/paddle/phi/infermeta/multiary.h b/paddle/phi/infermeta/multiary.h
@@ -538,7 +538,6 @@ void MoePermuteInferMeta(const MetaTensor& X,
                          const MetaTensor& XScale,
                          const MetaTensor& expert_routemap_topk,
                          const MetaTensor& expert_prob_topk,
-                         const int topk,
                          const int num_experts,
                          const std::vector<int>& tokens_per_expert,
                          const int padding_multiplex,
diff --git a/paddle/phi/kernels/CMakeLists.txt b/paddle/phi/kernels/CMakeLists.txt
@@ -89,7 +89,7 @@ foreach(arch ${NVCC_ARCH_BIN})
 endforeach()
 
 if(WITH_GPU
-   AND NOT has_arch_ge80
+   AND (NOT has_arch_ge80 OR (CUDA_VERSION VERSION_LESS 12.0))
    OR WIN32
    OR APPLE
    OR WITH_ROCM) # need support BF16
diff --git a/paddle/phi/kernels/gpu/moe_permute_kernel.cu b/paddle/phi/kernels/gpu/moe_permute_kernel.cu
@@ -221,7 +221,6 @@ void MoePermuteKernel(const Context &dev_ctx,
                       const paddle::optional<DenseTensor> &XScale,
                       const DenseTensor &expert_routemap_topk,
                       const DenseTensor &expert_prob_topk,
-                      const int topk,
                       const int num_experts,
                       const std::vector<int> &tokens_per_expert,
                       const int padding_multiplex,
diff --git a/paddle/phi/kernels/moe_permute.h b/paddle/phi/kernels/moe_permute.h
@@ -23,7 +23,6 @@ void MoePermuteKernel(const Context &dev_ctx,
                       const paddle::optional<DenseTensor> &XScale,
                       const DenseTensor &expert_routemap_topk,
                       const DenseTensor &expert_prob_topk,
-                      const int topk,
                       const int num_experts,
                       const std::vector<int> &tokens_per_expert,
                       const int padding_multiplex,
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3712,7 +3712,7 @@
   backward : moe_gate_dispatch_permute_grad
 
 - op : moe_permute
-  args : (Tensor hidden_states, Tensor scale, Tensor expert_routemap, Tensor expert_prob_topk, int topk, int num_experts, int[] tokens_per_expert, int padding_multiplex)
+  args : (Tensor hidden_states, Tensor scale, Tensor expert_routemap_topk, Tensor expert_prob_topk, int num_experts, int[] tokens_per_expert, int padding_alignment)
   output : Tensor(hidden_states_unzipped), Tensor(zipped_expertwise_rowmap), Tensor(token_prob_unzipped), Tensor(scale_unzipped)
   infer_meta:
     func : MoePermuteInferMeta
@@ -3722,13 +3722,13 @@
   optional : scale
 
 - op : moe_unpermute
-  args : (Tensor unzipped_tokens, Tensor zipped_expertwise_rowmap, Tensor expert_routemap_topk, Tensor unzipped_token_probs, int total_zipped_tokens_num, int num_experts, bool MP)
-  output : Tensor(zipped_tokens), Tensor(zipped_probs_topk)
+  args : (Tensor hidden_states_unzipped, Tensor zipped_expertwise_rowmap, Tensor expert_routemap_topk, Tensor token_prob_unzipped, int total_zipped_tokens_num, int num_experts, bool use_mix_precision)
+  output : Tensor(hidden_states), Tensor(expert_prob_topk)
   infer_meta :
     func : MoeUnpermuteInferMeta
   kernel :
     func : moe_unpermute
-    data_type : unzipped_tokens
+    data_type : hidden_states_unzipped
 
 - op : momentum_
   args : (Tensor param, Tensor grad, Tensor velocity, Tensor learning_rate, Tensor master_param, float mu, bool use_nesterov = false, str regularization_method = "", float regularization_coeff = 0.0f, bool multi_precision = false, float rescale_grad = 1.0f)
diff --git a/python/paddle/nn/functional/moe_permute.py b/python/paddle/nn/functional/moe_permute.py
@@ -24,32 +24,33 @@
 
 
 def moe_permute(
-    X: Tensor,
-    XScale: Tensor | None,
+    hidden_states: Tensor,
+    scale: Tensor | None,
     expert_routemap_topk: Tensor,
     expert_prob_topk: Tensor,
-    topk: int,
     num_experts: int,
     tokens_per_expert: list,
-    padding_multiplex: int,
+    padding_alignment: int,
     name: str | None = None,
-):
+) -> tuple[Tensor, Tensor, Tensor, Tensor]:
     if in_dynamic_or_pir_mode():
-        X_unzipped, zipped_experwise, token_prob_unzipped, XScale_unzipped = (
-            _C_ops.moe_permute(
-                X,
-                XScale,
-                expert_routemap_topk,
-                expert_prob_topk,
-                topk,
-                num_experts,
-                tokens_per_expert,
-                padding_multiplex,
-            )
+        (
+            hidden_states_unzipped,
+            zipped_expertwise_rowmap,
+            token_prob_unzipped,
+            scale_unzipped,
+        ) = _C_ops.moe_permute(
+            hidden_states,
+            scale,
+            expert_routemap_topk,
+            expert_prob_topk,
+            num_experts,
+            tokens_per_expert,
+            padding_alignment,
         )
         return (
-            X_unzipped,
-            zipped_experwise,
+            hidden_states_unzipped,
+            zipped_expertwise_rowmap,
             token_prob_unzipped,
-            XScale_unzipped,
+            scale_unzipped,
         )
diff --git a/python/paddle/nn/functional/moe_unpermute.py b/python/paddle/nn/functional/moe_unpermute.py
@@ -24,23 +24,23 @@
 
 
 def moe_unpermute(
-    unzipped_tokens: Tensor,
+    hidden_states_unzipped: Tensor,
     zipped_expertwise_rowmap: Tensor,
     expert_routemap_topk: Tensor,
-    unzipped_token_probs: Tensor,
+    token_prob_unzipped: Tensor,
     total_zipped_tokens: int,
     num_experts: int,
-    MP: bool = True,
+    use_mix_precision: bool = True,
     name: str | None = None,
-):
+) -> tuple[Tensor, Tensor]:
     if in_dynamic_or_pir_mode():
         zipped_tokens, zipped_probs_topk = _C_ops.moe_unpermute(
-            unzipped_tokens,
+            hidden_states_unzipped,
             zipped_expertwise_rowmap,
             expert_routemap_topk,
-            unzipped_token_probs,
+            token_prob_unzipped,
             total_zipped_tokens,
             num_experts,
-            MP,
+            use_mix_precision,
         )
         return (zipped_tokens, zipped_probs_topk)
diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt
@@ -523,9 +523,10 @@ foreach(arch ${NVCC_ARCH_BIN})
     break()
   endif()
 endforeach()
+
 if(NOT WITH_GPU
    OR WITH_GPU
-   AND NOT has_arch_ge80
+   AND (NOT has_arch_ge80 OR (CUDA_VERSION VERSION_LESS 12.0))
    OR WIN32
    OR APPLE
    OR WITH_ROCM)
diff --git a/test/legacy_test/test_moe_permute_unpermute.py b/test/legacy_test/test_moe_permute_unpermute.py
@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
 import unittest
 
+import numpy as np
+
 import paddle
 from paddle.nn.functional import moe_permute, moe_unpermute
 
@@ -27,12 +30,12 @@ def fabricate_dispatch_result(
     broadcast_ratio=0.5,
 ):
     """Helper function to generate test data."""
-    tokens = paddle.randn([seqlen, token_length], dtype=data_type)
+    hidden_states = paddle.randn([seqlen, token_length], dtype=data_type)
 
-    tokens_scale = paddle.empty([0])
+    scale = paddle.empty([0])
     if data_type == "float8_e4m3fn":
         scale_cols = (token_length + 127) // 128
-        tokens_scale = paddle.randn([seqlen, scale_cols], dtype="float32")
+        scale = paddle.randn([seqlen, scale_cols], dtype="float32")
 
     # Calculate expert counts with normal distribution
     expected_experts = max(1, min(broadcast_ratio * num_experts, topk))
@@ -44,35 +47,35 @@ def fabricate_dispatch_result(
     experts_count = paddle.cast(experts_count, "int32")
 
     # Preallocate results
-    dispatched_indices = paddle.full([seqlen, topk], -1, dtype="int32")
-    dispatched_probs = paddle.zeros([seqlen, topk], dtype="float32")
+    expert_routemap_topk = paddle.full([seqlen, topk], -1, dtype="int32")
+    expert_prob_topk = paddle.zeros([seqlen, topk], dtype="float32")
 
     # Batch generate expert indices and probabilities
     for i in range(seqlen):
         count = experts_count[i].item()
         indices = paddle.randperm(num_experts)[:count]
-        dispatched_indices[i, :count] = indices
+        expert_routemap_topk[i, :count] = indices
         prob_value = 1.0 / count
-        dispatched_probs[i, :count] = paddle.full(
+        expert_prob_topk[i, :count] = paddle.full(
             [count], prob_value, dtype=data_type
         )
 
     # Calculate expert token counts
-    valid_indices = dispatched_indices.reshape([-1])
+    valid_indices = expert_routemap_topk.reshape([-1])
     valid_mask = valid_indices >= 0
     valid_experts = valid_indices[valid_mask]
-    expert_counts = paddle.histogram(
+    tokens_per_expert = paddle.histogram(
         valid_experts, bins=num_experts, min=0, max=num_experts - 1
     )
-    expert_counts = paddle.cast(expert_counts, "int32")
-    expert_counts = list(expert_counts)
+    tokens_per_expert = paddle.cast(tokens_per_expert, "int32")
+    tokens_per_expert = list(tokens_per_expert)
 
     return (
-        tokens,
-        tokens_scale,
-        dispatched_indices,
-        dispatched_probs,
-        expert_counts,
+        hidden_states,
+        scale,
+        expert_routemap_topk,
+        expert_prob_topk,
+        tokens_per_expert,
     )
 
 
@@ -100,92 +103,72 @@ def setUp(self):
 
     def test_permute_unpermute_consistency(self):
         """Test that permute + unpermute recovers original tensors."""
-        for dt in self.DTYPES:
-            for expert_num in self.EXPERT_NUMS:
-                for topk in self.TOPKS:
-                    with self.subTest(
-                        dtype=dt, expert_num=expert_num, topk=topk
-                    ):
-                        print(
-                            f"Testing with {expert_num} experts, topk {topk}, dtype {dt}"
-                        )
-                        (
-                            tokens,
-                            tokens_scale,
-                            dispatched_indices,
-                            dispatched_probs,
-                            expert_tokens_count,
-                        ) = fabricate_dispatch_result(
-                            self.SEQLEN,
-                            self.TOKEN_LEN,
-                            topk,
-                            expert_num,
-                            data_type=dt,
-                            broadcast_ratio=0.5,
-                        )
-                        if dt == "bfloat16":
-                            tokens_scale = None
-
-                        # Permute step
-                        (
-                            unzipped_tokens,
-                            zipped_expertwise_rowmap,
-                            unzipped_probs,
-                            unzipped_scales,
-                        ) = moe_permute(
-                            tokens,
-                            tokens_scale,
-                            dispatched_indices,
-                            dispatched_probs,
-                            topk=topk,
-                            num_experts=expert_num,
-                            tokens_per_expert=expert_tokens_count,
-                            padding_multiplex=128,
-                        )
-
-                        # Unpermute step
-                        tokens_recovered, probs_recovered = moe_unpermute(
-                            (
-                                unzipped_tokens * unzipped_probs.unsqueeze(-1)
-                            ).astype("bfloat16"),
-                            zipped_expertwise_rowmap,
-                            dispatched_indices,
-                            unzipped_probs,
-                            total_zipped_tokens=self.SEQLEN,
-                            num_experts=expert_num,
-                        )
-
-                        # Check tensor recovery
-                        max_abs_err, max_rel_err = tensor_max_abs_rel_err(
-                            tokens, tokens_recovered
-                        )
-                        print(
-                            f"permute-unpermute tokens relative error: {max_rel_err}"
-                        )
-                        self.assertLess(
-                            max_rel_err, 1e-2, "Tokens relative error too large"
-                        )
-
-                        max_abs_err, max_rel_err = tensor_max_abs_rel_err(
-                            dispatched_probs, probs_recovered
-                        )
-                        print(
-                            f"ermute-unpermute probs max absolute error: {max_abs_err}, relative error: {max_rel_err}"
-                        )
-                        self.assertLess(
-                            max_abs_err, 1e-5, "Probs absolute error too large"
-                        )
-                        self.assertLess(
-                            max_rel_err, 1e-5, "Probs relative error too large"
-                        )
-
-    def assertLess(self, a, b, msg=None):
-        """Custom assert with better error message."""
-        if not a < b:
-            standard_msg = f"{a} not less than {b}"
-            if msg:
-                standard_msg = f"{msg}: {standard_msg}"
-            self.fail(standard_msg)
+        for dt, expert_num, topk in itertools.product(
+            self.DTYPES, self.EXPERT_NUMS, self.TOPKS
+        ):
+            with self.subTest(dtype=dt, expert_num=expert_num, topk=topk):
+                (
+                    hidden_states,
+                    scale,
+                    expert_routemap_topk,
+                    expert_prob_topk,
+                    tokens_per_expert,
+                ) = fabricate_dispatch_result(
+                    self.SEQLEN,
+                    self.TOKEN_LEN,
+                    topk,
+                    expert_num,
+                    data_type=dt,
+                    broadcast_ratio=0.5,
+                )
+                if dt == "bfloat16":
+                    scale = None
+
+                # Permute step
+                (
+                    unzipped_tokens,
+                    zipped_expertwise_rowmap,
+                    unzipped_probs,
+                    unzipped_scales,
+                ) = moe_permute(
+                    hidden_states,
+                    scale,
+                    expert_routemap_topk,
+                    expert_prob_topk,
+                    num_experts=expert_num,
+                    tokens_per_expert=tokens_per_expert,
+                    padding_alignment=128,
+                )
+
+                # Unpermute step
+                unzipped_tokens_recovered, expert_prob_topk_recovered = (
+                    moe_unpermute(
+                        (unzipped_tokens * unzipped_probs.unsqueeze(-1)).astype(
+                            "bfloat16"
+                        ),
+                        zipped_expertwise_rowmap,
+                        expert_routemap_topk,
+                        unzipped_probs,
+                        total_zipped_tokens=self.SEQLEN,
+                        num_experts=expert_num,
+                    )
+                )
+
+                # Check tensor recovery
+                max_abs_err, max_rel_err = tensor_max_abs_rel_err(
+                    hidden_states, unzipped_tokens_recovered
+                )
+                self.assertLess(
+                    max_rel_err,
+                    1e-2,
+                    f"Tokens relative error too large, permute-unpermute tokens max relative error: {max_rel_err}",
+                )
+
+                np.testing.assert_equal(
+                    expert_prob_topk._md5sum(),
+                    expert_prob_topk_recovered._md5sum(),
+                    err_msg="moe_permute_unpermute probs do not match",
+                )
 
 
 if __name__ == "__main__":
