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 9, 06:15 UTC
[perf] v1/spec_decode: skip softmax for all-greedy rejection sampling (#32852)
Signed-off-by: hdj <1293066020@qq.com>
caozuoba committed
Jan 31, 2026 at 17:51 UTC
8980001c9333c617f2bc3b2c60926b02434067a7
1 file changed
+11
-10
vllm/v1/sample/rejection_sampler.py
+11
-10
@@ -136,8 +136,6 @@ class RejectionSampler(nn.Module):
136
metadata.cu_num_draft_tokens,
137
sampling_metadata,
138
)
139
- # Compute probability distribution from target logits.
140
- target_probs = target_logits.softmax(dim=-1, dtype=torch.float32)
139
140
output_token_ids = rejection_sample(
141
metadata.draft_token_ids,
@@ -145,7 +143,7 @@ class RejectionSampler(nn.Module):
143
metadata.max_spec_len,
144
metadata.cu_num_draft_tokens,
145
draft_probs,
148
- target_probs,
146
+ target_logits,
147
bonus_token_ids,
148
sampling_metadata,
149
)
@@ -353,7 +351,7 @@ def rejection_sample(
351
# [num_tokens, vocab_size]
352
draft_probs: torch.Tensor | None,
353
# [num_tokens, vocab_size]
356
- target_probs: torch.Tensor,
354
+ target_logits: torch.Tensor,
355
# [batch_size, 1]
356
bonus_token_ids: torch.Tensor,
357
sampling_metadata: SamplingMetadata,
@@ -361,17 +359,16 @@ def rejection_sample(
359
assert draft_token_ids.ndim == 1
360
assert draft_probs is None or draft_probs.ndim == 2
361
assert cu_num_draft_tokens.ndim == 1
364
- assert target_probs.ndim == 2
362
+ assert target_logits.ndim == 2
363
364
batch_size = len(num_draft_tokens)
365
num_tokens = draft_token_ids.shape[0]
368
- vocab_size = target_probs.shape[-1]
369
- device = target_probs.device
366
+ vocab_size = target_logits.shape[-1]
367
+ device = target_logits.device
368
assert draft_token_ids.is_contiguous()
369
assert draft_probs is None or draft_probs.is_contiguous()
372
- assert target_probs.is_contiguous()
370
assert bonus_token_ids.is_contiguous()
374
- assert target_probs.shape == (num_tokens, vocab_size)
371
+ assert target_logits.shape == (num_tokens, vocab_size)
372
373
# Create output buffer.
374
output_token_ids = torch.full(
@@ -387,7 +384,7 @@ def rejection_sample(
384
is_greedy = sampling_metadata.temperature == GREEDY_TEMPERATURE
385
if not sampling_metadata.all_random:
386
# Rejection sampling for greedy sampling requests.
390
- target_argmax = target_probs.argmax(dim=-1)
387
+ target_argmax = target_logits.argmax(dim=-1)
388
rejection_greedy_sample_kernel[(batch_size,)](
389
output_token_ids,
390
cu_num_draft_tokens,
@@ -400,6 +397,10 @@ def rejection_sample(
397
if sampling_metadata.all_greedy:
398
return output_token_ids
399
400
+ # Compute probability distribution from target logits.
401
+ target_probs = target_logits.softmax(dim=-1, dtype=torch.float32)
402
+ assert target_probs.is_contiguous()
403
+
404
# Generate uniform probabilities for rejection sampling.
405
# [num_tokens]
406
uniform_probs = generate_uniform_probs(