wehub-resource-sync
265 lines
8.6 KiB
Python
265 lines
8.6 KiB
Python
# 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 unittest
|
|
|
|
from op_test import is_custom_device
|
|
|
|
import paddle
|
|
from paddle.incubate.nn.functional import fp8
|
|
|
|
|
|
class TestFP8GemmBlockwise(unittest.TestCase):
|
|
"""Test cases for FP8 GEMM blockwise operations"""
|
|
|
|
def setUp(self):
|
|
"""Set up test environment"""
|
|
# Skip tests if FP8 is not supported
|
|
if not (paddle.device.is_compiled_with_cuda() or is_custom_device()):
|
|
self.skipTest("CUDA is required for FP8 operations")
|
|
|
|
def cal_rmse(self, y_pred, y_true):
|
|
"""Calculate Root Mean Square Error"""
|
|
return paddle.sqrt(paddle.mean((y_pred - y_true) ** 2))
|
|
|
|
def pyop_ref_1x128_128x128(self, a, b, a_descales, b_descales, C):
|
|
"""Reference implementation for 1x128 @ 128x128 pattern"""
|
|
self.assertEqual(a.dtype, paddle.float8_e4m3fn)
|
|
self.assertEqual(b.dtype, paddle.float8_e4m3fn)
|
|
self.assertEqual(a_descales.dtype, paddle.float32)
|
|
self.assertEqual(b_descales.dtype, paddle.float32)
|
|
|
|
M, N, K = a.shape[0], b.shape[0], a.shape[1]
|
|
self.assertEqual(K, b.shape[1])
|
|
|
|
a_scales_m = a_descales.shape[1]
|
|
a_scales_k = a_descales.shape[0]
|
|
b_scales_k = b_descales.shape[1]
|
|
b_scales_n = b_descales.shape[0]
|
|
|
|
self.assertEqual(a_scales_m, M)
|
|
self.assertEqual(a_scales_k * 128, K)
|
|
self.assertEqual(b_scales_n * 128, N)
|
|
self.assertEqual(b_scales_k * 128, K)
|
|
|
|
a = a.astype(paddle.float32)
|
|
b = b.astype(paddle.float32)
|
|
|
|
out = paddle.zeros((M, N), dtype=paddle.float32)
|
|
|
|
for i in range(0, M):
|
|
for j in range(0, N, 128):
|
|
for k in range(0, K, 128):
|
|
out[i, j : j + 128] += (
|
|
(a[i, k : k + 128] @ b[j : j + 128, k : k + 128].t())
|
|
* a_descales[k // 128, i]
|
|
* b_descales[j // 128, k // 128]
|
|
)
|
|
out = out + C.t()
|
|
return out
|
|
|
|
def pyop_ref_128x128_1x128(self, a, b, a_descales, b_descales, C):
|
|
"""Reference implementation for 128x128 @ 1x128 pattern"""
|
|
self.assertEqual(a.dtype, paddle.float8_e4m3fn)
|
|
self.assertEqual(b.dtype, paddle.float8_e4m3fn)
|
|
self.assertEqual(a_descales.dtype, paddle.float32)
|
|
self.assertEqual(b_descales.dtype, paddle.float32)
|
|
|
|
M, N, K = a.shape[0], b.shape[0], a.shape[1]
|
|
self.assertEqual(K, b.shape[1])
|
|
|
|
a_scales_m = a_descales.shape[0]
|
|
a_scales_k = a_descales.shape[1]
|
|
b_scales_k = b_descales.shape[0]
|
|
b_scales_n = b_descales.shape[1]
|
|
|
|
self.assertEqual(a_scales_m * 128, M)
|
|
self.assertEqual(a_scales_k * 128, K)
|
|
self.assertEqual(b_scales_n, N)
|
|
self.assertEqual(b_scales_k * 128, K)
|
|
|
|
a = a.astype(paddle.float32)
|
|
b = b.astype(paddle.float32)
|
|
|
|
out = paddle.zeros((M, N), dtype=paddle.float32)
|
|
|
|
for i in range(0, M, 128):
|
|
for j in range(0, N):
|
|
for k in range(0, K, 128):
|
|
out[i : i + 128, j] += (
|
|
(a[i : i + 128, k : k + 128] @ b[j, k : k + 128].t())
|
|
* a_descales[i // 128, k // 128]
|
|
* b_descales[k // 128, j]
|
|
)
|
|
|
|
out = out + C
|
|
|
|
return out
|
|
|
|
def pyop_ref_1x128_1x128(self, a, b, a_descales, b_descales):
|
|
"""Reference implementation for 1x128 @ 1x128 pattern"""
|
|
self.assertEqual(a.dtype, paddle.float8_e4m3fn)
|
|
self.assertEqual(b.dtype, paddle.float8_e4m3fn)
|
|
self.assertEqual(a_descales.dtype, paddle.float32)
|
|
self.assertEqual(b_descales.dtype, paddle.float32)
|
|
|
|
M, N, K = a.shape[0], b.shape[0], a.shape[1]
|
|
self.assertEqual(K, b.shape[1])
|
|
|
|
a_scales_m = a_descales.shape[1]
|
|
a_scales_k = a_descales.shape[0]
|
|
b_scales_k = b_descales.shape[0]
|
|
b_scales_n = b_descales.shape[1]
|
|
|
|
self.assertEqual(a_scales_m, M)
|
|
self.assertEqual(a_scales_k * 128, K)
|
|
self.assertEqual(b_scales_n, N)
|
|
self.assertEqual(b_scales_k * 128, K)
|
|
|
|
a = a.astype(paddle.float32)
|
|
b = b.astype(paddle.float32)
|
|
|
|
out = paddle.zeros((M, N), dtype=paddle.float32)
|
|
|
|
for i in range(0, M):
|
|
for j in range(0, N):
|
|
for k in range(0, K, 128):
|
|
out[i, j] += (
|
|
(a[i, k : k + 128] @ b[j, k : k + 128].t())
|
|
* a_descales[k // 128, i]
|
|
* b_descales[k // 128, j]
|
|
)
|
|
|
|
return out
|
|
|
|
def test_1x128_128x128_bfloat16(self):
|
|
"""Test 1x128 @ 128x128 pattern with bfloat16 output"""
|
|
out_dtype = paddle.bfloat16
|
|
M, N, K = 256, 384, 512
|
|
seed = 0
|
|
paddle.seed(seed)
|
|
|
|
A = paddle.randn((M, K), dtype=paddle.bfloat16)
|
|
B = paddle.randn((N, K), dtype=paddle.bfloat16)
|
|
|
|
# Quantize A using fp8
|
|
data_A, scale_A = fp8.fp8_quant_blockwise(
|
|
A,
|
|
quant_method="1x128",
|
|
input_transpose=False,
|
|
output_scale_transpose=True,
|
|
using_pow2_scale=False,
|
|
)
|
|
qA, sA = data_A, scale_A
|
|
|
|
# Quantize B using fp8
|
|
data_B, scale_B = fp8.fp8_quant_blockwise(
|
|
B,
|
|
quant_method="128x128",
|
|
input_transpose=False,
|
|
output_scale_transpose=False,
|
|
using_pow2_scale=False,
|
|
)
|
|
qB, sB = data_B, scale_B
|
|
|
|
gold_matmul_result = A @ B.t()
|
|
|
|
C = paddle.ones([N, M], dtype=paddle.bfloat16)
|
|
# Test reference implementation
|
|
pyop_result = self.pyop_ref_1x128_128x128(qA, qB, sA, sB, C)
|
|
ref_rmse = self.cal_rmse(pyop_result, gold_matmul_result)
|
|
|
|
# Test fp8_gemm_blockwise
|
|
fp8_gemm_result = fp8.fp8_gemm_blockwise(
|
|
qB,
|
|
sB,
|
|
qA,
|
|
sA,
|
|
out_dtype,
|
|
C,
|
|
accumulate=True,
|
|
is_a_1d_scaled=False,
|
|
is_b_1d_scaled=True,
|
|
)
|
|
fp8_gemm_result = fp8_gemm_result.t()
|
|
|
|
rmse = self.cal_rmse(fp8_gemm_result, pyop_result)
|
|
|
|
# Assertions
|
|
self.assertLess(rmse, 0.06, f"RMSE {rmse} exceeds threshold 0.06")
|
|
|
|
def test_128x128_1x128_bfloat16(self):
|
|
"""Test 128x128 @ 1x128 pattern with bfloat16 output"""
|
|
out_dtype = paddle.bfloat16
|
|
M, N, K = 256, 384, 1024
|
|
seed = 0
|
|
paddle.seed(seed)
|
|
|
|
A = paddle.randn((M, K), dtype=paddle.bfloat16)
|
|
B = paddle.randn((N, K), dtype=paddle.bfloat16)
|
|
|
|
# Quantize A using fp8
|
|
data_A, scale_A = fp8.fp8_quant_blockwise(
|
|
A,
|
|
quant_method="128x128",
|
|
input_transpose=False,
|
|
output_scale_transpose=False,
|
|
using_pow2_scale=False,
|
|
)
|
|
qA, sA = data_A, scale_A
|
|
|
|
# Quantize B using fp8
|
|
data_B, scale_B = fp8.fp8_quant_blockwise(
|
|
B,
|
|
quant_method="1x128",
|
|
input_transpose=False,
|
|
output_scale_transpose=True,
|
|
using_pow2_scale=False,
|
|
)
|
|
qB, sB = data_B, scale_B
|
|
|
|
gold_matmul_result = A @ B.t()
|
|
|
|
C = paddle.ones([M, N], dtype=paddle.bfloat16)
|
|
# Test reference implementation
|
|
pyop_result = self.pyop_ref_128x128_1x128(qA, qB, sA, sB, C)
|
|
|
|
# Test fp8_gemm_blockwise
|
|
fp8_gemm_result = fp8.fp8_gemm_blockwise(
|
|
qA,
|
|
sA,
|
|
qB,
|
|
sB,
|
|
out_dtype,
|
|
C,
|
|
accumulate=True,
|
|
is_a_1d_scaled=False,
|
|
is_b_1d_scaled=True,
|
|
)
|
|
|
|
rmse = self.cal_rmse(fp8_gemm_result, pyop_result)
|
|
# Assertions
|
|
self.assertLess(rmse, 0.06, f"RMSE {rmse} exceeds threshold 0.06")
|
|
|
|
def test_1x128_1x128_bfloat16(self):
|
|
"""Test 1x128 @ 1x128 pattern with bfloat16 output"""
|
|
self._test_1x128_1x128(paddle.bfloat16)
|
|
|
|
def _test_1x128_1x128(self, out_dtype):
|
|
"""Helper method for 1x128 @ 1x128 pattern testing"""
|
|
pass
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|