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, 17:15 UTC
[V1][Spec Decode] Update target_logits in place for rejection sampling (#15427)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Mar 24, 2025 at 21:04 UTC
25f560a62c4f955672e2c6080b17ab3a48f96201
2 files changed
+12
-4
vllm/v1/sample/rejection_sampler.py
+5
-2
@@ -67,6 +67,7 @@ class RejectionSampler(nn.Module):
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
+ NOTE: `target_logits` can be updated in place to save memory.
71
bonus_token_ids_tensor (torch.Tensor):
72
A tensor containing bonus tokens. Shape is [batch_size, 1].
73
Bonus tokens are added to the end of the sequence if all
@@ -83,6 +84,8 @@ class RejectionSampler(nn.Module):
84
'''
85
assert metadata.max_spec_len <= MAX_SPEC_LEN
86
# [num_tokens, vocab_size]
87
+ # NOTE(woosuk): `target_logits` can be updated in place inside the
88
+ # `compute_probs` function.
89
target_probs = compute_probs(
90
target_logits,
91
metadata.cu_num_draft_tokens,
@@ -252,8 +255,8 @@ def compute_probs(
255
replace_from=GREEDY_TEMPERATURE,
256
replace_to=1,
257
)
255
- # TODO(woosuk): Consider using in-place op to reduce memory usage.
256
- logits = logits / temperature.unsqueeze(-1)
258
+ # NOTE(woosuk): Update `logits` in place to avoid allocating a new tensor.
259
+ logits.div_(temperature.unsqueeze(-1))
260
261
# Get expanded top_k and top_p tensors.
262
top_k = None
vllm/v1/worker/gpu_model_runner.py
+7
-2
@@ -1059,7 +1059,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1059
sampling_metadata=sampling_metadata,
1060
)
1061
else:
1062
- # TODO(woosuk): Optimize the memory usage.
1062
+ # When indexing with a tensor (bonus_logits_indices), PyTorch
1063
+ # creates a new tensor with separate storage from the original
1064
+ # logits tensor. This means any in-place operations on bonus_logits
1065
+ # won't affect the original logits tensor.
1066
bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
1067
sampler_output = self.model.sample(
1068
logits=bonus_logits,
@@ -1067,7 +1070,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1070
)
1071
bonus_token_ids = sampler_output.sampled_token_ids
1072
1070
- # TODO(woosuk): Optimize the memory usage.
1073
+ # Just like `bonus_logits`, `target_logits` is a new tensor with
1074
+ # separate storage from the original `logits` tensor. Therefore,
1075
+ # it is safe to update `target_logits` in place.
1076
target_logits = logits[spec_decode_metadata.target_logits_indices]
1077
output_token_ids = self.rejection_sampler(
1078
spec_decode_metadata,