@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e71b8e210d
Fork 0

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 14, 05:30 UTC

[Spec Decode] Add Batch Parallel Ngram. Upto 8x lower overhead. (#24986)

Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Co-authored-by: Nick Hill <nhill@redhat.com>

Ekagra Ranjan committed Sep 25, 2025 at 18:22 UTC e71b8e210db4b98ffa4398d25f8bdf280686ad78
5 files changed +381 -107
benchmarks/benchmark_ngram_proposer.py
+104 -3
@@ -1,17 +1,31 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import gc
4 +import time
5 +from unittest import mock
6
7 import numpy as np
8 from tabulate import tabulate
9
10 from benchmark_utils import TimeCollector
9 -from vllm.config import ModelConfig, SpeculativeConfig, VllmConfig
11 +from vllm.config import (
12 + CacheConfig,
13 + DeviceConfig,
14 + LoadConfig,
15 + ModelConfig,
16 + ParallelConfig,
17 + SchedulerConfig,
18 + SpeculativeConfig,
19 + VllmConfig,
20 +)
21 +from vllm.platforms import current_platform
22 from vllm.utils import FlexibleArgumentParser
23 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
24 +from vllm.v1.worker.gpu_input_batch import InputBatch
25 +from vllm.v1.worker.gpu_model_runner import GPUModelRunner
26
27
14 -def main(args):
28 +def benchmark_propose(args):
29 rows = []
30 for max_ngram in args.max_ngram:
31 collector = TimeCollector(TimeCollector.US)
@@ -69,10 +83,88 @@ def main(args):
83 )
84
85
86 +def benchmark_batched_propose(args):
87 + NUM_SPECULATIVE_TOKENS_NGRAM = 10
88 + PROMPT_LOOKUP_MIN = 5
89 + PROMPT_LOOKUP_MAX = 15
90 + MAX_MODEL_LEN = int(1e7)
91 + DEVICE = current_platform.device_type
92 +
93 + model_config = ModelConfig(model="facebook/opt-125m", runner="generate")
94 +
95 + speculative_config = SpeculativeConfig(
96 + target_model_config=model_config,
97 + target_parallel_config=ParallelConfig(),
98 + method="ngram",
99 + num_speculative_tokens=NUM_SPECULATIVE_TOKENS_NGRAM,
100 + prompt_lookup_max=PROMPT_LOOKUP_MAX,
101 + prompt_lookup_min=PROMPT_LOOKUP_MIN,
102 + )
103 +
104 + vllm_config = VllmConfig(
105 + model_config=model_config,
106 + cache_config=CacheConfig(),
107 + speculative_config=speculative_config,
108 + device_config=DeviceConfig(device=current_platform.device_type),
109 + parallel_config=ParallelConfig(),
110 + load_config=LoadConfig(),
111 + scheduler_config=SchedulerConfig(),
112 + )
113 +
114 + # monkey patch vllm.v1.worker.gpu_model_runner.get_pp_group
115 + mock_pp_group = mock.MagicMock()
116 + mock_pp_group.world_size = 1
117 + with mock.patch(
118 + "vllm.v1.worker.gpu_model_runner.get_pp_group", return_value=mock_pp_group
119 + ):
120 + runner = GPUModelRunner(vllm_config, DEVICE)
121 +
122 + # hack max model len
123 + runner.max_model_len = MAX_MODEL_LEN
124 + runner.drafter.max_model_len = MAX_MODEL_LEN
125 +
126 + dummy_input_batch = InputBatch(
127 + max_num_reqs=args.num_req,
128 + max_model_len=MAX_MODEL_LEN,
129 + max_num_batched_tokens=args.num_req * args.num_token,
130 + device=DEVICE,
131 + pin_memory=False,
132 + vocab_size=256000,
133 + block_sizes=[16],
134 + )
135 + dummy_input_batch._req_ids = list(str(id) for id in range(args.num_req))
136 + dummy_input_batch.spec_decode_unsupported_reqs = ()
137 + dummy_input_batch.num_tokens_no_spec = [args.num_token] * args.num_req
138 + dummy_input_batch.token_ids_cpu = np.random.randint(
139 + 0, 20, (args.num_req, args.num_token)
140 + )
141 +
142 + runner.input_batch = dummy_input_batch
143 +
144 + sampled_token_ids = [[0]] * args.num_req
145 +
146 + print("Starting benchmark")
147 + # first run is warmup so ignore it
148 + for _ in range(args.num_iteration):
149 + start = time.time()
150 + runner.drafter.propose(
151 + sampled_token_ids,
152 + dummy_input_batch.req_ids,
153 + dummy_input_batch.num_tokens_no_spec,
154 + dummy_input_batch.token_ids_cpu,
155 + dummy_input_batch.spec_decode_unsupported_reqs,
156 + )
157 + end = time.time()
158 + print(f"Iteration time (s): {end - start}")
159 +
160 +
161 def invoke_main() -> None:
162 parser = FlexibleArgumentParser(
163 description="Benchmark the performance of N-gram speculative decode drafting"
164 )
165 + parser.add_argument(
166 + "--batched", action="store_true", help="consider time to prepare batch"
167 + ) # noqa: E501
168 parser.add_argument(
169 "--num-iteration",
170 type=int,
@@ -105,8 +197,17 @@ def invoke_main() -> None:
197 help="Number of speculative tokens to generate",
198 )
199 args = parser.parse_args()
108 - main(args)
200 +
201 + if not args.batched:
202 + benchmark_propose(args)
203 + else:
204 + benchmark_batched_propose(args)
205
206
207 +"""
208 +# Example command lines:
209 +# time python3 benchmarks/benchmark_ngram_proposer.py
210 +# time python3 benchmarks/benchmark_ngram_proposer.py --batched --num-iteration 4 --num-token 1000000 --num-req 128
211 +""" # noqa: E501
212 if __name__ == "__main__":
213 invoke_main() # pragma: no cover
tests/v1/spec_decode/test_ngram.py
+114 -28
@@ -9,11 +9,13 @@ from vllm.v1.spec_decode.ngram_proposer import (
9
10 def test_find_longest_matched_ngram_and_propose_tokens():
11 tokens = np.array([1, 2, 3, 4, 1, 2, 3, 5, 6])
12 - assert _find_longest_matched_ngram_and_propose_tokens(origin_tokens=tokens,
13 - min_ngram=2,
14 - max_ngram=2,
15 - max_model_len=1024,
16 - k=2) is None
12 + result = _find_longest_matched_ngram_and_propose_tokens(
13 + origin_tokens=tokens,
14 + min_ngram=2,
15 + max_ngram=2,
16 + max_model_len=1024,
17 + k=2)
18 + assert len(result) == 0
19
20 tokens = np.array([1, 2, 3, 4, 1, 2, 3])
21 np.testing.assert_array_equal(
@@ -62,7 +64,7 @@ def test_find_longest_matched_ngram_and_propose_tokens():
64
65 def test_ngram_proposer():
66
65 - def ngram_proposer(min_n: int, max_n: int, k: int) -> NgramProposer:
67 + def get_ngram_proposer(min_n: int, max_n: int, k: int) -> NgramProposer:
68 # Dummy model config. Just to set max_model_len.
69 model_config = ModelConfig(model="facebook/opt-125m")
70 return NgramProposer(
@@ -75,36 +77,120 @@ def test_ngram_proposer():
77 )))
78
79 # No match.
78 - result = ngram_proposer(
79 - min_n=2, max_n=2,
80 - k=2).propose(context_token_ids=np.array([1, 2, 3, 4, 5]))
81 - assert result is None
80 + token_ids_cpu = np.array([[1, 2, 3, 4, 5]])
81 + result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(
82 + sampled_token_ids=[[0]],
83 + req_ids=["0"],
84 + num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
85 + token_ids_cpu=token_ids_cpu,
86 + spec_decode_unsupported_reqs=(),
87 + )
88 + assert len(result[0]) == 0
89
90 # No match for 4-gram.
84 - result = ngram_proposer(
85 - min_n=4, max_n=4,
86 - k=2).propose(context_token_ids=np.array([1, 2, 3, 4, 1, 2, 3]))
87 - assert result is None
91 + token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]])
92 + result = get_ngram_proposer(min_n=4, max_n=4, k=2).propose(
93 + sampled_token_ids=[[0]],
94 + req_ids=["0"],
95 + num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
96 + token_ids_cpu=token_ids_cpu,
97 + spec_decode_unsupported_reqs=(),
98 + )
99 + assert len(result[0]) == 0
100
101 # No match for 4-gram but match for 3-gram.
90 - result = ngram_proposer(
91 - min_n=3, max_n=4,
92 - k=2).propose(context_token_ids=np.array([1, 2, 3, 4, 1, 2, 3]))
93 - assert np.array_equal(result, np.array([4, 1]))
102 + token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]])
103 + result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose(
104 + sampled_token_ids=[[0]],
105 + req_ids=["0"],
106 + num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
107 + token_ids_cpu=token_ids_cpu,
108 + spec_decode_unsupported_reqs=(),
109 + )
110 + assert np.array_equal(result, np.array([[4, 1]]))
111
112 # Match for both 4-gram and 3-gram.
113 # In this case, the proposer should return the 4-gram match.
97 - result = ngram_proposer(min_n=3, max_n=4, k=2).propose(
98 - context_token_ids=np.array([2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]))
99 - assert np.array_equal(result, np.array([1, 2])) # Not [5, 1]
114 + token_ids_cpu = np.array([[2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]])
115 + result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose(
116 + sampled_token_ids=[[0]],
117 + req_ids=["0"],
118 + num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
119 + token_ids_cpu=token_ids_cpu,
120 + spec_decode_unsupported_reqs=(),
121 + )
122 + assert np.array_equal(result, np.array([[1, 2]])) # Not [5, 1]]
123
124 # Match for 2-gram and 3-gram, but not 4-gram.
102 - result = ngram_proposer(min_n=2, max_n=4, k=2).propose(
103 - context_token_ids=np.array([3, 4, 5, 2, 3, 4, 1, 2, 3, 4]))
104 - assert np.array_equal(result, np.array([1, 2])) # Not [5, 2]
125 + token_ids_cpu = np.array([[3, 4, 5, 2, 3, 4, 1, 2, 3, 4]])
126 + result = get_ngram_proposer(min_n=2, max_n=4, k=2).propose(
127 + sampled_token_ids=[[0]],
128 + req_ids=["0"],
129 + num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
130 + token_ids_cpu=token_ids_cpu,
131 + spec_decode_unsupported_reqs=(),
132 + )
133 + assert np.array_equal(result, np.array([[1, 2]])) # Not [5, 2]]
134
135 # Multiple 3-gram matched, but always pick the first one.
107 - result = ngram_proposer(
108 - min_n=3, max_n=3, k=2).propose(context_token_ids=np.array(
109 - [1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]))
110 - assert np.array_equal(result, np.array([100, 1]))
136 + token_ids_cpu = np.array(
137 + [[1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]])
138 + result = get_ngram_proposer(min_n=3, max_n=3, k=2).propose(
139 + sampled_token_ids=[[0]],
140 + req_ids=["0"],
141 + num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
142 + token_ids_cpu=token_ids_cpu,
143 + spec_decode_unsupported_reqs=(),
144 + )
145 + assert np.array_equal(result, np.array([[100, 1]]))
146 +
147 + # check empty input
148 + token_ids_cpu = np.array([[]])
149 + result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(
150 + sampled_token_ids=[[0]],
151 + req_ids=["0"],
152 + num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
153 + token_ids_cpu=token_ids_cpu,
154 + spec_decode_unsupported_reqs=(),
155 + )
156 + assert len(result[0]) == 0
157 +
158 + # check multibatch input
159 + # first request has 5 tokens and a match
160 + # second request has 3 tokens and no match. Padded with -1 for max len 5
161 + token_ids_cpu = np.array([[1, 2, 3, 1, 2], [4, 5, 6, -1, -1]])
162 + result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(
163 + sampled_token_ids=[[0], [1]],
164 + req_ids=["0", "1"],
165 + num_tokens_no_spec=np.array([5, 3]),
166 + token_ids_cpu=token_ids_cpu,
167 + spec_decode_unsupported_reqs=(),
168 + )
169 + assert len(result[0]) == 2
170 + assert np.array_equal(result[0], np.array([3, 1]))
171 + assert np.array_equal(result[1], np.array([]))
172 +
173 + # test if 0 threads available: can happen if TP size > CPU count
174 + ngram_proposer = get_ngram_proposer(min_n=2, max_n=2, k=2)
175 + ngram_proposer.num_numba_thread_available = 0
176 + # set max_model_len to 2 * threshold to ensure multithread is used
177 + num_tokens_threshold = ngram_proposer.num_tokens_threshold
178 + ngram_proposer.max_model_len = 2 * num_tokens_threshold
179 + # using multibatch test
180 + middle_integer = num_tokens_threshold // 2
181 + input_1 = [_ for _ in range(num_tokens_threshold)]
182 + input_1 += [middle_integer, middle_integer + 1]
183 + input_2 = [-1] * len(input_1)
184 + input_2[:3] = [4, 5, 6]
185 + token_ids_cpu = np.array([input_1, input_2])
186 + result = ngram_proposer.propose(
187 + sampled_token_ids=[[0], [1]],
188 + req_ids=["0", "1"],
189 + num_tokens_no_spec=np.array([len(input_1), 3]),
190 + token_ids_cpu=token_ids_cpu,
191 + spec_decode_unsupported_reqs=(),
192 + )
193 + assert len(result[0]) == 2
194 + assert np.array_equal(result[0],
195 + np.array([middle_integer + 2, middle_integer + 3]))
196 + assert np.array_equal(result[1], np.array([]))
vllm/v1/sample/rejection_sampler.py
+1 -1
@@ -17,7 +17,7 @@ PLACEHOLDER_TOKEN_ID: tl.constexpr = -1
17 GREEDY_TEMPERATURE: tl.constexpr = -1
18 # Maximum number of speculative draft tokens allowed per request in a single
19 # step. This value is chosen to be large enough to handle typical use cases.
20 -MAX_SPEC_LEN = 32
20 +MAX_SPEC_LEN = 128
21
22
23 class RejectionSampler(nn.Module):
vllm/v1/spec_decode/ngram_proposer.py
+157 -38
@@ -1,9 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional
3 +import os
4
5 import numpy as np
6 -from numba import jit
6 +from numba import get_num_threads, jit, njit, prange, set_num_threads
7
8 from vllm.config import VllmConfig
9
@@ -26,55 +26,174 @@ class NgramProposer:
26 # Maximum length of the model.
27 self.max_model_len = vllm_config.model_config.max_model_len
28
29 + # Pre-allocate buffers for numba batch propose.
30 + max_num_seqs = vllm_config.scheduler_config.max_num_seqs
31 + self.valid_ngram_draft = np.zeros((max_num_seqs, self.k),
32 + dtype=np.int32)
33 + self.valid_ngram_num_drafts = np.zeros((max_num_seqs), dtype=np.int32)
34 +
35 + # Threshold of total number of tokens in the batch to enable
36 + # multi-threading in numba batch propose.
37 + self.num_tokens_threshold = 8192
38 + tp_size = vllm_config.parallel_config.tensor_parallel_size
39 + cpu_count = os.cpu_count()
40 + # Max number of threads for numba parallel processing.
41 + if cpu_count:
42 + # Divide by 2 to use physical cores
43 + # and not logical cores (hyper-threading).
44 + # Cap the number of threads to 8 to avoid using too many threads
45 + # since other components like frontend (incl tokenization)
46 + # and Structured Outputs also use multiple threads.
47 + # TODO(ekagra-ranjan): bump up the cap from 1 to 8
48 + # when TP parallelization for ngram is implemented.
49 + self.num_numba_thread_available = min(1, (cpu_count // 2))
50 + # Divide by tp_size to ensure each tensor parallel rank
51 + # has some threads since all ranks will run this.
52 + self.num_numba_thread_available //= tp_size
53 + else:
54 + self.num_numba_thread_available = 1
55 +
56 # Trigger Numba JIT compilation for N-gram proposer.
57 # This usually takes less than 1 second.
31 - self.propose(np.zeros(1024, dtype=np.int32))
58 + self.propose([[]] * 1024, [""] * 1024, np.zeros(1024, dtype=np.int32),
59 + np.zeros((1024, self.max_model_len), dtype=np.int32),
60 + set())
61
33 - def propose(
62 + def batch_propose(
63 self,
35 - context_token_ids: np.ndarray,
36 - ) -> Optional[np.ndarray]:
37 - """Proposes the next sequence of tokens based on n-gram pattern
38 - matching in the context. The function finds matches of the last n
39 - tokens in the previous context, and returns k tokens that followed
40 - that match.
64 + num_requests: int,
65 + valid_ngram_requests: list,
66 + num_tokens_no_spec: np.ndarray,
67 + token_ids_cpu: np.ndarray,
68 + ) -> list[list[int]]:
69 + """Batch version of ngram proposer using numba for acceleration.
70
71 Args:
43 - context_token_ids: Numpy array of token IDs representing the
44 - context sequence.
72 + valid_ngram_requests:
73 + Set of indices of requests that need ngram proposals.
74 + num_tokens_no_spec:
75 + Numpy array of shape (batch_size,) representing the number
76 + of tokens without speculative tokens for each request.
77 + token_ids_cpu:
78 + Numpy array of shape (batch_size, max_model_len)
79 + representing the token IDs for each request.
80
81 Returns:
47 - np.ndarray: The sequence of tokens that followed
48 - the matched n-gram in the context.
49 - None: If no matching n-gram pattern is found.
50 -
51 - Example:
52 - If context_token_ids = [1,2,3,4,2,3], min_n = 2, max_n = 3, and
53 - k = 4:
54 - - The last 3 (= max_n) tokens [4,2,3] cannot find a match.
55 - - The last 2 tokens [2,3] will be matched against the previous
56 - 4 tokens [1,2,3,4].
57 - - Finding a match of [2,3] would return the tokens that
58 - followed that pattern. Here we will return [4,2,3] because
59 - we only have three tokens after the match.
82 + list[list[int]]:
83 + A list where each element is a list of proposed
84 + token IDs for the corresponding request.
85 """
61 - # TODO(woosuk): Optimize this.
62 - return _find_longest_matched_ngram_and_propose_tokens(
63 - origin_tokens=context_token_ids,
64 - min_ngram=self.min_n,
65 - max_ngram=self.max_n,
66 - max_model_len=self.max_model_len,
67 - k=self.k)
86 + draft_token_ids: list[list[int]] = []
87 +
88 + # Only run batch propose if there are requests needing ngram proposals.
89 + # avoid calling numba function with empty list which causes error
90 + # ValueError: cannot compute fingerprint of empty list
91 + if num_ngram_requests := len(valid_ngram_requests):
92 + original_num_numba_threads = get_num_threads()
93 + # Ensure we use at least one thread.
94 + # If total tokens is small, using multiple threads
95 + # may slow down due to overhead.
96 + total_tokens = np.sum(num_tokens_no_spec)
97 + if total_tokens >= self.num_tokens_threshold:
98 + final_num_threads = max(
99 + 1, min(self.num_numba_thread_available,
100 + num_ngram_requests))
101 + set_num_threads(final_num_threads)
102 + else:
103 + set_num_threads(1)
104 +
105 + batch_propose_numba(valid_ngram_requests, num_tokens_no_spec,
106 + token_ids_cpu, self.min_n, self.max_n,
107 + self.max_model_len, self.k,
108 + self.valid_ngram_draft,
109 + self.valid_ngram_num_drafts)
110 +
111 + # Restore original number of threads.
112 + set_num_threads(original_num_numba_threads)
113 +
114 + for i in range(num_requests):
115 + if i in valid_ngram_requests and \
116 + self.valid_ngram_num_drafts[i] > 0:
117 + draft_token_ids.append(self.valid_ngram_draft[
118 + i, :self.valid_ngram_num_drafts[i]].tolist())
119 + else:
120 + draft_token_ids.append([])
121 +
122 + return draft_token_ids
123 +
124 + def propose(
125 + self,
126 + sampled_token_ids: list[list[int]],
127 + req_ids: list[str],
128 + num_tokens_no_spec: np.ndarray,
129 + token_ids_cpu: np.ndarray,
130 + spec_decode_unsupported_reqs: set,
131 + ) -> list[list[int]]:
132 +
133 + # find which requests need ngram proposals
134 + valid_ngram_requests = []
135 + for i, sampled_ids in enumerate(sampled_token_ids):
136 + num_sampled_ids = len(sampled_ids)
137 + if not num_sampled_ids:
138 + # Skip speculative decoding.
139 + continue
140 +
141 + # Skip requests that require sampling parameters that are not
142 + # supported with speculative decoding.
143 + req_id = req_ids[i]
144 + if req_id in spec_decode_unsupported_reqs:
145 + continue
146 +
147 + num_tokens = num_tokens_no_spec[i]
148 + if num_tokens >= self.max_model_len:
149 + # Skip requests that have already reached the max model length.
150 + continue
151 +
152 + valid_ngram_requests.append(i)
153 +
154 + draft_token_ids = self.batch_propose(
155 + len(sampled_token_ids),
156 + valid_ngram_requests,
157 + num_tokens_no_spec,
158 + token_ids_cpu,
159 + )
160 +
161 + return draft_token_ids
162
163 def load_model(self, *args, **kwargs):
164 # No model to load.
165 pass
166
167
168 +@njit(parallel=True)
169 +def batch_propose_numba(valid_ngram_requests: list,
170 + num_tokens_no_spec: np.ndarray,
171 + token_ids_cpu: np.ndarray, min_n: int, max_n: int,
172 + max_model_len: int, k: int,
173 + valid_ngram_draft: np.ndarray,
174 + valid_ngram_num_drafts: np.ndarray):
175 + for i in prange(len(valid_ngram_requests)):
176 + idx = valid_ngram_requests[i]
177 + num_tokens = num_tokens_no_spec[idx]
178 + context_token_ids = token_ids_cpu[idx, :num_tokens]
179 + drafter_output = _find_longest_matched_ngram_and_propose_tokens(
180 + origin_tokens=context_token_ids,
181 + min_ngram=min_n,
182 + max_ngram=max_n,
183 + max_model_len=max_model_len,
184 + k=k)
185 +
186 + valid_ngram_num_drafts[i] = drafter_output.shape[0]
187 + if len(drafter_output):
188 + valid_ngram_draft[i, :drafter_output.shape[0]] = drafter_output
189 +
190 +
191 @jit(nopython=True)
75 -def _find_longest_matched_ngram_and_propose_tokens(
76 - origin_tokens: np.ndarray, min_ngram: int, max_ngram: int,
77 - max_model_len: int, k: int) -> Optional[np.ndarray]:
192 +def _find_longest_matched_ngram_and_propose_tokens(origin_tokens: np.ndarray,
193 + min_ngram: int,
194 + max_ngram: int,
195 + max_model_len: int,
196 + k: int) -> np.ndarray:
197 """
198 Find the longest n-gram which matches the suffix of the given tokens
199 whose length is within [min_ngram, max_ngram] (inclusive).
@@ -84,12 +203,12 @@ def _find_longest_matched_ngram_and_propose_tokens(
203 # Do not generate draft tokens is context is shorter than minimum n-gram
204 total_token = origin_tokens.shape[0]
205 if total_token < min_ngram:
87 - return None
206 + return np.empty((0, ), dtype=origin_tokens.dtype)
207
208 # Do not generate draft tokens beyond the max model length.
209 k = min(k, max_model_len - total_token)
210 if k <= 0:
92 - return None
211 + return np.empty((0, ), dtype=origin_tokens.dtype)
212
213 # Flip tokens, and the goal become to find longest ngram
214 # on the rightmost position which matches the prefix with
@@ -146,7 +265,7 @@ def _find_longest_matched_ngram_and_propose_tokens(
265
266 if longest_ngram < min_ngram:
267 # No valid ngram is found
149 - return None
268 + return np.empty((0, ), dtype=origin_tokens.dtype)
269
270 # Flip the position back, so in origin_tokens,
271 # origin_tokens[total_token-1-position:total_token-1-position+longest_ngram]
vllm/v1/worker/gpu_model_runner.py
+5 -37
@@ -2404,8 +2404,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2404 if self.speculative_config.method == "ngram":
2405 assert isinstance(sampled_token_ids, list)
2406 assert isinstance(self.drafter, NgramProposer)
2407 - draft_token_ids = self.propose_ngram_draft_token_ids(
2408 - sampled_token_ids)
2407 + draft_token_ids = self.drafter.propose(
2408 + sampled_token_ids, self.input_batch.req_ids,
2409 + self.input_batch.num_tokens_no_spec,
2410 + self.input_batch.token_ids_cpu,
2411 + self.input_batch.spec_decode_unsupported_reqs)
2412 elif self.speculative_config.method == "medusa":
2413 assert isinstance(sampled_token_ids, list)
2414 assert isinstance(self.drafter, MedusaProposer)
@@ -2515,41 +2518,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2518 )
2519 return draft_token_ids
2520
2518 - def propose_ngram_draft_token_ids(
2519 - self,
2520 - sampled_token_ids: list[list[int]],
2521 - ) -> list[list[int]]:
2522 - # TODO(woosuk): Optimize.
2523 - req_ids = self.input_batch.req_ids
2524 - draft_token_ids: list[list[int]] = []
2525 - for i, sampled_ids in enumerate(sampled_token_ids):
2526 - num_sampled_ids = len(sampled_ids)
2527 - if not num_sampled_ids:
2528 - # Skip speculative decoding.
2529 - draft_token_ids.append([])
2530 - continue
2531 -
2532 - # Skip requests that require sampling parameters that are not
2533 - # supported with speculative decoding.
2534 - req_id = req_ids[i]
2535 - if req_id in self.input_batch.spec_decode_unsupported_reqs:
2536 - draft_token_ids.append([])
2537 - continue
2538 -
2539 - num_tokens = self.input_batch.num_tokens_no_spec[i]
2540 - if num_tokens >= self.max_model_len:
2541 - # Skip requests that have already reached the max model length.
2542 - draft_token_ids.append([])
2543 - continue
2544 -
2545 - drafter_output = self.drafter.propose(
2546 - self.input_batch.token_ids_cpu[i, :num_tokens])
2547 - if drafter_output is None or len(drafter_output) == 0:
2548 - draft_token_ids.append([])
2549 - else:
2550 - draft_token_ids.append(drafter_output.tolist())
2551 - return draft_token_ids
2552 -
2521 def update_config(self, overrides: dict[str, Any]) -> None:
2522 allowed_config_names = {"load_config", "model_config"}
2523 for config_name, config_overrides in overrides.items():