A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 13:30 UTC
[Bugfix] Fix v1/spec_decode/test_ngram.py (#16895)
Signed-off-by: qizixi <qizixi@meta.com>
qizixi committed
Apr 20, 2025 at 20:54 UTC
bb3605db85c2d986bf20e0fdd0a50ec27502c367
2 files changed
+30
-38
tests/v1/spec_decode/test_ngram.py
+22
-31
@@ -2,6 +2,7 @@
2
3
import numpy as np
4
5
+from vllm.config import SpeculativeConfig, VllmConfig
6
from vllm.v1.spec_decode.ngram_proposer import (NgramProposer,
7
_find_subarray_kmp,
8
_kmp_lps_array)
@@ -39,50 +40,40 @@ def test_find_subarray_kmp():
40
41
42
def test_ngram_proposer():
42
- proposer = NgramProposer()
43
+
44
+ def ngram_proposer(min_n: int, max_n: int, k: int) -> NgramProposer:
45
+ return NgramProposer(vllm_config=VllmConfig(
46
+ speculative_config=SpeculativeConfig.from_dict(
47
+ {
48
+ "prompt_lookup_min": min_n,
49
+ "prompt_lookup_max": max_n,
50
+ "num_speculative_tokens": k,
51
+ "method": "ngram",
52
+ })))
53
54
# No match.
45
- result = proposer.propose(
46
- context_token_ids=np.array([1, 2, 3, 4, 5]),
47
- min_n=2,
48
- max_n=2,
49
- k=2,
50
- )
55
+ result = ngram_proposer(
56
+ 2, 2, 2).propose(context_token_ids=np.array([1, 2, 3, 4, 5]))
57
assert result is None
58
59
# No match for 4-gram.
54
- result = proposer.propose(
55
- context_token_ids=np.array([1, 2, 3, 4, 1, 2, 3]),
56
- min_n=4,
57
- max_n=4,
58
- k=2,
59
- )
60
+ result = ngram_proposer(
61
+ 4, 4, 2).propose(context_token_ids=np.array([1, 2, 3, 4, 1, 2, 3]))
62
assert result is None
63
64
# No match for 4-gram but match for 3-gram.
63
- result = proposer.propose(
64
- context_token_ids=np.array([1, 2, 3, 4, 1, 2, 3]),
65
- min_n=3,
66
- max_n=4,
67
- k=2,
68
- )
65
+ result = ngram_proposer(
66
+ 3, 4, 2).propose(context_token_ids=np.array([1, 2, 3, 4, 1, 2, 3]))
67
assert np.array_equal(result, np.array([4, 1]))
68
69
# Match for both 4-gram and 3-gram.
70
# In this case, the proposer should return the 4-gram match.
73
- result = proposer.propose(
74
- context_token_ids=np.array([2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]),
75
- min_n=3,
76
- max_n=4,
77
- k=2,
78
- )
71
+ result = ngram_proposer(3, 4, 2).propose(
72
+ context_token_ids=np.array([2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]))
73
assert np.array_equal(result, np.array([1, 2])) # Not [5, 1]
74
75
# Match for 2-gram and 3-gram, but not 4-gram.
82
- result = proposer.propose(
83
- context_token_ids=np.array([3, 4, 5, 2, 3, 4, 1, 2, 3, 4]),
84
- min_n=2,
85
- max_n=4,
86
- k=2,
87
- )
76
+ result = ngram_proposer(
77
+ 2, 4,
78
+ 2).propose(context_token_ids=np.array([3, 4, 5, 2, 3, 4, 1, 2, 3, 4]))
79
assert np.array_equal(result, np.array([1, 2])) # Not [5, 2]
vllm/config.py
+8
-7
@@ -120,7 +120,7 @@ def get_attr_docs(cls: type[Any]) -> dict[str, str]:
120
def pairwise(iterable):
121
"""
122
Manually implement https://docs.python.org/3/library/itertools.html#itertools.pairwise
123
-
123
+
124
Can be removed when Python 3.9 support is dropped.
125
"""
126
iterator = iter(iterable)
@@ -266,7 +266,7 @@ class ModelConfig:
266
config_format: The config format which shall be loaded.
267
Defaults to 'auto' which defaults to 'hf'.
268
hf_token: The token to use as HTTP bearer authorization for remote files
269
- . If `True`, will use the token generated when running
269
+ . If `True`, will use the token generated when running
270
`huggingface-cli login` (stored in `~/.huggingface`).
271
hf_overrides: If a dictionary, contains arguments to be forwarded to the
272
HuggingFace config. If a callable, it is called to update the
@@ -1624,7 +1624,7 @@ class ParallelConfig:
1624
"""The full name of the worker class to use. If "auto", the worker class
1625
will be determined based on the platform."""
1626
sd_worker_cls: str = "auto"
1627
- """The full name of the worker class to use for speculative decofing.
1627
+ """The full name of the worker class to use for speculative decofing.
1628
If "auto", the worker class will be determined based on the platform."""
1629
worker_extension_cls: str = ""
1630
"""The full name of the worker extension class to use. The worker extension
@@ -1815,13 +1815,13 @@ class SchedulerConfig:
1815
1816
max_num_batched_tokens: int = None # type: ignore
1817
"""Maximum number of tokens to be processed in a single iteration.
1818
-
1818
+
1819
This config has no static default. If left unspecified by the user, it will
1820
be set in `EngineArgs.create_engine_config` based on the usage context."""
1821
1822
max_num_seqs: int = None # type: ignore
1823
"""Maximum number of sequences to be processed in a single iteration.
1824
-
1824
+
1825
This config has no static default. If left unspecified by the user, it will
1826
be set in `EngineArgs.create_engine_config` based on the usage context."""
1827
@@ -1867,7 +1867,7 @@ class SchedulerConfig:
1867
# TODO (ywang96): Make this configurable.
1868
max_num_encoder_input_tokens: int = field(init=False)
1869
"""Multimodal encoder compute budget, only used in V1.
1870
-
1870
+
1871
NOTE: This is not currently configurable. It will be overridden by
1872
max_num_batched_tokens in case max multimodal embedding size is larger."""
1873
@@ -2306,7 +2306,8 @@ class SpeculativeConfig:
2306
if self.model is None and self.num_speculative_tokens is not None:
2307
# TODO(Shangming): Refactor mtp configuration logic when supporting
2308
# mtp acceleration for more models besides deepseek_v3
2309
- if self.target_model_config.hf_text_config.model_type \
2309
+ if self.target_model_config and \
2310
+ self.target_model_config.hf_text_config.model_type \
2311
== "deepseek_v3":
2312
# use the draft model from the same model:
2313
self.model = self.target_model_config.model