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 12, 16:00 UTC
[Model Runner V2] Add apply_temperature option to gumbel_sample (#29276)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 23, 2025 at 14:13 UTC
3e1ad406559c3b520eeda0e681ea68d33daf1be1
1 file changed
+15
-6
vllm/v1/worker/gpu/sampler.py
+15
-6
@@ -68,9 +68,10 @@ class Sampler:
68
69
sampled = gumbel_sample(
70
logits,
71
- is_greedy,
71
+ sampling_metadata.temperature,
72
sampling_metadata.seeds,
73
sampling_metadata.pos,
74
+ apply_temperature=False,
75
)
76
return sampled, logits if return_logits else None
77
@@ -85,9 +86,10 @@ def _gumbel_sample_kernel(
86
logits_stride,
87
seeds_ptr,
88
pos_ptr,
88
- is_greedy_ptr,
89
+ temp_ptr,
90
vocab_size,
91
BLOCK_SIZE: tl.constexpr,
92
+ APPLY_TEMPERATURE: tl.constexpr,
93
):
94
req_idx = tl.program_id(0)
95
block_idx = tl.program_id(1)
@@ -99,8 +101,8 @@ def _gumbel_sample_kernel(
101
other=float("-inf"),
102
)
103
102
- is_greedy = tl.load(is_greedy_ptr + req_idx)
103
- if not is_greedy:
104
+ temp = tl.load(temp_ptr + req_idx)
105
+ if temp != 0.0:
106
# Calculate the seed for gumbel noise.
107
seed = tl.load(seeds_ptr + req_idx)
108
pos = tl.load(pos_ptr + req_idx)
@@ -111,6 +113,11 @@ def _gumbel_sample_kernel(
113
gumbel_noise = -tl.log(-tl.log(r + 1e-20) + 1e-20)
114
gumbel_noise = gumbel_noise.to(tl.float32)
115
116
+ # Apply temperature.
117
+ if APPLY_TEMPERATURE:
118
+ # NOTE(woosuk): Use div_rn to match the behavior of torch.
119
+ logits = tl.div_rn(logits, temp.to(tl.float32))
120
+
121
# Apply gumbel noise.
122
logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
123
@@ -123,9 +130,10 @@ def _gumbel_sample_kernel(
130
131
def gumbel_sample(
132
logits: torch.Tensor, # [num_reqs, vocab_size]
126
- is_greedy: torch.Tensor, # [num_reqs]
133
+ temperature: torch.Tensor, # [num_reqs]
134
seed: torch.Tensor, # [num_reqs]
135
pos: torch.Tensor, # [num_reqs]
136
+ apply_temperature: bool,
137
) -> torch.Tensor:
138
num_reqs, vocab_size = logits.shape
139
BLOCK_SIZE = 1024
@@ -151,9 +159,10 @@ def gumbel_sample(
159
logits.stride(0),
160
seed,
161
pos,
154
- is_greedy,
162
+ temperature,
163
vocab_size,
164
BLOCK_SIZE=BLOCK_SIZE,
165
+ APPLY_TEMPERATURE=apply_temperature,
166
)
167
# NOTE(woosuk): Use int64 for later indexing.
168
max_block_idx = local_max.argmax(dim=-1, keepdim=True)