add speculate_update_v3 test

Author: co63oc

tests/operators/test_speculate_update_v3.py

@@ -0,0 +1,171 @@
+# 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
+
+import numpy as np
+import paddle
+
+from fastdeploy.model_executor.ops.gpu import speculate_update
+
+
+def speculate_update_v3_np(
+    seq_lens_encoder,
+    seq_lens_decoder,
+    not_need_stop,
+    draft_tokens,
+    actual_draft_token_nums,
+    accept_tokens,
+    accept_num,
+    stop_flags,
+    seq_lens_this_time,
+    is_block_step,
+    stop_nums,
+):
+    stop_sum = 0
+    real_bsz = seq_lens_this_time.shape[0]
+    max_bsz = stop_flags.shape[0]
+    max_draft_tokens = draft_tokens.shape[1]
+
+    for bid in range(max_bsz):
+        stop_flag_now_int = 0
+        inactive = bid >= real_bsz
+        block_step = (not inactive) and is_block_step[bid]
+
+        if (not block_step) and (not inactive):
+
+            if stop_flags[bid]:
+                stop_flag_now_int = 1
+
+            if seq_lens_encoder[bid] == 0:
+                seq_lens_decoder[bid] += accept_num[bid]
+
+            if (seq_lens_encoder[bid] == 0) and (seq_lens_this_time[bid] > 1):
+                cur_len = actual_draft_token_nums[bid]
+                if accept_num[bid] - 1 == cur_len:
+                    if cur_len + 2 <= max_draft_tokens - 1:
+                        cur_len += 2
+                    elif cur_len + 1 <= max_draft_tokens - 1:
+                        cur_len += 1
+                    else:
+                        cur_len = max_draft_tokens - 1
+                else:
+                    cur_len = max(1, cur_len - 1)
+                actual_draft_token_nums[bid] = cur_len
+
+            if seq_lens_encoder[bid] != 0:
+                seq_lens_decoder[bid] += seq_lens_encoder[bid]
+                seq_lens_encoder[bid] = 0
+
+            draft_tokens[bid, 0] = accept_tokens[bid, accept_num[bid] - 1]
+
+            if stop_flag_now_int:
+                seq_lens_decoder[bid] = 0
+
+        elif inactive:
+            stop_flag_now_int = 1
+
+        stop_sum += stop_flag_now_int
+    not_need_stop[0] = stop_sum < stop_nums[0]
+
+    return (
+        seq_lens_encoder,
+        seq_lens_decoder,
+        not_need_stop,
+        draft_tokens,
+        actual_draft_token_nums,
+    )
+
+
+def gen_inputs(
+    max_bsz=512,
+    max_draft_tokens=16,
+    real_bsz=123,
+    seed=2022,
+):
+    rng = np.random.default_rng(seed)
+
+    seq_lens_encoder = rng.integers(0, 3, size=max_bsz, dtype=np.int32)
+    seq_lens_decoder = rng.integers(0, 20, size=max_bsz, dtype=np.int32)
+    not_need_stop = rng.integers(0, 1, size=1, dtype=np.bool_)
+    draft_tokens = rng.integers(0, 1000, size=(max_bsz, max_draft_tokens), dtype=np.int64)
+    actual_draft_nums = rng.integers(1, max_draft_tokens, size=max_bsz, dtype=np.int32)
+    accept_tokens = rng.integers(0, 1000, size=(max_bsz, max_draft_tokens), dtype=np.int64)
+    accept_num = rng.integers(1, max_draft_tokens, size=max_bsz, dtype=np.int32)
+    stop_flags = rng.integers(0, 2, size=max_bsz, dtype=np.bool_)
+    is_block_step = rng.integers(0, 2, size=max_bsz, dtype=np.bool_)
+    stop_nums = np.array([5], dtype=np.int64)
+
+    seq_lens_this_time = rng.integers(1, max_draft_tokens, size=real_bsz, dtype=np.int32)
+
+    return {
+        "seq_lens_encoder": seq_lens_encoder,
+        "seq_lens_decoder": seq_lens_decoder,
+        "not_need_stop": not_need_stop,
+        "draft_tokens": draft_tokens,
+        "actual_draft_token_nums": actual_draft_nums,
+        "accept_tokens": accept_tokens,
+        "accept_num": accept_num,
+        "stop_flags": stop_flags,
+        "seq_lens_this_time": seq_lens_this_time,
+        "is_block_step": is_block_step,
+        "stop_nums": stop_nums,
+    }
+
+
+class TestSpeculateUpdateV3(unittest.TestCase):
+    def test_speculate_update_v3(self):
+        inputs = gen_inputs(max_bsz=512, max_draft_tokens=32, real_bsz=201)
+
+        paddle_inputs = {}
+        for k, v in inputs.items():
+            paddle_inputs[k] = paddle.to_tensor(v)
+
+        np_inputs = {
+            k: (paddle_inputs[k].numpy().copy() if isinstance(paddle_inputs[k], paddle.Tensor) else paddle_inputs[k])
+            for k in paddle_inputs
+        }
+
+        out_pd = speculate_update(*(paddle_inputs.values()))
+        (
+            seq_lens_encoder_pd,
+            seq_lens_decoder_pd,
+            not_need_stop_pd,
+            draft_tokens_pd,
+            actual_draft_nums_pd,
+        ) = out_pd
+
+        out_np = speculate_update_v3_np(**np_inputs)
+
+        names = [
+            "seq_lens_encoder",
+            "seq_lens_decoder",
+            "not_need_stop",
+            "draft_tokens",
+            "actual_draft_token_nums",
+        ]
+        pd_tensors = [
+            seq_lens_encoder_pd,
+            seq_lens_decoder_pd,
+            not_need_stop_pd,
+            draft_tokens_pd,
+            actual_draft_nums_pd,
+        ]
+
+        for name, pd_val, np_val in zip(names, pd_tensors, out_np):
+            np.testing.assert_allclose(pd_val.numpy(), np_val)
+
+
+if __name__ == "__main__":
+    unittest.main()