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, 22:15 UTC
[V1][Spec Decode] Change Spec Decode Rejection Sampling API (#13729)
Lily Liu committed
Feb 25, 2025 at 18:14 UTC
5629f26df7c0fe5126d90463e14292f4a0552a65
8 files changed
+104
-111
tests/v1/sample/test_rejection_sampler.py
+8
-9
@@ -29,7 +29,6 @@ def create_sampling_metadata(spec_tokens: List[List[int]]) -> SamplingMetadata:
29
temperature=torch.tensor([]),
30
all_greedy=True,
31
all_random=False,
32
- spec_token_ids=spec_tokens,
32
top_p=None,
33
top_k=None,
34
min_p=torch.empty(batch_size, ),
@@ -55,7 +54,7 @@ def test_perfect_match(sampler):
54
metadata = create_sampling_metadata(spec_tokens)
55
logits = create_logits_tensor(output_tokens)
56
58
- output = sampler(logits, metadata)
57
+ output = sampler(spec_tokens, logits, metadata)
58
expected = torch.tensor([[1, 2, 3, 4]],
59
dtype=torch.int,
60
device=logits.device)
@@ -70,7 +69,7 @@ def test_early_mismatch(sampler):
69
metadata = create_sampling_metadata(spec_tokens)
70
logits = create_logits_tensor(output_tokens)
71
73
- output = sampler(logits, metadata)
72
+ output = sampler(spec_tokens, logits, metadata)
73
expected = torch.tensor([[1, 5, INVALID_TOKEN_ID, INVALID_TOKEN_ID]],
74
dtype=torch.int,
75
device=logits.device)
@@ -85,7 +84,7 @@ def test_multiple_sequences(sampler):
84
metadata = create_sampling_metadata(spec_tokens)
85
logits = create_logits_tensor(output_tokens)
86
88
- output = sampler(logits, metadata)
87
+ output = sampler(spec_tokens, logits, metadata)
88
expected = torch.tensor([[1, 2, 5], [3, 4, INVALID_TOKEN_ID]],
89
dtype=torch.int,
90
device=logits.device)
@@ -100,7 +99,7 @@ def test_single_token_sequence(sampler):
99
metadata = create_sampling_metadata(spec_tokens)
100
logits = create_logits_tensor(output_tokens)
101
103
- output = sampler(logits, metadata)
102
+ output = sampler(spec_tokens, logits, metadata)
103
expected = torch.tensor([[1, 2]], dtype=torch.int, device=logits.device)
104
assert torch.equal(output.sampled_token_ids, expected)
105
@@ -113,7 +112,7 @@ def test_empty_sequence(sampler):
112
metadata = create_sampling_metadata(spec_tokens)
113
logits = create_logits_tensor(output_tokens)
114
116
- output = sampler(logits, metadata)
115
+ output = sampler(spec_tokens, logits, metadata)
116
expected = torch.tensor([[5]], dtype=torch.int, device=logits.device)
117
assert torch.equal(output.sampled_token_ids, expected)
118
@@ -126,7 +125,7 @@ def test_multiple_mismatches(sampler):
125
metadata = create_sampling_metadata(spec_tokens)
126
logits = create_logits_tensor(output_tokens)
127
129
- output = sampler(logits, metadata)
128
+ output = sampler(spec_tokens, logits, metadata)
129
expected = torch.tensor([[1, 2, 7, INVALID_TOKEN_ID],
130
[4, 8, INVALID_TOKEN_ID, INVALID_TOKEN_ID]],
131
dtype=torch.int,
@@ -147,7 +146,7 @@ def test_parametrized_cases(sampler, spec_tokens, output_tokens, expected):
146
metadata = create_sampling_metadata(spec_tokens)
147
logits = create_logits_tensor(output_tokens)
148
150
- output = sampler(logits, metadata)
149
+ output = sampler(spec_tokens, logits, metadata)
150
expected_tensor = torch.tensor(expected,
151
dtype=torch.int,
152
device=logits.device)
@@ -163,7 +162,7 @@ def test_logits_shape_handling(sampler):
162
metadata = create_sampling_metadata(spec_tokens)
163
logits = create_logits_tensor(output_tokens, vocab_size)
164
166
- output = sampler(logits, metadata)
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
tests/v1/sample/test_sampler.py
-1
@@ -105,7 +105,6 @@ def _create_default_sampling_metadata(
105
prompt_token_ids=_create_prompt_tokens_tensor(prompt_token_ids,
106
vocab_size, device),
107
output_token_ids=output_token_ids,
108
- spec_token_ids=None,
108
frequency_penalties=_create_penalty_tensor(batch_size, 0.0, device),
109
presence_penalties=_create_penalty_tensor(batch_size, 0.0, device),
110
repetition_penalties=_create_penalty_tensor(batch_size, 1.0, device),
tests/v1/worker/test_gpu_input_batch.py
-1
@@ -123,7 +123,6 @@ def _construct_expected_sampling_metadata(
123
dtype=torch.float,
124
device=device),
125
output_token_ids=output_token_ids,
126
- spec_token_ids=None,
126
min_tokens=min_tokens,
127
no_penalties=(all(x == 0 for x in presence_penalties)
128
and all(x == 0 for x in frequency_penalties)
vllm/v1/sample/metadata.py
-3
@@ -13,9 +13,6 @@ class SamplingMetadata:
13
all_greedy: bool
14
all_random: bool
15
16
- # None when there are no speculated tokens.
17
- spec_token_ids: Optional[List[List[int]]]
18
-
16
top_p: Optional[torch.Tensor]
17
top_k: Optional[torch.Tensor]
18
min_p: Optional[torch.Tensor]
vllm/v1/sample/rejection_sampler.py
+68
-66
@@ -1,4 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+from typing import List
3
+
4
import torch
5
import torch.nn as nn
6
from torch.nn.utils.rnn import pad_sequence
@@ -52,62 +54,62 @@ class RejectionSampler(nn.Module):
54
else:
55
self.forward_method = self.forward_native
56
55
- def forward(self, logits: torch.Tensor,
57
+ def forward(self, draft_token_ids: List[List[int]],
58
+ target_probs: torch.Tensor,
59
sampling_metadata: SamplingMetadata) -> SamplerOutput:
60
if not sampling_metadata.all_greedy:
61
raise NotImplementedError(
62
"Currently, only greedy sampling is supported by "
63
"rejection sampler.")
61
- return self.forward_method(logits, sampling_metadata)
64
+ return self.forward_method(draft_token_ids, target_probs,
65
+ sampling_metadata)
66
67
def flashinfer_sample(
68
self,
65
- logits: torch.Tensor,
69
+ draft_token_ids: List[List[int]],
70
+ target_probs: torch.Tensor,
71
sampling_metadata: SamplingMetadata,
72
) -> SamplerOutput:
73
# NOTE: The following input preparationg can be moved
74
# to the model runner with a persistent manner for better
75
# performance.
71
- assert sampling_metadata.spec_token_ids is not None
72
- spec_token_ids = sampling_metadata.spec_token_ids
73
- max_spec_len = max(len(s) for s in spec_token_ids)
74
- batch_size = len(spec_token_ids)
75
- draft_token_ids = torch.full((batch_size, max_spec_len),
76
- INVALID_TOKEN_ID,
77
- device="cpu",
78
- dtype=torch.long)
79
-
80
- target_token_ids = torch.full((batch_size, max_spec_len + 1),
81
- fill_value=INVALID_TOKEN_ID,
82
- device=logits.device,
83
- dtype=torch.long)
84
-
85
- # TODO: Vectorize the following loop for better performance.
86
- start_loc = 0
87
- for i in range(batch_size):
88
- num_spec_tokens = len(spec_token_ids[i])
89
- draft_token_ids[i, :num_spec_tokens] = torch.tensor(
90
- spec_token_ids[i], device="cpu", dtype=torch.long)
91
- end_loc = start_loc + num_spec_tokens + 1
92
- # Assume greedy sampling.
93
- target_token_ids[i, :num_spec_tokens + 1] = torch.argmax(
94
- logits[start_loc:end_loc], dim=-1)
95
- start_loc = end_loc
96
-
97
- vocab_size = logits.size(-1)
98
- # NOTE: CPU <-> GPU synchronization happens here.
99
- draft_token_ids = draft_token_ids.to(logits.device)
100
- draft_probs = _create_greedy_token_probs(draft_token_ids, vocab_size,
101
- logits.device)
102
- target_probs = _create_greedy_token_probs(target_token_ids, vocab_size,
103
- logits.device)
104
- uniform_samples = torch.zeros(batch_size,
105
- max_spec_len + 1,
106
- device=logits.device)
76
+ sample_lens = [len(x) + 1 for x in draft_token_ids]
77
+ # Convert draft token IDs to a tensor, split by sample_lens, then pad.
78
+ draft_token_ids = [
79
+ torch.tensor(x, dtype=int, device='cpu') for x in draft_token_ids
80
+ ]
81
+ draft_token_ids_tensor = pad_sequence(draft_token_ids,
82
+ batch_first=True,
83
+ padding_value=INVALID_TOKEN_ID)
84
+
85
+ if sampling_metadata.all_greedy:
86
+ target_token_ids = target_probs.argmax(dim=-1).view(-1)
87
+ target_token_ids = target_token_ids.split(sample_lens)
88
+ target_token_ids = pad_sequence(target_token_ids,
89
+ batch_first=True,
90
+ padding_value=INVALID_TOKEN_ID)
91
+
92
+ vocab_size = target_probs.size(-1)
93
+ # NOTE: CPU <-> GPU synchronization happens here.
94
+ draft_token_ids_tensor = draft_token_ids_tensor.to(
95
+ target_probs.device)
96
+ draft_probs = _create_greedy_token_probs(draft_token_ids_tensor,
97
+ vocab_size,
98
+ target_probs.device)
99
+ target_probs = _create_greedy_token_probs(target_token_ids,
100
+ vocab_size,
101
+ target_probs.device)
102
+ uniform_samples = torch.zeros(draft_token_ids_tensor.size(0),
103
+ draft_token_ids_tensor.size(1) + 1,
104
+ device=target_probs.device)
105
+ else:
106
+ raise NotImplementedError(
107
+ "Currently, only greedy sampling is supported by "
108
+ "rejection sampler.")
109
110
sampled_token_ids, _, _ = fs.chain_speculative_sampling(
111
draft_probs,
110
- draft_token_ids,
112
+ draft_token_ids_tensor,
113
uniform_samples,
114
target_probs,
115
)
@@ -117,35 +119,35 @@ class RejectionSampler(nn.Module):
119
# TODO: The following method can be optimized for better performance.
120
def forward_native(
121
self,
120
- logits: torch.Tensor,
122
+ draft_token_ids: List[List[int]],
123
+ target_probs: torch.Tensor,
124
sampling_metadata: SamplingMetadata,
125
) -> SamplerOutput:
123
- assert sampling_metadata.spec_token_ids is not None
124
- spec_lens = [len(x) for x in sampling_metadata.spec_token_ids]
125
- # Add 1 to include the 'bonus' token.
126
- sample_lens = [x + 1 for x in spec_lens]
127
-
128
- output_token_ids = logits.argmax(dim=-1).view(-1)
129
- output_token_ids = output_token_ids.split(sample_lens)
130
- output_token_ids = pad_sequence(output_token_ids,
131
- batch_first=True,
132
- padding_value=INVALID_TOKEN_ID)
133
-
134
- # Convert spec token IDs to a tensor, split by sample_lens, then pad.
135
- spec_token_ids = [
136
- torch.tensor(x,
137
- dtype=output_token_ids.dtype,
138
- device=output_token_ids.device)
139
- for x in sampling_metadata.spec_token_ids
126
+ sample_lens = [len(x) + 1 for x in draft_token_ids]
127
+ # Convert draft token IDs to a tensor, split by sample_lens, then pad.
128
+ draft_token_ids = [
129
+ torch.tensor(x, dtype=int, device='cpu') for x in draft_token_ids
130
]
141
- spec_token_ids = pad_sequence(spec_token_ids,
142
- batch_first=True,
143
- padding_value=INVALID_TOKEN_ID)
144
-
145
- # Produce a mask that remains 1 (True) until the first
146
- # mismatch (cumprod turns 0 after a mismatch).
147
- accept_mask = (output_token_ids[:, :-1] == spec_token_ids).cumprod(
148
- dim=1)
131
+ draft_token_ids_tensor = pad_sequence(draft_token_ids,
132
+ batch_first=True,
133
+ padding_value=INVALID_TOKEN_ID)
134
+ draft_token_ids_tensor = draft_token_ids_tensor.to(target_probs.device)
135
+ # Add 1 to include the 'bonus' token.
136
+ if sampling_metadata.all_greedy:
137
+ output_token_ids = target_probs.argmax(dim=-1).view(-1)
138
+ output_token_ids = output_token_ids.split(sample_lens)
139
+ output_token_ids = pad_sequence(output_token_ids,
140
+ batch_first=True,
141
+ padding_value=INVALID_TOKEN_ID)
142
+ # Produce a mask that remains 1 (True) until the first
143
+ # mismatch (cumprod turns 0 after a mismatch).
144
+ accept_mask = (
145
+ output_token_ids[:, :-1] == draft_token_ids_tensor).cumprod(
146
+ dim=1)
147
+ else:
148
+ raise NotImplementedError(
149
+ "Currently, only greedy sampling is supported by "
150
+ "rejection sampler.")
151
# Identify valid positions (non-padding).
152
valid_mask = output_token_ids != INVALID_TOKEN_ID
153
# Generate mask with bonus token.
vllm/v1/sample/sampler.py
+8
-11
@@ -9,7 +9,6 @@ from vllm.v1.sample.metadata import SamplingMetadata
9
from vllm.v1.sample.ops.penalties import (apply_all_penalties,
10
apply_min_token_penalties)
11
from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
12
-from vllm.v1.sample.rejection_sampler import RejectionSampler
12
13
_SAMPLING_EPS = 1e-5
14
@@ -19,22 +18,12 @@ class Sampler(nn.Module):
18
def __init__(self):
19
super().__init__()
20
self.topk_topp_sampler = TopKTopPSampler()
22
- self.rejection_sampler = RejectionSampler()
21
22
def forward(
23
self,
24
logits: torch.Tensor,
25
sampling_metadata: SamplingMetadata,
26
) -> SamplerOutput:
29
- if sampling_metadata.spec_token_ids:
30
- if sampling_metadata.max_num_logprobs:
31
- raise NotImplementedError(
32
- "Rejection sampling does not support logprobs.")
33
- return self.rejection_sampler(
34
- logits,
35
- sampling_metadata,
36
- )
37
-
27
# NOTE(woosuk): Use the original logits (before any penalties or
28
# temperature scaling) for the top-k logprobs.
29
# This is different from the V0 sampler, which uses the logits that
@@ -127,6 +116,14 @@ class Sampler(nn.Module):
116
)
117
return sampled
118
119
+ def compute_probs(self, logits: torch.Tensor,
120
+ sampling_metadata: SamplingMetadata) -> torch.Tensor:
121
+ if sampling_metadata.all_greedy:
122
+ return logits
123
+ # Apply temperature. This is an in-place op changing logits.
124
+ logits = self.apply_temperature(logits, sampling_metadata.temperature)
125
+ return logits.softmax(dim=-1, dtype=torch.float32)
126
+
127
def compute_logprobs(self, logits: torch.Tensor) -> torch.Tensor:
128
return logits.log_softmax(dim=-1, dtype=torch.float32)
129
vllm/v1/worker/gpu_input_batch.py
-11
@@ -490,23 +490,12 @@ class InputBatch:
490
presence_penalties=self.presence_penalties[:num_reqs],
491
repetition_penalties=self.repetition_penalties[:num_reqs],
492
output_token_ids=cast(List[List[int]], self.req_output_token_ids),
493
- spec_token_ids=None,
493
min_tokens=self.min_tokens,
494
no_penalties=self.no_penalties,
495
logit_bias=self.logit_bias[:num_reqs],
496
allowed_token_ids_mask=allowed_token_ids_mask,
497
)
498
500
- def get_sampling_metadata(
501
- self,
502
- req_id_to_spec_token_ids: Dict[str, List[int]],
503
- ) -> SamplingMetadata:
504
- # Set the new spec token ids in the cached sampling metadata.
505
- self.sampling_metadata.spec_token_ids = [
506
- req_id_to_spec_token_ids.get(req_id, []) for req_id in self.req_ids
507
- ] if req_id_to_spec_token_ids else None
508
- return self.sampling_metadata
509
-
499
def _make_prompt_token_ids_tensor(self) -> torch.Tensor:
500
max_prompt_len = self.num_prompt_tokens[:self.num_reqs].max()
501
prompt_token_ids_cpu_tensor = torch.empty(
vllm/v1/worker/gpu_model_runner.py
+20
-9
@@ -32,7 +32,7 @@ from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
32
KVCacheSpec)
33
from vllm.v1.outputs import LogprobsTensors, ModelRunnerOutput
34
from vllm.v1.sample.metadata import SamplingMetadata
35
-from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID
35
+from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID, RejectionSampler
36
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
37
from vllm.v1.utils import bind_kv_cache
38
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -122,7 +122,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
122
self.use_spec_decode = False
123
if self.speculative_config:
124
self.use_spec_decode = True
125
-
125
+ self.rejection_sampler = RejectionSampler()
126
# TODO: find a better way to check if we are using ngram.
127
assert self.speculative_config.ngram_prompt_lookup_min, \
128
"Currently, only ngram spec decode is supported in V1."
@@ -951,12 +951,24 @@ class GPUModelRunner(LoRAModelRunnerMixin):
951
logits = self.model.compute_logits(sample_hidden_states, None)
952
953
# Sample the next token and get logprobs if needed.
954
- sampling_metadata = self.input_batch.get_sampling_metadata(
955
- scheduler_output.scheduled_spec_decode_tokens)
956
- sampler_output = self.model.sample(
957
- logits=logits,
958
- sampling_metadata=sampling_metadata,
959
- )
954
+ sampling_metadata = self.input_batch.sampling_metadata
955
+ if not self.use_spec_decode:
956
+ sampler_output = self.model.sample(
957
+ logits=logits,
958
+ sampling_metadata=sampling_metadata,
959
+ )
960
+ else:
961
+ target_probs = self.model.sampler.compute_probs(
962
+ logits, sampling_metadata)
963
+ scheduled_request_ids = scheduler_output.num_scheduled_tokens.keys(
964
+ )
965
+ draft_token_ids = [
966
+ scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])
967
+ for req_id in scheduled_request_ids
968
+ ]
969
+ sampler_output = self.rejection_sampler(draft_token_ids,
970
+ target_probs,
971
+ sampling_metadata)
972
973
# TODO(woosuk): The following loop can be slow since it iterates over
974
# the requests one by one. Optimize.
@@ -1293,7 +1305,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1305
temperature=dummy_tensors(0.5),
1306
all_greedy=False,
1307
all_random=False,
1296
- spec_token_ids=None,
1308
top_p=dummy_tensors(0.9),
1309
top_k=dummy_tensors(logits.size(1) - 1),
1310
min_p=None,