chore: import upstream snapshot with attribution

test/legacy_test/test_fused_transpose_split_quant_op.py

@@ -0,0 +1,124 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+import paddle
+
+
+def dequant_ref(
+    fp8_tensor: paddle.Tensor, scale: paddle.Tensor, block_size: int = 128
+) -> paddle.Tensor:
+    """Helper function to dequantize fp8 tensor to bf16"""
+    expanded_scale = paddle.repeat_interleave(scale, repeats=128, axis=-1)
+    # Handle non-aligned cases by truncating
+    expanded_scale = expanded_scale[:, : fp8_tensor.shape[-1]]
+    return (fp8_tensor.astype('float32') * expanded_scale).astype('bfloat16')
+
+
+def fused_transpose_split_quant_ref(x, xscale, tokens_per_expert, pow_2_scales):
+    shape = x.shape
+    if x.dtype == paddle.float8_e4m3fn:
+        x = dequant_ref(x, xscale)
+    x = x.reshape([shape[0] // 128, 128, shape[1]])
+    amax = x.astype('float32').abs().max(axis=1)
+
+    scale = 448.0 / amax
+    if pow_2_scales:
+        _, exp = paddle.frexp(scale)
+        scale = paddle.ldexp(paddle.to_tensor([1.0]), exp - 1)
+    scale = paddle.where(amax == 0, 1.0, scale)
+
+    out = x * scale.unsqueeze(1)
+    out = out.reshape(shape).astype('float8_e4m3fn')
+    out = out.transpose([1, 0]).split(tokens_per_expert, axis=1)
+
+    scale = paddle.reciprocal(scale)
+    scale = scale.split([t // 128 for t in tokens_per_expert], axis=0)
+    return out, scale
+
+
+def test_fused_transpose_split_quant(
+    tokens_per_expert, seq_len, pow_2_scales, using_fp8=False
+):
+    x = paddle.randn([sum(tokens_per_expert), seq_len], dtype='bfloat16')
+    if using_fp8:
+        x = x.cast('float8_e4m3fn')
+    xscale = (
+        paddle.randn(
+            [sum(tokens_per_expert), (seq_len + 127) // 128], dtype='float32'
+        )
+        if using_fp8
+        else None
+    )
+    # x = paddle.clip(x, min=-50, max=50)
+
+    out, scale = paddle.incubate.nn.functional.fused_transpose_split_quant(
+        x, xscale, tokens_per_expert, pow_2_scales
+    )
+
+    out_ref, scale_ref = fused_transpose_split_quant_ref(
+        x, xscale, tokens_per_expert, pow_2_scales
+    )
+
+    for t, t_ref in zip(out, out_ref):
+        try:
+            np.testing.assert_allclose(
+                t.astype('float32'), t_ref.astype('float32')
+            )
+        except AssertionError as e:
+            print("AssertionError", e)
+
+    for t, t_ref in zip(scale, scale_ref):
+        try:
+            np.testing.assert_allclose(t, t_ref)
+        except AssertionError as e:
+            print("AssertionError", e)
+
+
+def run():
+    fp8_choice = [True, False]
+    for using_fp8 in fp8_choice:
+        test_fused_transpose_split_quant(
+            [0, 0], 1024, False, using_fp8=using_fp8
+        )
+        test_fused_transpose_split_quant(
+            [128, 2 * 128], 0, True, using_fp8=using_fp8
+        )
+        test_fused_transpose_split_quant([128], 1, False, using_fp8=using_fp8)
+        test_fused_transpose_split_quant(
+            [0, 128, 0, 2 * 128], 127, True, using_fp8=using_fp8
+        )
+        test_fused_transpose_split_quant(
+            [3 * 128, 4 * 128, 5 * 128], 233, False, using_fp8=using_fp8
+        )
+        test_fused_transpose_split_quant(
+            [24 * 128, 128, 50 * 128, 16 * 128], 2162, True, using_fp8=using_fp8
+        )
+        test_fused_transpose_split_quant(
+            [7 * 128, 29 * 128, 3 * 128, 128 * 128, 13 * 128],
+            4000,
+            False,
+            using_fp8=using_fp8,
+        )
+        test_fused_transpose_split_quant(
+            [18 * 128, 5 * 128, 24 * 128, 128, 6 * 128, 0, 27 * 128, 7 * 128],
+            7168,
+            True,
+            using_fp8=using_fp8,
+        )
+
+
+if __name__ == '__main__':
+    run()