@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 80f63a3966
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 15, 23:30 UTC

[V1][Spec Decode] Ngram Spec Decode (#12193)

Signed-off-by: LiuXiaoxuanPKU <lilyliupku@gmail.com>

Lily Liu committed Feb 15, 2025 at 18:05 UTC 80f63a3966a6dbfd492ccd74004da5929cdba2bb
21 files changed +1025 -84
tests/v1/core/test_scheduler.py
+196 -6
@@ -4,10 +4,12 @@ from typing import List, Optional
4 from vllm.config import CacheConfig, ModelConfig, SchedulerConfig
5 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
6 from vllm.sampling_params import SamplingParams
7 -from vllm.v1.core.scheduler import Scheduler
7 +from vllm.v1.core.scheduler import Scheduler, SchedulerOutput
8 from vllm.v1.outputs import ModelRunnerOutput
9 from vllm.v1.request import Request, RequestStatus
10
11 +EOS_TOKEN_ID = 50256
12 +
13
14 def create_scheduler(
15 model: str = "facebook/opt-125m",
@@ -38,6 +40,7 @@ def create_scheduler(
40 return Scheduler(scheduler_config,
41 model_config,
42 cache_config,
43 + speculative_config=None,
44 lora_config=None,
45 log_stats=True)
46
@@ -46,8 +49,12 @@ def create_requests(
49 num_requests: int,
50 num_tokens: int = 10,
51 mm_positions: Optional[List[PlaceholderRange]] = None,
52 + max_tokens: int = 16,
53 + stop_token_ids: Optional[List[int]] = None,
54 ):
50 - sampling_params = SamplingParams()
55 + sampling_params = SamplingParams(ignore_eos=False,
56 + max_tokens=max_tokens,
57 + stop_token_ids=stop_token_ids)
58 requests = []
59 for i in range(num_requests):
60 if mm_positions is not None:
@@ -64,7 +71,7 @@ def create_requests(
71 multi_modal_inputs=mm_inputs,
72 multi_modal_placeholders=mm_position,
73 multi_modal_hashes=None,
67 - eos_token_id=None,
74 + eos_token_id=EOS_TOKEN_ID,
75 arrival_time=0,
76 )
77 requests.append(request)
@@ -195,7 +202,7 @@ def test_schedule_partial_requests():
202 model_runner_output = ModelRunnerOutput(
203 req_ids=[request.request_id for request in requests],
204 req_id_to_index=req_to_index,
198 - sampled_token_ids=[0] * len(requests),
205 + sampled_token_ids=[[0] for _ in range(len(requests))],
206 logprobs=None,
207 prompt_logprobs_dict={},
208 )
@@ -215,6 +222,189 @@ def test_schedule_partial_requests():
222 assert requests[2].request_id not in output.num_scheduled_tokens
223
224
225 +def test_stop_via_update_from_output():
226 + """Test stopping behavior through update_from_output"""
227 + scheduler = create_scheduler()
228 +
229 + # Test case 1: Stop on EOS token
230 + requests = create_requests(num_requests=2, max_tokens=10)
231 + for req in requests:
232 + req.num_computed_tokens = req.num_tokens
233 + scheduler.requests[req.request_id] = req
234 + scheduler.running.append(req)
235 + scheduler.scheduled_req_ids.add(req.request_id)
236 +
237 + scheduler_output = SchedulerOutput(scheduled_new_reqs=[],
238 + scheduled_cached_reqs=[],
239 + num_scheduled_tokens={
240 + requests[0].request_id: 1,
241 + requests[1].request_id: 2
242 + },
243 + total_num_scheduled_tokens=3,
244 + scheduled_encoder_inputs={},
245 + scheduled_spec_decode_tokens={
246 + requests[0].request_id: [],
247 + requests[1].request_id: [10]
248 + },
249 + num_common_prefix_blocks=0,
250 + finished_req_ids=set(),
251 + free_encoder_input_ids=[])
252 +
253 + model_output = ModelRunnerOutput(
254 + req_ids=[req.request_id for req in requests],
255 + req_id_to_index={
256 + req.request_id: i
257 + for i, req in enumerate(requests)
258 + },
259 + sampled_token_ids=[[EOS_TOKEN_ID],
260 + [10,
261 + 11]], # First request hits EOS, second continues
262 + logprobs=None,
263 + prompt_logprobs_dict={})
264 +
265 + scheduler.update_from_output(scheduler_output, model_output)
266 +
267 + # Verify first request stopped, second continues
268 + assert len(scheduler.running) == 1
269 + assert scheduler.running[0].request_id == requests[1].request_id
270 + assert requests[0].status == RequestStatus.FINISHED_STOPPED
271 + assert requests[0].request_id in scheduler.finished_req_ids
272 + assert list(requests[0].output_token_ids) == [EOS_TOKEN_ID]
273 + assert list(requests[1].output_token_ids) == [10, 11]
274 +
275 + # Test case 2: Stop on custom stop token
276 + scheduler = create_scheduler()
277 + requests = create_requests(num_requests=2,
278 + max_tokens=10,
279 + stop_token_ids=[42, 43])
280 + for req in requests:
281 + req.num_computed_tokens = req.num_tokens
282 + scheduler.requests[req.request_id] = req
283 + scheduler.running.append(req)
284 + scheduler.scheduled_req_ids.add(req.request_id)
285 +
286 + scheduler_output = SchedulerOutput(scheduled_new_reqs=[],
287 + scheduled_cached_reqs=[],
288 + num_scheduled_tokens={
289 + requests[0].request_id: 3,
290 + requests[1].request_id: 2
291 + },
292 + total_num_scheduled_tokens=5,
293 + scheduled_encoder_inputs={},
294 + scheduled_spec_decode_tokens={
295 + requests[0].request_id: [10, 42],
296 + requests[1].request_id: [13]
297 + },
298 + num_common_prefix_blocks=0,
299 + finished_req_ids=set(),
300 + free_encoder_input_ids=[])
301 +
302 + model_output = ModelRunnerOutput(
303 + req_ids=[req.request_id for req in requests],
304 + req_id_to_index={
305 + req.request_id: i
306 + for i, req in enumerate(requests)
307 + },
308 + sampled_token_ids=[[10, 42, 12],
309 + [13, 14]], # First request hits stop token
310 + logprobs=None,
311 + prompt_logprobs_dict={})
312 +
313 + scheduler.update_from_output(scheduler_output, model_output)
314 +
315 + # Verify first request stopped on custom token
316 + assert len(scheduler.running) == 1
317 + assert scheduler.running[0].request_id == requests[1].request_id
318 + assert requests[0].status == RequestStatus.FINISHED_STOPPED
319 + assert requests[0].stop_reason == 42
320 + assert requests[0].request_id in scheduler.finished_req_ids
321 + assert list(requests[0].output_token_ids) == [10, 42]
322 + assert list(requests[1].output_token_ids) == [13, 14]
323 +
324 + # Test case 3: Stop on max tokens
325 + scheduler = create_scheduler()
326 + requests = create_requests(num_requests=2, max_tokens=2)
327 + for req in requests:
328 + req.num_computed_tokens = req.num_tokens
329 + scheduler.requests[req.request_id] = req
330 + scheduler.running.append(req)
331 + scheduler.scheduled_req_ids.add(req.request_id)
332 +
333 + scheduler_output = SchedulerOutput(scheduled_new_reqs=[],
334 + scheduled_cached_reqs=[],
335 + num_scheduled_tokens={
336 + requests[0].request_id: 3,
337 + requests[1].request_id: 1
338 + },
339 + total_num_scheduled_tokens=4,
340 + scheduled_encoder_inputs={},
341 + scheduled_spec_decode_tokens={
342 + requests[0].request_id: [10, 11],
343 + requests[1].request_id: []
344 + },
345 + num_common_prefix_blocks=0,
346 + finished_req_ids=set(),
347 + free_encoder_input_ids=[])
348 +
349 + model_output = ModelRunnerOutput(
350 + req_ids=[req.request_id for req in requests],
351 + req_id_to_index={
352 + req.request_id: i
353 + for i, req in enumerate(requests)
354 + },
355 + sampled_token_ids=[[10, 11, 12],
356 + [13]], # First request exceeds max_tokens
357 + logprobs=None,
358 + prompt_logprobs_dict={})
359 +
360 + scheduler.update_from_output(scheduler_output, model_output)
361 +
362 + # Verify first request stopped due to length
363 + assert len(scheduler.running) == 1
364 + assert scheduler.running[0].request_id == requests[1].request_id
365 + assert requests[0].status == RequestStatus.FINISHED_LENGTH_CAPPED
366 + assert requests[0].request_id in scheduler.finished_req_ids
367 + assert list(requests[0].output_token_ids) == [10, 11
368 + ] # Truncated to max_tokens
369 + assert list(requests[1].output_token_ids) == [13]
370 +
371 + # Test case 4: Ignore EOS flag
372 + scheduler = create_scheduler()
373 + requests = create_requests(num_requests=1, max_tokens=10)
374 + requests[0].sampling_params.ignore_eos = True
375 + requests[0].num_computed_tokens = requests[0].num_tokens
376 + scheduler.requests[requests[0].request_id] = requests[0]
377 + scheduler.running.append(requests[0])
378 + scheduler.scheduled_req_ids.add(requests[0].request_id)
379 +
380 + scheduler_output = SchedulerOutput(
381 + scheduled_new_reqs=[],
382 + scheduled_cached_reqs=[],
383 + num_scheduled_tokens={requests[0].request_id: 3},
384 + total_num_scheduled_tokens=3,
385 + scheduled_encoder_inputs={},
386 + scheduled_spec_decode_tokens={
387 + requests[0].request_id: [EOS_TOKEN_ID, 10]
388 + },
389 + num_common_prefix_blocks=0,
390 + finished_req_ids=set(),
391 + free_encoder_input_ids=[])
392 +
393 + model_output = ModelRunnerOutput(
394 + req_ids=[requests[0].request_id],
395 + req_id_to_index={requests[0].request_id: 0},
396 + sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],
397 + logprobs=None,
398 + prompt_logprobs_dict={})
399 +
400 + scheduler.update_from_output(scheduler_output, model_output)
401 +
402 + # Verify request continues past EOS
403 + assert len(scheduler.running) == 1
404 + assert not requests[0].is_finished()
405 + assert list(requests[0].output_token_ids) == [EOS_TOKEN_ID, 10, 11]
406 +
407 +
408 def test_schedule_concurrent_batches():
409 scheduler = create_scheduler(
410 max_num_batched_tokens=1024,
@@ -243,7 +433,7 @@ def test_schedule_concurrent_batches():
433 model_runner_output = ModelRunnerOutput(
434 req_ids=[requests[0].request_id],
435 req_id_to_index={requests[0].request_id: 0},
246 - sampled_token_ids=[0],
436 + sampled_token_ids=[[0]],
437 logprobs=None,
438 prompt_logprobs_dict={},
439 )
@@ -259,7 +449,7 @@ def test_schedule_concurrent_batches():
449 model_runner_output = ModelRunnerOutput(
450 req_ids=[requests[1].request_id],
451 req_id_to_index={requests[1].request_id: 0},
262 - sampled_token_ids=[0],
452 + sampled_token_ids=[[0]],
453 logprobs=None,
454 prompt_logprobs_dict={},
455 )
tests/v1/e2e/test_ngram_spec_decode.py new
+49
@@ -0,0 +1,49 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import pytest
3 +
4 +from vllm import LLM, SamplingParams
5 +
6 +
7 +@pytest.fixture
8 +def test_prompts():
9 + return [
10 + "Can you repeat the sentence ten times, this is a sentence.",
11 + "Can you repeat the sentence ten times, this is a test.",
12 + ]
13 +
14 +
15 +@pytest.fixture
16 +def sampling_config():
17 + # Only support greedy for now
18 + return SamplingParams(temperature=0, max_tokens=30, ignore_eos=False)
19 +
20 +
21 +@pytest.fixture
22 +def model_name():
23 + return "meta-llama/Meta-Llama-3-8B-Instruct"
24 +
25 +
26 +def test_ngram_correctness(monkeypatch, test_prompts, sampling_config,
27 + model_name):
28 + '''
29 + Compare the outputs of a original LLM and a speculative LLM
30 + should be the same when using ngram speculative decoding.
31 + '''
32 + with monkeypatch.context() as m:
33 + m.setenv("VLLM_USE_V1", "1")
34 +
35 + ref_llm = LLM(model=model_name)
36 + ref_outputs = ref_llm.generate(test_prompts, sampling_config)
37 + del ref_llm
38 +
39 + spec_llm = LLM(model=model_name,
40 + speculative_model='[ngram]',
41 + ngram_prompt_lookup_max=5,
42 + ngram_prompt_lookup_min=3,
43 + num_speculative_tokens=3)
44 + spec_outputs = spec_llm.generate(test_prompts, sampling_config)
45 + for ref_output, spec_output in zip(ref_outputs, spec_outputs):
46 + assert ref_output.outputs[0].text == spec_output.outputs[0].text, \
47 + (f"ref_output: {ref_output.outputs[0].text},"
48 + f"spec_output: {spec_output.outputs[0].text}")
49 + del spec_llm
tests/v1/sample/test_rejection_sampler.py new
+173
@@ -0,0 +1,173 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import List
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.v1.sample.metadata import SamplingMetadata
8 +from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID, RejectionSampler
9 +
10 +
11 +@pytest.fixture
12 +def sampler():
13 + return RejectionSampler()
14 +
15 +
16 +def create_logits_tensor(token_ids: List[int],
17 + vocab_size: int = 100) -> torch.Tensor:
18 + """Helper function to create logits tensor that
19 + will produce desired token ids on argmax"""
20 + logits = torch.full((len(token_ids), vocab_size), -100.0).cuda()
21 + for i, token_id in enumerate(token_ids):
22 + logits[i, token_id] = 100.0
23 + return logits
24 +
25 +
26 +def create_sampling_metadata(spec_tokens: List[List[int]]) -> SamplingMetadata:
27 + batch_size = len(spec_tokens)
28 + return SamplingMetadata(
29 + temperature=0.0,
30 + all_greedy=True,
31 + all_random=False,
32 + rejection_sampling=True,
33 + spec_token_ids=spec_tokens,
34 + top_p=None,
35 + top_k=None,
36 + no_top_p=False,
37 + no_top_k=False,
38 + min_p=torch.empty(batch_size, ),
39 + no_min_p=True,
40 + generators={},
41 + max_num_logprobs=0,
42 + no_penalties=False,
43 + prompt_token_ids=None,
44 + frequency_penalties=torch.tensor([]),
45 + presence_penalties=torch.tensor([]),
46 + repetition_penalties=torch.tensor([]),
47 + output_token_ids=[],
48 + min_tokens=[],
49 + stop_token_ids=[],
50 + logit_bias=[None] * batch_size,
51 + )
52 +
53 +
54 +def test_perfect_match(sampler):
55 + """Test when output tokens perfectly match speculated tokens"""
56 + spec_tokens = [[1, 2, 3]]
57 + output_tokens = [1, 2, 3, 4] # 4 is the bonus token
58 +
59 + metadata = create_sampling_metadata(spec_tokens)
60 + logits = create_logits_tensor(output_tokens)
61 +
62 + output = sampler(logits, metadata)
63 + expected = torch.tensor([[1, 2, 3, 4]],
64 + dtype=torch.int,
65 + device=logits.device)
66 + assert torch.equal(output.sampled_token_ids, expected)
67 +
68 +
69 +def test_early_mismatch(sampler):
70 + """Test when there's an early mismatch in tokens"""
71 + spec_tokens = [[1, 2, 3]]
72 + output_tokens = [1, 5, 3, 4] # Mismatch at position 1
73 +
74 + metadata = create_sampling_metadata(spec_tokens)
75 + logits = create_logits_tensor(output_tokens)
76 +
77 + output = sampler(logits, metadata)
78 + expected = torch.tensor([[1, 5, INVALID_TOKEN_ID, INVALID_TOKEN_ID]],
79 + dtype=torch.int,
80 + device=logits.device)
81 + assert torch.equal(output.sampled_token_ids, expected)
82 +
83 +
84 +def test_multiple_sequences(sampler):
85 + """Test handling multiple sequences of speculated tokens"""
86 + spec_tokens = [[1, 2], [3]]
87 + output_tokens = [1, 2, 5, 3, 4] # Two sequences with bonus tokens 5 and 4
88 +
89 + metadata = create_sampling_metadata(spec_tokens)
90 + logits = create_logits_tensor(output_tokens)
91 +
92 + output = sampler(logits, metadata)
93 + expected = torch.tensor([[1, 2, 5], [3, 4, INVALID_TOKEN_ID]],
94 + dtype=torch.int,
95 + device=logits.device)
96 + assert torch.equal(output.sampled_token_ids, expected)
97 +
98 +
99 +def test_single_token_sequence(sampler):
100 + """Test handling sequences with single token"""
101 + spec_tokens = [[1]]
102 + output_tokens = [1, 2] # Single token with bonus token 2
103 +
104 + metadata = create_sampling_metadata(spec_tokens)
105 + logits = create_logits_tensor(output_tokens)
106 +
107 + output = sampler(logits, metadata)
108 + expected = torch.tensor([[1, 2]], dtype=torch.int, device=logits.device)
109 + assert torch.equal(output.sampled_token_ids, expected)
110 +
111 +
112 +def test_empty_sequence(sampler):
113 + """Test handling empty sequence of speculated tokens"""
114 + spec_tokens: List[List[int]] = [[]]
115 + output_tokens = [5] # Just the bonus token
116 +
117 + metadata = create_sampling_metadata(spec_tokens)
118 + logits = create_logits_tensor(output_tokens)
119 +
120 + output = sampler(logits, metadata)
121 + expected = torch.tensor([[5]], dtype=torch.int, device=logits.device)
122 + assert torch.equal(output.sampled_token_ids, expected)
123 +
124 +
125 +def test_multiple_mismatches(sampler):
126 + """Test handling multiple sequences with mismatches"""
127 + spec_tokens = [[1, 2, 3], [4, 5, 6]]
128 + output_tokens = [1, 2, 7, 6, 4, 8, 6, 9] # Mismatches in both sequences
129 +
130 + metadata = create_sampling_metadata(spec_tokens)
131 + logits = create_logits_tensor(output_tokens)
132 +
133 + output = sampler(logits, metadata)
134 + expected = torch.tensor([[1, 2, 7, INVALID_TOKEN_ID],
135 + [4, 8, INVALID_TOKEN_ID, INVALID_TOKEN_ID]],
136 + dtype=torch.int,
137 + device=logits.device)
138 + assert torch.equal(output.sampled_token_ids, expected)
139 +
140 +
141 +@pytest.mark.parametrize(
142 + "spec_tokens,output_tokens,expected",
143 + [
144 + ([[1, 2]], [1, 2, 3], [[1, 2, 3]]), # Perfect match with bonus
145 + ([[1]], [2, 3], [[2, INVALID_TOKEN_ID]]), # First mismatch
146 + ([[1, 2], [3, 4]], [1, 5, 6, 3, 4, 7], [[1, 5, INVALID_TOKEN_ID],
147 + [3, 4, 7]]), # Mixed matches
148 + ])
149 +def test_parametrized_cases(sampler, spec_tokens, output_tokens, expected):
150 + """Parametrized test for various matching scenarios"""
151 + metadata = create_sampling_metadata(spec_tokens)
152 + logits = create_logits_tensor(output_tokens)
153 +
154 + output = sampler(logits, metadata)
155 + expected_tensor = torch.tensor(expected,
156 + dtype=torch.int,
157 + device=logits.device)
158 + assert torch.equal(output.sampled_token_ids, expected_tensor)
159 +
160 +
161 +def test_logits_shape_handling(sampler):
162 + """Test handling of different logits tensor shapes"""
163 + spec_tokens = [[1, 2]]
164 + output_tokens = [1, 2, 3]
165 + vocab_size = 1000
166 +
167 + metadata = create_sampling_metadata(spec_tokens)
168 + logits = create_logits_tensor(output_tokens, vocab_size)
169 +
170 + output = sampler(logits, metadata)
171 + expected = torch.tensor([[1, 2, 3]], dtype=torch.int, device=logits.device)
172 + assert torch.equal(output.sampled_token_ids, expected)
173 + assert logits.shape[-1] == vocab_size
tests/v1/sample/test_sampler.py
+2
@@ -77,6 +77,7 @@ def _create_default_sampling_metadata(
77 temperature=torch.full((batch_size, ), 0.0),
78 all_greedy=True,
79 all_random=False,
80 + rejection_sampling=False,
81 top_p=torch.empty(batch_size, ),
82 top_k=torch.empty(batch_size, ),
83 no_top_p=True,
@@ -88,6 +89,7 @@ def _create_default_sampling_metadata(
89 prompt_token_ids=_create_prompt_tokens_tensor(prompt_token_ids,
90 vocab_size, device),
91 output_token_ids=output_token_ids,
92 + spec_token_ids=[],
93 frequency_penalties=_create_penalty_tensor(batch_size, 0.0, device),
94 presence_penalties=_create_penalty_tensor(batch_size, 0.0, device),
95 repetition_penalties=_create_penalty_tensor(batch_size, 1.0, device),
tests/v1/spec_decode/test_ngram.py new
+32
@@ -0,0 +1,32 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import pytest
3 +
4 +from vllm.v1.spec_decode.ngram_proposer import NgramProposer
5 +from vllm.v1.utils import ConstantList
6 +
7 +
8 +@pytest.fixture
9 +def proposer():
10 + return NgramProposer()
11 +
12 +
13 +def test_kmp_lps_array(proposer):
14 + assert proposer._kmp_lps_array([]) == []
15 + assert proposer._kmp_lps_array([1]) == [0]
16 + assert proposer._kmp_lps_array([1, 1, 1]) == [0, 1, 2]
17 + assert proposer._kmp_lps_array([1, 2, 3, 4]) == [0, 0, 0, 0]
18 + assert proposer._kmp_lps_array([1, 2, 1, 2, 3]) == [0, 0, 1, 2, 0]
19 +
20 +
21 +def test_find_subarray_kmp(proposer):
22 + X = ConstantList([1, 2, 3, 4, 1, 2, 3, 5, 6])
23 + assert proposer._find_subarray_kmp(X, 2, 2) is None
24 + X = ConstantList([1, 2, 3, 4, 1, 2, 3])
25 + assert proposer._find_subarray_kmp(X, 2, 3) == [4, 1, 2]
26 + assert proposer._find_subarray_kmp(X, 2, 2) == [4, 1]
27 + assert proposer._find_subarray_kmp(X, 1, 3) == [4, 1, 2]
28 + assert proposer._find_subarray_kmp(X, 1, 2) == [4, 1]
29 + X = ConstantList([1, 3, 6, 2, 3, 4, 1, 2, 3])
30 + assert proposer._find_subarray_kmp(X, 2, 3) == [4, 1, 2]
31 + # Return on the first match
32 + assert proposer._find_subarray_kmp(X, 1, 3) == [6, 2, 3]
\ No newline at end of file
tests/v1/worker/test_gpu_input_batch.py
+3 -1
@@ -92,6 +92,7 @@ def _construct_expected_sampling_metadata(
92 device=device),
93 all_greedy=False,
94 all_random=True,
95 + rejection_sampling=False,
96 top_p=torch.tensor(top_p, dtype=torch.float, device=device),
97 top_k=torch.tensor(top_k, dtype=torch.int, device=device),
98 no_top_p=all(x == 1.0 for x in top_p),
@@ -116,6 +117,7 @@ def _construct_expected_sampling_metadata(
117 dtype=torch.float,
118 device=device),
119 output_token_ids=output_token_ids,
120 + spec_token_ids=[],
121 min_tokens=min_tokens,
122 stop_token_ids=stop_token_ids,
123 no_penalties=(all(x == 0 for x in presence_penalties)
@@ -205,7 +207,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
207
208 # Generate the sampling metadata
209 sampling_metadata = input_batch.make_sampling_metadata(
208 - req_id_output_token_ids, skip_copy=False)
210 + req_id_output_token_ids, req_id_to_spec_token_ids={}, skip_copy=False)
211
212 # Create expected output.
213 expected_sampling_metadata = _construct_expected_sampling_metadata(
tests/v1/worker/test_gpu_model_runner.py
+6
@@ -66,6 +66,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
66 scheduled_cached_reqs=[],
67 num_scheduled_tokens=num_scheduled_tokens,
68 total_num_scheduled_tokens=total_num_scheduled_tokens,
69 + scheduled_spec_decode_tokens={},
70 scheduled_encoder_inputs={},
71 num_common_prefix_blocks=0,
72 finished_req_ids=set(),
@@ -109,6 +110,7 @@ def test_update_states_request_finished(model_runner):
110 scheduled_cached_reqs=[],
111 num_scheduled_tokens={},
112 total_num_scheduled_tokens=0,
113 + scheduled_spec_decode_tokens={},
114 scheduled_encoder_inputs={},
115 num_common_prefix_blocks=0,
116 finished_req_ids={req_id},
@@ -137,6 +139,7 @@ def test_update_states_request_resumed(model_runner):
139 scheduled_cached_reqs=[],
140 num_scheduled_tokens={},
141 total_num_scheduled_tokens=0,
142 + scheduled_spec_decode_tokens={},
143 scheduled_encoder_inputs={},
144 num_common_prefix_blocks=0,
145 finished_req_ids={},
@@ -160,6 +163,7 @@ def test_update_states_request_resumed(model_runner):
163 scheduled_cached_reqs=[cached_req_data],
164 num_scheduled_tokens={req_id: 1},
165 total_num_scheduled_tokens=1,
166 + scheduled_spec_decode_tokens={},
167 scheduled_encoder_inputs={},
168 num_common_prefix_blocks=0,
169 finished_req_ids=set(),
@@ -188,6 +192,7 @@ def test_update_states_no_changes(model_runner):
192 scheduled_cached_reqs=[],
193 num_scheduled_tokens={req_id: 1},
194 total_num_scheduled_tokens=1,
195 + scheduled_spec_decode_tokens={},
196 scheduled_encoder_inputs={},
197 num_common_prefix_blocks=0,
198 finished_req_ids=set(),
@@ -220,6 +225,7 @@ def test_update_states_request_unscheduled(model_runner):
225 scheduled_cached_reqs=[],
226 num_scheduled_tokens={req_ids[0]: 1},
227 total_num_scheduled_tokens=1,
228 + scheduled_spec_decode_tokens={},
229 scheduled_encoder_inputs={},
230 num_common_prefix_blocks=0,
231 finished_req_ids=set(),
vllm/platforms/cuda.py
+2 -3
@@ -124,9 +124,8 @@ class CudaPlatformBase(Platform):
124 "vllm.worker.multi_step_worker.MultiStepWorker"
125 elif vllm_config.speculative_config:
126 if envs.VLLM_USE_V1:
127 - raise NotImplementedError(
128 - "Speculative decoding is not yet supported on VLLM V1."
129 - )
127 + parallel_config.worker_cls = \
128 + "vllm.v1.worker.gpu_worker.Worker"
129 else:
130 parallel_config.worker_cls = \
131 "vllm.spec_decode.spec_decode_worker.create_spec_worker"
vllm/v1/core/kv_cache_manager.py
+21 -12
@@ -82,6 +82,11 @@ class KVCacheManager:
82 self.req_to_block_hashes: DefaultDict[
83 str, List[BlockHashType]] = defaultdict(list)
84
85 + # {req_id: The number of cached blocks for this given request}
86 + # This is used to track the number of cached blocks for each request.
87 + # This is only used to track the RUNNING requests, we do not track the
88 + # data for reempted ones.
89 + self.num_cached_block: Dict[str, int] = defaultdict(int)
90 self.prefix_cache_stats = PrefixCacheStats()
91
92 @property
@@ -241,23 +246,25 @@ class KVCacheManager:
246 if not self.enable_caching:
247 return new_blocks
248
244 - # NOTE(rickyx): We are assuming the `num_tokens` are actual
245 - # tokens rather than lookahead slots (e.g. for speculative decoding).
246 - # TODO(rickyx): When supporting speculative decoding, we will need to
247 - # differentiate between them so that we can know how many blocks are
248 - # full after appending the actual tokens.
249 - num_full_blocks = (num_computed_tokens + num_tokens) // self.block_size
250 - num_computed_full_blocks = num_computed_tokens // self.block_size
251 - new_full_blocks = req_blocks[num_computed_full_blocks:num_full_blocks]
249 + num_cached_blocks = self.num_cached_block[request.request_id]
250 + # Speculated tokens might be rejected in the future, so we does
251 + # not cache any speculated tokens. We only cache blocks with
252 + # generated (accepted) tokens.
253 + num_full_blocks_after_append = (num_computed_tokens + num_tokens - len(
254 + request.spec_token_ids)) // self.block_size
255 + new_full_blocks = req_blocks[
256 + num_cached_blocks:num_full_blocks_after_append]
257 +
258 if new_full_blocks:
259 self._cache_full_blocks(
260 request=request,
255 - blk_start_idx=num_computed_full_blocks,
261 + blk_start_idx=num_cached_blocks,
262 # The new full blocks are the full blocks that are not computed.
263 full_blocks=new_full_blocks,
258 - prev_block=(req_blocks[num_computed_full_blocks - 1]
259 - if num_computed_full_blocks > 0 else None))
260 -
264 + prev_block=(req_blocks[num_cached_blocks -
265 + 1] if num_cached_blocks > 0 else None))
266 + self.num_cached_block[
267 + request.request_id] = num_full_blocks_after_append
268 return new_blocks
269
270 def free(self, request: Request) -> None:
@@ -281,6 +288,8 @@ class KVCacheManager:
288 if block.ref_cnt == 0:
289 self.free_block_queue.append(block)
290
291 + self.num_cached_block.pop(request.request_id, None)
292 +
293 def reset_prefix_cache(self) -> bool:
294 """Reset prefix cache. This function may be used in RLHF
295 flows to invalid prefix caching after the weights are updated,
vllm/v1/core/scheduler.py
+66 -31
@@ -4,7 +4,8 @@ import time
4 from collections import deque
5 from typing import Deque, Dict, Iterable, List, Optional, Set, Tuple, Union
6
7 -from vllm.config import CacheConfig, LoRAConfig, ModelConfig, SchedulerConfig
7 +from vllm.config import (CacheConfig, LoRAConfig, ModelConfig, SchedulerConfig,
8 + SpeculativeConfig)
9 from vllm.logger import init_logger
10 from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,
11 compute_encoder_budget)
@@ -28,11 +29,13 @@ class Scheduler:
29 model_config: ModelConfig,
30 cache_config: CacheConfig,
31 lora_config: Optional[LoRAConfig],
32 + speculative_config: Optional[SpeculativeConfig],
33 log_stats: bool,
34 ) -> None:
35 self.scheduler_config = scheduler_config
36 self.cache_config = cache_config
37 self.lora_config = lora_config
38 + self.speculative_config = speculative_config
39 self.log_stats = log_stats
40
41 # Scheduling constraints.
@@ -96,12 +99,14 @@ class Scheduler:
99 def schedule(self) -> "SchedulerOutput":
100 # NOTE(woosuk) on the scheduling algorithm:
101 # There's no "decoding phase" nor "prefill phase" in the scheduler.
99 - # Each request just has the num_computed_tokens and num_tokens,
100 - # which is equal to len(prompt_token_ids) + len(output_token_ids).
102 + # Each request just has the num_computed_tokens and
103 + # num_tokens_with_spec. num_tokens_with_spec =
104 + # len(prompt_token_ids) + len(output_token_ids) + len(spec_token_ids).
105 # At each step, the scheduler tries to assign tokens to the requests
106 # so that each request's num_computed_tokens can catch up its
103 - # num_tokens. This is general enough to cover chunked prefills,
104 - # prefix caching, and the "jump decoding" optimization in the future.
107 + # num_tokens_with_spec. This is general enough to cover
108 + # chunked prefills, prefix caching, speculative decoding,
109 + # and the "jump decoding" optimization in the future.
110
111 scheduled_new_reqs: List[Request] = []
112 scheduled_resumed_reqs: List[Request] = []
@@ -114,7 +119,8 @@ class Scheduler:
119 # Encoder-related.
120 scheduled_encoder_inputs: Dict[str, List[int]] = {}
121 encoder_budget = self.max_num_encoder_input_tokens
117 -
122 + # Spec decode-related.
123 + scheduled_spec_decode_tokens: Dict[str, List[int]] = {}
124 scheduled_timestamp = time.monotonic()
125
126 # First, schedule the RUNNING requests.
@@ -126,7 +132,8 @@ class Scheduler:
132 req_index += 1
133 continue
134
129 - num_new_tokens = request.num_tokens - request.num_computed_tokens
135 + num_new_tokens = (request.num_tokens_with_spec -
136 + request.num_computed_tokens)
137 num_new_tokens = min(num_new_tokens, token_budget)
138 assert num_new_tokens > 0
139
@@ -189,6 +196,11 @@ class Scheduler:
196 self.encoder_cache_manager.allocate(request, i)
197 encoder_budget = new_encoder_budget
198
199 + # Speculative decode related.
200 + if request.spec_token_ids:
201 + scheduled_spec_decode_tokens[
202 + request.request_id] = request.spec_token_ids
203 +
204 # Record the LoRAs in scheduled_running_reqs
205 requested_loras: Set[int] = set()
206 if self.lora_config:
@@ -338,6 +350,7 @@ class Scheduler:
350 num_scheduled_tokens=num_scheduled_tokens,
351 total_num_scheduled_tokens=total_num_scheduled_tokens,
352 scheduled_encoder_inputs=scheduled_encoder_inputs,
353 + scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
354 num_common_prefix_blocks=num_common_prefix_blocks,
355 # finished_req_ids is an existing state in the scheduler,
356 # instead of being newly scheduled in this step.
@@ -447,11 +460,11 @@ class Scheduler:
460 scheduler_output: "SchedulerOutput",
461 model_runner_output: "ModelRunnerOutput",
462 ) -> EngineCoreOutputs:
450 - # NOTE(woosuk): This method doesn't consider speculative decoding.
463 sampled_token_ids = model_runner_output.sampled_token_ids
464 logprobs = model_runner_output.logprobs
465 prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
466 num_scheduled_tokens = scheduler_output.num_scheduled_tokens
467 +
468 new_running: List[Request] = []
469 outputs: List[EngineCoreOutput] = []
470
@@ -466,11 +479,30 @@ class Scheduler:
479 new_running.append(request)
480 continue
481
469 - request.num_computed_tokens += num_tokens_scheduled
470 - # When the request's num_computed_tokens catches up its num_tokens,
471 - # the request generates output tokens. Otherwise, we ignore the
472 - # sampler output for the request.
473 - assert request.num_computed_tokens <= request.num_tokens
482 + req_index = model_runner_output.req_id_to_index[req_id]
483 + generated_token_ids = sampled_token_ids[req_index]
484 + if req_id not in scheduler_output.scheduled_spec_decode_tokens:
485 + # When the request's num_computed_tokens catches up
486 + # its num_tokens, the request generates output tokens.
487 + # Otherwise, we ignore the sampler output for the request.
488 + request.num_computed_tokens += num_tokens_scheduled
489 + assert request.num_computed_tokens <= request.num_tokens
490 + else:
491 + # num_computed_tokens_step represents the number of tokens
492 + # processed in the current step, considering scheduled
493 + # tokens and rejections.
494 + # It is calculated as:
495 + # num_computed_tokens_step = num_scheduled_tokens -
496 + # num_tokens_rejected,
497 + # where num_tokens_rejected is given by:
498 + # len(scheduled_spec_token_ids) + 1 - len(generated_token_ids).
499 + scheduled_spec_token_ids = (
500 + scheduler_output.scheduled_spec_decode_tokens[req_id])
501 +
502 + num_computed_tokens_step = num_scheduled_tokens[req_id] - (
503 + len(scheduled_spec_token_ids) + 1 -
504 + len(generated_token_ids))
505 + request.num_computed_tokens += num_computed_tokens_step
506
507 cached_encoder_input_ids = (
508 self.encoder_cache_manager.get_cached_input_ids(request))
@@ -485,27 +517,32 @@ class Scheduler:
517 self.encoder_cache_manager.free_encoder_input(
518 request, input_id)
519
520 + if request.num_computed_tokens >= request.num_tokens:
521 + # Clear the spec tokens as the request has generated
522 + # a new token. Here, We assume all spec tokens are verified
523 + # if we perform speculative decoding for this request.
524 + # Therefore, we can clear all spec tokens after
525 + # the generation step.
526 + request.clear_spec_tokens()
527 +
528 # Get prompt logprobs for this request.
529 prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
530
531 stopped = False
532 new_logprobs = None
493 - new_token_ids = None
494 -
495 - if request.num_computed_tokens == request.num_tokens:
496 - req_index = model_runner_output.req_id_to_index[req_id]
497 - # NOTE(woosuk): Currently, we assume that each request
498 - # generates at most one token at each step.
499 - token_id = sampled_token_ids[req_index]
500 - request.append_output_token_ids(token_id)
501 - num_new_tokens = 1
502 - # TODO: Update the KV cache manager for prefix caching.
503 -
504 - # Check for stop and update request state.
505 - # This must be called before we make the EngineCoreOutput.
506 - stopped = self._check_stop(request)
507 - if stopped:
508 - self._free_request(request)
533 + new_token_ids: List[int] = []
534 +
535 + if request.num_computed_tokens >= request.num_tokens:
536 + for output_token_id in generated_token_ids:
537 + request.append_output_token_ids(output_token_id)
538 + new_token_ids.append(output_token_id)
539 +
540 + # Check for stop and update request state.
541 + # This must be called before we make the EngineCoreOutput.
542 + stopped = self._check_stop(request)
543 + if stopped:
544 + self._free_request(request)
545 + break
546
547 # Extract sample logprobs if needed.
548 if request.sampling_params.logprobs is not None:
@@ -514,8 +551,6 @@ class Scheduler:
551 # the outer lists can be of length > 1.
552 new_logprobs = logprobs.slice(req_index, req_index + 1)
553
517 - new_token_ids = request.output_token_ids[-num_new_tokens:]
518 -
554 # Transmit partial if chunked prefill & prompt logprobs is enabled
555 if new_token_ids or prompt_logprobs_tensors is not None:
556 # Add EngineCoreOutput for this Request.
vllm/v1/core/scheduler_output.py
+4
@@ -91,6 +91,10 @@ class SchedulerOutput:
91 # Total number of tokens scheduled for all requests.
92 # Equal to sum(num_scheduled_tokens.values())
93 total_num_scheduled_tokens: int
94 + # req_id -> spec_decode_tokens
95 + # If a request does not have any spec decode tokens, it will
96 + # not be included in the dictionary.
97 + scheduled_spec_decode_tokens: Dict[str, List[int]]
98 # req_id -> encoder input indices that need processing.
99 # E.g., if a request has [0, 1], it could mean the vision encoder needs
100 # to process that the request's 0-th and 1-th images in the current step.
vllm/v1/engine/core.py
+31
@@ -27,6 +27,7 @@ from vllm.v1.executor.abstract import Executor
27 from vllm.v1.outputs import ModelRunnerOutput
28 from vllm.v1.request import Request, RequestStatus
29 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
30 +from vllm.v1.spec_decode.ngram_proposer import NgramProposer
31 from vllm.version import __version__ as VLLM_VERSION
32
33 logger = init_logger(__name__)
@@ -65,6 +66,7 @@ class EngineCore:
66 model_config=vllm_config.model_config,
67 cache_config=vllm_config.cache_config,
68 lora_config=vllm_config.lora_config,
69 + speculative_config=vllm_config.speculative_config,
70 log_stats=self.log_stats,
71 )
72
@@ -84,6 +86,15 @@ class EngineCore:
86 self.batch_queue_size)
87 self.batch_queue = queue.Queue(self.batch_queue_size)
88
89 + # Setup speculative decode.
90 + # TODO: find a better way to check if we are using ngram.
91 + self.use_spec_decode = False
92 + if self.scheduler.speculative_config:
93 + assert self.scheduler.speculative_config.ngram_prompt_lookup_min \
94 + , "Only ngram spec decode is supported in V1."
95 + self.proposer = NgramProposer()
96 + self.use_spec_decode = True
97 +
98 def _initialize_kv_caches(self,
99 vllm_config: VllmConfig) -> Tuple[int, int]:
100 start = time.time()
@@ -147,6 +158,9 @@ class EngineCore:
158 return EngineCoreOutputs(
159 outputs=[], scheduler_stats=self.scheduler.make_stats())
160
161 + if self.use_spec_decode:
162 + self.propose_tokens()
163 +
164 scheduler_output = self.scheduler.schedule()
165 output = self.model_executor.execute_model(scheduler_output)
166 engine_core_outputs = self.scheduler.update_from_output(
@@ -207,6 +221,23 @@ class EngineCore:
221 def profile(self, is_start: bool = True):
222 self.model_executor.profile(is_start)
223
224 + def propose_tokens(self):
225 + assert self.scheduler.speculative_config is not None
226 + for req in self.scheduler.running:
227 + # Ignore requests that are doing chunked prefill.
228 + if req.num_computed_tokens < req.num_tokens - 1:
229 + continue
230 + # Ignore requests that already have spec tokens.
231 + if req.spec_token_ids:
232 + continue
233 + spec_tokens = self.proposer.propose(
234 + req.all_token_ids,
235 + self.scheduler.speculative_config.ngram_prompt_lookup_min,
236 + self.scheduler.speculative_config.num_speculative_tokens,
237 + )
238 + if spec_tokens:
239 + req.append_spec_token_ids(spec_tokens)
240 +
241 def reset_prefix_cache(self):
242 self.scheduler.reset_prefix_cache()
243
vllm/v1/outputs.py
+9 -3
@@ -43,7 +43,10 @@ class LogprobsTensors(NamedTuple):
43 @dataclass
44 class SamplerOutput:
45
46 - # [num_reqs]
46 + # [num_reqs, max_num_generated_tokens]
47 + # Different requests can have different number of generated tokens.
48 + # All requests are padded to max_num_generated_tokens.
49 + # INVALID_TOKEN_ID (-1 by default) is used for padding.
50 sampled_token_ids: torch.Tensor
51 logprobs_tensors: Optional[LogprobsTensors]
52
@@ -58,8 +61,11 @@ class ModelRunnerOutput:
61 # req_id -> index
62 req_id_to_index: Dict[str, int]
63
61 - # [num_reqs]
62 - sampled_token_ids: List[int]
64 + # num_reqs x num_generated_tokens
65 + # num_generated_tokens is the number of tokens
66 + # generated in the current step. It can be different for
67 + # each request due to speculative/jump decoding.
68 + sampled_token_ids: List[List[int]]
69
70 # [num_reqs, max_num_logprobs + 1]
71 # [num_reqs, max_num_logprobs + 1]
vllm/v1/request.py
+17
@@ -46,6 +46,7 @@ class Request:
46 self.num_prompt_tokens = len(self.prompt_token_ids)
47 self._output_token_ids: List[int] = []
48 self._all_token_ids: List[int] = self.prompt_token_ids.copy()
49 + self.spec_token_ids: List[int] = []
50 self.num_computed_tokens = 0
51
52 # Multi-modal related
@@ -103,10 +104,26 @@ class Request:
104 self._output_token_ids.extend(token_ids)
105 self._all_token_ids.extend(token_ids)
106
107 + def append_spec_token_ids(
108 + self,
109 + token_ids: Union[int, List[int]],
110 + ) -> None:
111 + if isinstance(token_ids, int):
112 + self.spec_token_ids.append(token_ids)
113 + else:
114 + self.spec_token_ids.extend(token_ids)
115 +
116 + def clear_spec_tokens(self) -> None:
117 + self.spec_token_ids.clear()
118 +
119 @property
120 def num_tokens(self) -> int:
121 return len(self._all_token_ids)
122
123 + @property
124 + def num_tokens_with_spec(self) -> int:
125 + return len(self._all_token_ids) + len(self.spec_token_ids)
126 +
127 @property
128 def num_output_tokens(self) -> int:
129 return len(self._output_token_ids)
vllm/v1/sample/metadata.py
+2
@@ -12,6 +12,8 @@ class SamplingMetadata:
12 temperature: torch.Tensor
13 all_greedy: bool
14 all_random: bool
15 + rejection_sampling: bool
16 + spec_token_ids: List[List[int]]
17
18 top_p: torch.Tensor
19 top_k: torch.Tensor
vllm/v1/sample/rejection_sampler.py new
+160
@@ -0,0 +1,160 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import torch
3 +import torch.nn as nn
4 +from torch.nn.utils.rnn import pad_sequence
5 +
6 +from vllm.logger import init_logger
7 +from vllm.v1.outputs import SamplerOutput
8 +from vllm.v1.sample.metadata import SamplingMetadata
9 +
10 +try:
11 + import flashinfer.sampling as fs
12 + is_flashinfer_available = True
13 +except ImportError:
14 + is_flashinfer_available = False
15 +
16 +logger = init_logger(__name__)
17 +INVALID_TOKEN_ID = -1
18 +
19 +
20 +class RejectionSampler(nn.Module):
21 +
22 + def forward(self, logits: torch.Tensor,
23 + sampling_metadata: SamplingMetadata) -> SamplerOutput:
24 + if not sampling_metadata.all_greedy:
25 + raise NotImplementedError(
26 + "Only greedy sampling is supported by rejection sampler.")
27 +
28 + if is_flashinfer_available:
29 + logger.info("User FlashInfer for rejection sampling.")
30 + return RejectionSampler.flashinfer_sample(logits,
31 + sampling_metadata)
32 + else:
33 + logger.warning(
34 + "FlashInfer is not available. Falling back to the PyTorch-"
35 + "native implementation of rejection sampling.")
36 + return RejectionSampler.greedy_sample_native(
37 + logits, sampling_metadata)
38 +
39 + @staticmethod
40 + def flashinfer_sample(
41 + logits: torch.Tensor,
42 + sampling_metadata: SamplingMetadata) -> SamplerOutput:
43 + # NOTE: The following input preparationg can be moved
44 + # to the model runner with a persistent manner for better
45 + # performance.
46 + spec_token_ids = sampling_metadata.spec_token_ids
47 + max_spec_len = max(len(s) for s in spec_token_ids)
48 + batch_size = len(spec_token_ids)
49 + draft_token_ids = torch.full((batch_size, max_spec_len),
50 + INVALID_TOKEN_ID,
51 + device="cpu",
52 + dtype=torch.long)
53 +
54 + target_token_ids = torch.full((batch_size, max_spec_len + 1),
55 + fill_value=INVALID_TOKEN_ID,
56 + device=logits.device,
57 + dtype=torch.long)
58 +
59 + # TODO: Vectorize the following loop for better performance.
60 + start_loc = 0
61 + for i in range(batch_size):
62 + num_spec_tokens = len(spec_token_ids[i])
63 + draft_token_ids[i, :num_spec_tokens] = torch.tensor(
64 + spec_token_ids[i], device="cpu", dtype=torch.long)
65 + end_loc = start_loc + num_spec_tokens + 1
66 + # Assume greedy sampling.
67 + target_token_ids[i, :num_spec_tokens + 1] = torch.argmax(
68 + logits[start_loc:end_loc], dim=-1)
69 + start_loc = end_loc
70 +
71 + vocab_size = logits.size(-1)
72 + # NOTE: CPU <-> GPU synchronization happens here.
73 + draft_token_ids = draft_token_ids.to(logits.device)
74 + draft_probs = RejectionSampler._create_greedy_token_probs(
75 + draft_token_ids, vocab_size, logits.device)
76 + target_probs = RejectionSampler._create_greedy_token_probs(
77 + target_token_ids, vocab_size, logits.device)
78 + uniform_samples = torch.zeros(batch_size,
79 + max_spec_len + 1,
80 + device=logits.device)
81 +
82 + sampled_token_ids, _, _ = fs.chain_speculative_sampling(
83 + draft_probs,
84 + draft_token_ids,
85 + uniform_samples,
86 + target_probs,
87 + )
88 + return SamplerOutput(sampled_token_ids=sampled_token_ids,
89 + logprobs_tensors=None)
90 +
91 + # TODO: The following method can be optimized for better performance.
92 + @staticmethod
93 + def greedy_sample_native(
94 + logits: torch.Tensor,
95 + sampling_metadata: SamplingMetadata) -> SamplerOutput:
96 + spec_lens = [len(x) for x in sampling_metadata.spec_token_ids]
97 + # Add 1 to include the 'bonus' token.
98 + sample_lens = [x + 1 for x in spec_lens]
99 +
100 + output_token_ids = logits.argmax(dim=-1).view(-1)
101 + output_token_ids = output_token_ids.split(sample_lens)
102 + output_token_ids = pad_sequence(output_token_ids,
103 + batch_first=True,
104 + padding_value=INVALID_TOKEN_ID)
105 +
106 + # Convert spec token IDs to a tensor, split by sample_lens, then pad.
107 + spec_token_ids = [
108 + torch.tensor(x,
109 + dtype=output_token_ids.dtype,
110 + device=output_token_ids.device)
111 + for x in sampling_metadata.spec_token_ids
112 + ]
113 + spec_token_ids = pad_sequence(spec_token_ids,
114 + batch_first=True,
115 + padding_value=INVALID_TOKEN_ID)
116 +
117 + # Produce a mask that remains 1 (True) until the first
118 + # mismatch (cumprod turns 0 after a mismatch).
119 + accept_mask = (output_token_ids[:, :-1] == spec_token_ids).cumprod(
120 + dim=1)
121 + # Identify valid positions (non-padding).
122 + valid_mask = output_token_ids != INVALID_TOKEN_ID
123 + # Generate mask with bonus token.
124 + generate_mask = torch.cat([
125 + accept_mask,
126 + torch.zeros(accept_mask.size(0), 1, device=accept_mask.device)
127 + ],
128 + dim=1).to(torch.bool) & valid_mask
129 + zeros_mask = (generate_mask == 0)
130 + first_zero_idx = zeros_mask.float().argmax(dim=1)
131 + # Figure out which rows actually contain at least one zero.
132 + rows_with_zero = zeros_mask.any(dim=1)
133 + # Use indexing to set the first zero in each of those rows to 1.
134 + generate_mask[rows_with_zero, first_zero_idx[rows_with_zero]] = 1
135 +
136 + output_token_ids[~generate_mask] = INVALID_TOKEN_ID
137 + return SamplerOutput(sampled_token_ids=output_token_ids,
138 + logprobs_tensors=None)
139 +
140 + @staticmethod
141 + def _create_greedy_token_probs(token_ids: torch.Tensor, vocab_size: int,
142 + out_device: torch.device) -> torch.Tensor:
143 + batch_size, num_tokens = token_ids.shape
144 +
145 + token_probs = torch.zeros(batch_size,
146 + num_tokens,
147 + vocab_size,
148 + dtype=torch.float,
149 + device=out_device)
150 +
151 + # Ignore INVALID_TOKEN_ID.
152 + valid_mask = (token_ids != INVALID_TOKEN_ID)
153 + valid_indices = token_ids.clone()
154 + valid_indices[~valid_mask] = 0
155 +
156 + token_probs.scatter_(dim=2,
157 + index=valid_indices.unsqueeze(-1),
158 + src=valid_mask.unsqueeze(-1).float())
159 +
160 + return token_probs
vllm/v1/sample/sampler.py
+14 -1
@@ -9,6 +9,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
9 from vllm.v1.sample.ops.penalties import (apply_all_penalties,
10 apply_min_token_penalties)
11 from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
12 +from vllm.v1.sample.rejection_sampler import RejectionSampler
13
14 _SAMPLING_EPS = 1e-5
15
@@ -18,12 +19,21 @@ class Sampler(nn.Module):
19 def __init__(self):
20 super().__init__()
21 self.topk_topp_sampler = TopKTopPSampler()
22 + self.rejection_sampler = RejectionSampler()
23
24 def forward(
25 self,
26 logits: torch.Tensor,
27 sampling_metadata: SamplingMetadata,
28 ) -> SamplerOutput:
29 + if sampling_metadata.rejection_sampling:
30 + if sampling_metadata.max_num_logprobs:
31 + raise NotImplementedError(
32 + "Rejection sampling does not support logprobs.")
33 + return self.rejection_sampler(
34 + logits,
35 + sampling_metadata,
36 + )
37
38 # NOTE(woosuk): Use the original logits (before any penalties or
39 # temperature scaling) for the top-k logprobs.
@@ -54,7 +64,10 @@ class Sampler(nn.Module):
64
65 # These are GPU tensors.
66 sampler_output = SamplerOutput(
57 - sampled_token_ids=sampled,
67 + # The sampled tokens are expanded to 2D tensor with shape
68 + # [num_requests, 1], where each row represents one generated
69 + # token per request.
70 + sampled_token_ids=sampled.unsqueeze(-1),
71 logprobs_tensors=logprobs_tensors,
72 )
73 return sampler_output
vllm/v1/spec_decode/ngram_proposer.py new
+99
@@ -0,0 +1,99 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import List, Optional
3 +
4 +from vllm.v1.utils import ConstantList
5 +
6 +
7 +class NgramProposer:
8 +
9 + def __init__(self):
10 + pass
11 +
12 + def propose(self, context_token_ids: ConstantList[int], n: int,
13 + k: int) -> Optional[List[int]]:
14 + """Proposes the next sequence of tokens based on n-gram pattern
15 + matching in the context. The function finds matches of the last n
16 + tokens in the previous context, and returns k tokens that followed
17 + that match.
18 +
19 + Args:
20 + context_token_ids: List of token IDs representing the
21 + context sequence.
22 + n: Length of the n-gram to match.
23 + k: Number of tokens follow the match. If there are less
24 + than k tokens follow the match, we will return
25 + the maximum amount of tokens until the end.
26 +
27 + Returns:
28 + List[int]: The sequence of tokens that followed
29 + the matched n-gram in the context.
30 + None: If no matching n-gram pattern is found.
31 +
32 + Example:
33 + If context_token_ids = [1,2,3,4,2,3], n = 2, and k = 4:
34 + - The last 2 tokens [2,3] will be matched against the previous
35 + 4 tokens [1,2,3,4].
36 + - Finding a match of [2,3] would return the tokens that
37 + followed that pattern. Here we will return [4,2,3] because
38 + we only have three tokens after the match.
39 + """
40 + # TODO: Use c++ to implement the _find_subarray_kmp to
41 + # improve the efficiency
42 + return self._find_subarray_kmp(context_token_ids, n, k)
43 +
44 + @staticmethod
45 + def _kmp_lps_array(pattern: List[int]) -> List[int]:
46 + """
47 + Build the lps (longest proper prefix which is also suffix)
48 + array for the pattern.
49 + """
50 + lps = [0] * len(pattern)
51 + prev_lps = 0 # length of the previous longest prefix suffix
52 + i = 1
53 +
54 + while i < len(pattern):
55 + if pattern[i] == pattern[prev_lps]:
56 + prev_lps += 1
57 + lps[i] = prev_lps
58 + i += 1
59 + else:
60 + if prev_lps != 0:
61 + prev_lps = lps[prev_lps - 1]
62 + else:
63 + lps[i] = 0
64 + i += 1
65 +
66 + return lps
67 +
68 + @staticmethod
69 + def _find_subarray_kmp(context_token_ids: ConstantList[int], n: int,
70 + k: int) -> Optional[List[int]]:
71 + context_len = len(context_token_ids)
72 + assert n > 0
73 +
74 + pattern = context_token_ids[-n:]
75 + # Precompute lps array for Y
76 + lps = NgramProposer._kmp_lps_array(pattern)
77 +
78 + i = 0
79 + j = 0
80 + # -n because the last n tokens are used as pattern
81 + while i < context_len - n:
82 + if context_token_ids[i] == pattern[j]:
83 + i += 1
84 + j += 1
85 +
86 + # If we have matched the entire Y
87 + if j == n:
88 + # Found pattern in context, gather the next K elements
89 + return context_token_ids[i:i + k]
90 + else:
91 + # Mismatch
92 + if j != 0:
93 + # Use the lps array to avoid re-checking elements
94 + j = lps[j - 1]
95 + else:
96 + i += 1
97 +
98 + # Y not found
99 + return None
vllm/v1/worker/gpu_input_batch.py
+11 -1
@@ -390,6 +390,7 @@ class InputBatch:
390 def make_sampling_metadata(
391 self,
392 req_id_output_token_ids: Dict[str, List[int]],
393 + req_id_to_spec_token_ids: Dict[str, List[int]],
394 skip_copy: bool = False,
395 ) -> SamplingMetadata:
396 if not skip_copy:
@@ -423,7 +424,8 @@ class InputBatch:
424 self.prompt_token_ids = self._make_prompt_token_ids_tensor()
425
426 output_token_ids: List[List[int]] = []
426 -
427 + spec_token_ids: List[List[int]] = []
428 + rejection_sampling = False
429 for req_id in self.req_ids[:self.num_reqs]:
430 assert req_id is not None
431 # Currently we create a tensor for output_token_ids from scratch
@@ -434,11 +436,18 @@ class InputBatch:
436 # TODO - Replace this with incremental update to output token
437 # statistics.
438 output_token_ids.append(req_id_output_token_ids[req_id])
439 + req_spec_token_ids = req_id_to_spec_token_ids.get(req_id, [])
440 + spec_token_ids.append(req_spec_token_ids)
441 + if req_spec_token_ids:
442 + # If any of the requests require speculative decoding, set the
443 + # flag to True.
444 + rejection_sampling = True
445
446 return SamplingMetadata(
447 temperature=self.temperature[:self.num_reqs],
448 all_greedy=self.all_greedy,
449 all_random=self.all_random,
450 + rejection_sampling=rejection_sampling,
451 top_p=self.top_p[:self.num_reqs],
452 top_k=self.top_k[:self.num_reqs],
453 min_p=self.min_p[:self.num_reqs],
@@ -452,6 +461,7 @@ class InputBatch:
461 presence_penalties=self.presence_penalties[:self.num_reqs],
462 repetition_penalties=self.repetition_penalties[:self.num_reqs],
463 output_token_ids=output_token_ids,
464 + spec_token_ids=spec_token_ids,
465 min_tokens=self.min_tokens[:self.num_reqs],
466 stop_token_ids=self.stop_token_ids[:self.num_reqs],
467 no_penalties=self.no_penalties,
vllm/v1/worker/gpu_model_runner.py
+127 -25
@@ -32,6 +32,7 @@ from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
32 KVCacheSpec)
33 from vllm.v1.outputs import LogprobsTensors, ModelRunnerOutput
34 from vllm.v1.sample.metadata import SamplingMetadata
35 +from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID
36 from vllm.v1.utils import bind_kv_cache
37 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
38 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
@@ -180,6 +181,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
181 self.max_model_len,
182 self.max_num_tokens),
183 dtype=np.int32)
184 + self.arange_cpu = torch.from_numpy(self.arange_np)
185 # NOTE(woosuk): These tensors are "stateless", i.e., they are literally
186 # a faster version of creating a new tensor every time. Thus, we should
187 # not make any assumptions about the values in these tensors.
@@ -368,7 +370,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
370
371 return batch_changed
372
371 - def _prepare_inputs(self, scheduler_output: "SchedulerOutput"):
373 + def _prepare_inputs(
374 + self, scheduler_output: "SchedulerOutput"
375 + ) -> Tuple[FlashAttentionMetadata, torch.Tensor]:
376 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
377 assert total_num_scheduled_tokens > 0
378 num_reqs = self.input_batch.num_reqs
@@ -382,12 +386,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
386 # TODO: The Python loop can be slow. Optimize.
387 num_scheduled_tokens_list: List[int] = []
388 max_num_scheduled_tokens = 0
385 - for req_id in self.input_batch.req_ids[:num_reqs]:
389 + all_spec_token_ids: List[int] = []
390 + num_spec_tokens_list: List[int] = []
391 + for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):
392 assert req_id is not None
393 num_tokens = scheduler_output.num_scheduled_tokens[req_id]
394 num_scheduled_tokens_list.append(num_tokens)
395 max_num_scheduled_tokens = max(max_num_scheduled_tokens,
396 num_tokens)
397 + spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
398 + req_id, [])
399 + all_spec_token_ids.extend(spec_token_ids)
400 + num_spec_tokens_list.append(len(spec_token_ids))
401 +
402 num_scheduled_tokens: np.ndarray = np.array(num_scheduled_tokens_list,
403 dtype=np.int32)
404 assert max_num_scheduled_tokens > 0
@@ -426,6 +437,79 @@ class GPUModelRunner(LoRAModelRunnerMixin):
437 # where M is the max_model_len.
438 token_indices = (positions_np +
439 req_indices * self.input_batch.token_ids_cpu.shape[1])
440 +
441 + use_spec_decode = len(all_spec_token_ids) > 0
442 + if use_spec_decode:
443 +
444 + # 1. Write spec_token_ids to input batch.
445 + # Step 1. Get req indices that perform spec decode and repeat
446 + # the req indices by the number of spec tokens. Note
447 + # for requests that don't perform spec decode, the
448 + # number of spec tokens is 0 and the req index is
449 + # repeated 0 times.
450 + # E.g., num_spec_tokens_list: [3, 0, 2, 0, 1]
451 + # spec_req_indices: [0, 0, 0, 2, 2, 4]
452 + spec_req_indices = np.repeat(self.arange_np[:num_reqs],
453 + num_spec_tokens_list)
454 + # spec_offsets: offsets within each spec token list.
455 + # E.g., [1, 2, 3, 1, 2, 1], TODO: avoid the for loop here
456 + spec_offsets = np.concatenate(
457 + [self.arange_np[1:val + 1] for val in num_spec_tokens_list])
458 + # spec_seq_offsets: offsets within each sequence.
459 + # E.g., num_computed_tokens_cpu: [1, 4, 3, 6, 2]
460 + # after repeating: [1, 1, 1, 3, 3, 2]
461 + # spec_seq_offsets: [1, 1, 1, 3, 3, 2] + [1, 2, 3, 1, 2, 1]
462 + # = [2, 3, 4, 4, 5, 3]
463 + spec_seq_offsets = np.repeat(
464 + self.input_batch.num_computed_tokens_cpu[:num_reqs],
465 + num_spec_tokens_list) + spec_offsets
466 + # cumsums_spec_offsets: [0, 0, 0, 2M, 2M, 4M] + [2, 3, 4, 4, 5, 3]
467 + cumsums_spec_offsets = (
468 + spec_seq_offsets +
469 + spec_req_indices * self.input_batch.token_ids_cpu.shape[1])
470 + cumsums_spec_offsets = torch.from_numpy(cumsums_spec_offsets).to(
471 + torch.int64)
472 + all_spec_token_ids = torch.tensor(all_spec_token_ids,
473 + device="cpu",
474 + dtype=self.input_ids_cpu.dtype)
475 +
476 + # Step 2. Write spec token ids to input_ids_cpu.
477 + self.input_batch.token_ids_cpu_tensor.flatten().scatter_(
478 + 0, cumsums_spec_offsets, all_spec_token_ids)
479 +
480 + # 2. Get spec decode logits indices.
481 + # E.g., num_scheduled_tokens: [4, 100, 3, 100, 2]
482 + # cu_num_tokens: [4, 104, 107, 207, 209]
483 + # num_spec_tokens_list: [3, 0, 2, 0, 1]
484 + # num_sampled_tokens: [4, 1, 3, 1, 2]
485 + # spec_decode_logits_indices:
486 + # [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]
487 + num_spec_tokens_np = np.array(num_spec_tokens_list, dtype=np.int32)
488 + num_sampled_tokens = num_spec_tokens_np + 1
489 + # logits_start_loc: [0, 103, 104, 206, 207]
490 + logits_start_loc = cu_num_tokens - num_sampled_tokens
491 + # [0, 103, 104, 206, 207] ->
492 + # [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207]
493 + logits_start_loc = np.repeat(logits_start_loc, num_sampled_tokens)
494 + # The following three lines:
495 + # [4, 1, 3, 1, 2] -> [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]
496 + # Step 1. [4, 1, 3, 1, 2] -> [4, 5, 8, 9, 11]
497 + cu_num_sampled_tokens = np.cumsum(num_sampled_tokens)
498 + # Step 2. [4, 5, 8, 9, 11] -> [0, 4, 5, 8, 9]
499 + # -> [0, 0, 0, 0, 4, 5, 5, 5, 8, 9, 9]
500 + cumsums_sampled_offsets = np.repeat(
501 + cu_num_sampled_tokens - num_sampled_tokens, num_sampled_tokens)
502 + # Step 3. [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
503 + # - [0, 0, 0, 0, 4, 5, 5, 5, 8, 9, 9]
504 + # -> [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]
505 + total_num_sampled_tokens = num_sampled_tokens.sum()
506 + sampled_arange = (self.arange_np[:total_num_sampled_tokens] -
507 + cumsums_sampled_offsets)
508 +
509 + # [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207] ->
510 + # [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]
511 + spec_decode_logits_indices = logits_start_loc + sampled_arange
512 +
513 # NOTE(woosuk): We use torch.index_select instead of np.take here
514 # because torch.index_select is much faster than np.take for large
515 # tensors.
@@ -519,16 +603,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
603 suffix_kv_lens=suffix_kv_lens,
604 )
605
606 + if use_spec_decode:
607 + logits_indices = torch.from_numpy(spec_decode_logits_indices).to(
608 + self.device, non_blocking=True)
609 + else:
610 + # NOTE(woosuk): Due to chunked prefills, the batch may contain
611 + # partial requests. While we should not sample any token
612 + # from these partial requests, we do so for simplicity.
613 + # We will ignore the sampled tokens from the partial requests.
614 + # TODO: Support prompt logprobs.
615 + logits_indices = query_start_loc[1:] - 1
616 +
617 # Hot-Swap lora model
618 if self.lora_config:
619 self.set_active_loras(self.input_batch, num_scheduled_tokens)
620
526 - # NOTE(woosuk): Due to chunked prefills, the batch may contain partial
527 - # requests. While we should not sample any token from these partial
528 - # requests, we do so for simplicity. We will ignore the sampled
529 - # tokens from the partial requests.
530 - # TODO: Support prompt logprobs.
531 - logits_indices = query_start_loc[1:] - 1
621 return attn_metadata, logits_indices
622
623 def _compute_cascade_attn_prefix_len(
@@ -673,6 +762,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
762 def _prepare_sampling(
763 self,
764 batch_changed: bool,
765 + req_to_spec_token_ids: Dict[str, List[int]],
766 ) -> SamplingMetadata:
767 # Create the sampling metadata.
768 req_id_output_token_ids: Dict[str, List[int]] = \
@@ -680,7 +770,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
770 for req_id, req in self.requests.items()}
771
772 sampling_metadata = self.input_batch.make_sampling_metadata(
683 - req_id_output_token_ids, skip_copy=not batch_changed)
773 + req_id_output_token_ids, req_to_spec_token_ids, not batch_changed)
774 return sampling_metadata
775
776 def _execute_encoder(self, scheduler_output: "SchedulerOutput"):
@@ -847,7 +937,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
937 logits = self.model.compute_logits(sample_hidden_states, None)
938
939 # Sample the next token and get logprobs if needed.
850 - sampling_metadata = self._prepare_sampling(batch_changed)
940 + sampling_metadata = self._prepare_sampling(
941 + batch_changed, scheduler_output.scheduled_spec_decode_tokens)
942 sampler_output = self.model.sample(
943 logits=logits,
944 sampling_metadata=sampling_metadata,
@@ -857,18 +948,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
948 # the requests one by one. Optimize.
949 num_reqs = self.input_batch.num_reqs
950 request_seq_lens: List[Tuple[int, CachedRequestState, int]] = []
860 - for i, req_id in enumerate( # type: ignore[assignment]
861 - self.input_batch.req_ids[:num_reqs]):
951 + for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):
952 assert req_id is not None
953 req_state = self.requests[req_id]
954 seq_len = (req_state.num_computed_tokens +
955 scheduler_output.num_scheduled_tokens[req_id])
866 - assert seq_len <= req_state.num_tokens
867 - if seq_len == req_state.num_tokens:
868 - # Append the sampled token to the output token ids.
869 - self.input_batch.num_tokens[i] += 1
870 - # OPTIMIZATION: Priming the state updates for later updates.
871 - req_state.output_token_ids.append(0)
956 + if seq_len >= req_state.num_tokens:
957 request_seq_lens.append((i, req_state, seq_len))
958 else:
959 # Ignore the sampled token from the partial request.
@@ -886,7 +971,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
971
972 # NOTE: GPU -> CPU Sync happens here.
973 # Move as many CPU operations as possible before this sync point.
889 - sampled_token_ids = sampler_output.sampled_token_ids.tolist()
974 logprobs_tensors = sampler_output.logprobs_tensors
975 logprobs_lists = logprobs_tensors.tolists() \
976 if logprobs_tensors is not None else None
@@ -897,16 +981,34 @@ class GPUModelRunner(LoRAModelRunnerMixin):
981 scheduler_output,
982 )
983
900 - # Update with the actual token ids
901 - for i, req_state, seq_len in request_seq_lens:
902 - token_id = sampled_token_ids[i]
903 - self.input_batch.token_ids_cpu[i, seq_len] = token_id
904 - req_state.output_token_ids[-1] = token_id
984 + # Update batch with the valid generated tokens.
985 + sampled_token_ids = sampler_output.sampled_token_ids
986 + max_gen_len = sampled_token_ids.shape[-1]
987 + if max_gen_len == 1:
988 + valid_sampled_token_ids = sampled_token_ids.tolist()
989 + for i, req_state, seq_len in request_seq_lens:
990 + token_id = valid_sampled_token_ids[i][0]
991 + self.input_batch.token_ids_cpu[i, seq_len] = token_id
992 + req_state.output_token_ids.append(token_id)
993 + self.input_batch.num_tokens[i] += 1
994 + else:
995 + valid_mask = sampled_token_ids != INVALID_TOKEN_ID
996 + gen_lens = valid_mask.sum(dim=1).tolist()
997 + valid_sampled_token_ids = [
998 + seq.tolist()
999 + for seq in sampled_token_ids[valid_mask].split(gen_lens)
1000 + ]
1001 + self.input_batch.num_tokens[:num_reqs] += gen_lens
1002 + for i, req_state, seq_len in request_seq_lens:
1003 + target_slice = slice(seq_len - gen_lens[i] + 1, seq_len + 1)
1004 + self.input_batch.token_ids_cpu[
1005 + i, target_slice] = valid_sampled_token_ids[i]
1006 + req_state.output_token_ids.extend(valid_sampled_token_ids[i])
1007
1008 model_runner_output = ModelRunnerOutput(
1009 req_ids=req_ids,
1010 req_id_to_index=self.input_batch.req_id_to_index,
909 - sampled_token_ids=sampled_token_ids,
1011 + sampled_token_ids=valid_sampled_token_ids,
1012 logprobs=logprobs_lists,
1013 prompt_logprobs_dict=prompt_logprobs_dict,
1014 )
vllm/v1/worker/tpu_model_runner.py
+1 -1
@@ -695,7 +695,7 @@ class TPUModelRunner:
695 model_runner_output = ModelRunnerOutput(
696 req_ids=all_req_ids,
697 req_id_to_index=self.input_batch.req_id_to_index,
698 - sampled_token_ids=sampled_token_ids,
698 + sampled_token_ids=[[token_id] for token_id in sampled_token_ids],
699 logprobs=None,
700 prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type]
701 )