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
[V1][Spec Decode] Fix greedy temperature detection after sampler refactor (#27077)
Signed-off-by: Pradyun Ramadorai <pradyunr@amazon.com> Co-authored-by: Pradyun Ramadorai <pradyunr@amazon.com>
Pradyun92 committed
Oct 17, 2025 at 16:27 UTC
acedc74b1a82ac75144868894bc43ab4cb723252
5 files changed
+22
-6
vllm/v1/sample/rejection_sampler.py
+1
-1
@@ -15,7 +15,7 @@ from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
15
logger = init_logger(__name__)
16
17
PLACEHOLDER_TOKEN_ID: tl.constexpr = -1
18
-GREEDY_TEMPERATURE: tl.constexpr = -1
18
+GREEDY_TEMPERATURE: tl.constexpr = 0
19
# Maximum number of speculative draft tokens allowed per request in a single
20
# step. This value is chosen to be large enough to handle typical use cases.
21
MAX_SPEC_LEN = 128
vllm/v1/sample/tpu/metadata.py
+2
@@ -30,6 +30,7 @@ class TPUSupportedSamplingMetadata:
30
top_p: torch.Tensor = None
31
32
all_greedy: bool = True
33
+ all_random: bool = False
34
35
# Whether logprobs are to be gathered in this batch of request. To balance
36
# out compile time and runtime, a fixed `max_number_logprobs` value is used
@@ -110,6 +111,7 @@ class TPUSupportedSamplingMetadata:
111
xla_device
112
),
113
all_greedy=input_batch.all_greedy,
114
+ all_random=input_batch.all_random,
115
# TODO enable more and avoid returning None values
116
top_p=input_batch.top_p_cpu_tensor[:padded_num_reqs].to(xla_device),
117
top_k=input_batch.top_k_cpu_tensor[:padded_num_reqs].to(xla_device),
vllm/v1/sample/tpu/sampler.py
+7
-1
@@ -40,7 +40,11 @@ class Sampler(nn.Module):
40
self,
41
logits: torch.Tensor,
42
temp: torch.Tensor,
43
+ all_random: bool = False,
44
) -> torch.Tensor:
45
+ # Avoid division by zero for greedy sampling (temperature ~ 0.0).
46
+ if not all_random:
47
+ temp = torch.where(temp < _SAMPLING_EPS, 1.0, temp)
48
return logits.div_(temp.unsqueeze(dim=1))
49
50
def greedy_sample(self, logits: torch.Tensor) -> torch.Tensor:
@@ -56,7 +60,9 @@ class Sampler(nn.Module):
60
assert sampling_metadata.temperature is not None
61
62
# Apply temperature.
59
- logits = self.apply_temperature(logits, sampling_metadata.temperature)
63
+ logits = self.apply_temperature(
64
+ logits, sampling_metadata.temperature, sampling_metadata.all_random
65
+ )
66
67
# Apply min_p.
68
if sampling_metadata.min_p is not None:
vllm/v1/spec_decode/eagle.py
+10
-2
@@ -37,6 +37,7 @@ from vllm.v1.attention.backends.utils import (
37
)
38
from vllm.v1.kv_cache_interface import KVCacheConfig
39
from vllm.v1.sample.metadata import SamplingMetadata
40
+from vllm.v1.sample.sampler import _SAMPLING_EPS
41
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
42
from vllm.v1.utils import CpuGpuBuffer
43
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -1140,8 +1141,15 @@ def compute_probs_and_sample_next_token(
1141
next_token_ids = logits.argmax(dim=-1)
1142
return next_token_ids, probs
1143
1143
- is_greedy = sampling_metadata.temperature == -1
1144
- temperature = torch.where(is_greedy, 1.0, sampling_metadata.temperature)
1144
+ assert sampling_metadata.temperature is not None
1145
+
1146
+ # Use epsilon comparison to detect greedy sampling (temperature ~ 0.0)
1147
+ # consistent with sampler.py's _SAMPLING_EPS threshold
1148
+ temperature = sampling_metadata.temperature
1149
+ # Avoid division by zero if there are greedy requests.
1150
+ if not sampling_metadata.all_random:
1151
+ is_greedy = temperature < _SAMPLING_EPS
1152
+ temperature = torch.where(is_greedy, 1.0, temperature)
1153
logits.div_(temperature.view(-1, 1))
1154
probs = logits.softmax(dim=-1, dtype=torch.float32)
1155
vllm/v1/worker/tpu_input_batch.py
+2
-2
@@ -215,8 +215,8 @@ class InputBatch:
215
sampling_params = request.sampling_params
216
assert sampling_params is not None, "pooling requests not supported yet"
217
if sampling_params.sampling_type == SamplingType.GREEDY:
218
- # Avoid later division by zero.
219
- self.temperature_cpu[req_index] = -1.0
218
+ # Should avoid division by zero later when apply_temperature.
219
+ self.temperature_cpu[req_index] = 0.0
220
self.greedy_reqs.add(req_id)
221
else:
222
self.temperature_cpu[req_index] = sampling_params.temperature