add rebuild_padding tests

Author: co63oc

tests/operators/test_rebuild_padding.py

@@ -0,0 +1,200 @@
+import unittest
+
+import numpy as np
+import paddle
+
+from fastdeploy.model_executor.ops.gpu import rebuild_padding
+
+
+def RebuildPaddingKernel(
+    out,
+    tmp_out,
+    cu_seqlens_q,
+    seq_len_this_time,
+    seq_lens_decoder,
+    seq_lens_encoder,
+    bsz,
+):
+    for bi in range(bsz):
+        seq_id = 0
+        if seq_len_this_time[bi] == 0:
+            continue
+        if seq_lens_decoder[bi] == 0 and seq_lens_encoder[bi] == 0:
+            continue
+        if seq_lens_encoder[bi] > 0:
+            seq_id = seq_lens_encoder[bi] - 1
+        out[bi] = tmp_out[cu_seqlens_q[bi] + seq_id][:]
+
+
+def RebuildAppendPaddingKernel(
+    out,
+    tmp_out,
+    cu_seqlens_q,
+    seq_len_this_time,
+    seq_lens_decoder,
+    seq_lens_encoder,
+    output_padding_offset,
+    max_input_length,
+    token_num,
+    need_delete_token_num,
+):
+    for token_id in range(token_num - need_delete_token_num):
+        bi = int(token_id / max_input_length)
+        if seq_len_this_time[bi] == 0 or (seq_lens_decoder[bi] == 0 and seq_lens_encoder[bi] == 0):
+            continue
+        ori_token_id = token_id + output_padding_offset[token_id]
+        seq_id = 0
+        if seq_lens_encoder[bi] > 0:
+            seq_id = seq_lens_encoder[bi] - 1
+        cum_offset_bi = bi * max_input_length - cu_seqlens_q[bi]
+        input_token_id = ori_token_id - cum_offset_bi + seq_id
+        out[token_id] = tmp_out[input_token_id][:]
+
+
+def rebuild_padding_ref(
+    tmp_out,  # [token_num, dim_embed]
+    cu_seqlens_q,  # [bsz+1, 1]
+    seq_len_this_time,
+    seq_lens_decoder,
+    seq_lens_encoder,
+    output_padding_offset,
+    max_input_length,
+):
+
+    tmp_out_shape = tmp_out.shape
+    token_num = tmp_out_shape[0]
+    dim_embed = tmp_out_shape[1]
+    bsz = cu_seqlens_q.shape[0] - 1
+
+    out = np.zeros([bsz, dim_embed])
+    if output_padding_offset is not None:
+        need_delete_token_num = 0
+        for i in range(bsz):
+            if seq_lens_encoder[i] > 0:
+                need_delete_token_num += seq_lens_encoder[i] - 1
+        out = np.zeros([token_num - need_delete_token_num, dim_embed])
+    else:
+        out = np.zeros([bsz, dim_embed])
+
+    if output_padding_offset is not None:
+        RebuildAppendPaddingKernel(
+            out,
+            tmp_out,
+            cu_seqlens_q,
+            seq_len_this_time,
+            seq_lens_decoder,
+            seq_lens_encoder,
+            output_padding_offset,
+            max_input_length,
+            token_num,
+            need_delete_token_num,
+        )
+    else:
+        RebuildPaddingKernel(
+            out,
+            tmp_out,
+            cu_seqlens_q,
+            seq_len_this_time,
+            seq_lens_decoder,
+            seq_lens_encoder,
+            bsz,
+        )
+    return out
+
+
+class TestRebuildPadding(unittest.TestCase):
+    # test no offset
+    def test_rebuild_padding_no_offset(self):
+        token_num = 100
+        dim_embed = 256
+        # bsz = 4
+        max_input_length = 512
+        # tmp_out: [token_num, dim_embed]
+        tmp_out = np.random.randn(token_num, dim_embed).astype(np.float32)
+        # cu_seqlens_q: [bsz + 1]，accumulate the number of tokens for each batch.
+        cu_seqlens_q = np.array(
+            [0, 1, 21, 22, 42, 43, 63, 64, 84], dtype=np.int32
+        )  # Assume there are 4 batches, and the total token_num = 100.
+
+        # Simulated sequence length information
+        seq_len_this_time = np.array([1, 20, 1, 20, 1, 20, 1, 20], dtype=np.int32)
+        seq_lens_encoder = np.array([0, 20, 0, 20, 0, 20, 0, 20], dtype=np.int32)
+        seq_lens_decoder = np.array([21, 0, 21, 0, 21, 0, 21, 0], dtype=np.int32)
+        out_no_offset_ref = rebuild_padding_ref(
+            tmp_out=tmp_out,
+            cu_seqlens_q=cu_seqlens_q,
+            seq_len_this_time=seq_len_this_time,
+            seq_lens_decoder=seq_lens_decoder,
+            seq_lens_encoder=seq_lens_encoder,
+            output_padding_offset=None,
+            max_input_length=max_input_length,
+        )
+
+        tmp_out = paddle.to_tensor(tmp_out)
+        cu_seqlens_q = paddle.to_tensor(cu_seqlens_q)
+        seq_len_this_time = paddle.to_tensor(seq_len_this_time)
+        seq_lens_decoder = paddle.to_tensor(seq_lens_decoder)
+        seq_lens_encoder = paddle.to_tensor(seq_lens_encoder)
+
+        out_no_offset = rebuild_padding(
+            tmp_out,
+            cu_seqlens_q,
+            seq_len_this_time,
+            seq_lens_decoder,
+            seq_lens_encoder,
+            None,
+            max_input_length,
+        )
+        np.testing.assert_allclose(out_no_offset.numpy(), out_no_offset_ref)
+
+    # test with offset
+    def test_rebuild_padding_with_offset(self):
+        paddle.seed(42)
+        token_num = 100
+        dim_embed = 256
+        # bsz = 4
+        max_input_length = 512
+        # tmp_out: [token_num, dim_embed]
+        tmp_out = np.random.randn(token_num, dim_embed).astype(np.float32)
+        # cu_seqlens_q: [bsz + 1]，accumulate the number of tokens for each batch.
+        cu_seqlens_q = np.array(
+            [0, 1, 21, 22, 42, 43, 63, 64, 84], dtype=np.int32
+        )  # Assume there are 4 batches, and the total token_num = 100.
+
+        # Simulated sequence length information
+        seq_len_this_time = np.array([1, 20, 1, 20, 1, 20, 1, 20], dtype=np.int32)
+        seq_lens_encoder = np.array([0, 20, 0, 20, 0, 20, 0, 20], dtype=np.int32)
+        seq_lens_decoder = np.array([21, 0, 21, 0, 21, 0, 21, 0], dtype=np.int32)
+
+        num_output_tokens = 80
+        output_padding_offset = np.random.randint(0, 10, [num_output_tokens], dtype=np.int32)
+        out_with_offset_ref = rebuild_padding_ref(
+            tmp_out=tmp_out,
+            cu_seqlens_q=cu_seqlens_q,
+            seq_len_this_time=seq_len_this_time,
+            seq_lens_decoder=seq_lens_decoder,
+            seq_lens_encoder=seq_lens_encoder,
+            output_padding_offset=output_padding_offset,
+            max_input_length=max_input_length,
+        )
+
+        tmp_out = paddle.to_tensor(tmp_out)
+        cu_seqlens_q = paddle.to_tensor(cu_seqlens_q)
+        seq_len_this_time = paddle.to_tensor(seq_len_this_time)
+        seq_lens_decoder = paddle.to_tensor(seq_lens_decoder)
+        seq_lens_encoder = paddle.to_tensor(seq_lens_encoder)
+        output_padding_offset = paddle.to_tensor(output_padding_offset)
+        out_with_offset = rebuild_padding(
+            tmp_out,
+            cu_seqlens_q,
+            seq_len_this_time,
+            seq_lens_decoder,
+            seq_lens_encoder,
+            output_padding_offset,
+            max_input_length,
+        )
+        np.testing.assert_allclose(out_with_offset.numpy(), out_with_offset_ref)
+
+
+if __name__ == "__main__":
+    unittest.main()