frozenleaves/vllm
1
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2025-10-20 14:53:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Lily Liu
2025-03-16 00:29:22 -07:00
committed by GitHub
parent b82662d952
commit d1ad2a57af
1 changed files with 30 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
tests/v1/spec_decode/test_ngram.py
View File

@ -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]))