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, 13:30 UTC
[V1][Spec Decode] Always use argmax for sampling draft tokens (#16899)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Apr 23, 2025 at 14:57 UTC
41fb013d299941486cb026a0e607d2f77a7ec1e8
3 files changed
+18
-23
vllm/v1/sample/rejection_sampler.py
+7
-7
@@ -226,7 +226,7 @@ def rejection_sample(
226
is_greedy,
227
max_spec_len,
228
vocab_size,
229
- IS_NGRAM=draft_probs is None,
229
+ NO_DRAFT_PROBS=draft_probs is None,
230
num_warps=1,
231
)
232
return output_token_ids
@@ -423,7 +423,7 @@ def sample_recovered_tokens(
423
q,
424
vocab_size,
425
triton.next_power_of_2(vocab_size),
426
- IS_NGRAM=draft_probs is None,
426
+ NO_DRAFT_PROBS=draft_probs is None,
427
)
428
return recovered_token_ids
429
@@ -490,7 +490,7 @@ def rejection_random_sample_kernel(
490
is_greedy_ptr, # [batch_size]
491
max_spec_len,
492
vocab_size,
493
- IS_NGRAM: tl.constexpr,
493
+ NO_DRAFT_PROBS: tl.constexpr,
494
):
495
req_idx = tl.program_id(0)
496
is_greedy = tl.load(is_greedy_ptr + req_idx)
@@ -509,7 +509,7 @@ def rejection_random_sample_kernel(
509
for pos in range(num_draft_tokens):
510
if not rejected:
511
draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
512
- if IS_NGRAM:
512
+ if NO_DRAFT_PROBS:
513
draft_prob = 1
514
else:
515
draft_prob = tl.load(draft_probs_ptr +
@@ -575,7 +575,7 @@ def sample_recovered_tokens_kernel(
575
q_ptr, # [batch_size, vocab_size]
576
vocab_size,
577
PADDED_VOCAB_SIZE: tl.constexpr,
578
- IS_NGRAM: tl.constexpr,
578
+ NO_DRAFT_PROBS: tl.constexpr,
579
):
580
req_idx = tl.program_id(0)
581
if req_idx == 0:
@@ -591,7 +591,7 @@ def sample_recovered_tokens_kernel(
591
return
592
593
vocab_offset = tl.arange(0, PADDED_VOCAB_SIZE)
594
- if IS_NGRAM:
594
+ if NO_DRAFT_PROBS:
595
draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
596
orig_prob = tl.load(target_probs_ptr + (start_idx + pos) * vocab_size +
597
draft_token_id)
@@ -624,7 +624,7 @@ def sample_recovered_tokens_kernel(
624
recovered_id = tl.argmax(prob / q, axis=-1)
625
tl.store(output_token_ids_ptr + start_idx + pos, recovered_id)
626
627
- if IS_NGRAM:
627
+ if NO_DRAFT_PROBS:
628
# Restore the original probability.
629
tl.store(
630
target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id,
vllm/v1/spec_decode/eagle.py
+10
-12
@@ -51,7 +51,7 @@ class EagleProposer:
51
# [batch_size, max_num_blocks_per_req]
52
block_table: torch.Tensor,
53
sampling_metadata: SamplingMetadata,
54
- ) -> tuple[torch.Tensor, torch.Tensor]:
54
+ ) -> torch.Tensor:
55
num_tokens = target_token_ids.shape[0]
56
batch_size = next_token_ids.shape[0]
57
last_token_indices = cu_num_tokens[1:] - 1
@@ -94,17 +94,15 @@ class EagleProposer:
94
)
95
sample_hidden_states = hidden_states[last_token_indices]
96
logits = self.model.compute_logits(sample_hidden_states, None)
97
- draft_token_ids, draft_probs = compute_probs_and_sample_next_token(
98
- logits, sampling_metadata)
97
+ draft_token_ids = logits.argmax(dim=-1)
98
99
# Early exit if there is only one draft token to be generated.
100
if self.num_speculative_tokens == 1:
102
- # [batch_size, 1] and [batch_size, 1, vocab_size]
103
- return draft_token_ids.view(-1, 1), draft_probs.unsqueeze(dim=1)
101
+ # [batch_size, 1]
102
+ return draft_token_ids.view(-1, 1)
103
104
# Generate the remaining draft tokens.
105
draft_token_ids_list = [draft_token_ids]
107
- draft_probs_list = [draft_probs]
106
107
positions = target_positions[last_token_indices]
108
hidden_states = sample_hidden_states
@@ -159,16 +157,12 @@ class EagleProposer:
157
positions=clamped_positions,
158
)
159
logits = self.model.compute_logits(hidden_states, None)
162
- draft_token_ids, probs = compute_probs_and_sample_next_token(
163
- logits, sampling_metadata)
160
+ draft_token_ids = logits.argmax(dim=-1)
161
draft_token_ids_list.append(draft_token_ids)
165
- draft_probs_list.append(probs)
162
163
# [batch_size, num_speculative_tokens]
164
draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
169
- # [batch_size, num_speculative_tokens, vocab_size]
170
- draft_probs = torch.stack(draft_probs_list, dim=1)
171
- return draft_token_ids, draft_probs
165
+ return draft_token_ids
166
167
@staticmethod
168
def prepare_inputs(
@@ -238,6 +232,10 @@ class EagleProposer:
232
self.model.lm_head = target_model.lm_head
233
234
235
+# NOTE(woosuk): Currently, the below code is not used and we always use argmax
236
+# to sample the draft tokens. We will use this after we find a way to manage
237
+# the draft prob tensor.
238
+# Refer to https://github.com/vllm-project/vllm/pull/16899 for the details.
239
# FIXME(woosuk): The logic here is duplicated with the main sampling code.
240
# We should refactor this to reuse the same sampling implementation.
241
def compute_probs_and_sample_next_token(
vllm/v1/worker/gpu_model_runner.py
+1
-4
@@ -1230,7 +1230,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1230
target_hidden_states = hidden_states[token_indices]
1231
target_slot_mapping = attn_metadata.slot_mapping[token_indices]
1232
1233
- draft_token_ids, draft_probs = self.drafter.propose(
1233
+ draft_token_ids = self.drafter.propose(
1234
target_token_ids=target_token_ids,
1235
target_positions=target_positions,
1236
target_hidden_states=target_hidden_states,
@@ -1241,9 +1241,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1241
sampling_metadata=sampling_metadata,
1242
)
1243
spec_token_ids = draft_token_ids.tolist()
1244
- # TODO(woosuk): Cache draft_probs and use it for rejection sampling
1245
- # in the next step.
1246
- del draft_probs
1244
1245
# Clear KVConnector state after all KVs are generated.
1246
if has_kv_transfer_group():