@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 99abb8b650
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, 18:30 UTC

[V1][Spec Decode] Optimize Rejection Sampler with Triton Kernels (#14930)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Mar 18, 2025 at 14:31 UTC 99abb8b650c66664cdc84d815b7f306f33bd9881
8 files changed +898 -431
tests/v1/sample/test_rejection_sampler.py
+166 -65
@@ -6,20 +6,23 @@ import torch
6 import torch.nn.functional as F
7
8 from vllm.v1.sample.metadata import SamplingMetadata
9 -from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID, RejectionSampler
9 +from vllm.v1.sample.rejection_sampler import (PLACEHOLDER_TOKEN_ID,
10 + RejectionSampler)
11 +from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
12
11 -DEVICE = "cpu"
13 +DEVICE = "cuda"
14
15
16 @pytest.fixture
15 -def sampler():
17 +def rejection_sampler():
18 return RejectionSampler()
19
20
19 -def create_logits_tensor(token_ids: list[list[int]],
21 +def create_logits_tensor(output_token_ids: list[list[int]],
22 vocab_size: int = 100) -> torch.Tensor:
23 """Helper function to create logits tensor that
24 will produce desired token ids on argmax"""
25 + token_ids = [tokens[:-1] for tokens in output_token_ids]
26 num_total_tokens = sum(len(tokens) for tokens in token_ids)
27 logits = torch.full((num_total_tokens, vocab_size), -100.0, device=DEVICE)
28 start_loc = 0
@@ -31,15 +34,22 @@ def create_logits_tensor(token_ids: list[list[int]],
34
35
36 def create_sampling_metadata(
34 - all_greedy: bool,
35 - generators: Optional[dict[int, Any]] = None) -> SamplingMetadata:
37 + all_greedy: bool,
38 + temperature: Optional[torch.Tensor] = None,
39 + generators: Optional[dict[int, Any]] = None,
40 +) -> SamplingMetadata:
41 """Create a v1 sampling metadata object with all_greedy set
42 to the given value. Either all greedy or all random sampling
43 is used.
44 """
45 generators = generators or {}
46 + if all_greedy:
47 + temperature = None
48 + else:
49 + assert temperature is not None
50 +
51 return SamplingMetadata(
42 - temperature=torch.tensor([]),
52 + temperature=temperature,
53 all_greedy=all_greedy,
54 all_random=not all_greedy,
55 top_p=None,
@@ -61,7 +71,7 @@ def create_sampling_metadata(
71
72
73 ########################### Tests for Greedy Sampling ###################
64 -def test_perfect_match(sampler):
74 +def test_perfect_match(rejection_sampler):
75 """Test when output tokens perfectly match speculated tokens"""
76 spec_tokens = [[1, 2, 3]]
77 output_tokens = [[1, 2, 3, 4]] # 4 is the bonus token
@@ -70,15 +80,23 @@ def test_perfect_match(sampler):
80 logits = create_logits_tensor(output_tokens)
81 bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
82 device=logits.device)
73 -
74 - output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
83 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
84 + device=logits.device)
85 +
86 + output = rejection_sampler(
87 + spec_decode_metadata,
88 + draft_probs=None,
89 + target_logits=logits,
90 + bonus_token_ids=bonus_token_tensor,
91 + sampling_metadata=metadata,
92 + )
93 expected = torch.tensor([[1, 2, 3, 4]],
94 dtype=torch.int,
95 device=logits.device)
96 assert torch.equal(output, expected)
97
98
81 -def test_early_mismatch(sampler):
99 +def test_early_mismatch(rejection_sampler):
100 """Test when there's an early mismatch in tokens"""
101 spec_tokens = [[1, 2, 3]]
102 output_tokens = [[1, 5, 3, 4]] # Mismatch at position 1
@@ -87,15 +105,25 @@ def test_early_mismatch(sampler):
105 logits = create_logits_tensor(output_tokens)
106 bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
107 device=logits.device)
90 -
91 - output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
92 - expected = torch.tensor([[1, 5, INVALID_TOKEN_ID, INVALID_TOKEN_ID]],
93 - dtype=torch.int,
94 - device=logits.device)
108 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
109 + device=logits.device)
110 +
111 + output = rejection_sampler(
112 + spec_decode_metadata,
113 + draft_probs=None,
114 + target_logits=logits,
115 + bonus_token_ids=bonus_token_tensor,
116 + sampling_metadata=metadata,
117 + )
118 + expected = torch.tensor(
119 + [[1, 5, PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID]],
120 + dtype=torch.int,
121 + device=logits.device,
122 + )
123 assert torch.equal(output, expected)
124
125
98 -def test_multiple_sequences(sampler):
126 +def test_multiple_sequences(rejection_sampler):
127 """Test handling multiple sequences of speculated tokens"""
128 spec_tokens = [[1, 2], [3]]
129 output_tokens = [[1, 2, 5], [3,
@@ -105,15 +133,23 @@ def test_multiple_sequences(sampler):
133 logits = create_logits_tensor(output_tokens)
134 bonus_token_tensor = torch.tensor(
135 [output_tokens[0][-1], output_tokens[1][-1]], device=logits.device)
108 -
109 - output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
110 - expected = torch.tensor([[1, 2, 5], [3, 4, INVALID_TOKEN_ID]],
136 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
137 + device=logits.device)
138 +
139 + output = rejection_sampler(
140 + spec_decode_metadata,
141 + draft_probs=None,
142 + target_logits=logits,
143 + bonus_token_ids=bonus_token_tensor,
144 + sampling_metadata=metadata,
145 + )
146 + expected = torch.tensor([[1, 2, 5], [3, 4, PLACEHOLDER_TOKEN_ID]],
147 dtype=torch.int,
148 device=logits.device)
149 assert torch.equal(output, expected)
150
151
116 -def test_single_token_sequence(sampler):
152 +def test_single_token_sequence(rejection_sampler):
153 """Test handling sequences with single token"""
154 spec_tokens = [[1]]
155 output_tokens = [[1, 2]] # Single token with bonus token 2
@@ -122,13 +158,21 @@ def test_single_token_sequence(sampler):
158 logits = create_logits_tensor(output_tokens)
159 bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
160 device=logits.device)
125 -
126 - output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
161 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
162 + device=logits.device)
163 +
164 + output = rejection_sampler(
165 + spec_decode_metadata,
166 + draft_probs=None,
167 + target_logits=logits,
168 + bonus_token_ids=bonus_token_tensor,
169 + sampling_metadata=metadata,
170 + )
171 expected = torch.tensor([[1, 2]], dtype=torch.int, device=logits.device)
172 assert torch.equal(output, expected)
173
174
131 -def test_empty_sequence(sampler):
175 +def test_empty_sequence(rejection_sampler):
176 """Test handling empty sequence of speculated tokens"""
177 spec_tokens: list[list[int]] = [[]]
178 output_tokens = [[5]] # Just the bonus token
@@ -137,13 +181,21 @@ def test_empty_sequence(sampler):
181 logits = create_logits_tensor(output_tokens)
182 bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
183 device=logits.device)
140 -
141 - output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
184 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
185 + device=logits.device)
186 +
187 + output = rejection_sampler(
188 + spec_decode_metadata,
189 + draft_probs=None,
190 + target_logits=logits,
191 + bonus_token_ids=bonus_token_tensor,
192 + sampling_metadata=metadata,
193 + )
194 expected = torch.tensor([[5]], dtype=torch.int, device=logits.device)
195 assert torch.equal(output, expected)
196
197
146 -def test_multiple_mismatches(sampler):
198 +def test_multiple_mismatches(rejection_sampler):
199 """Test handling multiple sequences with mismatches"""
200 spec_tokens = [[1, 2, 3], [4, 5, 6]]
201 output_tokens = [[1, 2, 7, 6], [4, 8, 6,
@@ -153,12 +205,22 @@ def test_multiple_mismatches(sampler):
205 logits = create_logits_tensor(output_tokens)
206 bonus_token_tensor = torch.tensor(
207 [output_tokens[0][-1], output_tokens[1][-1]], device=logits.device)
156 -
157 - output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
158 - expected = torch.tensor([[1, 2, 7, INVALID_TOKEN_ID],
159 - [4, 8, INVALID_TOKEN_ID, INVALID_TOKEN_ID]],
160 - dtype=torch.int,
161 - device=logits.device)
208 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
209 + device=logits.device)
210 +
211 + output = rejection_sampler(
212 + spec_decode_metadata,
213 + draft_probs=None,
214 + target_logits=logits,
215 + bonus_token_ids=bonus_token_tensor,
216 + sampling_metadata=metadata,
217 + )
218 + expected = torch.tensor(
219 + [[1, 2, 7, PLACEHOLDER_TOKEN_ID],
220 + [4, 8, PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID]],
221 + dtype=torch.int,
222 + device=logits.device,
223 + )
224 assert torch.equal(output, expected)
225
226
@@ -166,18 +228,27 @@ def test_multiple_mismatches(sampler):
228 "spec_tokens,output_tokens,expected",
229 [
230 ([[1, 2]], [[1, 2, 3]], [[1, 2, 3]]), # Perfect match with bonus
169 - ([[1]], [[2, 3]], [[2, INVALID_TOKEN_ID]]), # First mismatch
231 + ([[1]], [[2, 3]], [[2, PLACEHOLDER_TOKEN_ID]]), # First mismatch
232 ([[1, 2], [3, 4]], [[1, 5, 6], [3, 4, 7]],
171 - [[1, 5, INVALID_TOKEN_ID], [3, 4, 7]]), # Mixed matches
233 + [[1, 5, PLACEHOLDER_TOKEN_ID], [3, 4, 7]]), # Mixed matches
234 ])
173 -def test_parametrized_cases(sampler, spec_tokens, output_tokens, expected):
235 +def test_parametrized_cases(rejection_sampler, spec_tokens, output_tokens,
236 + expected):
237 """Parametrized test for various matching scenarios"""
238 metadata = create_sampling_metadata(all_greedy=True)
239 logits = create_logits_tensor(output_tokens)
240 bonus_token_tensor = torch.tensor([tokens[-1] for tokens in output_tokens],
241 device=logits.device)
179 -
180 - output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
242 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
243 + device=logits.device)
244 +
245 + output = rejection_sampler(
246 + spec_decode_metadata,
247 + draft_probs=None,
248 + target_logits=logits,
249 + bonus_token_ids=bonus_token_tensor,
250 + sampling_metadata=metadata,
251 + )
252 expected_tensor = torch.tensor(expected,
253 dtype=torch.int,
254 device=logits.device)
@@ -190,21 +261,31 @@ def test_parametrized_cases(sampler, spec_tokens, output_tokens, expected):
261 @pytest.mark.parametrize("batch_size", [1, 4, 8])
262 @pytest.mark.parametrize("frac_seeded", [0.0, 0.5])
263 @pytest.mark.parametrize("n_rep", [20])
193 -def test_deterministic_when_seeded(sampler, k: int, vocab_size: int,
194 - batch_size: int, frac_seeded: float,
195 - n_rep: int):
196 - draft_probs = torch.rand(batch_size, k, vocab_size, dtype=torch.float32)
197 - target_probs = torch.rand(batch_size * (k + 1),
198 - vocab_size,
199 - dtype=torch.float32)
264 +def test_deterministic_when_seeded(
265 + rejection_sampler,
266 + k: int,
267 + vocab_size: int,
268 + batch_size: int,
269 + frac_seeded: float,
270 + n_rep: int,
271 +):
272 + num_tokens = batch_size * k
273 + draft_probs = torch.rand(num_tokens,
274 + vocab_size,
275 + dtype=torch.float32,
276 + device=DEVICE)
277 + draft_probs = F.softmax(draft_probs, dim=-1)
278 + target_logits = torch.rand_like(draft_probs)
279 bonus_token_ids = torch.randint(low=0,
280 high=vocab_size,
281 size=(batch_size, 1),
203 - dtype=torch.int64)
282 + dtype=torch.int64,
283 + device=DEVICE)
284 draft_token_ids = torch.randint(low=0,
285 high=vocab_size,
286 size=(batch_size, k),
207 - dtype=torch.int64)
287 + dtype=torch.int64,
288 + device=DEVICE)
289
290 seeded_mask = torch.rand(batch_size, dtype=torch.float32) <= frac_seeded
291
@@ -215,10 +296,21 @@ def test_deterministic_when_seeded(sampler, k: int, vocab_size: int,
296 for i in range(batch_size) if seeded_mask[i]
297 }
298
299 + temperature = torch.ones(batch_size,
300 + dtype=torch.float32,
301 + device=DEVICE)
302 sampling_metadata = create_sampling_metadata(all_greedy=False,
303 + temperature=temperature,
304 generators=seeded_seqs)
220 - rep_result = sampler(draft_token_ids.tolist(), draft_probs,
221 - bonus_token_ids, target_probs, sampling_metadata)
305 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(
306 + draft_token_ids.tolist(), device=DEVICE)
307 + rep_result = rejection_sampler(
308 + spec_decode_metadata,
309 + draft_probs=draft_probs,
310 + target_logits=target_logits,
311 + bonus_token_ids=bonus_token_ids,
312 + sampling_metadata=sampling_metadata,
313 + )
314
315 results.append(rep_result)
316
@@ -257,10 +349,10 @@ def test_rejection_sampling_approximates_target_distribution():
349 num_reference_probs = 100
350
351 # Prepare draft, target, and reference probability distributions
260 - draft_probs, target_probs = (F.softmax(
261 - torch.rand(vocab_size, dtype=torch.float32),
262 - dim=-1,
263 - ) for _ in range(2))
352 + draft_probs = F.softmax(torch.rand(vocab_size, dtype=torch.float32),
353 + dim=-1)
354 + target_logits = torch.rand(vocab_size, dtype=torch.float32)
355 + target_probs = F.softmax(target_logits, dim=-1)
356 reference_probs = F.softmax(
357 torch.rand(num_reference_probs, vocab_size, dtype=torch.float32),
358 dim=-1,
@@ -273,7 +365,7 @@ def test_rejection_sampling_approximates_target_distribution():
365 for num_samples in sample_sizes:
366 # Sample using rejection sampling.
367 rej_sample_probs = estimate_rejection_sampling_pdf(
276 - draft_probs, target_probs, k, vocab_size, num_samples)
368 + draft_probs, target_logits, k, vocab_size, num_samples)
369 rej_sample_probs = rej_sample_probs.to(DEVICE)
370
371 # Average distance from reference probs.
@@ -313,7 +405,7 @@ def get_ratio_first_to_last(elements: list[float]) -> float:
405
406 def estimate_rejection_sampling_pdf(
407 draft_probs: torch.Tensor,
316 - target_probs: torch.Tensor,
408 + target_logits: torch.Tensor,
409 k: int,
410 vocab_size: int,
411 num_samples: int,
@@ -323,35 +415,44 @@ def estimate_rejection_sampling_pdf(
415
416 Args:
417 draft_probs: Draft probability distribution.
326 - target_probs: Target probability distribution.
418 + target_logits: Target logits.
419 num_samples: Number of samples to draw.
420
421 Returns:
422 Estimated probability distribution of the output tokens.
423 """
332 - sampler = RejectionSampler()
333 - # Repeat draft probs num_samples times.
424 + rejection_sampler = RejectionSampler()
425 + num_tokens = num_samples * k
426 + # Repeat draft probs num_samples * k times.
427 draft_probs = draft_probs.reshape(1, 1,
428 vocab_size).repeat(num_samples, k, 1)
429
337 - # Repeat target probs num_samples * (k + 1) times.
338 - target_probs = target_probs.reshape(1, 1, vocab_size).repeat(
339 - num_samples, k + 1, 1).reshape(num_samples * (k + 1), vocab_size)
430 + # Repeat target probs num_tokens times.
431 + target_logits = target_logits.reshape(1, vocab_size).repeat(num_tokens, 1)
432
433 # Randomly sample draft token ids from draft probs.
434 draft_token_ids = torch.multinomial(draft_probs[:, 0, :],
435 num_samples=k,
436 replacement=True).reshape(
437 num_samples, k)
438 + draft_probs = draft_probs.view(num_tokens, vocab_size)
439
440 # Bonus tokens not used but required.
441 bonus_token_ids = torch.zeros((1, 1), dtype=torch.int64,
442 device=DEVICE).repeat(num_samples, 1)
443
351 - sampling_metadata = create_sampling_metadata(all_greedy=False)
352 - output_token_ids = sampler(draft_token_ids.tolist(), draft_probs,
353 - bonus_token_ids, target_probs,
354 - sampling_metadata)
444 + temperature = torch.ones(num_samples, dtype=torch.float32, device=DEVICE)
445 + sampling_metadata = create_sampling_metadata(all_greedy=False,
446 + temperature=temperature)
447 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(
448 + draft_token_ids.tolist(), device=bonus_token_ids.device)
449 + output_token_ids = rejection_sampler(
450 + spec_decode_metadata,
451 + draft_probs=draft_probs,
452 + target_logits=target_logits,
453 + bonus_token_ids=bonus_token_ids,
454 + sampling_metadata=sampling_metadata,
455 + )
456 output_token_ids = output_token_ids[:, :-1].flatten()
457
458 hist = torch.histogram(output_token_ids.to(dtype=torch.float,
vllm/envs.py
-1
@@ -35,7 +35,6 @@ if TYPE_CHECKING:
35 VLLM_TRACE_FUNCTION: int = 0
36 VLLM_ATTENTION_BACKEND: Optional[str] = None
37 VLLM_USE_FLASHINFER_SAMPLER: Optional[bool] = None
38 - VLLM_USE_FLASHINFER_REJECTION_SAMPLER: bool = False
38 VLLM_FLASHINFER_FORCE_TENSOR_CORES: bool = False
39 VLLM_PP_LAYER_PARTITION: Optional[str] = None
40 VLLM_CPU_KVCACHE_SPACE: int = 0
vllm/v1/outputs.py
+1 -1
@@ -46,7 +46,7 @@ class SamplerOutput:
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.
49 + # PLACEHOLDER_TOKEN_ID (-1 by default) is used for padding.
50 sampled_token_ids: torch.Tensor
51 logprobs_tensors: Optional[LogprobsTensors]
52
vllm/v1/sample/ops/utils.py new
+30
@@ -0,0 +1,30 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Union
3 +
4 +import torch
5 +
6 +
7 +def compiled_softmax(
8 + logits: torch.Tensor,
9 + temperature: Union[float, torch.Tensor] = 1.0,
10 +) -> torch.Tensor:
11 + """Faster softmax kernel generated by torch.compile.
12 +
13 + Args:
14 + logits: [n, vocab_size]
15 + temperature: [n] or float
16 + """
17 + # NOTE(woosuk): Avoid recompilation by marking the first dim as dynamic.
18 + torch._dynamo.mark_dynamic(logits, index=0)
19 + if isinstance(temperature, torch.Tensor):
20 + torch._dynamo.mark_dynamic(temperature, index=0)
21 + return _softmax(logits, temperature)
22 +
23 +
24 +@torch.compile
25 +def _softmax(
26 + logits: torch.Tensor,
27 + temperature: Union[float, torch.Tensor],
28 +) -> torch.Tensor:
29 + logits = logits / temperature
30 + return torch.softmax(logits, dim=-1, dtype=torch.float32)
vllm/v1/sample/rejection_sampler.py
+506 -292
@@ -3,25 +3,32 @@ from typing import Optional
3
4 import torch
5 import torch.nn as nn
6 -from torch.nn.utils.rnn import pad_sequence
6 +import triton
7 +import triton.language as tl
8
9 from vllm.logger import init_logger
10 from vllm.v1.sample.metadata import SamplingMetadata
10 -from vllm.v1.spec_decode.utils import random_sample
11 +from vllm.v1.sample.ops.utils import compiled_softmax
12 +from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
13
14 logger = init_logger(__name__)
13 -INVALID_TOKEN_ID = -1
15 +
16 +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
21
22
23 class RejectionSampler(nn.Module):
24 """
18 - The implementation strictly follows the algorithm described in
25 + The implementation strictly follows the algorithm described in
26 https://arxiv.org/abs/2211.17192.
27 However, we want to clarify the terminology used in the implementation:
21 - accepted tokens: tokens that are accepted based on the relationship
28 + accepted tokens: tokens that are accepted based on the relationship
29 between the "raw" draft and target probabilities.
30 recovered tokens: tokens that are sampled based on the adjusted probability
24 - distribution, which is derived from both the draft and target
31 + distribution, which is derived from both the draft and target
32 probabilities.
33 bonus tokens:
34 If all proposed tokens are accepted, the bonus token is added to the
@@ -31,48 +38,42 @@ class RejectionSampler(nn.Module):
38 sampling process. For example, we can use top_p, top_k sampling for
39 bonus tokens, while spec decode does not support these sampling
40 strategies.
34 - output tokens:
35 - Tokens are finally generated with the rejection sampler.
41 + output tokens:
42 + Tokens are finally generated with the rejection sampler.
43 output tokens = accepted tokens + recovered tokens + bonus tokens
44 """
45
39 - def __init__(self):
40 - super().__init__()
41 -
46 def forward(
47 self,
44 - draft_token_ids: list[list[int]],
48 + metadata: SpecDecodeMetadata,
49 + # [num_tokens, vocab_size]
50 draft_probs: Optional[torch.Tensor],
46 - bonus_token_ids_tensor: torch.Tensor, # [batch_size, 1]
47 - target_probs: torch.Tensor, # [num_total_tokens, vocab_size]
51 + # [num_tokens, vocab_size]
52 + target_logits: torch.Tensor,
53 + # [batch_size, 1]
54 + bonus_token_ids: torch.Tensor,
55 sampling_metadata: SamplingMetadata,
56 ) -> torch.Tensor:
57 '''
58 Args:
52 - draft_token_ids (List[List[int]]):
53 - A 2D list of token IDs for each request in the batch.
54 - Each request might have different number of draft tokens.
55 - It may also contain empty lists for requests that have
56 - no draft tokens.
59 + metadata:
60 + Metadata for spec decoding.
61 draft_probs (Optional[torch.Tensor]):
62 Probability distribution for the draft tokens. Shape is
59 - [batch_size, max_spec_len, vocab_size]. Can be None if
60 - probabilities are not provided, which is the case for
61 - ngram spec decode.
63 + [num_tokens, vocab_size]. Can be None if probabilities are
64 + not provided, which is the case for ngram spec decode.
65 + target_logits (torch.Tensor):
66 + Target model's logits probability distribution.
67 + Shape is [num_tokens, vocab_size]. Here, probabilities from
68 + different requests are flattened into a single tensor because
69 + this is the shape of the output logits.
70 bonus_token_ids_tensor (torch.Tensor):
63 - A tensor containing bonus tokens. Shape is [batch_size, 1].
64 - Bonus tokens are added to the end of the sequence if all
65 - proposed tokens are accepted. We generate the bonus tokens
66 - outside of the rejection sampler with the default sampling
67 - strategy. It allows for more flexibility in the sampling
71 + A tensor containing bonus tokens. Shape is [batch_size, 1].
72 + Bonus tokens are added to the end of the sequence if all
73 + proposed tokens are accepted. We generate the bonus tokens
74 + outside of the rejection sampler with the default sampling
75 + strategy. It allows for more flexibility in the sampling
76 process such as top_p, top_k sampling.
69 - target_probs (torch.Tensor):
70 - Target model probability distribution.
71 - Shape is [num_total_tokens, vocab_size]. num_total_tokens
72 - is the total number of tokens from all requests. Here,
73 - probabilities from different requests are flattened into
74 - a single tensor because this is the shape of the output
75 - logits.
77 sampling_metadata (SamplingMetadata):
78 Additional metadata needed for sampling, such as temperature,
79 top-k/top-p parameters, or other relevant information.
@@ -80,268 +81,481 @@ class RejectionSampler(nn.Module):
81 output_token_ids (torch.Tensor):
82 A tensor containing the final output token IDs.
83 '''
83 -
84 - # NOTE: The following input preparationg can be moved
85 - # to the model runner with a persistent manner for better
86 - # performance.
87 - # Convert draft token IDs to a tensor, split by sample_lens, then pad.
88 - draft_token_ids = [
89 - torch.tensor(x, dtype=int, device='cpu') for x in draft_token_ids
90 - ]
91 - draft_token_ids_tensor = pad_sequence(draft_token_ids,
92 - batch_first=True,
93 - padding_value=INVALID_TOKEN_ID)
94 -
95 - # NOTE: CPU <-> GPU synchronization happens here.
96 - draft_token_ids_tensor = draft_token_ids_tensor.to(target_probs.device)
97 -
98 - # Create one-hot tensor for draft token ids.
99 - # This is used for ngram where we don't have draft_probs.
100 - if draft_probs is None and not sampling_metadata.all_greedy:
101 - vocab_size = target_probs.size(-1)
102 - draft_probs = _create_greedy_token_probs(draft_token_ids_tensor,
103 - vocab_size,
104 - target_probs.device)
105 - sample_lens = [len(x) + 1 for x in draft_token_ids]
106 - target_probs = _convert_2d_probs(target_probs, sample_lens)
107 -
108 - return self.forward_native(draft_token_ids_tensor, draft_probs,
109 - bonus_token_ids_tensor, target_probs,
110 - sampling_metadata)
111 -
112 - # TODO: The following method can be optimized for better performance.
113 - def forward_native(
114 - self,
115 - draft_token_ids_tensor: torch.Tensor,
116 - # [batch_size, max_spec_len, vocab_size]
117 - draft_probs: Optional[torch.Tensor],
118 - bonus_token_ids_tensor: torch.Tensor,
119 - # [batch_size, max_spec_len + 1, vocab_size]
120 - target_probs: torch.Tensor,
121 - sampling_metadata: SamplingMetadata,
122 - ) -> torch.Tensor:
123 - # Add 1 to include the 'bonus' token.
124 - if sampling_metadata.all_greedy:
125 - # Produce a mask that remains 1 (True) until the first
126 - # mismatch (cumprod turns 0 after a mismatch).
127 - target_token_ids_tensor = target_probs.argmax(dim=-1)
128 - accept_mask = (target_token_ids_tensor[:, :-1] ==
129 - draft_token_ids_tensor).cumprod(dim=1)
130 -
131 - # Identify valid positions (non-padding).
132 - valid_mask = target_token_ids_tensor != INVALID_TOKEN_ID
133 - # Generate mask with bonus token.
134 - generate_mask = torch.cat([
135 - accept_mask,
136 - torch.zeros(accept_mask.size(0), 1, device=accept_mask.device)
137 - ],
138 - dim=1).to(torch.bool) & valid_mask
139 - zeros_mask = (generate_mask == 0)
140 - first_zero_idx = zeros_mask.float().argmax(dim=1)
141 - # Figure out which rows actually contain at least one zero.
142 - rows_with_zero = zeros_mask.any(dim=1)
143 - # Use indexing to set the first zero in each of those rows to 1.
144 - generate_mask[rows_with_zero, first_zero_idx[rows_with_zero]] = 1
145 -
146 - output_token_ids = target_token_ids_tensor
147 - output_token_ids[~generate_mask] = INVALID_TOKEN_ID
148 - else:
149 - # Reference: https://arxiv.org/pdf/2211.17192
150 - # 1. Extract the probabilities of the draft tokens.
151 - # [batch_size, max_spec_len]
152 - batch_size = draft_token_ids_tensor.size(0)
153 - max_spec_len = draft_token_ids_tensor.size(1)
154 - invalid_idx = draft_token_ids_tensor == INVALID_TOKEN_ID
155 - draft_token_ids_tensor[invalid_idx] = 0
156 - assert draft_probs is not None
157 - draft_token_probs = draft_probs.gather(
158 - dim=-1, index=draft_token_ids_tensor.unsqueeze(-1)).squeeze(-1)
159 - target_token_probs = target_probs.gather(
160 - dim=-1, index=draft_token_ids_tensor.unsqueeze(-1)).squeeze(-1)
161 - # Force the probabilities of invalid tokens to inf
162 - # so that they are not accepted.
163 - draft_token_probs[invalid_idx] = float('inf')
164 -
165 - # 2. Generate uniform samples.
166 - # [batch_size, max_spec_len + 1]
167 - uniform_samples = _create_uniform_samples(
168 - sampling_metadata.generators, batch_size, max_spec_len,
169 - target_probs.device)
170 -
171 - # 3. Accept or reject the samples.
172 - # [batch_size, max_spec_len]
173 - # If the draft token probabilities are 0, set them to the smallest
174 - # positive normal value representable by float32.
175 - safe_draft_probs = torch.where(draft_token_probs > 0,
176 - draft_token_probs,
177 - torch.finfo(torch.float32).tiny)
178 - accepted = uniform_samples <= target_token_probs / safe_draft_probs
179 - accept_mask = accepted.cumprod(dim=1)
180 - # Set the token ids to the draft token ids if accepted, otherwise
181 - # set them to INVALID_TOKEN_ID.
182 - accepted_token_ids = (draft_token_ids_tensor * accept_mask +
183 - INVALID_TOKEN_ID * (1 - accept_mask))
184 -
185 - # 4. Adjust the distribution for the recovered tokens.
186 - # Clamp the bonus probabilities to the smallest positive normal
187 - # value representable by float32.
188 - bonus_prob = torch.clamp(target_probs[:, :-1, :] - draft_probs,
189 - min=torch.finfo(torch.float32).tiny)
190 - normalized_bonus_prob = bonus_prob / bonus_prob.sum(dim=-1,
191 - keepdim=True)
192 -
193 - # 5. Sample recovered token ids.
194 - recovered_token_ids = random_sample(
195 - normalized_bonus_prob,
196 - sampling_metadata.generators).reshape(batch_size, max_spec_len)
197 -
198 - # 6. Get the final output token ids.
199 - # output_token_ids = accepted_token_ids +
200 - # recovered_token_ids +
201 - # bonus_token_id
202 - recovered_bonus_token_ids = torch.cat(
203 - [recovered_token_ids, bonus_token_ids_tensor], dim=1)
204 - # Generate mask with bonus tokens.
205 - generate_mask = torch.cat([
206 - accept_mask,
207 - torch.zeros(batch_size, 1, device=accept_mask.device)
208 - ],
209 - dim=1).to(torch.bool)
210 - zeros_mask = (generate_mask == 0)
211 - first_zero_idx = zeros_mask.float().argmax(dim=1)
212 - output_token_ids = torch.cat([
213 - accepted_token_ids,
214 - torch.full((batch_size, 1),
215 - fill_value=INVALID_TOKEN_ID,
216 - device=accept_mask.device)
217 - ],
218 - dim=1)
219 - output_token_ids[torch.arange(batch_size),
220 - first_zero_idx] = recovered_bonus_token_ids[
221 - torch.arange(batch_size), first_zero_idx]
222 -
84 + assert metadata.max_spec_len <= MAX_SPEC_LEN
85 + # [num_tokens, vocab_size]
86 + target_probs = compute_probs(
87 + target_logits,
88 + metadata.cu_num_draft_tokens,
89 + sampling_metadata,
90 + )
91 +
92 + output_token_ids = rejection_sample(
93 + metadata.draft_token_ids,
94 + metadata.num_draft_tokens,
95 + metadata.max_spec_len,
96 + metadata.cu_num_draft_tokens,
97 + draft_probs,
98 + target_probs,
99 + bonus_token_ids,
100 + sampling_metadata,
101 + )
102 return output_token_ids
103
225 - def compute_probs(self, logits: torch.Tensor,
226 - sampling_metadata: SamplingMetadata,
227 - sample_lens: list[int]) -> torch.Tensor:
228 - """
229 - Compute probability distribution from logits based on sampling metadata.
230 -
231 - This function applies temperature scaling to the logits and converts
232 - them to probabilities using softmax. Note that division by
233 - temperature is not performed inplace to preserve the original logits
234 - tensor, which will be used by the original sampler to get bonus tokens.
235 -
236 - Args:
237 - logits: Input logits tensor to be converted to probabilities
238 - sampling_metadata: Metadata containing sampling parameters such
239 - as temperature and whether greedy sampling is used
240 - sample_lens: List of sample lengths used for repeating
241 - temperature values
242 -
243 - Returns:
244 - torch.Tensor: Probability distribution (softmax of scaled logits)
245 - if non-greedy sampling is used, otherwise returns the
246 - original logits
247 - """
104 + @staticmethod
105 + def parse_output(
106 + output_token_ids: torch.Tensor,
107 + vocab_size: int,
108 + ) -> list[list[int]]:
109 + output_token_ids_np = output_token_ids.cpu().numpy()
110 + # Create mask for valid tokens.
111 + valid_mask = ((output_token_ids_np != PLACEHOLDER_TOKEN_ID) &
112 + (output_token_ids_np < vocab_size))
113 + outputs = [
114 + row[valid_mask[i]].tolist()
115 + for i, row in enumerate(output_token_ids_np)
116 + ]
117 + return outputs
118 +
119 +
120 +def rejection_sample(
121 + # [num_tokens]
122 + draft_token_ids: torch.Tensor,
123 + # [batch_size]
124 + num_draft_tokens: list[int],
125 + max_spec_len: int,
126 + # [batch_size]
127 + cu_num_draft_tokens: torch.Tensor,
128 + # [num_tokens, vocab_size]
129 + draft_probs: Optional[torch.Tensor],
130 + # [num_tokens, vocab_size]
131 + target_probs: torch.Tensor,
132 + # [batch_size, 1]
133 + bonus_token_ids: torch.Tensor,
134 + sampling_metadata: SamplingMetadata,
135 +) -> torch.Tensor:
136 + assert draft_token_ids.ndim == 1
137 + assert draft_probs is None or draft_probs.ndim == 2
138 + assert cu_num_draft_tokens.ndim == 1
139 + assert target_probs.ndim == 2
140 +
141 + batch_size = len(num_draft_tokens)
142 + num_tokens = draft_token_ids.shape[0]
143 + vocab_size = target_probs.shape[-1]
144 + device = target_probs.device
145 + assert draft_token_ids.is_contiguous()
146 + assert draft_probs is None or draft_probs.is_contiguous()
147 + assert target_probs.is_contiguous()
148 + assert bonus_token_ids.is_contiguous()
149 + assert target_probs.shape == (num_tokens, vocab_size)
150 +
151 + # Create output buffer.
152 + output_token_ids = torch.empty(
153 + (batch_size, max_spec_len + 1),
154 + dtype=torch.int32, # Consistent with SamplerOutput.sampled_token_ids.
155 + device=device,
156 + )
157 + output_token_ids.fill_(PLACEHOLDER_TOKEN_ID)
158 +
159 + if sampling_metadata.all_greedy:
160 + is_greedy = None
161 + else:
162 + is_greedy = sampling_metadata.temperature == GREEDY_TEMPERATURE
163 + if not sampling_metadata.all_random:
164 + # Rejection sampling for greedy sampling requests.
165 + target_argmax = target_probs.argmax(dim=-1)
166 + rejection_greedy_sample_kernel[(batch_size, )](
167 + output_token_ids,
168 + cu_num_draft_tokens,
169 + draft_token_ids,
170 + target_argmax,
171 + bonus_token_ids,
172 + is_greedy,
173 + max_spec_len,
174 + num_warps=1,
175 + )
176 if sampling_metadata.all_greedy:
249 - return logits
250 - assert sampling_metadata.temperature is not None
251 - # We should optimize the following code as
252 - # it will cause CPU -> GPU synchronization.
253 - temperature = torch.repeat_interleave(
254 - sampling_metadata.temperature,
255 - torch.tensor(sample_lens,
256 - device=sampling_metadata.temperature.device))
257 - temperature = temperature.unsqueeze(dim=1)
258 - logits = logits / temperature
259 - return logits.softmax(dim=-1, dtype=torch.float32)
260 -
261 -
262 -def _create_greedy_token_probs(
263 - token_ids: torch.Tensor,
264 - vocab_size: int,
265 - out_device: torch.device,
177 + return output_token_ids
178 +
179 + # Generate uniform probabilities for rejection sampling.
180 + # [num_tokens]
181 + uniform_probs = generate_uniform_probs(
182 + num_tokens,
183 + num_draft_tokens,
184 + sampling_metadata.generators,
185 + device,
186 + )
187 +
188 + # Sample recovered tokens for each position.
189 + # [num_tokens]
190 + recovered_token_ids = sample_recovered_tokens(
191 + max_spec_len,
192 + num_draft_tokens,
193 + cu_num_draft_tokens,
194 + draft_token_ids,
195 + draft_probs,
196 + target_probs,
197 + sampling_metadata,
198 + device,
199 + )
200 +
201 + # Rejection sampling for random sampling requests.
202 + rejection_random_sample_kernel[(batch_size, )](
203 + output_token_ids,
204 + cu_num_draft_tokens,
205 + draft_token_ids,
206 + draft_probs,
207 + target_probs,
208 + bonus_token_ids,
209 + recovered_token_ids,
210 + uniform_probs,
211 + is_greedy,
212 + max_spec_len,
213 + vocab_size,
214 + IS_NGRAM=draft_probs is None,
215 + num_warps=1,
216 + )
217 + return output_token_ids
218 +
219 +
220 +def compute_probs(
221 + logits: torch.Tensor, # [num_tokens, vocab_size]
222 + cu_num_draft_tokens: torch.Tensor, # [batch_size]
223 + sampling_metadata: SamplingMetadata,
224 ) -> torch.Tensor:
267 - batch_size, num_tokens = token_ids.shape
268 -
269 - token_probs = torch.zeros(batch_size,
270 - num_tokens,
271 - vocab_size,
272 - dtype=torch.float,
273 - device=out_device)
274 -
275 - # Ignore INVALID_TOKEN_ID.
276 - valid_mask = (token_ids != INVALID_TOKEN_ID)
277 - valid_indices = token_ids.clone()
278 - valid_indices[~valid_mask] = 0
279 -
280 - token_probs.scatter_(dim=2,
281 - index=valid_indices.unsqueeze(-1),
282 - src=valid_mask.unsqueeze(-1).float())
283 -
284 - return token_probs
285 -
286 -
287 -def _convert_2d_probs(
288 - probs: torch.Tensor, # [num_total_tokens, vocab_size]
289 - sample_lens: list[int]) -> torch.Tensor:
225 + """Compute probability distribution from logits based on sampling metadata.
226 +
227 + This function applies temperature scaling to the logits and converts
228 + them to probabilities using softmax. For greedy decoding, it returns
229 + the original logits.
230 +
231 + Args:
232 + logits: Input logits tensor to be converted to probabilities.
233 + cu_num_draft_tokens: Cumulative number of draft tokens.
234 + sampling_metadata: Metadata containing sampling parameters such as
235 + temperature and whether greedy sampling is used.
236 +
237 + Returns:
238 + torch.Tensor: Probability distribution (softmax of scaled logits)
239 + if non-greedy sampling is used, otherwise returns the
240 + original logits.
241 """
291 - Converts a 2D tensor of probabilities to a 3D tensor with padding.
292 - [num_total_tokens, vocab_size] ->
293 - [batch_size, max_spec_len + 1, vocab_size]
242 + assert logits.ndim == 2
243 + assert cu_num_draft_tokens.ndim == 1
244 + if sampling_metadata.all_greedy:
245 + return logits
246 +
247 + num_tokens = logits.shape[0]
248 + batch_size = cu_num_draft_tokens.shape[0]
249 + expanded_temperature = torch.empty(
250 + (num_tokens, 1),
251 + dtype=torch.float32,
252 + device=logits.device,
253 + )
254 + expand_kernel[(batch_size, )](
255 + expanded_temperature,
256 + sampling_metadata.temperature,
257 + cu_num_draft_tokens,
258 + GREEDY_TEMPERATURE, # replace_from
259 + 1, # replace_to
260 + MAX_NUM_TOKENS=MAX_SPEC_LEN,
261 + num_warps=1,
262 + )
263 + output_prob = compiled_softmax(logits, expanded_temperature)
264 + return output_prob
265 +
266 +
267 +def generate_uniform_probs(
268 + num_tokens: int,
269 + num_draft_tokens: list[int],
270 + generators: dict[int, torch.Generator],
271 + device: torch.device,
272 +) -> torch.Tensor:
273 """
295 - cumulative_lens = torch.cumsum(torch.tensor(sample_lens,
296 - device=probs.device),
297 - dim=0)
298 - split_indices = cumulative_lens[:-1].tolist() # Exclude last index
299 -
300 - # Split into chunks without loops
301 - chunks = torch.tensor_split(probs, split_indices, dim=0)
302 -
303 - # Pad all sequences to maximum length
304 - padded_probs = pad_sequence(chunks, batch_first=True, padding_value=0.0)
305 - return padded_probs
306 -
307 -
308 -def _create_uniform_samples(seeded_seqs: dict[int, torch.Generator],
309 - batch_size: int, k: int,
310 - device: torch.device) -> torch.Tensor:
274 + Generates a batch of uniform random samples, with optional seeding
275 + if available.
276 +
277 + This method creates a tensor of shape `(num_tokens, )` filled
278 + with uniform random values in the range [0, 1). If `generators` is provided,
279 + the requests with their own seeds will use the provided `torch.Generator`
280 + for reproducibility. The samples for the other requests will be generated
281 + without a seed.
282 +
283 + Args:
284 + num_tokens : int
285 + Total number of tokens.
286 + num_draft_tokens : List[List[int]]
287 + Number of draft tokens per request.
288 + generators : Optional[Dict[int, torch.Generator]]
289 + A dictionary mapping indices in the batch to
290 + `torch.Generator` objects.
291 + device : torch.device
292 + The device on which to allocate the tensor.
293 + Returns:
294 + uniform_rand : torch.Tensor
295 + A tensor of shape `(num_tokens, )` containing uniform
296 + random values in the range [0, 1).
297 """
312 - Generates a batch of uniform random samples, with optional seeding
313 - for specific sequences.
314 -
315 - This method creates a tensor of shape `(batch_size, k)` filled
316 - with uniform random values in the range [0, 1). If `seeded_seqs`
317 - is provided, the sequences corresponding to specific indices
318 - will be generated using the provided `torch.Generator` for
319 - reproducibility. The other sequences will be generated without
320 - a seed.
321 -
322 - Args:
323 - seeded_seqs : Optional[Dict[int, torch.Generator]]
324 - A dictionary mapping indices in the batch to
325 - `torch.Generator` objects.
326 - batch_size : int
327 - The number of sequences to generate.
328 - k : int
329 - The number of random samples per sequence.
330 - device : torch.device
331 - The device on which to allocate the tensor.
332 -
333 - Returns:
334 - uniform_rand : torch.Tensor
335 - A tensor of shape `(batch_size, k)` containing uniform
336 - random values in the range [0, 1).
337 - """
338 -
339 - uniform_rand = torch.rand(batch_size,
340 - k,
341 - dtype=torch.float32,
342 - device=device)
343 - # Apply seeded generators only where needed
344 - if seeded_seqs:
345 - for idx, generator in seeded_seqs.items():
346 - uniform_rand[idx].uniform_(0, 1, generator=generator)
347 - return uniform_rand
298 + uniform_probs = torch.rand(
299 + (num_tokens, ),
300 + dtype=torch.float32,
301 + device=device,
302 + )
303 + start_idx = 0
304 + for req_idx, n in enumerate(num_draft_tokens):
305 + # Do not generate random numbers for requests with no draft tokens.
306 + # This can be important for reproducibility.
307 + if n == 0:
308 + continue
309 + end_idx = start_idx + n
310 + generator = generators.get(req_idx)
311 + if generator is not None:
312 + uniform_probs[start_idx:end_idx].uniform_(generator=generator)
313 + start_idx = end_idx
314 + return uniform_probs
315 +
316 +
317 +def sample_recovered_tokens(
318 + max_spec_len: int,
319 + num_draft_tokens: list[int],
320 + # [batch_size]
321 + cu_num_draft_tokens: torch.Tensor,
322 + # [num_tokens]
323 + draft_token_ids: torch.Tensor,
324 + # [num_tokens, vocab_size]
325 + draft_probs: Optional[torch.Tensor],
326 + # [num_tokens, vocab_size]
327 + target_probs: torch.Tensor,
328 + sampling_metadata: SamplingMetadata,
329 + device: torch.device,
330 +) -> torch.Tensor:
331 + # NOTE(woosuk): Create only one distribution for each request.
332 + batch_size = len(num_draft_tokens)
333 + vocab_size = target_probs.shape[-1]
334 + q = torch.empty(
335 + (batch_size, vocab_size),
336 + dtype=torch.float32,
337 + device=device,
338 + )
339 + q.exponential_()
340 + for i, generator in sampling_metadata.generators.items():
341 + # Do not generate random numbers for requests with no draft tokens.
342 + # This can be important for reproducibility.
343 + if num_draft_tokens[i] > 0:
344 + q[i].exponential_(generator=generator)
345 +
346 + recovered_token_ids = torch.empty_like(draft_token_ids)
347 + sample_recovered_tokens_kernel[(batch_size, max_spec_len)](
348 + recovered_token_ids,
349 + cu_num_draft_tokens,
350 + draft_token_ids,
351 + draft_probs,
352 + target_probs,
353 + q,
354 + vocab_size,
355 + triton.next_power_of_2(vocab_size),
356 + IS_NGRAM=draft_probs is None,
357 + )
358 + return recovered_token_ids
359 +
360 +
361 +# NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
362 +@triton.jit(do_not_specialize=["max_spec_len"])
363 +def rejection_greedy_sample_kernel(
364 + output_token_ids_ptr, # [batch_size, max_spec_len + 1]
365 + cu_num_draft_tokens_ptr, # [batch_size]
366 + draft_token_ids_ptr, # [num_tokens]
367 + target_argmax_ptr, # [num_tokens]
368 + bonus_token_ids_ptr, # [batch_size]
369 + is_greedy_ptr, # [batch_size] or None
370 + max_spec_len,
371 +):
372 + req_idx = tl.program_id(0)
373 + # FIXME(woosuk): Because is_greedy_ptr is not None at profiling run,
374 + # re-compilation may happen during runtime when is_greedy_ptr is None.
375 + if is_greedy_ptr is None:
376 + is_greedy = True
377 + else:
378 + is_greedy = tl.load(is_greedy_ptr + req_idx)
379 + if not is_greedy:
380 + # Early exit for non-greedy sampling requests.
381 + return
382 +
383 + if req_idx == 0:
384 + start_idx = 0
385 + else:
386 + start_idx = tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
387 + end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
388 + num_draft_tokens = end_idx - start_idx
389 +
390 + rejected = False
391 + for pos in range(num_draft_tokens):
392 + if not rejected:
393 + draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
394 + target_argmax_id = tl.load(target_argmax_ptr + start_idx + pos)
395 + tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos,
396 + target_argmax_id)
397 + if draft_token_id != target_argmax_id:
398 + # Reject.
399 + rejected = True
400 +
401 + if not rejected:
402 + # If all tokens are accepted, append the bonus token.
403 + bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx)
404 + tl.store(
405 + output_token_ids_ptr + req_idx * (max_spec_len + 1) +
406 + num_draft_tokens, bonus_token_id)
407 +
408 +
409 +# NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
410 +@triton.jit(do_not_specialize=["max_spec_len"])
411 +def rejection_random_sample_kernel(
412 + output_token_ids_ptr, # [batch_size, max_spec_len + 1]
413 + cu_num_draft_tokens_ptr, # [batch_size]
414 + draft_token_ids_ptr, # [num_tokens]
415 + draft_probs_ptr, # [num_tokens, vocab_size] or None
416 + target_probs_ptr, # [num_tokens, vocab_size]
417 + bonus_token_ids_ptr, # [batch_size]
418 + recovered_token_ids_ptr, # [num_tokens]
419 + uniform_probs_ptr, # [num_tokens]
420 + is_greedy_ptr, # [batch_size]
421 + max_spec_len,
422 + vocab_size,
423 + IS_NGRAM: tl.constexpr,
424 +):
425 + req_idx = tl.program_id(0)
426 + is_greedy = tl.load(is_greedy_ptr + req_idx)
427 + if is_greedy:
428 + # Early exit for greedy sampling requests.
429 + return
430 +
431 + if req_idx == 0:
432 + start_idx = 0
433 + else:
434 + start_idx = tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
435 + end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
436 + num_draft_tokens = end_idx - start_idx
437 +
438 + rejected = False
439 + for pos in range(num_draft_tokens):
440 + if not rejected:
441 + draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
442 + if IS_NGRAM:
443 + draft_prob = 1
444 + else:
445 + draft_prob = tl.load(draft_probs_ptr +
446 + (start_idx + pos) * vocab_size +
447 + draft_token_id)
448 + target_prob = tl.load(target_probs_ptr +
449 + (start_idx + pos) * vocab_size +
450 + draft_token_id)
451 + uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
452 + # NOTE(woosuk): While the draft probability should never be 0,
453 + # we check it to avoid NaNs. If it happens to be 0, we reject.
454 + if draft_prob > 0 and target_prob / draft_prob >= uniform_prob:
455 + # Accept.
456 + token_id = draft_token_id
457 + else:
458 + # Reject. Use recovered token.
459 + rejected = True
460 + token_id = tl.load(recovered_token_ids_ptr + start_idx + pos)
461 + tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos,
462 + token_id)
463 +
464 + if not rejected:
465 + # If all tokens are accepted, append the bonus token.
466 + bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx)
467 + tl.store(
468 + output_token_ids_ptr + req_idx * (max_spec_len + 1) +
469 + num_draft_tokens, bonus_token_id)
470 +
471 +
472 +# NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
473 +@triton.jit(do_not_specialize=["replace_from", "replace_to"])
474 +def expand_kernel(
475 + output_ptr, # [num_tokens]
476 + input_ptr, # [batch_size]
477 + cu_num_tokens_ptr, # [batch_size]
478 + replace_from,
479 + replace_to,
480 + MAX_NUM_TOKENS: tl.constexpr,
481 +):
482 + req_idx = tl.program_id(0)
483 + if req_idx == 0: # noqa: SIM108
484 + start_idx = 0
485 + else:
486 + start_idx = tl.load(cu_num_tokens_ptr + req_idx - 1)
487 + end_idx = tl.load(cu_num_tokens_ptr + req_idx)
488 + num_tokens = end_idx - start_idx
489 +
490 + src_val = tl.load(input_ptr + req_idx)
491 + src_val = tl.where(src_val == replace_from, replace_to, src_val)
492 + offset = tl.arange(0, MAX_NUM_TOKENS)
493 + tl.store(output_ptr + start_idx + offset,
494 + src_val,
495 + mask=offset < num_tokens)
496 +
497 +
498 +@triton.jit
499 +def sample_recovered_tokens_kernel(
500 + output_token_ids_ptr, # [num_tokens]
501 + cu_num_draft_tokens_ptr, # [batch_size]
502 + draft_token_ids_ptr, # [num_tokens]
503 + draft_probs_ptr, # [num_tokens, vocab_size] or None
504 + target_probs_ptr, # [num_tokens, vocab_size]
505 + q_ptr, # [batch_size, vocab_size]
506 + vocab_size,
507 + PADDED_VOCAB_SIZE: tl.constexpr,
508 + IS_NGRAM: tl.constexpr,
509 +):
510 + req_idx = tl.program_id(0)
511 + if req_idx == 0:
512 + start_idx = 0
513 + else:
514 + start_idx = tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
515 + end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
516 + num_draft_tokens = end_idx - start_idx
517 +
518 + # Early exit for out-of-range positions.
519 + pos = tl.program_id(1)
520 + if pos >= num_draft_tokens:
521 + return
522 +
523 + vocab_offset = tl.arange(0, PADDED_VOCAB_SIZE)
524 + if IS_NGRAM:
525 + draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
526 + orig_prob = tl.load(target_probs_ptr + (start_idx + pos) * vocab_size +
527 + draft_token_id)
528 + # Temporarily zero out the probability of the draft token.
529 + # This is essentially the same as target_prob - draft_prob, except that
530 + # n-gram does not have draft_prob. We regard it as 1.
531 + tl.store(
532 + target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id,
533 + 0)
534 + prob = tl.load(target_probs_ptr + (start_idx + pos) * vocab_size +
535 + vocab_offset,
536 + mask=vocab_offset < vocab_size,
537 + other=0)
538 + else:
539 + draft_prob = tl.load(draft_probs_ptr + (start_idx + pos) * vocab_size +
540 + vocab_offset,
541 + mask=vocab_offset < vocab_size,
542 + other=0)
543 + target_prob = tl.load(target_probs_ptr +
544 + (start_idx + pos) * vocab_size + vocab_offset,
545 + mask=vocab_offset < vocab_size,
546 + other=0)
547 + prob = tl.maximum(target_prob - draft_prob, 0)
548 + # NOTE(woosuk): We don't need `prob = prob / tl.sum(prob)` here because
549 + # `tl.argmax` will select the maximum value.
550 +
551 + q = tl.load(q_ptr + req_idx * vocab_size + vocab_offset,
552 + mask=vocab_offset < vocab_size,
553 + other=float("-inf"))
554 + recovered_id = tl.argmax(prob / q, axis=-1)
555 + tl.store(output_token_ids_ptr + start_idx + pos, recovered_id)
556 +
557 + if IS_NGRAM:
558 + # Restore the original probability.
559 + tl.store(
560 + target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id,
561 + orig_prob)
vllm/v1/spec_decode/metadata.py new
+61
@@ -0,0 +1,61 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from dataclasses import dataclass
3 +
4 +import numpy as np
5 +import torch
6 +
7 +
8 +@dataclass
9 +class SpecDecodeMetadata:
10 +
11 + # [num_tokens]
12 + draft_token_ids: torch.Tensor
13 + # [batch_size]
14 + num_draft_tokens: list[int]
15 + # [batch_size]
16 + cu_num_draft_tokens: torch.Tensor
17 + # [num_tokens]
18 + target_logits_indices: torch.Tensor
19 + # [batch_size]
20 + bonus_logits_indices: torch.Tensor
21 + # [num_tokens + batch_size]
22 + logits_indices: torch.Tensor
23 +
24 + def __post_init__(self):
25 + self.max_spec_len = max(self.num_draft_tokens)
26 +
27 + @classmethod
28 + def make_dummy(
29 + cls,
30 + draft_token_ids: list[list[int]],
31 + device: torch.device,
32 + ) -> "SpecDecodeMetadata":
33 + batch_size = len(draft_token_ids)
34 + num_draft_tokens = [len(ids) for ids in draft_token_ids]
35 + flattened_draft_token_ids = sum(draft_token_ids, [])
36 + num_tokens = len(flattened_draft_token_ids)
37 +
38 + draft_token_ids_tensor = torch.tensor(flattened_draft_token_ids,
39 + dtype=torch.int32,
40 + device=device)
41 + cu_num_draft_tokens = np.cumsum(num_draft_tokens, dtype=np.int32)
42 + cu_num_draft_tokens_tensor = torch.from_numpy(cu_num_draft_tokens).to(
43 + device)
44 +
45 + target_logits_indices = torch.zeros(num_tokens,
46 + dtype=torch.int32,
47 + device=device)
48 + bonus_logits_indices = torch.zeros(batch_size,
49 + dtype=torch.int32,
50 + device=device)
51 + logits_indices = torch.zeros(num_tokens + batch_size,
52 + dtype=torch.int32,
53 + device=device)
54 + return cls(
55 + draft_token_ids=draft_token_ids_tensor,
56 + num_draft_tokens=num_draft_tokens,
57 + cu_num_draft_tokens=cu_num_draft_tokens_tensor,
58 + target_logits_indices=target_logits_indices,
59 + bonus_logits_indices=bonus_logits_indices,
60 + logits_indices=logits_indices,
61 + )
vllm/v1/spec_decode/utils.py
-1
@@ -1,5 +1,4 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -from vllm.v1.sample.ops.topk_topp_sampler import random_sample # noqa
2 from vllm.v1.worker.gpu_input_batch import InputBatch
3
4
vllm/v1/worker/gpu_model_runner.py
+134 -71
@@ -34,7 +34,8 @@ from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
34 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
35 ModelRunnerOutput)
36 from vllm.v1.sample.metadata import SamplingMetadata
37 -from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID, RejectionSampler
37 +from vllm.v1.sample.rejection_sampler import RejectionSampler
38 +from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
39 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
40 from vllm.v1.spec_decode.utils import is_spec_decode_supported
41 from vllm.v1.utils import bind_kv_cache
@@ -149,7 +150,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
150 self.use_spec_decode = False
151 if self.speculative_config:
152 self.use_spec_decode = True
152 - self.rejection_sampler = RejectionSampler()
153 # TODO: find a better way to check if we are using ngram.
154 assert self.speculative_config.ngram_prompt_lookup_min, \
155 "Currently, only ngram spec decode is supported in V1."
@@ -162,6 +162,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
162 self.speculative_config.ngram_prompt_lookup_min,
163 self.speculative_config.num_speculative_tokens,
164 )
165 + self.rejection_sampler = RejectionSampler()
166
167 # Request states.
168 self.requests: dict[str, CachedRequestState] = {}
@@ -452,7 +453,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
453 def _prepare_inputs(
454 self,
455 scheduler_output: "SchedulerOutput",
455 - ) -> tuple[FlashAttentionMetadata, torch.Tensor]:
456 + ) -> tuple[FlashAttentionMetadata, torch.Tensor,
457 + Optional[SpecDecodeMetadata]]:
458 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
459 assert total_num_scheduled_tokens > 0
460 num_reqs = self.input_batch.num_reqs
@@ -577,22 +579,33 @@ class GPUModelRunner(LoRAModelRunnerMixin):
579
580 use_spec_decode = len(
581 scheduler_output.scheduled_spec_decode_tokens) > 0
580 - if use_spec_decode:
581 - logits_indices = self._calc_spec_decode_metadata(
582 - scheduler_output, cu_num_tokens)
583 - else:
582 + if not use_spec_decode:
583 # NOTE(woosuk): Due to chunked prefills, the batch may contain
584 # partial requests. While we should not sample any token
585 # from these partial requests, we do so for simplicity.
586 # We will ignore the sampled tokens from the partial requests.
587 # TODO: Support prompt logprobs.
588 logits_indices = attn_metadata.query_start_loc[1:] - 1
589 + spec_decode_metadata = None
590 + else:
591 + # Get the number of draft tokens for each request.
592 + # Iterate over the dictionary rather than all requests since not all
593 + # requests have draft tokens.
594 + num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
595 + for req_id, draft_token_ids in (
596 + scheduler_output.scheduled_spec_decode_tokens.items()):
597 + req_idx = self.input_batch.req_id_to_index[req_id]
598 + num_draft_tokens[req_idx] = len(draft_token_ids)
599 +
600 + spec_decode_metadata = self._calc_spec_decode_metadata(
601 + num_draft_tokens, cu_num_tokens)
602 + logits_indices = spec_decode_metadata.logits_indices
603
604 # Hot-Swap lora model
605 if self.lora_config:
606 self.set_active_loras(self.input_batch, num_scheduled_tokens)
607
595 - return attn_metadata, logits_indices
608 + return attn_metadata, logits_indices, spec_decode_metadata
609
610 def _compute_cascade_attn_prefix_len(
611 self,
@@ -732,50 +745,79 @@ class GPUModelRunner(LoRAModelRunnerMixin):
745
746 def _calc_spec_decode_metadata(
747 self,
735 - scheduler_output: "SchedulerOutput",
736 - cu_num_tokens: np.ndarray,
737 - ) -> torch.Tensor:
738 - # Get the number of spec decode tokens for each request.
739 - num_reqs = self.input_batch.num_reqs
740 - num_spec_decode_tokens = np.empty(num_reqs, dtype=np.int32)
741 - for i, req_id in enumerate(self.input_batch.req_ids):
742 - num_spec_decode_tokens[i] = len(
743 - scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()))
744 -
745 - # Get spec decode logits indices.
746 - # E.g., num_scheduled_tokens: [4, 100, 3, 100, 2]
747 - # cu_num_tokens: [4, 104, 107, 207, 209]
748 - # num_spec_tokens_list: [3, 0, 2, 0, 1]
749 - # num_sampled_tokens: [4, 1, 3, 1, 2]
750 - # spec_decode_logits_indices:
751 - # [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]
752 - num_sampled_tokens = num_spec_decode_tokens + 1
753 - # logits_start_loc: [0, 103, 104, 206, 207]
754 - logits_start_loc = cu_num_tokens - num_sampled_tokens
755 - # [0, 103, 104, 206, 207] ->
756 - # [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207]
757 - logits_start_loc = np.repeat(logits_start_loc, num_sampled_tokens)
758 - # The following three lines:
759 - # [4, 1, 3, 1, 2] -> [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]
760 - # Step 1. [4, 1, 3, 1, 2] -> [4, 5, 8, 9, 11]
761 - cu_num_sampled_tokens = np.cumsum(num_sampled_tokens)
762 - # Step 2. [4, 5, 8, 9, 11] -> [0, 4, 5, 8, 9]
763 - # -> [0, 0, 0, 0, 4, 5, 5, 5, 8, 9, 9]
764 - cumsums_sampled_offsets = np.repeat(
765 - cu_num_sampled_tokens - num_sampled_tokens, num_sampled_tokens)
766 - # Step 3. [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
767 - # - [0, 0, 0, 0, 4, 5, 5, 5, 8, 9, 9]
768 - # -> [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]
769 - total_num_sampled_tokens = num_sampled_tokens.sum()
770 - sampled_arange = (self.arange_np[:total_num_sampled_tokens] -
771 - cumsums_sampled_offsets)
772 -
773 - # [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207] ->
774 - # [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]
775 - spec_decode_logits_indices = logits_start_loc + sampled_arange
776 - return torch.from_numpy(spec_decode_logits_indices).to(
748 + num_draft_tokens: np.ndarray,
749 + cu_num_scheduled_tokens: np.ndarray,
750 + ) -> SpecDecodeMetadata:
751 + # Inputs:
752 + # cu_num_scheduled_tokens: [ 4, 104, 107, 207, 209]
753 + # num_draft_tokens: [ 3, 0, 2, 0, 1]
754 + # Outputs:
755 + # cu_num_draft_tokens: [ 3, 3, 5, 5, 6]
756 + # logits_indices: [ 0, 1, 2, 3, 103, 104, 105, 106,
757 + # 206, 207, 208]
758 + # target_logits_indices: [ 0, 1, 2, 5, 6, 9]
759 + # bonus_logits_indices: [ 3, 4, 7, 8, 10]
760 +
761 + # Compute the logits indices.
762 + # [4, 1, 3, 1, 2]
763 + num_sampled_tokens = num_draft_tokens + 1
764 + # Step 1. [4, 5, 8, 9, 11]
765 + cu_num_sampled_tokens = np.cumsum(num_sampled_tokens, dtype=np.int32)
766 + total_num_sampled_tokens = cu_num_sampled_tokens[-1]
767 + # Step 2. [0, 0, 0, 0, 4, 5, 5, 5, 8, 9, 9]
768 + cumsums_offsets = np.repeat(cu_num_sampled_tokens - num_sampled_tokens,
769 + num_sampled_tokens)
770 + # Step 3. [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]
771 + arange = self.arange_np[:total_num_sampled_tokens] - cumsums_offsets
772 + # Step 4. [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207]
773 + logits_indices = np.repeat(
774 + cu_num_scheduled_tokens - num_sampled_tokens, num_sampled_tokens)
775 + # Step 5. [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]
776 + logits_indices += arange
777 +
778 + # Compute the bonus logits indices.
779 + bonus_logits_indices = cu_num_sampled_tokens - 1
780 +
781 + # Compute the draft logits indices.
782 + # [3, 3, 5, 5, 6]
783 + cu_num_draft_tokens = np.cumsum(num_draft_tokens, dtype=np.int32)
784 + total_num_draft_tokens = cu_num_draft_tokens[-1]
785 + # [0, 0, 0, 3, 3, 5]
786 + cumsums_offsets = np.repeat(cu_num_draft_tokens - num_draft_tokens,
787 + num_draft_tokens)
788 + # [0, 1, 2, 0, 1, 0]
789 + arange = self.arange_np[:total_num_draft_tokens] - cumsums_offsets
790 + # [0, 0, 0, 5, 5, 9]
791 + target_logits_indices = np.repeat(
792 + cu_num_sampled_tokens - num_sampled_tokens, num_draft_tokens)
793 + # [0, 1, 2, 5, 6, 9]
794 + target_logits_indices += arange
795 +
796 + # TODO: Optimize the CPU -> GPU copy.
797 + cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to(
798 + self.device, non_blocking=True)
799 + logits_indices = torch.from_numpy(logits_indices).to(self.device,
800 + non_blocking=True)
801 + target_logits_indices = torch.from_numpy(target_logits_indices).to(
802 + self.device, non_blocking=True)
803 + bonus_logits_indices = torch.from_numpy(bonus_logits_indices).to(
804 self.device, non_blocking=True)
805
806 + # Compute the draft token ids.
807 + # draft_token_indices: [ 1, 2, 3, 105, 106, 208]
808 + draft_token_ids = self.input_ids[logits_indices]
809 + draft_token_ids = draft_token_ids[target_logits_indices + 1]
810 +
811 + metadata = SpecDecodeMetadata(
812 + draft_token_ids=draft_token_ids,
813 + num_draft_tokens=num_draft_tokens.tolist(),
814 + cu_num_draft_tokens=cu_num_draft_tokens,
815 + target_logits_indices=target_logits_indices,
816 + bonus_logits_indices=bonus_logits_indices,
817 + logits_indices=logits_indices,
818 + )
819 + return metadata
820 +
821 def _execute_encoder(self, scheduler_output: "SchedulerOutput"):
822 scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
823 if not scheduled_encoder_inputs:
@@ -931,7 +973,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
973 encoder_outputs = []
974
975 # Prepare the decoder inputs.
934 - attn_metadata, logits_indices = self._prepare_inputs(scheduler_output)
976 + attn_metadata, logits_indices, spec_decode_metadata = (
977 + self._prepare_inputs(scheduler_output))
978 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
979 if (self.use_cuda_graph
980 and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
@@ -1006,31 +1049,29 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1049
1050 # Sample the next token and get logprobs if needed.
1051 sampling_metadata = self.input_batch.sampling_metadata
1009 - if not self.use_spec_decode:
1052 + if spec_decode_metadata is None:
1053 sampler_output = self.model.sample(
1054 logits=logits,
1055 sampling_metadata=sampling_metadata,
1056 )
1057 else:
1015 - draft_token_ids = [
1016 - scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])
1017 - for req_id in self.input_batch.req_ids
1018 - ]
1019 - sample_lens = [len(tokens) + 1 for tokens in draft_token_ids]
1020 - recover_logits_idx = np.cumsum(sample_lens) - 1
1021 - target_probs = self.rejection_sampler.compute_probs(
1022 - logits, sampling_metadata, sample_lens)
1058 + # TODO(woosuk): Optimize the memory usage.
1059 + bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
1060 sampler_output = self.model.sample(
1024 - logits=logits[recover_logits_idx, :],
1061 + logits=bonus_logits,
1062 sampling_metadata=sampling_metadata,
1063 )
1064 bonus_token_ids = sampler_output.sampled_token_ids
1065 +
1066 + # TODO(woosuk): Optimize the memory usage.
1067 + target_logits = logits[spec_decode_metadata.target_logits_indices]
1068 output_token_ids = self.rejection_sampler(
1029 - draft_token_ids,
1069 + spec_decode_metadata,
1070 None, # draft_probs
1071 + target_logits,
1072 bonus_token_ids,
1032 - target_probs,
1033 - sampling_metadata)
1073 + sampling_metadata,
1074 + )
1075 sampler_output.sampled_token_ids = output_token_ids
1076
1077 # TODO(woosuk): The following loop can be slow since it iterates over
@@ -1066,13 +1107,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1107 valid_sampled_token_ids = sampled_token_ids.tolist()
1108 else:
1109 # Includes spec decode tokens.
1069 - valid_mask = sampled_token_ids != INVALID_TOKEN_ID
1070 - gen_lens = valid_mask.sum(dim=1).tolist()
1071 - # TODO(woosuk): Optimize this.
1072 - valid_sampled_token_ids = [
1073 - seq.tolist()
1074 - for seq in sampled_token_ids[valid_mask].split(gen_lens)
1075 - ]
1110 + valid_sampled_token_ids = self.rejection_sampler.parse_output(
1111 + sampled_token_ids, self.input_batch.vocab_size)
1112
1113 if not self.use_spec_decode:
1114 spec_token_ids = None
@@ -1316,6 +1352,33 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1352 "initializing the engine.") from e
1353 else:
1354 raise e
1355 + if self.use_spec_decode:
1356 + draft_token_ids = [[0] for _ in range(num_reqs)]
1357 + dummy_spec_decode_metadata = SpecDecodeMetadata.make_dummy(
1358 + draft_token_ids, self.device)
1359 +
1360 + num_tokens = sum(len(ids) for ids in draft_token_ids)
1361 + # draft_probs = torch.randn(
1362 + # num_tokens, logits.shape[-1], device=self.device,
1363 + # dtype=logits.dtype)
1364 + draft_probs = None
1365 + target_logits = torch.randn(num_tokens,
1366 + logits.shape[-1],
1367 + device=self.device,
1368 + dtype=logits.dtype)
1369 + # NOTE(woosuk): Here, we should use int32 because the sampler uses
1370 + # int32 for bonus_token_ids. If the dtype mismatches, re-compilation
1371 + # will occur at runtime.
1372 + bonus_token_ids = torch.zeros(num_reqs,
1373 + device=self.device,
1374 + dtype=torch.int32)
1375 + self.rejection_sampler(
1376 + dummy_spec_decode_metadata,
1377 + draft_probs,
1378 + target_logits,
1379 + bonus_token_ids,
1380 + dummy_metadata,
1381 + )
1382 return sampler_output
1383
1384 def profile_run(self) -> None: