[CI] Standard unittest (#3606)

* standard unittest

* fix bugs

* fix script

fastdeploy/model_executor/models/ernie4_5_vl/ernie4_5_vl_moe.py

@@ -647,6 +647,9 @@ class Ernie4_5_VLMoeForConditionalGeneration(ModelForCasualLM):
             model_sublayer_name = re.sub(r"\.(up_gate_proj_weight|down_proj_weight|weight)$", "", model_param_name)
             process_weights_after_loading_fn(model_sublayer_name, param)
         if self.tie_word_embeddings:
+            # because we use lazy guard and is not initialized by default
+            if not self.lm_head.linear.weight._is_initialized():
+                self.lm_head.linear.weight.initialize()
             self.lm_head.linear.weight.set_value(self.ernie.embed_tokens.embeddings.weight.transpose([1, 0]))
 
     @paddle.no_grad()

scripts/coverage_run.sh

@@ -11,40 +11,6 @@ cd "$run_path" || exit 1
 failed_tests_file="failed_tests.log"
 > "$failed_tests_file"
 
-##################################
-# 执行特殊单测case(不符合unittest/pytest格式)
-##################################
-special_tests=(
-    "graph_optimization/test_cuda_graph_dynamic_subgraph.py"
-    "graph_optimization/test_cuda_graph_spec_decode.py"
-    "layers/test_quant_layer.py"
-    "operators/test_token_penalty.py"
-    "operators/test_split_fuse.py"
-    "operators/test_flash_mask_attn.py"
-    "operators/test_w4afp8_gemm.py"
-    "model_loader/test_load_ernie_vl.py"
-    "operators/test_tree_mask.py"
-)
-
-failed_special=0
-success_special=0
-
-for test_file in "${special_tests[@]}"; do
-    if [ -f "$test_file" ]; then
-        echo "Running special test: $test_file"
-        python -m coverage run --parallel-mode "$test_file"
-        status=$?
-        if [ "$status" -ne 0 ]; then
-            echo "$test_file" >> "$failed_tests_file"
-            failed_special=$((failed_special+1))
-        else
-            success_special=$((success_special+1))
-        fi
-    else
-        echo "Warning: $test_file not found"
-        failed_special=$((failed_special+1))
-    fi
-done
 
 ##################################
 # 执行 pytest，每个文件单独跑
@@ -78,9 +44,8 @@ echo "Pytest failed: $failed_pytest"
 
 echo "Special tests total: ${#special_tests[@]}"
 echo "Special tests successful: $success_special"
-echo "Special tests failed: $failed_special"
 
-if [ "$failed_pytest" -ne 0 ] || [ "$failed_special" -ne 0 ]; then
+if [ "$failed_pytest" -ne 0 ]; then
     echo "Failed test cases are listed in $failed_tests_file"
     cat "$failed_tests_file"
     exit 8

tests/operators/test_flash_mask_attn.py

@@ -1,10 +1,22 @@
+import unittest
+
 import numpy as np
 import paddle
 
 from fastdeploy.model_executor.ops.gpu import flash_attention_mask
 
 
-def naive_attn(q_input, k_input, v_input, mask):
+class TestFlashMaskAttention(unittest.TestCase):
+    def setUp(self):
+        self.bsz = 1
+        self.num_head = 8
+        self.num_kv_head = 1
+        self.q_seq_len = 1024
+        self.k_seq_len = 1024
+        self.head_dim = 128
+        np.random.seed(self.q_seq_len)
+
+    def naive_attn(self, q_input, k_input, v_input, mask):
         gqa_group_size = q_input.shape[2] // k_input.shape[2]
 
         q_cur = q_input.transpose([0, 2, 1, 3])
@@ -28,8 +40,7 @@ def naive_attn(q_input, k_input, v_input, mask):
                 out[bsz, hi] = (np.matmul(qk, v_cur[bsz, hi // gqa_group_size]) * exp_sum_inv).astype(q_input.dtype)
         return out
 
-
+    def paddle_flash_attn_mask(self, q_input, k_input, v_input, mask):
-def paddle_flash_attn_mask(q_input, k_input, v_input, mask):
         bsz = q_input.shape[0]
         cu_seq_q = paddle.arange(bsz + 1) * q_input.shape[1]
         cu_seq_k = paddle.arange(bsz + 1) * k_input.shape[1]
@@ -61,33 +72,28 @@ def paddle_flash_attn_mask(q_input, k_input, v_input, mask):
         )
         return out
 
+    def test_flash_attention_mask(self):
+        q_input = np.random.normal(0, 0.5, size=(self.bsz, self.q_seq_len, self.num_head, self.head_dim))
+        k_input = np.random.normal(
+            0, 0.5, size=(self.bsz, self.q_seq_len + self.k_seq_len, self.num_kv_head, self.head_dim)
+        )
+        v_input = np.random.normal(
+            0, 0.5, size=(self.bsz, self.q_seq_len + self.k_seq_len, self.num_kv_head, self.head_dim)
+        )
 
-def test(bsz, num_head, num_kv_head, q_seq_len, k_seq_len):
+        random_len = np.random.randint(self.q_seq_len // 2, size=2)
-    head_dim = 128
-    q_input = np.random.normal(0, 0.5, size=(bsz, q_seq_len, num_head, head_dim))
-    k_input = np.random.normal(0, 0.5, size=(bsz, q_seq_len + k_seq_len, num_kv_head, head_dim))
-    v_input = np.random.normal(0, 0.5, size=(bsz, q_seq_len + k_seq_len, num_kv_head, head_dim))
-
-    random_len = np.random.randint(q_seq_len // 2, size=2)
-
         text_len = random_len[0]
         image_len = random_len[1]
 
-    mask = np.array([i + 1 for i in range(0, q_seq_len)]) + k_seq_len
+        mask = np.array([i + 1 for i in range(0, self.q_seq_len)]) + self.k_seq_len
+        mask[text_len : text_len + image_len] = text_len + image_len + self.k_seq_len
 
-    mask[text_len : text_len + image_len] = text_len + image_len + k_seq_len
+        naive_attn_out = self.naive_attn(q_input, k_input, v_input, mask)
+        paddle_attn_out = self.paddle_flash_attn_mask(q_input, k_input, v_input, mask)
 
-    naive_attn_out = naive_attn(q_input, k_input, v_input, mask)
+        max_diff = float((paddle_attn_out.reshape([-1]) - paddle.to_tensor(naive_attn_out).reshape([-1])).max())
-    paddle_attn_out = paddle_flash_attn_mask(q_input, k_input, v_input, mask)
+        self.assertLessEqual(max_diff, 0.05)
-
-    assert float((paddle_attn_out.reshape([-1]) - paddle.to_tensor(naive_attn_out).reshape([-1])).max()) <= 0.05
 
 
 if __name__ == "__main__":
-    bsz = 1
+    unittest.main()
-    num_head = 8
-    num_kv_head = 1
-    q_seq_len = 1024
-    k_seq_len = 1024
-    np.random.seed(q_seq_len)
-    test(bsz, num_head, num_kv_head, q_seq_len, k_seq_len)

tests/operators/test_perchannel_gemm.py

@@ -1,88 +0,0 @@
-# Copyright (c) 2024 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.
-
-"""UT for per_channel_fp8_fp8_half_gemm_fused kernel"""
-
-import os
-import unittest
-from itertools import product
-
-import numpy as np
-import paddle
-
-
-class Test(unittest.TestCase):
-    def setUp(self):
-        """
-            Initialize the test environment,
-        including setting random seeds and environment variables.
-        """
-        paddle.seed(2003)
-        os.environ["FLAGS_use_cutlass_device_best_config_path"] = "default"
-
-    def testcase1(self):
-        """
-        Check if the per_channel_fp8_fp8_half_gemm_fused function works properly.
-        """
-        prop = paddle.device.cuda.get_device_properties()
-        cc = prop.major * 10 + prop.minor
-        if cc < 89:
-            self.skipTest("per_channel_fp8_fp8_half_gemm_fused only support sm89+")
-
-        from fastdeploy.model_executor.ops.gpu import (
-            per_channel_fp8_fp8_half_gemm_fused,
-        )
-
-        nks = [[2048, 2048], [2048, 5504], [6144, 2048]]
-        nks = nks + [[4096, 4096], [4096, 12800], [6144, 4096]]
-        nks = nks + [[5120, 5120], [5120, 13824], [15360, 5120]]
-
-        m = [1, 32, 64, 128, 256, 512, 1024, 2048]
-
-        combinations = list(product(m, nks))
-        for m, (n, k) in combinations:
-            A_bf16 = paddle.rand(shape=[m, k], dtype="bfloat16")
-            A_fp8 = paddle.cast(A_bf16, "float8_e4m3fn")
-            B_bf16 = paddle.rand(shape=[n, k], dtype="bfloat16")
-            B_fp8 = B_bf16.astype("float8_e4m3fn")
-
-            scalar_scale = paddle.full([1], 0.5, dtype="float32")
-            channel_scale = paddle.rand(shape=[n], dtype="float32")
-            bias = paddle.rand(shape=[n], dtype="bfloat16")
-
-            result_bf16 = paddle.matmul(A_bf16, B_bf16, transpose_y=True) * scalar_scale * channel_scale + bias
-            result_fp8 = per_channel_fp8_fp8_half_gemm_fused(
-                A_fp8,
-                B_fp8,
-                bias=bias,
-                scalar_scale=scalar_scale,
-                channel_scale=channel_scale,
-                transpose_x=False,
-                transpose_y=True,
-                output_dtype="bfloat16",
-            )
-            # absolute_error = paddle.abs(result_bf16 - result_fp8)
-            # mean_absolute_error = paddle.mean(absolute_error)
-            relative_error = paddle.abs(result_bf16 - result_fp8) / (paddle.abs(result_bf16))
-            mean_relative_error = paddle.mean(relative_error)
-            np.testing.assert_allclose(
-                mean_relative_error.numpy(),
-                np.array([0.001]),
-                rtol=0.001,
-                atol=0.25,
-            )
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/operators/test_split_fuse.py

@@ -13,16 +13,22 @@
 # limitations under the License.
 
 """UT for set_stop_value"""
+import unittest
+
 import paddle
 
 from fastdeploy.model_executor.ops.gpu import get_mm_split_fuse
 
-input_ids = []
-image_type_ids = []
-grid_thw = []
 
+class TestSplitFuse(unittest.TestCase):
+    def setUp(self):
+        self.grid_thw = [[6, 20, 20], [6, 40, 20]]
+        self.split_fuse_img_size = 16
+        self.split_fuse_text_size = 384  # 1024
+        self.max_seq_len = 2048
+        self.image_token_id = 100295
 
-def split_grid(origin_grid_thw):
+    def split_grid(self, origin_grid_thw):
         # 划分grid_thw，该函数用于视频场景
         # origin_grid_thw = [6, 10, 12] ---> [2, 10, 12, 2, 10, 12, 2, 10, 12]
         grid_thw = []
@@ -38,28 +44,24 @@ def split_grid(origin_grid_thw):
                 grid_thw.extend([t, h, w])
         return grid_thw
 
-
+    def test_get_mm_split_fuse(self):
-if __name__ == "__main__":
+        grid_thw = self.split_grid(self.grid_thw)
-    grid_thw = [[6, 20, 20], [6, 40, 20]]
-    grid_thw = split_grid(grid_thw)
         image_bs = len(grid_thw) // 3
         image_type_ids = [0] * image_bs
+
         # 随机拼接input_ids: [txt0+img1+tx1+img2]
         input_ids = [2] * 19
-    img1 = [100295] * 100 * 3
+        img1 = [self.image_token_id] * 100 * 3
         txt1 = [3] * 19
-    img2 = [100295] * 200 * 3
+        img2 = [self.image_token_id] * 200 * 3
         input_ids.extend(img1)
         input_ids.extend(txt1)
         input_ids.extend(img2)
 
-    split_fuse_img_size = 16
-    split_fuse_text_size = 384  # 1024
-
         seq_len = len(input_ids)
         input_ids_tensor = paddle.to_tensor(input_ids, dtype="int64")
         image_type_ids_tensor = paddle.to_tensor(image_type_ids, dtype="int32")
-    is_image_token = paddle.where(input_ids_tensor == 100295, 1, 0)
+        is_image_token = paddle.where(input_ids_tensor == self.image_token_id, 1, 0)
         image_token_sum = paddle.cumsum(is_image_token)  # 前缀和
         image_token_sum = paddle.concat([paddle.zeros([1], dtype="int64"), image_token_sum])
 
@@ -69,16 +71,23 @@ if __name__ == "__main__":
             image_type_ids_tensor.cast("int32").cpu(),
             image_token_sum.cast("int32").cpu(),
             grid_thw_tensor.cpu(),
-        100295,
+            self.image_token_id,
             image_bs,
             0,
             seq_len,
-        split_fuse_img_size,
+            self.split_fuse_img_size,
-        split_fuse_text_size,
+            self.split_fuse_text_size,
-        2048,
+            self.max_seq_len,
         )
 
-    print("seq_len: ", seq_len)
+        # Verify the outputs are not None
-    print("grid_thw", grid_thw_tensor)
+        self.assertIsNotNone(image_chunk_selections)
-    print("image_chunk_selections: ", image_chunk_selections)
+        self.assertIsNotNone(split_fuse_cur_seq_lens)
-    print("split_fuse_cur_seq_lens: ", split_fuse_cur_seq_lens)
+
+        # Verify the shapes are as expected
+        self.assertEqual(len(image_chunk_selections.shape), 1)
+        self.assertEqual(len(split_fuse_cur_seq_lens.shape), 1)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/operators/test_token_penalty.py

@@ -13,40 +13,58 @@
 # limitations under the License.
 
 """UT for get_token_penalty"""
+import unittest
+
 import numpy as np
 import paddle
 
 from fastdeploy.model_executor.ops.gpu import get_token_penalty_once
 
+
+class TestTokenPenalty(unittest.TestCase):
+    def setUp(self):
         paddle.seed(2023)
+        self.pre_ids = paddle.randint(0, 10000, (8, 1000))
+        self.pre_ids[:, -1] = self.pre_ids[:, -2]
+        self.logits = paddle.rand(shape=[8, 10000], dtype="float16")
+        self.penalty_scores = np.array([1.2] * 8).astype(np.float16).reshape(-1, 1)
+        self.penalty_scores = paddle.to_tensor(self.penalty_scores)
 
-pre_ids = paddle.randint(0, 10000, (8, 1000))
+    def test_token_penalty_once(self):
-pre_ids[:, -1] = pre_ids[:, -2]
+        res = get_token_penalty_once(self.pre_ids, self.logits, self.penalty_scores)
-print(pre_ids)
-logits = paddle.rand(shape=[8, 10000], dtype="float16")
-penalty_scores = np.array([1.2] * 8).astype(np.float16).reshape(-1, 1)
-penalty_scores = paddle.to_tensor(penalty_scores)
 
-print("logits[0][pre_ids[0]]: ", logits[0][pre_ids[0]])
+        # 验证结果形状
-res = get_token_penalty_once(pre_ids, logits, penalty_scores)
+        self.assertEqual(res.shape, self.logits.shape)
+
+        # 验证惩罚逻辑
         for i in range(8):
-    print(f"res[{i}]:{res[i][pre_ids[i]]}")
+            original_values = self.logits[i][self.pre_ids[i]]
+            penalized_values = res[i][self.pre_ids[i]]
+            # 检查是否应用了惩罚
+            for orig, penal in zip(original_values.numpy(), penalized_values.numpy()):
+                if orig < 0:
+                    self.assertLess(penal, orig, "负值应该乘以惩罚因子")
+                else:
+                    self.assertLess(penal, orig, "正值应该除以惩罚因子")
+
+    def test_compare_with_naive_implementation(self):
+        res = get_token_penalty_once(self.pre_ids, self.logits, self.penalty_scores)
+
+        # 朴素实现
+        score = paddle.index_sample(self.logits, self.pre_ids)
+        score = paddle.where(score < 0, score * self.penalty_scores, score / self.penalty_scores)
+
+        bsz = paddle.shape(self.logits)[0]
+        bsz_range = paddle.arange(start=bsz * 0, end=bsz, step=bsz / bsz, name="bsz_range", dtype="int64").unsqueeze(
+            -1
+        )
+        input_ids = self.pre_ids + bsz_range * self.logits.shape[-1]
+        res2 = paddle.scatter(self.logits.flatten(), input_ids.flatten(), score.flatten()).reshape(self.logits.shape)
+
+        # 比较两种实现的结果差异
+        max_diff = (res - res2).abs().max().item()
+        self.assertLess(max_diff, 1e-5)
 
 
-input_ids = pre_ids
+if __name__ == "__main__":
-score = paddle.index_sample(logits, input_ids)
+    unittest.main()
-score = paddle.where(score < 0, score * penalty_scores, score / penalty_scores)
-
-bsz = paddle.shape(logits)[0]  # TODO: Bsz as input for inference with dynamic batch_size
-bsz_range = paddle.arange(start=bsz * 0, end=bsz, step=bsz / bsz, name="bsz_range", dtype="int64").unsqueeze(-1)
-input_ids = input_ids + bsz_range * logits.shape[-1]
-res2 = paddle.scatter(logits.flatten(), input_ids.flatten(), score.flatten()).reshape(logits.shape)
-print("-------------------------------------------")
-for i in range(8):
-    print(res2[i][pre_ids[i]])
-
-print("res_sub:")
-for i in range(8):
-    print(res2[i][pre_ids[i]] - res[i][pre_ids[i]])
-
-print((res.numpy() - res2.numpy()).sum())

tests/operators/test_tree_mask.py

@@ -1,5 +1,6 @@
 import math
 import time
+import unittest
 
 import numpy as np
 import paddle
@@ -10,46 +11,70 @@ from fastdeploy.model_executor.layers.attention.ops import (
     get_block_shape_and_split_kv_block,
 )
 
+
+class TestTreeMask(unittest.TestCase):
+    def setUp(self):
         paddle.seed(0)
+        self.max_seq_len = 32768
+        self.encoder_max_partition_size = self.max_seq_len
+        self.max_partition_size = self.max_seq_len
 
-max_seq_len = 32768
+        self.max_dec_len = 1024
-encoder_max_partition_size = max_seq_len
+        self.bsz = 64
-max_partition_size = max_seq_len
+        self.run_time = 10
+        self.warm_up = 2
+        self.block_size = 64
+        self.head_dim = 128
+        self.num_q_head = 20
+        self.num_kv_head = 4
+        self.dtype = "bfloat16"
 
-max_dec_len = 1024
+        self.rope_3d = False
-bsz = 64
+        self.use_neox_rotary_style = False
-run_time = 10
+        self.CURRENT_Q = [None]
-warm_up = 2
+        self.TOTAL_K = []
-block_size = 64
+        self.TOTAL_V = []
-head_dim = 128
-num_q_head = 20
-num_kv_head = 4
-dtype = "bfloat16"
 
-rope_3d = False
+        # Initialize cache and block tables
-use_neox_rotary_style = False
+        block_num_per_seq = (self.max_seq_len + self.block_size - 1) // self.block_size
-CURRENT_Q = [None]
+        max_block_num = block_num_per_seq * self.bsz
-TOTAL_K = []
+        cache_shape = (
-TOTAL_V = []
+            max_block_num,
+            self.num_kv_head,
+            self.block_size,
+            self.head_dim,
+        )
 
+        self.cache_k = paddle.zeros(shape=cache_shape).astype(self.dtype)
+        self.cache_v = paddle.zeros(shape=cache_shape).astype(self.dtype)
 
-def split_qkv(qkv, bsz, seq_len, num_q_head, num_kv_head, head_dim):
+        self.block_tables = paddle.zeros(shape=(self.bsz, block_num_per_seq), dtype="int32")
-    # [token_num, (num_q_head + 2 * num_kv_head) * head_dim]
-    qkv = qkv.reshape([bsz, seq_len, -1, head_dim])
-    q = qkv[:, :, :num_q_head, :]
-    # [bsz,  seq_len, num_q_head, head_dim]
-    CURRENT_Q[0] = q
 
-    # [bsz,  seq_len, num_kv_head, head_dim]
+        free_list = list(range(max_block_num - 1, -1, -1))
-    k = qkv[:, :, num_q_head : num_q_head + num_kv_head, :]
-    TOTAL_K.append(k)
 
-    # [bsz,  seq_len, num_kv_head, head_dim]
+        for i in range(self.bsz):
-    v = qkv[:, :, num_q_head + num_kv_head :, :]
+            need_block_num = (self.max_seq_len + self.block_size - 1) // self.block_size
-    TOTAL_V.append(v)
+            for j in range(need_block_num):
+                block_id = free_list.pop()
+                self.block_tables[i, j] = block_id
 
+    def tearDown(self):
+        self.CURRENT_Q = [None]
+        self.TOTAL_K = []
+        self.TOTAL_V = []
 
-def get_padding_offset(bsz, seq_lens_this_time, seq_lens_decoder):
+    def split_qkv(self, qkv, bsz, seq_len):
+        qkv = qkv.reshape([bsz, seq_len, -1, self.head_dim])
+        q = qkv[:, :, : self.num_q_head, :]
+        self.CURRENT_Q[0] = q
+
+        k = qkv[:, :, self.num_q_head : self.num_q_head + self.num_kv_head, :]
+        self.TOTAL_K.append(k)
+
+        v = qkv[:, :, self.num_q_head + self.num_kv_head :, :]
+        self.TOTAL_V.append(v)
+
+    def get_padding_offset(self, bsz, seq_lens_this_time, seq_lens_decoder):
         batch_id_per_token = []
         cu_seqlens_q = paddle.zeros(shape=(bsz + 1), dtype="int32")
         cu_seqlens_k = paddle.zeros(shape=(bsz + 1), dtype="int32")
@@ -66,32 +91,7 @@ def get_padding_offset(bsz, seq_lens_this_time, seq_lens_decoder):
             cu_seqlens_k[i + 1] = cum_seq_len_k
         return paddle.to_tensor(batch_id_per_token, dtype="int32"), cu_seqlens_q, cu_seqlens_k
 
-
+    def ref_attention(self, q, k, v, mask):
-# block_table
-block_num_per_seq = (max_seq_len + block_size - 1) // block_size
-max_block_num = block_num_per_seq * bsz
-cache_shape = (
-    max_block_num,
-    num_kv_head,
-    block_size,
-    head_dim,
-)
-
-cache_k = paddle.zeros(shape=cache_shape).astype(dtype)
-cache_v = paddle.zeros(shape=cache_shape).astype(dtype)
-
-block_tables = paddle.zeros(shape=(bsz, block_num_per_seq), dtype="int32")
-
-free_list = list(range(max_block_num - 1, -1, -1))
-
-for i in range(bsz):
-    need_block_num = (max_seq_len + block_size - 1) // block_size
-    for j in range(need_block_num):
-        block_id = free_list.pop()
-        block_tables[i, j] = block_id
-
-
-def ref_attention(q, k, v, num_q_head, num_kv_head, head_dim, mask):
         q = q.transpose([0, 2, 1, 3])
         if len(k) > 1:
             k = paddle.concat(k, axis=1)
@@ -105,44 +105,44 @@ def ref_attention(q, k, v, num_q_head, num_kv_head, head_dim, mask):
         v = v.transpose([0, 2, 1, 3])
         total_len = k.shape[2]
 
-    scores = q.reshape([bsz, num_kv_head, -1, head_dim]) @ k.transpose([0, 1, 3, 2]) * (1.0 / math.sqrt(head_dim))
+        scores = (
-    scores = scores.reshape([bsz, num_q_head, -1, total_len])
+            q.reshape([self.bsz, self.num_kv_head, -1, self.head_dim])
+            @ k.transpose([0, 1, 3, 2])
+            * (1.0 / math.sqrt(self.head_dim))
+        )
+        scores = scores.reshape([self.bsz, self.num_q_head, -1, total_len])
 
         if mask is not None:
             if mask.ndim == 2:
-            mask = mask.unsqueeze(0).unsqueeze(0)  # [1,1,q_len,kv_len]
+                mask = mask.unsqueeze(0).unsqueeze(0)
             elif mask.ndim == 3:
-            mask = mask.unsqueeze(1)  # [bsz,1,q_len,kv_len]
+                mask = mask.unsqueeze(1)
             scores = paddle.add(scores, mask)
         weights = F.softmax(scores, axis=-1)
 
-    o = weights.reshape([bsz, num_kv_head, -1, total_len]) @ v
+        o = weights.reshape([self.bsz, self.num_kv_head, -1, total_len]) @ v
-    return o.reshape([bsz, num_q_head, -1, head_dim]).transpose([0, 2, 1, 3]).reshape([-1, num_q_head, head_dim])
+        return (
+            o.reshape([self.bsz, self.num_q_head, -1, self.head_dim])
+            .transpose([0, 2, 1, 3])
+            .reshape([-1, self.num_q_head, self.head_dim])
+        )
 
-
+    def run_append_c16_attention(self, q_len, kv_len, prefill=False, attn_mask=None):
-def clear_param():
-    global CURRENT_Q, TOTAL_K, TOTAL_V
-    CURRENT_Q = [None]
-    TOTAL_K = []
-    TOTAL_V = []
-
-
-def test_append_c16_attention(q_len, kv_len, prefill=False, attn_mask=None):
         if prefill:
             seq_lens_enc = [
                 q_len,
-        ] * bsz
+            ] * self.bsz
         else:
             seq_lens_enc = [
                 0,
-        ] * bsz
+            ] * self.bsz
 
         seq_lens_dec = [
             kv_len,
-    ] * bsz
+        ] * self.bsz
         seq_lens_cur = [
             q_len,
-    ] * bsz
+        ] * self.bsz
         token_num = sum(seq_lens_cur)
         decoder_step_token_num = 1 if prefill else q_len
 
@@ -150,39 +150,37 @@ def test_append_c16_attention(q_len, kv_len, prefill=False, attn_mask=None):
         seq_lens_this_time = paddle.to_tensor(seq_lens_cur, "int32")
         seq_lens_decoder = paddle.to_tensor(seq_lens_dec, "int32")
 
-    batch_id_per_token, cu_seqlens_q, cu_seqlens_k = get_padding_offset(bsz, seq_lens_this_time, seq_lens_decoder)
+        batch_id_per_token, cu_seqlens_q, cu_seqlens_k = self.get_padding_offset(
+            self.bsz, seq_lens_this_time, seq_lens_decoder
+        )
 
-    # random data
+        qkv_varlen_shape = [token_num, (self.num_q_head + 2 * self.num_kv_head) * self.head_dim]
-    qkv_varlen_shape = [token_num, (num_q_head + 2 * num_kv_head) * head_dim]
+        rotary_embs_shape = [
+            2,
+            1,
+            self.max_seq_len,
+            1,
+            self.head_dim if self.use_neox_rotary_style else self.head_dim // 2,
+        ]
 
-    rotary_embs_shape = [2, 1, max_seq_len, 1, head_dim if use_neox_rotary_style else head_dim // 2]
+        qkv = paddle.randn(shape=qkv_varlen_shape).astype(self.dtype)
-    # qkv_bias_shape = [num_q_head + 2 * num_kv_head, head_dim]
+        self.split_qkv(qkv, self.bsz, q_len)
-
-    qkv = paddle.randn(shape=qkv_varlen_shape).astype(dtype)
-
-    # save q, k, v for ref
-    split_qkv(qkv, bsz, q_len, num_q_head, num_kv_head, head_dim)
 
         rotary_embs = paddle.randn(shape=rotary_embs_shape).astype("float32")
         rotary_embs[0, :, :, :, :] = 1
         rotary_embs[1, :, :, :, :] = 0
 
-    # qkv_scale = None
-    # qkv_bias = None
-
         cache_k_scale = None
         cache_v_scale = None
         cache_k_out_scale = None
         cache_v_out_scale = None
-    # shift_bias = None
-    # smooth_weight = None
 
         encoder_block_shape_q = 64
         decoder_block_shape_q = 16
 
         decode_max_tile_size = (
-        bsz
+            self.bsz
-        * (decoder_step_token_num * (num_q_head // num_kv_head) + decoder_block_shape_q - 1)
+            * (decoder_step_token_num * (self.num_q_head // self.num_kv_head) + decoder_block_shape_q - 1)
             / decoder_block_shape_q
         )
         decoder_batch_ids = paddle.full([int(decode_max_tile_size)], 0, dtype="int32")
@@ -208,24 +206,24 @@ def test_append_c16_attention(q_len, kv_len, prefill=False, attn_mask=None):
             max_len_tensor_cpu,
             encoder_block_shape_q,
             decoder_block_shape_q,
-        num_q_head // num_kv_head,
+            self.num_q_head // self.num_kv_head,
-        block_size,
+            self.block_size,
             decoder_step_token_num,
         )
         s_time = 0
-    for i in range(run_time + warm_up):
+        for i in range(self.run_time + self.warm_up):
-        if i == warm_up:
+            if i == self.warm_up:
                 s_time = time.time()
             out = append_attention(
                 qkv,
-            cache_k,
+                self.cache_k,
-            cache_v,
+                self.cache_v,
                 seq_lens_encoder,
                 seq_lens_decoder,
                 seq_lens_this_time,
                 batch_id_per_token,
                 cu_seqlens_q,
-            block_tables,
+                self.block_tables,
                 encoder_batch_ids,
                 encoder_tile_ids_per_batch,
                 encoder_num_blocks,
@@ -238,15 +236,15 @@ def test_append_c16_attention(q_len, kv_len, prefill=False, attn_mask=None):
                 max_len_tensor_cpu,
                 max_len_kv,
                 rotary_embs,
-            attn_mask,  # attn_mask
+                attn_mask,
                 None,
                 None,
                 cache_k_scale,
                 cache_v_scale,
                 cache_k_out_scale,
                 cache_v_out_scale,
-            None,  # cache_k_zp
+                None,
-            None,  # cache_v_zp
+                None,
                 None,
                 None,
                 None,
@@ -255,47 +253,45 @@ def test_append_c16_attention(q_len, kv_len, prefill=False, attn_mask=None):
                 None,
                 1e-6,
                 "bf16",
-            "none",  # cache_quant_type
+                "none",
-            use_neox_rotary_style,
+                self.use_neox_rotary_style,
-            rope_3d,
+                self.rope_3d,
-            max_seq_len,
+                self.max_seq_len,
                 0.0,
                 0.0,
-            -1.0,  # out_linear_in_scale
+                -1.0,
-            encoder_block_shape_q,  # encoder_block_shape_q
+                encoder_block_shape_q,
-            decoder_block_shape_q,  # decoder_block_shape_q
+                decoder_block_shape_q,
-            max_partition_size,  # max_partition_size
+                self.max_partition_size,
-            encoder_max_partition_size,  # encoder_max_partition_size
+                self.encoder_max_partition_size,
-            decoder_step_token_num,  # speculate_max_draft_token_num
+                decoder_step_token_num,
-            True,  # causal
+                True,
-            decoder_step_token_num > 1,  # speculate_decoder
+                decoder_step_token_num > 1,
             )
             paddle.device.synchronize()
         e_time = time.time()
-    print(f"mean infer time: {np.mean((e_time - s_time) * 1000 / run_time):.2f}")
+        print(f"mean infer time: {np.mean((e_time - s_time) * 1000 / self.run_time):.2f}")
-    return out[0].reshape([token_num, num_q_head, head_dim])
+        return out[0].reshape([token_num, self.num_q_head, self.head_dim])
 
-
+    def test_naive_speculative_decoding(self):
-def test_naive_speculative_decoding(num_q_head, num_kv_head, head_dim):
         prefill_len = 8192
         dec_len_q = 5
         total_len = prefill_len + dec_len_q
-    mask = paddle.tril(paddle.ones((bsz, dec_len_q, total_len), dtype="float32"), diagonal=prefill_len)
+        mask = paddle.tril(paddle.ones((self.bsz, dec_len_q, total_len), dtype="float32"), diagonal=prefill_len)
         mask = paddle.where(mask == 1, paddle.zeros_like(mask), paddle.full_like(mask, fill_value=float("-inf")))
-    test_append_c16_attention(prefill_len, 0, True)
+        self.run_append_c16_attention(prefill_len, 0, True)
-    dec_out = test_append_c16_attention(dec_len_q, prefill_len, False)
+        dec_out = self.run_append_c16_attention(dec_len_q, prefill_len, False)
 
-    ref_out = ref_attention(CURRENT_Q[0], TOTAL_K, TOTAL_V, num_q_head, num_kv_head, head_dim, mask)
+        ref_out = self.ref_attention(self.CURRENT_Q[0], self.TOTAL_K, self.TOTAL_V, mask)
         np.testing.assert_allclose(
             ref_out.astype("float32").numpy(), dec_out.astype("float32").numpy(), rtol=1e-03, atol=5e-03
         )
 
-
+    def test_mask(self):
-def test_mask(num_q_head, num_kv_head, head_dim):
         prefill_len = 8192
         dec_len_q = 5
         total_len = prefill_len + dec_len_q
-    mask = paddle.tril(paddle.ones((bsz, dec_len_q, total_len), dtype="float32"), diagonal=prefill_len)
+        mask = paddle.tril(paddle.ones((self.bsz, dec_len_q, total_len), dtype="float32"), diagonal=prefill_len)
         mask_ref = paddle.where(mask == 1, paddle.zeros_like(mask), paddle.full_like(mask, fill_value=float("-inf")))
 
         mask_append_attn = mask[:, :, prefill_len:]
@@ -305,28 +301,20 @@ def test_mask(num_q_head, num_kv_head, head_dim):
             paddle.full_like(mask_append_attn, fill_value=True, dtype=bool),
         )
 
-    test_append_c16_attention(prefill_len, 0, True)
+        self.run_append_c16_attention(prefill_len, 0, True)
-    dec_out = test_append_c16_attention(dec_len_q, prefill_len, False, mask_append_attn)
+        dec_out = self.run_append_c16_attention(dec_len_q, prefill_len, False, mask_append_attn)
 
-    ref_out = ref_attention(CURRENT_Q[0], TOTAL_K, TOTAL_V, num_q_head, num_kv_head, head_dim, mask_ref)
+        ref_out = self.ref_attention(self.CURRENT_Q[0], self.TOTAL_K, self.TOTAL_V, mask_ref)
 
         np.testing.assert_allclose(
             ref_out.astype("float32").numpy(), dec_out.astype("float32").numpy(), rtol=1e-03, atol=5e-03
         )
 
-
+    def test_tree_mask(self):
-def test_tree_mask(num_q_head, num_kv_head, head_dim):
-    # tree
-    #       [N,   N+1,    N+1,    N+2,    N+2]
-    # N     [0,   -inf,   -inf,   -inf,   -inf]
-    # N+1   [0,   0,      -inf,   -inf,   -inf]
-    # N+1   [0,   -inf,   0,      -inf,   -inf]
-    # N+2   [0,   0,      -inf,   0,      -inf]
-    # N+2   [0,   -inf,   0,      -inf,   0]
         prefill_len = 8192
         dec_len_q = 5
         total_len = prefill_len + dec_len_q
-    mask = paddle.tril(paddle.ones((bsz, dec_len_q, total_len), dtype="float32"), diagonal=prefill_len)
+        mask = paddle.tril(paddle.ones((self.bsz, dec_len_q, total_len), dtype="float32"), diagonal=prefill_len)
         mask[:, 2, prefill_len + 1] = 0
         mask[:, 3, prefill_len + 2] = 0
         mask[:, 4, prefill_len + 1] = 0
@@ -341,20 +329,13 @@ def test_tree_mask(num_q_head, num_kv_head, head_dim):
             paddle.full_like(mask_append_attn, fill_value=True, dtype=bool),
         )
 
-    test_append_c16_attention(prefill_len, 0, True)
+        self.run_append_c16_attention(prefill_len, 0, True)
-    dec_out = test_append_c16_attention(dec_len_q, prefill_len, False, mask_append_attn)
+        dec_out = self.run_append_c16_attention(dec_len_q, prefill_len, False, mask_append_attn)
-    ref_out = ref_attention(CURRENT_Q[0], TOTAL_K, TOTAL_V, num_q_head, num_kv_head, head_dim, mask_ref)
+        ref_out = self.ref_attention(self.CURRENT_Q[0], self.TOTAL_K, self.TOTAL_V, mask_ref)
         np.testing.assert_allclose(
             ref_out.astype("float32").numpy(), dec_out.astype("float32").numpy(), rtol=1e-03, atol=5e-03
         )
 
 
 if __name__ == "__main__":
-
+    unittest.main()
-    test_naive_speculative_decoding(num_q_head, num_kv_head, head_dim)
-    clear_param()
-
-    test_mask(num_q_head, num_kv_head, head_dim)
-    clear_param()
-
-    test_tree_mask(num_q_head, num_kv_head, head_dim)

tests/operators/test_w4afp8_gemm.py

@@ -12,17 +12,47 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import unittest
+
 import numpy as np
 import paddle
 
 from fastdeploy.model_executor.ops.gpu import w4afp8_gemm, w4afp8_gemm_weight_convert
 
 
-def w4afp8_gemm_naive(input_bf16, weight_quant, tokens, weight_dequant_scale, BATCH, N):
+class TestW4AFP8GEMM(unittest.TestCase):
+    def setUp(self):
+        paddle.seed(0)
+        self.tokens_per_group = 256
+        self.N = 256
+        self.K = 256
+        self.BATCH = 1
+        self.TokenPadding = 0
+
+        tokens = [self.tokens_per_group] * self.BATCH
+        self.tokens_perfix_sum = np.cumsum(tokens)
+
+        self.tokens = paddle.to_tensor(tokens, dtype="int64")
+        self.tokens_perfix_sum = paddle.to_tensor(self.tokens_perfix_sum, dtype="int64")
+        self.all_tokens = int(self.tokens.sum())
+
+        self.input_fp8 = paddle.randn([self.all_tokens, self.K], dtype="bfloat16").astype(paddle.float8_e4m3fn)
+        self.input_bf16 = self.input_fp8.astype("bfloat16")
+        self.weight = paddle.randn([self.BATCH, self.N, self.K], dtype="bfloat16") / 10
+
+        self.weight_scale = 7 / self.weight.abs().max(axis=-1).reshape([self.BATCH, self.N, 1])
+        self.weight_quant = (self.weight * self.weight_scale).astype("int") + 7
+        self.weight_quant = paddle.clip(self.weight_quant, 0, 14)
+        self.weight_quant = self.weight_quant.astype("bfloat16")
+        self.weight_dequant_scale = 1 / self.weight_scale.astype("float32")
+        self.input_row_sum = self.input_bf16.sum(axis=1) * -7 / 512
+        self.max_tokens = int(self.tokens.max())
+
+    def w4afp8_gemm_naive(self, input_bf16, weight_quant, tokens, weight_dequant_scale):
         all_tokens = int(tokens.sum())
-    out = paddle.zeros([all_tokens, N], dtype="bfloat16")
+        out = paddle.zeros([all_tokens, self.N], dtype="bfloat16")
         pre_fix_token = 0
-    for i in range(BATCH):
+        for i in range(self.BATCH):
             input = input_bf16[pre_fix_token : pre_fix_token + tokens[i], :]
             weight = (weight_quant[i] - 7.0) * weight_dequant_scale[i]
             out_i = paddle.matmul(input, weight.astype("bfloat16"), transpose_y=True)
@@ -30,74 +60,49 @@ def w4afp8_gemm_naive(input_bf16, weight_quant, tokens, weight_dequant_scale, BA
             pre_fix_token += tokens[i]
         return out
 
-
+    def permute_scale(self, weight_scale):
-def permute_scale(weight_scale):
+        weight_scale = weight_scale.reshape([self.BATCH, self.N])
-    weight_scale = weight_scale.reshape([BATCH, N])
         temp = paddle.zeros([16])
-    for b in range(BATCH):
+        for b in range(self.BATCH):
-        for n in range(0, N, 16):
+            for n in range(0, self.N, 16):
                 temp[:] = weight_scale[b, n : n + 16]
                 for j in range(0, 16, 2):
                     weight_scale[b, n + j] = temp[j // 2]
                     weight_scale[b, n + j + 1] = temp[j // 2 + 8]
         return weight_scale
 
+    def test_w4afp8_gemm(self):
+        out_naive = self.w4afp8_gemm_naive(self.input_bf16, self.weight_quant, self.tokens, self.weight_dequant_scale)
 
-paddle.seed(0)
+        weight_dequant_scale = paddle.to_tensor(self.permute_scale(self.weight_dequant_scale) * 512)
-tokens_per_group = 256
+        weight_int4 = w4afp8_gemm_weight_convert(self.weight_quant.astype("uint8").cpu())
-N = 256
-K = 256
-BATCH = 1
-TokenPadding = 0
 
-tokens = [tokens_per_group] * BATCH
+        if self.TokenPadding == 0:
-tokens_perfix_sum = np.cumsum(tokens)
-
-
-tokens = paddle.to_tensor(tokens, dtype="int64")
-tokens_perfix_sum = paddle.to_tensor(tokens_perfix_sum, dtype="int64")
-
-all_tokens = int(tokens.sum())
-
-input_fp8 = paddle.randn([all_tokens, K], dtype="bfloat16").astype(paddle.float8_e4m3fn)
-input_bf16 = input_fp8.astype("bfloat16")
-weight = paddle.randn([BATCH, N, K], dtype="bfloat16") / 10
-
-weight_scale = 7 / weight.abs().max(axis=-1).reshape([BATCH, N, 1])
-weight_quant = (weight * weight_scale).astype("int") + 7
-weight_quant = paddle.clip(weight_quant, 0, 14)
-weight_quant = weight_quant.astype("bfloat16")
-weight_dequant_scale = 1 / weight_scale.astype("float32")
-input_row_sum = input_bf16.sum(axis=1) * -7 / 512
-max_tokens = int(tokens.max())
-
-out_naive = w4afp8_gemm_naive(input_bf16, weight_quant, tokens, weight_dequant_scale, BATCH, N)
-weight_dequant_scale = paddle.to_tensor(permute_scale(weight_dequant_scale) * 512)
-
-weight_int4 = w4afp8_gemm_weight_convert(weight_quant.astype("uint8").cpu())
-
-if TokenPadding == 0:
             out_cuda = w4afp8_gemm(
-        input_fp8,
+                self.input_fp8,
                 weight_int4.cuda(),
-        tokens_perfix_sum,
+                self.tokens_perfix_sum,
-        input_row_sum.astype("float32"),
+                self.input_row_sum.astype("float32"),
                 weight_dequant_scale.astype("float32"),
-        int(TokenPadding),
+                int(self.TokenPadding),
-        max_tokens,
+                self.max_tokens,
                 True,
             )
         else:
             out_cuda = w4afp8_gemm(
-        input_fp8,
+                self.input_fp8,
                 weight_int4.cuda(),
-        tokens,
+                self.tokens,
-        input_row_sum.astype("float32"),
+                self.input_row_sum.astype("float32"),
                 weight_dequant_scale.astype("float32"),
-        int(TokenPadding),
+                int(self.TokenPadding),
-        max_tokens,
+                self.max_tokens,
                 True,
             )
 
         gap = (out_cuda - out_naive).abs()
-assert float(gap.mean()) < 0.07
+        self.assertLess(float(gap.mean()), 0.07)
+
+
+if __name__ == "__main__":
+    unittest.main()