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][Sampler] Avoid an operation during temperature application (#13587)
Nick Hill committed
Feb 20, 2025 at 22:05 UTC
31aa045c11c996a78dc10fdda60a60fda6788c2a
4 files changed
+18
-10
vllm/v1/sample/metadata.py
+1
-1
@@ -9,7 +9,7 @@ import torch
9
@dataclass
10
class SamplingMetadata:
11
12
- temperature: torch.Tensor
12
+ temperature: Optional[torch.Tensor]
13
all_greedy: bool
14
all_random: bool
15
vllm/v1/sample/sampler.py
+4
-4
@@ -77,11 +77,8 @@ class Sampler(nn.Module):
77
logits: torch.Tensor,
78
temp: torch.Tensor,
79
) -> torch.Tensor:
80
- # Avoid division by zero.
81
- temp = torch.where(temp < _SAMPLING_EPS, 1.0, temp)
80
# Use in-place division to avoid creating a new tensor.
83
- logits.div_(temp.unsqueeze(dim=1))
84
- return logits
81
+ return logits.div_(temp.unsqueeze(dim=1))
82
83
def greedy_sample(self, logits: torch.Tensor) -> torch.Tensor:
84
return logits.argmax(dim=-1).view(-1)
@@ -100,6 +97,8 @@ class Sampler(nn.Module):
97
if sampling_metadata.all_greedy:
98
return greedy_sampled
99
100
+ assert sampling_metadata.temperature is not None
101
+
102
# Apply temperature.
103
logits = self.apply_temperature(logits, sampling_metadata.temperature)
104
@@ -122,6 +121,7 @@ class Sampler(nn.Module):
121
sampling_metadata.temperature < _SAMPLING_EPS,
122
greedy_sampled,
123
random_sampled,
124
+ out=greedy_sampled, # Reuse tensor
125
)
126
return sampled
127
vllm/v1/utils.py
+4
-2
@@ -191,11 +191,13 @@ def bind_kv_cache(
191
192
193
def copy_slice(from_tensor: torch.Tensor, to_tensor: torch.Tensor,
194
- length: int) -> None:
194
+ length: int) -> torch.Tensor:
195
"""
196
Copy the first length elements of a tensor into another tensor in a
197
non-blocking manner.
198
199
Used to copy pinned CPU tensor data to pre-allocated GPU tensors.
200
+
201
+ Returns the sliced target tensor.
202
"""
201
- to_tensor[:length].copy_(from_tensor[:length], non_blocking=True)
203
+ return to_tensor[:length].copy_(from_tensor[:length], non_blocking=True)
vllm/v1/worker/gpu_input_batch.py
+9
-3
@@ -242,10 +242,12 @@ class InputBatch:
242
self.block_table.add_row(req_index, request.block_ids)
243
244
sampling_params = request.sampling_params
245
- self.temperature_cpu[req_index] = sampling_params.temperature
245
if sampling_params.sampling_type == SamplingType.GREEDY:
246
+ # Avoid later division by zero.
247
+ self.temperature_cpu[req_index] = -1.0
248
self.greedy_reqs.add(req_id)
249
else:
250
+ self.temperature_cpu[req_index] = sampling_params.temperature
251
self.random_reqs.add(req_id)
252
253
self.top_p_cpu[req_index] = sampling_params.top_p
@@ -410,7 +412,11 @@ class InputBatch:
412
413
def _make_sampling_metadata(self) -> SamplingMetadata:
414
num_reqs = self.num_reqs
413
- copy_slice(self.temperature_cpu_tensor, self.temperature, num_reqs)
415
+ if not self.all_greedy:
416
+ temperature = copy_slice(self.temperature_cpu_tensor,
417
+ self.temperature, num_reqs)
418
+ else:
419
+ temperature = None
420
if not self.no_top_p:
421
copy_slice(self.top_p_cpu_tensor, self.top_p, num_reqs)
422
if not self.no_top_k:
@@ -437,7 +443,7 @@ class InputBatch:
443
prompt_token_ids = None
444
445
return SamplingMetadata(
440
- temperature=self.temperature[:num_reqs],
446
+ temperature=temperature,
447
all_greedy=self.all_greedy,
448
all_random=self.all_random,
449
top_p=None if self.no_top_p else self.top_p[:num_reqs],