[V1] [Spec Decode] Fix ngram tests (vllm-project#14878)

Signed-off-by: DefTruth <[email protected]>

Author: LiuXiaoxuanPKU

tests/v1/spec_decode/test_ngram.py

@@ -1,32 +1,37 @@
 # SPDX-License-Identifier: Apache-2.0
-import pytest
 
-from vllm.v1.spec_decode.ngram_proposer import NgramProposer
-from vllm.v1.utils import ConstantList
+import numpy as np
 
+from vllm.v1.spec_decode.ngram_proposer import (_find_subarray_kmp,
+                                                _kmp_lps_array)
 
-@pytest.fixture
-def proposer():
-    return NgramProposer()
 
+def test_kmp_lps_array():
+    np.testing.assert_array_equal(_kmp_lps_array(np.array([])), np.array([]))
+    np.testing.assert_array_equal(_kmp_lps_array(np.array([1])), np.array([0]))
+    np.testing.assert_array_equal(_kmp_lps_array(np.array([1, 1, 1])),
+                                  np.array([0, 1, 2]))
+    np.testing.assert_array_equal(_kmp_lps_array(np.array([1, 2, 3, 4])),
+                                  np.array([0, 0, 0, 0]))
+    np.testing.assert_array_equal(_kmp_lps_array(np.array([1, 2, 1, 2, 3])),
+                                  np.array([0, 0, 1, 2, 0]))
 
-def test_kmp_lps_array(proposer):
-    assert proposer._kmp_lps_array([]) == []
-    assert proposer._kmp_lps_array([1]) == [0]
-    assert proposer._kmp_lps_array([1, 1, 1]) == [0, 1, 2]
-    assert proposer._kmp_lps_array([1, 2, 3, 4]) == [0, 0, 0, 0]
-    assert proposer._kmp_lps_array([1, 2, 1, 2, 3]) == [0, 0, 1, 2, 0]
 
-
-def test_find_subarray_kmp(proposer):
-    X = ConstantList([1, 2, 3, 4, 1, 2, 3, 5, 6])
-    assert proposer._find_subarray_kmp(X, 2, 2) is None
-    X = ConstantList([1, 2, 3, 4, 1, 2, 3])
-    assert proposer._find_subarray_kmp(X, 2, 3) == [4, 1, 2]
-    assert proposer._find_subarray_kmp(X, 2, 2) == [4, 1]
-    assert proposer._find_subarray_kmp(X, 1, 3) == [4, 1, 2]
-    assert proposer._find_subarray_kmp(X, 1, 2) == [4, 1]
-    X = ConstantList([1, 3, 6, 2, 3, 4, 1, 2, 3])
-    assert proposer._find_subarray_kmp(X, 2, 3) == [4, 1, 2]
+def test_find_subarray_kmp():
+    X = np.array([1, 2, 3, 4, 1, 2, 3, 5, 6])
+    assert _find_subarray_kmp(X, 2, 2) is None
+    X = np.array([1, 2, 3, 4, 1, 2, 3])
+    np.testing.assert_array_equal(_find_subarray_kmp(X, 2, 3),
+                                  np.array([4, 1, 2]))
+    np.testing.assert_array_equal(_find_subarray_kmp(X, 2, 2), np.array([4,
+                                                                         1]))
+    np.testing.assert_array_equal(_find_subarray_kmp(X, 1, 3),
+                                  np.array([4, 1, 2]))
+    np.testing.assert_array_equal(_find_subarray_kmp(X, 1, 2), np.array([4,
+                                                                         1]))
+    X = np.array([1, 3, 6, 2, 3, 4, 1, 2, 3])
+    np.testing.assert_array_equal(_find_subarray_kmp(X, 2, 3),
+                                  np.array([4, 1, 2]))
     # Return on the first match
-    assert proposer._find_subarray_kmp(X, 1, 3) == [6, 2, 3]
+    np.testing.assert_array_equal(_find_subarray_kmp(X, 1, 3),
+                                  np.array([6, 2, 3]))