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 9, 23:00 UTC
[Model Runner V2] Decouple temperature from penalties (#32629)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 19, 2026 at 19:13 UTC
6c01ffb89712fc2d18c39034b86b0991ef0076e5
3 files changed
+82
-56
vllm/v1/worker/gpu/sample/gumbel.py
+45
-1
@@ -5,6 +5,50 @@ import torch
5
from vllm.triton_utils import tl, triton
6
7
8
+@triton.jit
9
+def _temperature_kernel(
10
+ logits_ptr,
11
+ logits_stride,
12
+ idx_mapping_ptr,
13
+ temperature_ptr,
14
+ vocab_size,
15
+ BLOCK_SIZE: tl.constexpr,
16
+):
17
+ batch_idx = tl.program_id(0)
18
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
19
+ temperature = tl.load(temperature_ptr + req_state_idx).to(tl.float32)
20
+ if temperature == 0.0 or temperature == 1.0:
21
+ # Early return to avoid loading logits.
22
+ return
23
+
24
+ block_idx = tl.program_id(1)
25
+ block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
26
+ mask = block < vocab_size
27
+
28
+ logits = tl.load(logits_ptr + batch_idx * logits_stride + block, mask=mask)
29
+ logits = logits.to(tl.float32)
30
+ logits = logits / temperature
31
+ tl.store(logits_ptr + batch_idx * logits_stride + block, logits, mask=mask)
32
+
33
+
34
+def apply_temperature(
35
+ logits: torch.Tensor,
36
+ idx_mapping: torch.Tensor,
37
+ temperature: torch.Tensor,
38
+) -> None:
39
+ num_reqs, vocab_size = logits.shape
40
+ BLOCK_SIZE = 8192
41
+ num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
42
+ _temperature_kernel[(num_reqs, num_blocks)](
43
+ logits,
44
+ logits.stride(0),
45
+ idx_mapping,
46
+ temperature,
47
+ vocab_size,
48
+ BLOCK_SIZE=BLOCK_SIZE,
49
+ )
50
+
51
+
52
@triton.jit
53
def _gumbel_sample_kernel(
54
local_argmax_ptr,
@@ -48,7 +92,7 @@ def _gumbel_sample_kernel(
92
93
# Apply temperature.
94
if APPLY_TEMPERATURE:
51
- # NOTE(woosuk): Match the behavior of _penalties_and_temperature_kernel.
95
+ # NOTE(woosuk): Match the behavior of _temperature_kernel.
96
# E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too.
97
logits = logits / temp
98
vllm/v1/worker/gpu/sample/penalties.py
+31
-50
@@ -66,16 +66,10 @@ class PenaltiesState:
66
self.frequency_penalty.copy_to_uva()
67
self.presence_penalty.copy_to_uva()
68
69
- def apply_penalties_and_temperature(
70
- self,
71
- logits: torch.Tensor,
72
- idx_mapping: torch.Tensor,
73
- temperature: torch.Tensor,
74
- ) -> None:
75
- apply_penalties_and_temperature(
69
+ def apply_penalties(self, logits: torch.Tensor, idx_mapping: torch.Tensor) -> None:
70
+ apply_penalties(
71
logits,
72
idx_mapping,
78
- temperature,
73
self.repetition_penalty.gpu,
74
self.frequency_penalty.gpu,
75
self.presence_penalty.gpu,
@@ -85,14 +79,13 @@ class PenaltiesState:
79
80
81
@triton.jit
88
-def _penalties_and_temperature_kernel(
82
+def _penalties_kernel(
83
logits_ptr,
84
logits_stride,
85
idx_mapping_ptr,
86
repetition_penalty_ptr,
87
frequency_penalty_ptr,
88
presence_penalty_ptr,
95
- temperature_ptr,
89
prompt_bin_mask_ptr,
90
prompt_bin_mask_stride,
91
output_bin_counts_ptr,
@@ -105,15 +98,12 @@ def _penalties_and_temperature_kernel(
98
rep_penalty = tl.load(repetition_penalty_ptr + req_state_idx)
99
freq_penalty = tl.load(frequency_penalty_ptr + req_state_idx)
100
pres_penalty = tl.load(presence_penalty_ptr + req_state_idx)
108
- temperature = tl.load(temperature_ptr + req_state_idx)
109
- temperature = tl.where(temperature == 0.0, 1.0, temperature)
101
102
use_rep_penalty = rep_penalty != 1.0
103
use_freq_penalty = freq_penalty != 0.0
104
use_pres_penalty = pres_penalty != 0.0
105
use_penalty = use_rep_penalty or use_freq_penalty or use_pres_penalty
115
- use_temperature = temperature != 1.0
116
- if not (use_penalty or use_temperature):
106
+ if not use_penalty:
107
# Early return to avoid loading logits.
108
return
109
@@ -123,47 +113,39 @@ def _penalties_and_temperature_kernel(
113
logits = tl.load(logits_ptr + batch_idx * logits_stride + block, mask=mask)
114
logits = logits.to(tl.float32)
115
126
- if use_penalty:
127
- output_bin_counts = tl.load(
128
- output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block,
129
- mask=mask,
116
+ output_bin_counts = tl.load(
117
+ output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block,
118
+ mask=mask,
119
+ )
120
+ output_bin_mask = output_bin_counts > 0
121
+
122
+ # Apply repetition penalties.
123
+ if use_rep_penalty:
124
+ packed_block = block_idx * BLOCK_SIZE // 32 + tl.arange(0, BLOCK_SIZE // 32)
125
+ packed_mask = tl.load(
126
+ prompt_bin_mask_ptr + req_state_idx * prompt_bin_mask_stride + packed_block,
127
+ mask=packed_block < tl.cdiv(vocab_size, 32),
128
)
131
- output_bin_mask = output_bin_counts > 0
132
-
133
- # Apply repetition penalties.
134
- if use_rep_penalty:
135
- packed_block = block_idx * BLOCK_SIZE // 32 + tl.arange(0, BLOCK_SIZE // 32)
136
- packed_mask = tl.load(
137
- prompt_bin_mask_ptr
138
- + req_state_idx * prompt_bin_mask_stride
139
- + packed_block,
140
- mask=packed_block < tl.cdiv(vocab_size, 32),
141
- )
142
- prompt_bin_mask = (packed_mask[:, None] >> (tl.arange(0, 32)[None, :])) & 1
143
- prompt_bin_mask = prompt_bin_mask.to(tl.int1)
144
- prompt_bin_mask = prompt_bin_mask.reshape(BLOCK_SIZE)
145
-
146
- # If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
147
- scale = tl.where(prompt_bin_mask | output_bin_mask, rep_penalty, 1.0)
148
- # If logits are positive, divide by penalty, otherwise multiply by penalty.
149
- logits *= tl.where(logits > 0, 1.0 / scale, scale)
150
-
151
- # Apply frequency penalties.
152
- logits -= freq_penalty * output_bin_counts
153
- # Apply presence penalties.
154
- logits -= pres_penalty * output_bin_mask
155
-
156
- # Apply temperature.
157
- logits = logits / temperature
158
-
129
+ prompt_bin_mask = (packed_mask[:, None] >> (tl.arange(0, 32)[None, :])) & 1
130
+ prompt_bin_mask = prompt_bin_mask.to(tl.int1)
131
+ prompt_bin_mask = prompt_bin_mask.reshape(BLOCK_SIZE)
132
+
133
+ # If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
134
+ scale = tl.where(prompt_bin_mask | output_bin_mask, rep_penalty, 1.0)
135
+ # If logits are positive, divide by penalty, otherwise multiply by penalty.
136
+ logits *= tl.where(logits > 0, 1.0 / scale, scale)
137
+
138
+ # Apply frequency penalties.
139
+ logits -= freq_penalty * output_bin_counts
140
+ # Apply presence penalties.
141
+ logits -= pres_penalty * output_bin_mask
142
# Store back to logits.
143
tl.store(logits_ptr + batch_idx * logits_stride + block, logits, mask=mask)
144
145
163
-def apply_penalties_and_temperature(
146
+def apply_penalties(
147
logits: torch.Tensor,
148
idx_mapping: torch.Tensor,
166
- temperature: torch.Tensor,
149
repetition_penalty: torch.Tensor,
150
frequency_penalty: torch.Tensor,
151
presence_penalty: torch.Tensor,
@@ -173,14 +155,13 @@ def apply_penalties_and_temperature(
155
num_reqs, vocab_size = logits.shape
156
BLOCK_SIZE = 8192
157
num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
176
- _penalties_and_temperature_kernel[(num_reqs, num_blocks)](
158
+ _penalties_kernel[(num_reqs, num_blocks)](
159
logits,
160
logits.stride(0),
161
idx_mapping,
162
repetition_penalty,
163
frequency_penalty,
164
presence_penalty,
183
- temperature,
165
prompt_bin_mask,
166
prompt_bin_mask.stride(0),
167
output_bin_counts,
vllm/v1/worker/gpu/sample/sampler.py
+6
-5
@@ -9,7 +9,7 @@ from vllm.config.model import LogprobsMode
9
from vllm.sampling_params import SamplingParams
10
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
11
from vllm.v1.worker.gpu.metrics.logits import get_num_nans
12
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
12
+from vllm.v1.worker.gpu.sample.gumbel import apply_temperature, gumbel_sample
13
from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
14
from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
15
from vllm.v1.worker.gpu.sample.min_p import apply_min_p
@@ -106,10 +106,11 @@ class Sampler:
106
# Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
107
self.logit_bias_state.apply_logit_bias(logits, idx_mapping, pos)
108
109
- # Apply penalties and temperature in place.
110
- self.penalties_state.apply_penalties_and_temperature(
111
- logits, idx_mapping, self.sampling_states.temperature.gpu
112
- )
109
+ # Apply penalties in place.
110
+ self.penalties_state.apply_penalties(logits, idx_mapping)
111
+
112
+ # Apply temperature in place.
113
+ apply_temperature(logits, idx_mapping, self.sampling_states.temperature.gpu)
114
115
# Apply min_p in place if any request has a non-zero min_p.
116
do_min_p = self.sampling_states.do_min_p(idx_mapping_np)