diff --git a/tests/v1/spec_decode/test_ngram.py b/tests/v1/spec_decode/test_ngram.py index ec663c84d0d2a..2c2e125ade48c 100644 --- a/tests/v1/spec_decode/test_ngram.py +++ b/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] \ No newline at end of file + np.testing.assert_array_equal(_find_subarray_kmp(X, 1, 3), + np.array([6, 2, 3]))