@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8d6cf89526
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] Support random sampling for spec decode (#13933)

Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Lily Liu committed Mar 16, 2025 at 22:00 UTC 8d6cf89526ff983b7eb74aad3903138004ae95cd
5 files changed +572 -198
tests/v1/sample/test_rejection_sampler.py
+251 -56
@@ -1,37 +1,51 @@
1 # SPDX-License-Identifier: Apache-2.0
2 +from typing import Any, Optional
3
4 import pytest
5 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
10
11 +DEVICE = "cpu"
12 +
13
14 @pytest.fixture
15 def sampler():
16 return RejectionSampler()
17
18
15 -def create_logits_tensor(token_ids: list[int],
19 +def create_logits_tensor(token_ids: list[list[int]],
20 vocab_size: int = 100) -> torch.Tensor:
21 """Helper function to create logits tensor that
22 will produce desired token ids on argmax"""
19 - logits = torch.full((len(token_ids), vocab_size), -100.0).cuda()
20 - for i, token_id in enumerate(token_ids):
21 - logits[i, token_id] = 100.0
23 + num_total_tokens = sum(len(tokens) for tokens in token_ids)
24 + logits = torch.full((num_total_tokens, vocab_size), -100.0, device=DEVICE)
25 + start_loc = 0
26 + for tokens in token_ids:
27 + for j, token_id in enumerate(tokens):
28 + logits[start_loc + j, token_id] = 100.0
29 + start_loc += len(tokens)
30 return logits
31
32
25 -def create_sampling_metadata(spec_tokens: list[list[int]]) -> SamplingMetadata:
26 - batch_size = len(spec_tokens)
33 +def create_sampling_metadata(
34 + all_greedy: bool,
35 + generators: Optional[dict[int, Any]] = None) -> SamplingMetadata:
36 + """Create a v1 sampling metadata object with all_greedy set
37 + to the given value. Either all greedy or all random sampling
38 + is used.
39 + """
40 + generators = generators or {}
41 return SamplingMetadata(
42 temperature=torch.tensor([]),
29 - all_greedy=True,
30 - all_random=False,
43 + all_greedy=all_greedy,
44 + all_random=not all_greedy,
45 top_p=None,
46 top_k=None,
33 - min_p=torch.empty(batch_size, ),
34 - generators={},
47 + min_p=torch.empty(1, ),
48 + generators=generators,
49 max_num_logprobs=0,
50 no_penalties=False,
51 prompt_token_ids=None,
@@ -40,129 +54,310 @@ def create_sampling_metadata(spec_tokens: list[list[int]]) -> SamplingMetadata:
54 repetition_penalties=torch.tensor([]),
55 output_token_ids=[],
56 min_tokens={},
43 - logit_bias=[None] * batch_size,
57 + logit_bias=[None],
58 allowed_token_ids_mask=None,
59 bad_words_token_ids={},
60 )
61
62
63 +########################### Tests for Greedy Sampling ###################
64 def test_perfect_match(sampler):
65 """Test when output tokens perfectly match speculated tokens"""
66 spec_tokens = [[1, 2, 3]]
52 - output_tokens = [1, 2, 3, 4] # 4 is the bonus token
67 + output_tokens = [[1, 2, 3, 4]] # 4 is the bonus token
68
54 - metadata = create_sampling_metadata(spec_tokens)
69 + metadata = create_sampling_metadata(all_greedy=True)
70 logits = create_logits_tensor(output_tokens)
71 + bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
72 + device=logits.device)
73
57 - output = sampler(spec_tokens, logits, metadata)
74 + output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
75 expected = torch.tensor([[1, 2, 3, 4]],
76 dtype=torch.int,
77 device=logits.device)
61 - assert torch.equal(output.sampled_token_ids, expected)
78 + assert torch.equal(output, expected)
79
80
81 def test_early_mismatch(sampler):
82 """Test when there's an early mismatch in tokens"""
83 spec_tokens = [[1, 2, 3]]
67 - output_tokens = [1, 5, 3, 4] # Mismatch at position 1
84 + output_tokens = [[1, 5, 3, 4]] # Mismatch at position 1
85
69 - metadata = create_sampling_metadata(spec_tokens)
86 + metadata = create_sampling_metadata(all_greedy=True)
87 logits = create_logits_tensor(output_tokens)
88 + bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
89 + device=logits.device)
90
72 - output = sampler(spec_tokens, logits, metadata)
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)
76 - assert torch.equal(output.sampled_token_ids, expected)
95 + assert torch.equal(output, expected)
96
97
98 def test_multiple_sequences(sampler):
99 """Test handling multiple sequences of speculated tokens"""
100 spec_tokens = [[1, 2], [3]]
82 - output_tokens = [1, 2, 5, 3, 4] # Two sequences with bonus tokens 5 and 4
101 + output_tokens = [[1, 2, 5], [3,
102 + 4]] # Two sequences with bonus tokens 5 and 4
103
84 - metadata = create_sampling_metadata(spec_tokens)
104 + metadata = create_sampling_metadata(all_greedy=True)
105 logits = create_logits_tensor(output_tokens)
106 + bonus_token_tensor = torch.tensor(
107 + [output_tokens[0][-1], output_tokens[1][-1]], device=logits.device)
108
87 - output = sampler(spec_tokens, logits, metadata)
109 + output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
110 expected = torch.tensor([[1, 2, 5], [3, 4, INVALID_TOKEN_ID]],
111 dtype=torch.int,
112 device=logits.device)
91 - assert torch.equal(output.sampled_token_ids, expected)
113 + assert torch.equal(output, expected)
114
115
116 def test_single_token_sequence(sampler):
117 """Test handling sequences with single token"""
118 spec_tokens = [[1]]
97 - output_tokens = [1, 2] # Single token with bonus token 2
119 + output_tokens = [[1, 2]] # Single token with bonus token 2
120
99 - metadata = create_sampling_metadata(spec_tokens)
121 + metadata = create_sampling_metadata(all_greedy=True)
122 logits = create_logits_tensor(output_tokens)
123 + bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
124 + device=logits.device)
125
102 - output = sampler(spec_tokens, logits, metadata)
126 + output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
127 expected = torch.tensor([[1, 2]], dtype=torch.int, device=logits.device)
104 - assert torch.equal(output.sampled_token_ids, expected)
128 + assert torch.equal(output, expected)
129
130
131 def test_empty_sequence(sampler):
132 """Test handling empty sequence of speculated tokens"""
133 spec_tokens: list[list[int]] = [[]]
110 - output_tokens = [5] # Just the bonus token
134 + output_tokens = [[5]] # Just the bonus token
135
112 - metadata = create_sampling_metadata(spec_tokens)
136 + metadata = create_sampling_metadata(all_greedy=True)
137 logits = create_logits_tensor(output_tokens)
138 + bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
139 + device=logits.device)
140
115 - output = sampler(spec_tokens, logits, metadata)
141 + output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
142 expected = torch.tensor([[5]], dtype=torch.int, device=logits.device)
117 - assert torch.equal(output.sampled_token_ids, expected)
143 + assert torch.equal(output, expected)
144
145
146 def test_multiple_mismatches(sampler):
147 """Test handling multiple sequences with mismatches"""
148 spec_tokens = [[1, 2, 3], [4, 5, 6]]
123 - output_tokens = [1, 2, 7, 6, 4, 8, 6, 9] # Mismatches in both sequences
149 + output_tokens = [[1, 2, 7, 6], [4, 8, 6,
150 + 9]] # Mismatches in both sequences
151
125 - metadata = create_sampling_metadata(spec_tokens)
152 + metadata = create_sampling_metadata(all_greedy=True)
153 logits = create_logits_tensor(output_tokens)
154 + bonus_token_tensor = torch.tensor(
155 + [output_tokens[0][-1], output_tokens[1][-1]], device=logits.device)
156
128 - output = sampler(spec_tokens, logits, metadata)
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)
133 - assert torch.equal(output.sampled_token_ids, expected)
162 + assert torch.equal(output, expected)
163
164
165 @pytest.mark.parametrize(
166 "spec_tokens,output_tokens,expected",
167 [
139 - ([[1, 2]], [1, 2, 3], [[1, 2, 3]]), # Perfect match with bonus
140 - ([[1]], [2, 3], [[2, INVALID_TOKEN_ID]]), # First mismatch
141 - ([[1, 2], [3, 4]], [1, 5, 6, 3, 4, 7], [[1, 5, INVALID_TOKEN_ID],
142 - [3, 4, 7]]), # Mixed matches
168 + ([[1, 2]], [[1, 2, 3]], [[1, 2, 3]]), # Perfect match with bonus
169 + ([[1]], [[2, 3]], [[2, INVALID_TOKEN_ID]]), # First mismatch
170 + ([[1, 2], [3, 4]], [[1, 5, 6], [3, 4, 7]],
171 + [[1, 5, INVALID_TOKEN_ID], [3, 4, 7]]), # Mixed matches
172 ])
173 def test_parametrized_cases(sampler, spec_tokens, output_tokens, expected):
174 """Parametrized test for various matching scenarios"""
146 - metadata = create_sampling_metadata(spec_tokens)
175 + metadata = create_sampling_metadata(all_greedy=True)
176 logits = create_logits_tensor(output_tokens)
177 + bonus_token_tensor = torch.tensor([tokens[-1] for tokens in output_tokens],
178 + device=logits.device)
179
149 - output = sampler(spec_tokens, logits, metadata)
180 + output = sampler(spec_tokens, None, bonus_token_tensor, logits, metadata)
181 expected_tensor = torch.tensor(expected,
182 dtype=torch.int,
183 device=logits.device)
153 - assert torch.equal(output.sampled_token_ids, expected_tensor)
154 -
155 -
156 -def test_logits_shape_handling(sampler):
157 - """Test handling of different logits tensor shapes"""
158 - spec_tokens = [[1, 2]]
159 - output_tokens = [1, 2, 3]
160 - vocab_size = 1000
161 -
162 - metadata = create_sampling_metadata(spec_tokens)
163 - logits = create_logits_tensor(output_tokens, vocab_size)
184 + assert torch.equal(output, expected_tensor)
185 +
186 +
187 +########################### Tests for Random Sampling ###################
188 +@pytest.mark.parametrize("k", [1, 3, 5])
189 +@pytest.mark.parametrize("vocab_size", [1000])
190 +@pytest.mark.parametrize("batch_size", [1, 4, 8])
191 +@pytest.mark.parametrize("frac_seeded", [0.0, 0.5])
192 +@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)
200 + bonus_token_ids = torch.randint(low=0,
201 + high=vocab_size,
202 + size=(batch_size, 1),
203 + dtype=torch.int64)
204 + draft_token_ids = torch.randint(low=0,
205 + high=vocab_size,
206 + size=(batch_size, k),
207 + dtype=torch.int64)
208 +
209 + seeded_mask = torch.rand(batch_size, dtype=torch.float32) <= frac_seeded
210 +
211 + results = []
212 + for _ in range(n_rep):
213 + seeded_seqs = {
214 + i: torch.Generator(device=DEVICE).manual_seed(i)
215 + for i in range(batch_size) if seeded_mask[i]
216 + }
217 +
218 + sampling_metadata = create_sampling_metadata(all_greedy=False,
219 + generators=seeded_seqs)
220 + rep_result = sampler(draft_token_ids.tolist(), draft_probs,
221 + bonus_token_ids, target_probs, sampling_metadata)
222 +
223 + results.append(rep_result)
224 +
225 + for i in range(batch_size):
226 + if seeded_mask[i]:
227 + for j in range(1, n_rep):
228 + assert torch.equal(results[j][i], results[0][i])
229 +
230 +
231 +def test_rejection_sampling_approximates_target_distribution():
232 + """Verify rejection sampling approximates target distribution,
233 + despite sampling from a potentially distinct draft distribution.
234 +
235 + This is done by first creating a random target probability
236 + distribution and a random draft probability distribution. We then
237 + sample token ids from the rejection sampler using these draft
238 + and target distributions. The samples are used to estimate
239 + the output probability distribution, which we expect to approximate
240 + the target distribution.
241 +
242 + A basic distance metric is used to determine similarity between
243 + distributions.
244 +
245 + We expect that as we increase the number of samples,
246 + the distance between the observed distribution and the target
247 + distribution decreases. To measure this, we compare the distance
248 + of the observed distribution against both the target distribution
249 + and a uniform random distribution. We expect the distance between
250 + the observed distribution and the target distribution to improve
251 + much more than the distance improvement between the observed
252 + distribution and the random distribution.
253 + """
254 + torch.set_default_device(DEVICE)
255 + vocab_size = 10
256 + k = 2
257 + num_reference_probs = 100
258 +
259 + # 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))
264 + reference_probs = F.softmax(
265 + torch.rand(num_reference_probs, vocab_size, dtype=torch.float32),
266 + dim=-1,
267 + )
268
165 - output = sampler(spec_tokens, logits, metadata)
166 - expected = torch.tensor([[1, 2, 3]], dtype=torch.int, device=logits.device)
167 - assert torch.equal(output.sampled_token_ids, expected)
168 - assert logits.shape[-1] == vocab_size
269 + sample_sizes = [10, 100, 1_000, 10_000, 100_000]
270 + distance_wrt_reference: list[float] = []
271 + distance_wrt_target: list[float] = []
272 +
273 + for num_samples in sample_sizes:
274 + # Sample using rejection sampling.
275 + rej_sample_probs = estimate_rejection_sampling_pdf(
276 + draft_probs, target_probs, k, vocab_size, num_samples)
277 + rej_sample_probs = rej_sample_probs.to(DEVICE)
278 +
279 + # Average distance from reference probs.
280 + reference_vs_rejsample_dist = torch.dist(
281 + reference_probs,
282 + rej_sample_probs).item() / reference_probs.shape[0]
283 + target_vs_rejsample_dist = torch.dist(target_probs,
284 + rej_sample_probs).item()
285 +
286 + distance_wrt_reference.append(reference_vs_rejsample_dist)
287 + distance_wrt_target.append(target_vs_rejsample_dist)
288 +
289 + relative_change_in_distance_wrt_target = get_ratio_first_to_last(
290 + distance_wrt_target)
291 + relative_change_in_distance_wrt_reference = get_ratio_first_to_last(
292 + distance_wrt_reference)
293 +
294 + print(f"{num_samples=} {target_vs_rejsample_dist=:.05f} "
295 + f"{reference_vs_rejsample_dist=:.05f}")
296 + print(f"{num_samples=} {relative_change_in_distance_wrt_target=:.02f} "
297 + f"{relative_change_in_distance_wrt_reference=:.02f}")
298 +
299 + relative_change_in_distance_wrt_target = get_ratio_first_to_last(
300 + distance_wrt_target)
301 + relative_change_in_distance_wrt_reference = get_ratio_first_to_last(
302 + distance_wrt_reference)
303 +
304 + expected_improvement_multiplier = 20
305 + assert (relative_change_in_distance_wrt_target
306 + > relative_change_in_distance_wrt_reference *
307 + expected_improvement_multiplier)
308 +
309 +
310 +def get_ratio_first_to_last(elements: list[float]) -> float:
311 + return elements[0] / elements[-1]
312 +
313 +
314 +def estimate_rejection_sampling_pdf(
315 + draft_probs: torch.Tensor,
316 + target_probs: torch.Tensor,
317 + k: int,
318 + vocab_size: int,
319 + num_samples: int,
320 +) -> torch.Tensor:
321 + """Estimate the probability distribution of the output tokens
322 + using rejection sampling.
323 +
324 + Args:
325 + draft_probs: Draft probability distribution.
326 + target_probs: Target probability distribution.
327 + num_samples: Number of samples to draw.
328 +
329 + Returns:
330 + Estimated probability distribution of the output tokens.
331 + """
332 + sampler = RejectionSampler()
333 + # Repeat draft probs num_samples times.
334 + draft_probs = draft_probs.reshape(1, 1,
335 + vocab_size).repeat(num_samples, k, 1)
336 +
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)
340 +
341 + # Randomly sample draft token ids from draft probs.
342 + draft_token_ids = torch.multinomial(draft_probs[:, 0, :],
343 + num_samples=k,
344 + replacement=True).reshape(
345 + num_samples, k)
346 +
347 + # Bonus tokens not used but required.
348 + bonus_token_ids = torch.zeros((1, 1), dtype=torch.int64,
349 + device=DEVICE).repeat(num_samples, 1)
350 +
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)
355 + output_token_ids = output_token_ids[:, :-1].flatten()
356 +
357 + hist = torch.histogram(output_token_ids.to(dtype=torch.float,
358 + device="cpu"),
359 + bins=vocab_size,
360 + range=(0, vocab_size),
361 + density=True)
362 +
363 + return hist.hist
vllm/v1/sample/rejection_sampler.py
+274 -128
@@ -1,87 +1,89 @@
1 # SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional
3
4 import torch
5 import torch.nn as nn
6 from torch.nn.utils.rnn import pad_sequence
7
7 -from vllm import envs
8 from vllm.logger import init_logger
9 -from vllm.platforms import current_platform
10 -from vllm.v1.outputs import SamplerOutput
9 from vllm.v1.sample.metadata import SamplingMetadata
12 -
13 -try:
14 - import flashinfer.sampling as fs
15 - is_flashinfer_available = True
16 -except ImportError:
17 - is_flashinfer_available = False
10 +from vllm.v1.spec_decode.utils import random_sample
11
12 logger = init_logger(__name__)
13 INVALID_TOKEN_ID = -1
14
15
16 class RejectionSampler(nn.Module):
17 + """
18 + The implementation strictly follows the algorithm described in
19 + https://arxiv.org/abs/2211.17192.
20 + However, we want to clarify the terminology used in the implementation:
21 + accepted tokens: tokens that are accepted based on the relationship
22 + between the "raw" draft and target probabilities.
23 + recovered tokens: tokens that are sampled based on the adjusted probability
24 + distribution, which is derived from both the draft and target
25 + probabilities.
26 + bonus tokens:
27 + If all proposed tokens are accepted, the bonus token is added to the
28 + end of the sequence. The bonus token is only sampled from the target
29 + probabilities. We pass in the bonus tokens instead of sampling them
30 + in the rejection sampler to allow for more flexibility in the
31 + sampling process. For example, we can use top_p, top_k sampling for
32 + bonus tokens, while spec decode does not support these sampling
33 + strategies.
34 + output tokens:
35 + Tokens are finally generated with the rejection sampler.
36 + output tokens = accepted tokens + recovered tokens + bonus tokens
37 + """
38
39 def __init__(self):
40 super().__init__()
27 - if current_platform.is_cuda():
28 - if is_flashinfer_available:
29 - if envs.VLLM_USE_FLASHINFER_SAMPLER is not False:
30 - # FIXME(woosuk): Currently, we have errors when using
31 - # FlashInfer for rejection sampling. As a workaround, we
32 - # disable FlashInfer for rejection sampling by default.
33 - logger.info("Currently, FlashInfer rejection sampler is "
34 - "disabled because of a bug. Falling back to "
35 - "the PyTorch-native implementation of "
36 - "rejection sampling.")
37 - self.forward_method = self.forward_native
38 -
39 - # NOTE(woosuk): The V0 sampler doesn't use FlashInfer for
40 - # sampling unless VLLM_USE_FLASHINFER_SAMPLER=1 (i.e., by
41 - # default it is unused). For backward compatibility, we set
42 - # `VLLM_USE_FLASHINFER_SAMPLER` as None by default and
43 - # interpret it differently in V0 and V1 samplers: In V0,
44 - # None means False, while in V1, None means True. This is
45 - # why we use the condition
46 - # `envs.VLLM_USE_FLASHINFER_SAMPLER is not False` here.
47 - # logger.info("Using FlashInfer for rejection sampling.")
48 - # self.forward_method = self.flashinfer_sample
49 - else:
50 - logger.warning(
51 - "FlashInfer is available, but it is not enabled. "
52 - "Falling back to the PyTorch-native implementation of "
53 - "rejection sampling. For the best performance, "
54 - "please set VLLM_USE_FLASHINFER_SAMPLER=1.")
55 - self.forward_method = self.forward_native
56 - else:
57 - logger.warning(
58 - "FlashInfer is not available. Falling back to the PyTorch-"
59 - "native implementation of rejection sampling. For the "
60 - "best performance, please install FlashInfer.")
61 - self.forward_method = self.forward_native
62 - else:
63 - self.forward_method = self.forward_native
64 -
65 - def forward(self, draft_token_ids: list[list[int]],
66 - target_probs: torch.Tensor,
67 - sampling_metadata: SamplingMetadata) -> SamplerOutput:
68 - if not sampling_metadata.all_greedy:
69 - raise NotImplementedError(
70 - "Currently, only greedy sampling is supported by "
71 - "rejection sampler.")
72 - return self.forward_method(draft_token_ids, target_probs,
73 - sampling_metadata)
41
75 - def flashinfer_sample(
42 + def forward(
43 self,
44 draft_token_ids: list[list[int]],
78 - target_probs: torch.Tensor,
45 + 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]
48 sampling_metadata: SamplingMetadata,
80 - ) -> SamplerOutput:
49 + ) -> torch.Tensor:
50 + '''
51 + 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.
57 + draft_probs (Optional[torch.Tensor]):
58 + 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.
62 + 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
68 + 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.
76 + sampling_metadata (SamplingMetadata):
77 + Additional metadata needed for sampling, such as temperature,
78 + top-k/top-p parameters, or other relevant information.
79 + Returns:
80 + output_token_ids (torch.Tensor):
81 + A tensor containing the final output token IDs.
82 + '''
83 +
84 # NOTE: The following input preparationg can be moved
85 # to the model runner with a persistent manner for better
86 # performance.
84 - sample_lens = [len(x) + 1 for x in draft_token_ids]
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,90 +92,171 @@ class RejectionSampler(nn.Module):
92 batch_first=True,
93 padding_value=INVALID_TOKEN_ID)
94
93 - if sampling_metadata.all_greedy:
94 - target_token_ids = target_probs.argmax(dim=-1).view(-1)
95 - target_token_ids = target_token_ids.split(sample_lens)
96 - target_token_ids = pad_sequence(target_token_ids,
97 - batch_first=True,
98 - padding_value=INVALID_TOKEN_ID)
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)
101 - # NOTE: CPU <-> GPU synchronization happens here.
102 - draft_token_ids_tensor = draft_token_ids_tensor.to(
103 - target_probs.device)
102 draft_probs = _create_greedy_token_probs(draft_token_ids_tensor,
103 vocab_size,
104 target_probs.device)
107 - target_probs = _create_greedy_token_probs(target_token_ids,
108 - vocab_size,
109 - target_probs.device)
110 - uniform_samples = torch.zeros(draft_token_ids_tensor.size(0),
111 - draft_token_ids_tensor.size(1) + 1,
112 - device=target_probs.device)
113 - else:
114 - raise NotImplementedError(
115 - "Currently, only greedy sampling is supported by "
116 - "rejection sampler.")
117 -
118 - sampled_token_ids, _, _ = fs.chain_speculative_sampling(
119 - draft_probs,
120 - draft_token_ids_tensor,
121 - uniform_samples,
122 - target_probs,
123 - )
124 - return SamplerOutput(sampled_token_ids=sampled_token_ids,
125 - logprobs_tensors=None)
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,
130 - draft_token_ids: list[list[int]],
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,
133 - ) -> SamplerOutput:
134 - sample_lens = [len(x) + 1 for x in draft_token_ids]
135 - # Convert draft token IDs to a tensor, split by sample_lens, then pad.
136 - draft_token_ids = [
137 - torch.tensor(x, dtype=int, device='cpu') for x in draft_token_ids
138 - ]
139 - draft_token_ids_tensor = pad_sequence(draft_token_ids,
140 - batch_first=True,
141 - padding_value=INVALID_TOKEN_ID)
142 - draft_token_ids_tensor = draft_token_ids_tensor.to(target_probs.device)
122 + ) -> torch.Tensor:
123 # Add 1 to include the 'bonus' token.
124 if sampling_metadata.all_greedy:
145 - output_token_ids = target_probs.argmax(dim=-1).view(-1)
146 - output_token_ids = output_token_ids.split(sample_lens)
147 - output_token_ids = pad_sequence(output_token_ids,
148 - batch_first=True,
149 - padding_value=INVALID_TOKEN_ID)
125 # Produce a mask that remains 1 (True) until the first
126 # mismatch (cumprod turns 0 after a mismatch).
152 - accept_mask = (
153 - output_token_ids[:, :-1] == draft_token_ids_tensor).cumprod(
154 - dim=1)
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:
156 - raise NotImplementedError(
157 - "Currently, only greedy sampling is supported by "
158 - "rejection sampler.")
159 - # Identify valid positions (non-padding).
160 - valid_mask = output_token_ids != INVALID_TOKEN_ID
161 - # Generate mask with bonus token.
162 - generate_mask = torch.cat([
163 - accept_mask,
164 - torch.zeros(accept_mask.size(0), 1, device=accept_mask.device)
165 - ],
166 - dim=1).to(torch.bool) & valid_mask
167 - zeros_mask = (generate_mask == 0)
168 - first_zero_idx = zeros_mask.float().argmax(dim=1)
169 - # Figure out which rows actually contain at least one zero.
170 - rows_with_zero = zeros_mask.any(dim=1)
171 - # Use indexing to set the first zero in each of those rows to 1.
172 - generate_mask[rows_with_zero, first_zero_idx[rows_with_zero]] = 1
173 -
174 - output_token_ids[~generate_mask] = INVALID_TOKEN_ID
175 - return SamplerOutput(sampled_token_ids=output_token_ids,
176 - logprobs_tensors=None)
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 +
223 + return output_token_ids
224 +
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 + """
248 + 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(
@@ -199,3 +282,66 @@ def _create_greedy_token_probs(
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:
290 + """
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]
294 + """
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:
311 + """
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
vllm/v1/sample/sampler.py
-8
@@ -119,14 +119,6 @@ class Sampler(nn.Module):
119 )
120 return sampled
121
122 - def compute_probs(self, logits: torch.Tensor,
123 - sampling_metadata: SamplingMetadata) -> torch.Tensor:
124 - if sampling_metadata.all_greedy:
125 - return logits
126 - # Apply temperature. This is an in-place op changing logits.
127 - logits = self.apply_temperature(logits, sampling_metadata.temperature)
128 - return logits.softmax(dim=-1, dtype=torch.float32)
129 -
122 def compute_logprobs(self, logits: torch.Tensor) -> torch.Tensor:
123 return logits.log_softmax(dim=-1, dtype=torch.float32)
124
vllm/v1/spec_decode/utils.py new
+22
@@ -0,0 +1,22 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from vllm.v1.sample.ops.topk_topp_sampler import random_sample # noqa
3 +from vllm.v1.worker.gpu_input_batch import InputBatch
4 +
5 +
6 +def is_spec_decode_supported(req_id: str, input_batch: InputBatch) -> bool:
7 + if req_id in input_batch.top_k_reqs or req_id in input_batch.top_p_reqs:
8 + # Spec decode doesn't support top_p/top_k sampling.
9 + return False
10 + elif req_id in input_batch.min_p_reqs:
11 + # Spec decode doesn't support min_p sampling.
12 + return False
13 + elif (req_id in input_batch.frequency_penalties_reqs
14 + or req_id in input_batch.presence_penalties_reqs
15 + or req_id in input_batch.repetition_penalties_reqs):
16 + # Spec decode doesn't support penalties.
17 + return False
18 + elif req_id in input_batch.num_logprobs:
19 + # Spec decode doesn't support logprobs.
20 + return False
21 +
22 + return True
vllm/v1/worker/gpu_model_runner.py
+25 -6
@@ -37,6 +37,7 @@ from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
37 from vllm.v1.sample.metadata import SamplingMetadata
38 from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID, RejectionSampler
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
42 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
43 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
@@ -1020,15 +1021,26 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1021 sampling_metadata=sampling_metadata,
1022 )
1023 else:
1023 - target_probs = self.model.sampler.compute_probs(
1024 - logits, sampling_metadata)
1024 draft_token_ids = [
1025 scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])
1026 for req_id in self.input_batch.req_ids
1027 ]
1029 - sampler_output = self.rejection_sampler(draft_token_ids,
1030 - target_probs,
1031 - sampling_metadata)
1028 + sample_lens = [len(tokens) + 1 for tokens in draft_token_ids]
1029 + recover_logits_idx = np.cumsum(sample_lens) - 1
1030 + target_probs = self.rejection_sampler.compute_probs(
1031 + logits, sampling_metadata, sample_lens)
1032 + sampler_output = self.model.sample(
1033 + logits=logits[recover_logits_idx, :],
1034 + sampling_metadata=sampling_metadata,
1035 + )
1036 + bonus_token_ids = sampler_output.sampled_token_ids
1037 + output_token_ids = self.rejection_sampler(
1038 + draft_token_ids,
1039 + None, # draft_probs
1040 + bonus_token_ids,
1041 + target_probs,
1042 + sampling_metadata)
1043 + sampler_output.sampled_token_ids = output_token_ids
1044
1045 # TODO(woosuk): The following loop can be slow since it iterates over
1046 # the requests one by one. Optimize.
@@ -1075,7 +1087,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1087 spec_token_ids = None
1088 else:
1089 spec_token_ids = self.generate_draft_token_ids(
1078 - valid_sampled_token_ids)
1090 + valid_sampled_token_ids, sampling_metadata)
1091
1092 return ModelRunnerOutput(
1093 req_ids=self.input_batch.req_ids,
@@ -1089,6 +1101,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1101 def generate_draft_token_ids(
1102 self,
1103 sampled_token_ids: list[list[int]],
1104 + sampling_metadata: SamplingMetadata,
1105 ) -> list[list[int]]:
1106 # TODO(woosuk): Optimize.
1107 draft_token_ids: list[list[int]] = []
@@ -1099,6 +1112,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1112 draft_token_ids.append([])
1113 continue
1114
1115 + # Skip requests that require top-p, top-k, etc.
1116 + req_id = self.input_batch.req_ids[i]
1117 + if not is_spec_decode_supported(req_id, self.input_batch):
1118 + draft_token_ids.append([])
1119 + continue
1120 +
1121 # Add sampled_token_ids to token_ids_cpu.
1122 start_idx = self.input_batch.num_tokens_no_spec[i]
1123 end_idx = start_idx + num_sampled_ids