support mtp in v1_scheduler mode (#3695)

tests/operators/test_speculative_schedule_cache.py

@@ -0,0 +1,239 @@
+import unittest
+
+import numpy as np
+import paddle
+
+from fastdeploy.model_executor.ops.gpu import speculate_schedule_cache
+
+
+def cpu_reference(
+    draft_tokens,
+    block_tables,
+    stop_flags,
+    seq_lens_this_time,
+    seq_lens_decoder,
+    step_seq_lens_decoder,
+    step_draft_tokens,
+    step_seq_lens_this_time,
+    accept_num,
+    accept_tokens,
+    is_block_step,
+    not_need_stop,
+    stop_nums,
+    block_size,
+    max_draft_tokens,
+):
+    """Pure-NumPy mirror of the CUDA kernel's logic (single block of 512 threads).
+
+    Shapes are the same as inputs to the custom op. This mutates the provided
+    NumPy arrays in-place, exactly like the kernel does.
+    """
+    real_bsz = seq_lens_this_time.shape[0]
+    max_bsz = stop_flags.shape[0]
+    draft_tokens_len = draft_tokens.shape[1]
+    block_num_per_seq = block_tables.shape[1]
+
+    max_next_step_tokens = 2 * max_draft_tokens + 2
+
+    # Block-local reduction input per thread (threadIdx.x -> bid)
+    stop_flag_now_int = np.zeros(512, dtype=np.int64)  # THREADBLOCK_SIZE = 512
+
+    for bid in range(512):
+        if bid < real_bsz:
+            if not stop_flags[bid]:
+                max_possible_block_idx = (seq_lens_decoder[bid] + max_next_step_tokens) // block_size
+                if max_possible_block_idx < block_num_per_seq and block_tables[bid, max_possible_block_idx] == -1:
+                    is_block_step[bid] = True
+                    step_seq_lens_this_time[bid] = seq_lens_this_time[bid]
+                    seq_lens_this_time[bid] = 0
+                    stop_flags[bid] = True
+                    step_seq_lens_decoder[bid] = seq_lens_decoder[bid]
+                    seq_lens_decoder[bid] = 0
+                    accept_num[bid] = 0
+                    accept_tokens[bid, :] = -1
+                    step_draft_tokens[bid, :draft_tokens_len] = draft_tokens[bid, :draft_tokens_len]
+                    stop_flag_now_int[bid] = 1
+                else:
+                    stop_flag_now_int[bid] = 0
+            else:
+                stop_flag_now_int[bid] = 1
+        elif bid < max_bsz:
+            # Threads in [real_bsz, max_bsz) contribute 1 to reduction
+            stop_flag_now_int[bid] = 1
+        else:
+            stop_flag_now_int[bid] = 0
+
+    stop_sum = int(stop_flag_now_int.sum())
+    not_need_stop[0] = stop_sum < int(stop_nums[0])
+
+
+class TestSpeculateScheduleCache(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not paddle.is_compiled_with_cuda():
+            raise unittest.SkipTest("Paddle is not compiled with CUDA; skipping GPU op test.")
+        paddle.device.set_device("gpu")
+
+    def setUp(self):
+        # --- Construct a deterministic case that exercises all branches ---
+        # real_bsz < max_bsz to test the padding logic in the CUB reduction
+        self.real_bsz = 3
+        self.max_bsz = 5  # only stop_flags has length max_bsz
+
+        self.draft_tokens_len = 6
+        self.accept_tokens_len = 5
+        self.block_size = 4
+        self.block_num_per_seq = 3
+        self.max_draft_tokens = 2  # -> max_next_step_tokens = 6
+
+        # Inputs that will trigger for bid 0, not trigger for bid 2, and bid 1 is already stopped
+        # seq_lens_decoder + 6 // 4 -> indices: [1, 1, 4]. Index 4 is out of range -> no trigger on bid 2
+        self.draft_tokens = paddle.to_tensor(
+            np.array(
+                [
+                    [1, 1, 1, 1, 1, 1],
+                    [2, 2, 2, 2, 2, 2],
+                    [3, 3, 3, 3, 3, 3],
+                ],
+                dtype=np.int64,
+            )
+        )
+        self.block_tables = paddle.to_tensor(np.full((self.real_bsz, self.block_num_per_seq), -1, dtype=np.int32))
+        # stop_flags length is max_bsz, others are real_bsz
+        self.stop_flags = paddle.to_tensor(np.array([False, True, False, False, False], dtype=np.bool_))
+        self.seq_lens_this_time = paddle.to_tensor(np.array([5, 6, 7], dtype=np.int32))
+        self.seq_lens_decoder = paddle.to_tensor(np.array([1, 1, 10], dtype=np.int32))
+
+        # Will be filled by kernel for the triggering bids only
+        self.step_seq_lens_decoder = paddle.zeros((self.real_bsz,), dtype="int32")
+        self.step_draft_tokens = paddle.zeros((self.real_bsz, self.draft_tokens_len), dtype="int64")
+        self.step_seq_lens_this_time = paddle.zeros((self.real_bsz,), dtype="int32")
+
+        # Intentionally non-zero so we can verify in-place zeroing only where triggered
+        self.accept_num = paddle.to_tensor(np.array([9, 8, 7], dtype=np.int32))
+        self.accept_tokens = paddle.to_tensor(
+            np.arange(self.real_bsz * self.accept_tokens_len, dtype=np.int64).reshape(
+                self.real_bsz, self.accept_tokens_len
+            )
+        )
+        self.is_block_step = paddle.zeros((self.real_bsz,), dtype=paddle.bool)
+
+        # not_need_stop lives on CPU in the caller; the kernel copies to device internally
+        self.not_need_stop = paddle.zeros((1,), dtype=paddle.bool).cpu()
+
+        # Choose threshold so with: bid0 triggers, bid1 already stopped, padding (5-3)=2 -> stop_sum = 1+1+2 = 4
+        # Set stop_nums to 5 so not_need_stop = (4 < 5) = True
+        self.stop_nums = paddle.to_tensor([5], dtype=paddle.int64)
+
+        # Keep NumPy copies for CPU reference
+        self.np_draft_tokens = self.draft_tokens.numpy().copy()
+        self.np_block_tables = self.block_tables.numpy().copy()
+        self.np_stop_flags = self.stop_flags.numpy().copy()
+        self.np_seq_lens_this_time = self.seq_lens_this_time.numpy().copy()
+        self.np_seq_lens_decoder = self.seq_lens_decoder.numpy().copy()
+        self.np_step_seq_lens_decoder = self.step_seq_lens_decoder.numpy().copy()
+        self.np_step_draft_tokens = self.step_draft_tokens.numpy().copy()
+        self.np_step_seq_lens_this_time = self.step_seq_lens_this_time.numpy().copy()
+        self.np_accept_num = self.accept_num.numpy().copy()
+        self.np_accept_tokens = self.accept_tokens.numpy().copy()
+        self.np_is_block_step = self.is_block_step.numpy().copy()
+        self.np_not_need_stop = self.not_need_stop.numpy().copy()
+        self.np_stop_nums = self.stop_nums.numpy().copy()
+
+    def test_correctness_against_cpu_reference(self):
+        # Run GPU kernel (in-place)
+        speculate_schedule_cache(
+            self.draft_tokens,
+            self.block_tables,
+            self.stop_flags,
+            self.seq_lens_this_time,
+            self.seq_lens_decoder,
+            self.step_seq_lens_decoder,
+            self.step_draft_tokens,
+            self.step_seq_lens_this_time,
+            self.accept_num,
+            self.accept_tokens,
+            self.is_block_step,
+            self.not_need_stop,
+            self.stop_nums,
+            self.block_size,
+            self.max_draft_tokens,
+        )
+
+        # Compute CPU reference (in-place on NumPy copies)
+        cpu_reference(
+            self.np_draft_tokens,
+            self.np_block_tables,
+            self.np_stop_flags,
+            self.np_seq_lens_this_time,
+            self.np_seq_lens_decoder,
+            self.np_step_seq_lens_decoder,
+            self.np_step_draft_tokens,
+            self.np_step_seq_lens_this_time,
+            self.np_accept_num,
+            self.np_accept_tokens,
+            self.np_is_block_step,
+            self.np_not_need_stop,
+            self.np_stop_nums,
+            self.block_size,
+            self.max_draft_tokens,
+        )
+
+        # Compare all mutated tensors
+        np.testing.assert_array_equal(self.step_draft_tokens.numpy(), self.np_step_draft_tokens)
+        np.testing.assert_array_equal(self.accept_tokens.numpy(), self.np_accept_tokens)
+        np.testing.assert_array_equal(self.stop_flags.numpy(), self.np_stop_flags)
+        np.testing.assert_array_equal(self.is_block_step.numpy(), self.np_is_block_step)
+        np.testing.assert_array_equal(self.seq_lens_this_time.numpy(), self.np_seq_lens_this_time)
+        np.testing.assert_array_equal(self.seq_lens_decoder.numpy(), self.np_seq_lens_decoder)
+        np.testing.assert_array_equal(self.step_seq_lens_decoder.numpy(), self.np_step_seq_lens_decoder)
+        np.testing.assert_array_equal(self.step_seq_lens_this_time.numpy(), self.np_step_seq_lens_this_time)
+        np.testing.assert_array_equal(self.accept_num.numpy(), self.np_accept_num)
+        self.assertEqual(bool(self.not_need_stop.numpy()[0]), bool(self.np_not_need_stop[0]))
+
+    def test_no_trigger_path(self):
+        # Make block_tables at candidate index != -1 so nothing triggers
+        # Candidate index for bid 0/1 is 1, set it to 7
+        bt = self.block_tables.numpy()
+        bt[:, 1] = 7
+        self.block_tables = paddle.to_tensor(bt)
+
+        # Reset outputs to distinctive values
+        self.step_seq_lens_decoder[:] = 0
+        self.step_draft_tokens[:] = 0
+        self.step_seq_lens_this_time[:] = 0
+        self.accept_num[:] = -123
+        self.accept_tokens[:] = -777
+        self.is_block_step[:] = False
+        self.not_need_stop[:] = False
+
+        # For not_need_stop: stopped_in_real = (bid1 True) = 1, padding = 2 -> stop_sum=3
+        # With stop_nums=5 -> True
+        speculate_schedule_cache(
+            self.draft_tokens,
+            self.block_tables,
+            self.stop_flags,
+            self.seq_lens_this_time,
+            self.seq_lens_decoder,
+            self.step_seq_lens_decoder,
+            self.step_draft_tokens,
+            self.step_seq_lens_this_time,
+            self.accept_num,
+            self.accept_tokens,
+            self.is_block_step,
+            self.not_need_stop,
+            self.stop_nums,
+            self.block_size,
+            self.max_draft_tokens,
+        )
+
+        # Nothing should have changed except not_need_stop
+        np.testing.assert_array_equal(self.step_draft_tokens.numpy(), np.zeros_like(self.step_draft_tokens.numpy()))
+        np.testing.assert_array_equal(self.is_block_step.numpy(), np.zeros_like(self.is_block_step.numpy()))
+        np.testing.assert_array_equal(self.accept_tokens.numpy(), np.full_like(self.accept_tokens.numpy(), -777))
+        np.testing.assert_array_equal(self.accept_num.numpy(), np.full_like(self.accept_num.numpy(), -123))
+        self.assertTrue(bool(self.not_need_stop.numpy()[0]))
+
+
+if __name__ == "__main__":
+    unittest.main()