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 14, 05:30 UTC
[Bugfix] fix apply_temperature to avoid nan in probs (#24734)
Signed-off-by: courage17340 <courage17340@163.com>
courage17340 committed
Sep 25, 2025 at 13:32 UTC
a676e668ee10585099816b60a440ec089fd391c4
2 files changed
+8
-3
vllm/v1/sample/sampler.py
+6
-1
@@ -128,8 +128,12 @@ class Sampler(nn.Module):
128
self,
129
logits: torch.Tensor,
130
temp: torch.Tensor,
131
+ all_random: bool,
132
) -> torch.Tensor:
133
# Use in-place division to avoid creating a new tensor.
134
+ # Avoid division by zero if there are greedy requests.
135
+ if not all_random:
136
+ temp = torch.where(temp < _SAMPLING_EPS, 1.0, temp)
137
return logits.div_(temp.unsqueeze(dim=1))
138
139
def greedy_sample(self, logits: torch.Tensor) -> torch.Tensor:
@@ -164,7 +168,8 @@ class Sampler(nn.Module):
168
assert sampling_metadata.temperature is not None
169
170
# Apply temperature.
167
- logits = self.apply_temperature(logits, sampling_metadata.temperature)
171
+ logits = self.apply_temperature(logits, sampling_metadata.temperature,
172
+ sampling_metadata.all_random)
173
174
# Apply logits processors that only apply to random sampling
175
# (argmax invariant)
vllm/v1/worker/gpu_input_batch.py
+2
-2
@@ -354,8 +354,8 @@ class InputBatch:
354
and is_spec_decode_unsupported(sampling_params)):
355
self.spec_decode_unsupported_reqs.add(req_id)
356
if sampling_params.sampling_type == SamplingType.GREEDY:
357
- # Avoid later division by zero.
358
- self.temperature_cpu[req_index] = -1.0
357
+ # Should avoid division by zero later when apply_temperature.
358
+ self.temperature_cpu[req_index] = 0.0
359
self.greedy_reqs.add(req_id)
360
else:
361
self.temperature_cpu[req_index] = sampling_params.temperature