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 7, 10:15 UTC
[ModelRunnerV2] Rename sampler functions and variables for clarity (#35459)
Signed-off-by: Andy Lo <andy@mistral.ai>
Andy Lo committed
Mar 3, 2026 at 04:48 UTC
0a7165fd7196bb3111f87ae2a0b074dec8af4359
7 files changed
+99
-91
vllm/v1/worker/gpu/sample/bad_words.py
+9
-9
@@ -72,7 +72,7 @@ class BadWordsState:
72
def apply_bad_words(
73
self,
74
logits: torch.Tensor,
75
- idx_mapping: torch.Tensor,
75
+ expanded_idx_mapping: torch.Tensor,
76
idx_mapping_np: np.ndarray,
77
input_ids: torch.Tensor,
78
expanded_local_pos: torch.Tensor,
@@ -84,7 +84,7 @@ class BadWordsState:
84
85
apply_bad_words(
86
logits,
87
- idx_mapping,
87
+ expanded_idx_mapping,
88
self.bad_word_token_ids.gpu,
89
self.bad_word_offsets.gpu,
90
self.num_bad_words.gpu,
@@ -114,17 +114,17 @@ def _bad_words_kernel(
114
input_ids_ptr,
115
expanded_local_pos_ptr,
116
):
117
- logit_idx = tl.program_id(0)
117
+ token_idx = tl.program_id(0)
118
bw_idx = tl.program_id(1)
119
120
- req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx)
120
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
121
num_bad_words = tl.load(num_bad_words_ptr + req_state_idx)
122
123
if bw_idx >= num_bad_words:
124
return
125
126
- pos = tl.load(expanded_local_pos_ptr + logit_idx)
127
- cur_req_first_pos = logit_idx - pos
126
+ pos = tl.load(expanded_local_pos_ptr + token_idx)
127
+ cur_req_first_pos = token_idx - pos
128
129
prompt_len = tl.load(prompt_len_ptr + req_state_idx)
130
total_len = tl.load(total_len_ptr + req_state_idx)
@@ -159,7 +159,7 @@ def _bad_words_kernel(
159
match = match & (expected == actual)
160
161
if match:
162
- tl.store(logits_ptr + logit_idx * logits_stride + last_token, -float("inf"))
162
+ tl.store(logits_ptr + token_idx * logits_stride + last_token, -float("inf"))
163
164
165
def apply_bad_words(
@@ -175,8 +175,8 @@ def apply_bad_words(
175
expanded_local_pos: torch.Tensor,
176
max_num_bad_words: int,
177
) -> None:
178
- total_num_tokens = logits.shape[0]
179
- _bad_words_kernel[(total_num_tokens, max_num_bad_words)](
178
+ num_tokens = logits.shape[0]
179
+ _bad_words_kernel[(num_tokens, max_num_bad_words)](
180
logits,
181
logits.stride(0),
182
expanded_idx_mapping,
vllm/v1/worker/gpu/sample/gumbel.py
+24
-24
@@ -9,13 +9,13 @@ from vllm.triton_utils import tl, triton
9
def _temperature_kernel(
10
logits_ptr,
11
logits_stride,
12
- idx_mapping_ptr,
12
+ expanded_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)
17
+ token_idx = tl.program_id(0)
18
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + token_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.
@@ -25,24 +25,24 @@ def _temperature_kernel(
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)
28
+ logits = tl.load(logits_ptr + token_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)
31
+ tl.store(logits_ptr + token_idx * logits_stride + block, logits, mask=mask)
32
33
34
def apply_temperature(
35
logits: torch.Tensor,
36
- idx_mapping: torch.Tensor,
36
+ expanded_idx_mapping: torch.Tensor,
37
temperature: torch.Tensor,
38
) -> None:
39
- num_reqs, vocab_size = logits.shape
39
+ num_tokens, vocab_size = logits.shape
40
BLOCK_SIZE = 8192
41
num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
42
- _temperature_kernel[(num_reqs, num_blocks)](
42
+ _temperature_kernel[(num_tokens, num_blocks)](
43
logits,
44
logits.stride(0),
45
- idx_mapping,
45
+ expanded_idx_mapping,
46
temperature,
47
vocab_size,
48
BLOCK_SIZE=BLOCK_SIZE,
@@ -57,7 +57,7 @@ def _gumbel_sample_kernel(
57
local_max_stride,
58
logits_ptr,
59
logits_stride,
60
- idx_mapping_ptr,
60
+ expanded_idx_mapping_ptr,
61
seeds_ptr,
62
pos_ptr,
63
temp_ptr,
@@ -65,14 +65,14 @@ def _gumbel_sample_kernel(
65
BLOCK_SIZE: tl.constexpr,
66
APPLY_TEMPERATURE: tl.constexpr,
67
):
68
- batch_idx = tl.program_id(0)
69
- req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
68
+ token_idx = tl.program_id(0)
69
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
70
71
block_idx = tl.program_id(1)
72
block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
73
mask = block < vocab_size
74
logits = tl.load(
75
- logits_ptr + batch_idx * logits_stride + block,
75
+ logits_ptr + token_idx * logits_stride + block,
76
mask=mask,
77
other=float("-inf"),
78
)
@@ -82,7 +82,7 @@ def _gumbel_sample_kernel(
82
if temp != 0.0:
83
# Calculate the seed for gumbel noise.
84
seed = tl.load(seeds_ptr + req_state_idx)
85
- pos = tl.load(pos_ptr + batch_idx)
85
+ pos = tl.load(pos_ptr + token_idx)
86
gumbel_seed = tl.randint(seed, pos)
87
88
# Generate gumbel noise in FP32.
@@ -101,41 +101,41 @@ def _gumbel_sample_kernel(
101
102
value, idx = tl.max(logits, axis=0, return_indices=True)
103
token_id = block_idx * BLOCK_SIZE + idx
104
- tl.store(local_argmax_ptr + batch_idx * local_argmax_stride + block_idx, token_id)
105
- tl.store(local_max_ptr + batch_idx * local_max_stride + block_idx, value)
104
+ tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id)
105
+ tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, value)
106
107
108
def gumbel_sample(
109
- logits: torch.Tensor, # [num_reqs, vocab_size]
110
- idx_mapping: torch.Tensor, # [max_num_reqs]
109
+ logits: torch.Tensor, # [num_tokens, vocab_size]
110
+ expanded_idx_mapping: torch.Tensor, # [num_tokens]
111
temperature: torch.Tensor, # [max_num_reqs]
112
seed: torch.Tensor, # [max_num_reqs]
113
- pos: torch.Tensor, # [num_reqs]
113
+ pos: torch.Tensor, # [num_tokens]
114
apply_temperature: bool,
115
) -> torch.Tensor:
116
- num_reqs, vocab_size = logits.shape
116
+ num_tokens, vocab_size = logits.shape
117
BLOCK_SIZE = 1024
118
num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
119
local_argmax = torch.empty(
120
- num_reqs,
120
+ num_tokens,
121
num_blocks,
122
dtype=torch.int64,
123
device=logits.device,
124
)
125
local_max = torch.empty(
126
- num_reqs,
126
+ num_tokens,
127
num_blocks,
128
dtype=torch.float32,
129
device=logits.device,
130
)
131
- _gumbel_sample_kernel[(num_reqs, num_blocks)](
131
+ _gumbel_sample_kernel[(num_tokens, num_blocks)](
132
local_argmax,
133
local_argmax.stride(0),
134
local_max,
135
local_max.stride(0),
136
logits,
137
logits.stride(0),
138
- idx_mapping,
138
+ expanded_idx_mapping,
139
seed,
140
pos,
141
temperature,
vllm/v1/worker/gpu/sample/logit_bias.py
+16
-16
@@ -121,7 +121,7 @@ class LogitBiasState:
121
def apply_logit_bias(
122
self,
123
logits: torch.Tensor,
124
- idx_mapping: torch.Tensor,
124
+ expanded_idx_mapping: torch.Tensor,
125
idx_mapping_np: np.ndarray,
126
pos: torch.Tensor,
127
) -> None:
@@ -131,7 +131,7 @@ class LogitBiasState:
131
132
apply_logit_bias(
133
logits,
134
- idx_mapping,
134
+ expanded_idx_mapping,
135
pos,
136
self.num_allowed_token_ids.gpu,
137
self.allowed_token_ids.gpu,
@@ -149,7 +149,7 @@ def _bias_kernel(
149
logits_ptr,
150
logits_stride,
151
vocab_size,
152
- idx_mapping_ptr,
152
+ expanded_idx_mapping_ptr,
153
# Allowed token IDs.
154
num_allowed_token_ids_ptr,
155
allowed_token_ids_ptr,
@@ -169,8 +169,8 @@ def _bias_kernel(
169
BLOCK_SIZE: tl.constexpr,
170
LOGITS_BLOCK_SIZE: tl.constexpr,
171
):
172
- batch_idx = tl.program_id(0)
173
- req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
172
+ token_idx = tl.program_id(0)
173
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
174
175
block = tl.arange(0, BLOCK_SIZE)
176
@@ -186,21 +186,21 @@ def _bias_kernel(
186
mask=mask,
187
)
188
logits = tl.load(
189
- logits_ptr + batch_idx * logits_stride + allowed_token_ids, mask=mask
189
+ logits_ptr + token_idx * logits_stride + allowed_token_ids, mask=mask
190
)
191
192
# Set logits to -inf for all tokens.
193
for i in range(0, vocab_size, LOGITS_BLOCK_SIZE):
194
offset = i + tl.arange(0, LOGITS_BLOCK_SIZE)
195
tl.store(
196
- logits_ptr + batch_idx * logits_stride + offset,
196
+ logits_ptr + token_idx * logits_stride + offset,
197
-float("inf"),
198
mask=offset < vocab_size,
199
)
200
201
# Restore logits for allowed token IDs.
202
tl.store(
203
- logits_ptr + batch_idx * logits_stride + allowed_token_ids,
203
+ logits_ptr + token_idx * logits_stride + allowed_token_ids,
204
logits,
205
mask=mask,
206
)
@@ -214,13 +214,13 @@ def _bias_kernel(
214
mask=mask,
215
)
216
bias = tl.load(bias_ptr + req_state_idx * bias_stride + block, mask=mask)
217
- logits = tl.load(logits_ptr + batch_idx * logits_stride + token_ids, mask=mask)
217
+ logits = tl.load(logits_ptr + token_idx * logits_stride + token_ids, mask=mask)
218
logits += bias
219
- tl.store(logits_ptr + batch_idx * logits_stride + token_ids, logits, mask=mask)
219
+ tl.store(logits_ptr + token_idx * logits_stride + token_ids, logits, mask=mask)
220
221
# Apply min tokens.
222
num_stop_token_ids = tl.load(num_stop_token_ids_ptr + req_state_idx)
223
- pos = tl.load(pos_ptr + batch_idx)
223
+ pos = tl.load(pos_ptr + token_idx)
224
min_len = tl.load(min_lens_ptr + req_state_idx)
225
if num_stop_token_ids > 0 and pos < min_len:
226
mask = block < num_stop_token_ids
@@ -229,7 +229,7 @@ def _bias_kernel(
229
mask=mask,
230
)
231
tl.store(
232
- logits_ptr + batch_idx * logits_stride + stop_token_ids,
232
+ logits_ptr + token_idx * logits_stride + stop_token_ids,
233
-float("inf"),
234
mask=mask,
235
)
@@ -237,7 +237,7 @@ def _bias_kernel(
237
238
def apply_logit_bias(
239
logits: torch.Tensor,
240
- idx_mapping: torch.Tensor,
240
+ expanded_idx_mapping: torch.Tensor,
241
pos: torch.Tensor,
242
num_allowed_token_ids: torch.Tensor,
243
allowed_token_ids: torch.Tensor,
@@ -248,7 +248,7 @@ def apply_logit_bias(
248
num_stop_token_ids: torch.Tensor,
249
stop_token_ids: torch.Tensor,
250
) -> None:
251
- num_reqs, vocab_size = logits.shape
251
+ num_tokens, vocab_size = logits.shape
252
BLOCK_SIZE = triton.next_power_of_2(
253
max(
254
allowed_token_ids.shape[-1],
@@ -257,11 +257,11 @@ def apply_logit_bias(
257
)
258
)
259
LOGITS_BLOCK_SIZE = 8192
260
- _bias_kernel[(num_reqs,)](
260
+ _bias_kernel[(num_tokens,)](
261
logits,
262
logits.stride(0),
263
vocab_size,
264
- idx_mapping,
264
+ expanded_idx_mapping,
265
num_allowed_token_ids,
266
allowed_token_ids,
267
allowed_token_ids.stride(0),
vllm/v1/worker/gpu/sample/min_p.py
+14
-10
@@ -9,13 +9,13 @@ from vllm.triton_utils import tl, triton
9
def _min_p_kernel(
10
logits_ptr,
11
logits_stride,
12
- idx_mapping_ptr,
12
+ expanded_idx_mapping_ptr,
13
min_p_ptr,
14
vocab_size,
15
BLOCK_SIZE: tl.constexpr,
16
):
17
- req_idx = tl.program_id(0)
18
- req_state_idx = tl.load(idx_mapping_ptr + req_idx)
17
+ token_idx = tl.program_id(0)
18
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
19
min_p = tl.load(min_p_ptr + req_state_idx).to(tl.float32)
20
if min_p == 0.0:
21
return
@@ -25,7 +25,9 @@ def _min_p_kernel(
25
block = i + tl.arange(0, BLOCK_SIZE)
26
mask = block < vocab_size
27
logits = tl.load(
28
- logits_ptr + req_idx * logits_stride + block, mask=mask, other=float("-inf")
28
+ logits_ptr + token_idx * logits_stride + block,
29
+ mask=mask,
30
+ other=float("-inf"),
31
)
32
max_val = tl.max(tl.maximum(logits, max_val))
33
max_val = max_val.to(tl.float32) # type: ignore
@@ -35,21 +37,23 @@ def _min_p_kernel(
37
block = i + tl.arange(0, BLOCK_SIZE)
38
mask = block < vocab_size
39
logits = tl.load(
38
- logits_ptr + req_idx * logits_stride + block, mask=mask, other=float("-inf")
40
+ logits_ptr + token_idx * logits_stride + block,
41
+ mask=mask,
42
+ other=float("-inf"),
43
)
44
logits = tl.where(logits < threshold, float("-inf"), logits)
41
- tl.store(logits_ptr + req_idx * logits_stride + block, logits, mask=mask)
45
+ tl.store(logits_ptr + token_idx * logits_stride + block, logits, mask=mask)
46
47
48
def apply_min_p(
45
- logits: torch.Tensor, idx_mapping: torch.Tensor, min_p: torch.Tensor
49
+ logits: torch.Tensor, expanded_idx_mapping: torch.Tensor, min_p: torch.Tensor
50
) -> None:
47
- num_reqs, vocab_size = logits.shape
51
+ num_tokens, vocab_size = logits.shape
52
BLOCK_SIZE = 1024
49
- _min_p_kernel[(num_reqs,)](
53
+ _min_p_kernel[(num_tokens,)](
54
logits,
55
logits.stride(0),
52
- idx_mapping,
56
+ expanded_idx_mapping,
57
min_p,
58
vocab_size,
59
BLOCK_SIZE=BLOCK_SIZE,
vllm/v1/worker/gpu/sample/penalties.py
+15
-15
@@ -82,7 +82,7 @@ class PenaltiesState:
82
def apply_penalties(
83
self,
84
logits: torch.Tensor,
85
- idx_mapping: torch.Tensor,
85
+ expanded_idx_mapping: torch.Tensor,
86
idx_mapping_np: np.ndarray,
87
input_ids: torch.Tensor,
88
expanded_local_pos: torch.Tensor,
@@ -94,7 +94,7 @@ class PenaltiesState:
94
95
apply_penalties(
96
logits,
97
- idx_mapping,
97
+ expanded_idx_mapping,
98
input_ids,
99
expanded_local_pos,
100
self.repetition_penalty.gpu,
@@ -110,7 +110,7 @@ class PenaltiesState:
110
def _penalties_kernel(
111
logits_ptr,
112
logits_stride,
113
- idx_mapping_ptr,
113
+ expanded_idx_mapping_ptr,
114
token_ids_ptr,
115
expanded_local_pos_ptr,
116
repetition_penalty_ptr,
@@ -125,7 +125,7 @@ def _penalties_kernel(
125
MAX_SPEC_LEN: tl.constexpr,
126
):
127
token_idx = tl.program_id(0)
128
- req_state_idx = tl.load(idx_mapping_ptr + token_idx)
128
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
129
rep_penalty = tl.load(repetition_penalty_ptr + req_state_idx)
130
freq_penalty = tl.load(frequency_penalty_ptr + req_state_idx)
131
pres_penalty = tl.load(presence_penalty_ptr + req_state_idx)
@@ -191,7 +191,7 @@ def _penalties_kernel(
191
192
def apply_penalties(
193
logits: torch.Tensor,
194
- idx_mapping: torch.Tensor,
194
+ expanded_idx_mapping: torch.Tensor,
195
token_ids: torch.Tensor,
196
expanded_local_pos: torch.Tensor,
197
repetition_penalty: torch.Tensor,
@@ -207,7 +207,7 @@ def apply_penalties(
207
_penalties_kernel[(num_tokens, num_blocks)](
208
logits,
209
logits.stride(0),
210
- idx_mapping,
210
+ expanded_idx_mapping,
211
token_ids,
212
expanded_local_pos,
213
repetition_penalty,
@@ -225,7 +225,7 @@ def apply_penalties(
225
226
@triton.jit
227
def _bincount_kernel(
228
- idx_mapping_ptr,
228
+ expanded_idx_mapping_ptr,
229
all_token_ids_ptr,
230
all_token_ids_stride,
231
prompt_len_ptr,
@@ -236,9 +236,9 @@ def _bincount_kernel(
236
output_bin_counts_stride,
237
BLOCK_SIZE: tl.constexpr,
238
):
239
- batch_idx = tl.program_id(0)
239
+ token_idx = tl.program_id(0)
240
block_idx = tl.program_id(1)
241
- req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
241
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
242
243
prefill_len = tl.load(prefill_len_ptr + req_state_idx)
244
if block_idx * BLOCK_SIZE >= prefill_len:
@@ -276,7 +276,7 @@ def _bincount_kernel(
276
277
278
def bincount(
279
- idx_mapping: torch.Tensor,
279
+ expanded_idx_mapping: torch.Tensor,
280
all_token_ids: torch.Tensor,
281
prompt_len: torch.Tensor,
282
prefill_len: torch.Tensor,
@@ -284,13 +284,13 @@ def bincount(
284
output_bin_counts: torch.Tensor,
285
max_prefill_len: int,
286
) -> None:
287
- prompt_bin_mask[idx_mapping] = 0
288
- output_bin_counts[idx_mapping] = 0
289
- num_reqs = idx_mapping.shape[0]
287
+ prompt_bin_mask[expanded_idx_mapping] = 0
288
+ output_bin_counts[expanded_idx_mapping] = 0
289
+ num_tokens = expanded_idx_mapping.shape[0]
290
BLOCK_SIZE = 1024
291
num_blocks = triton.cdiv(max_prefill_len, BLOCK_SIZE)
292
- _bincount_kernel[(num_reqs, num_blocks)](
293
- idx_mapping,
292
+ _bincount_kernel[(num_tokens, num_blocks)](
293
+ expanded_idx_mapping,
294
all_token_ids,
295
all_token_ids.stride(0),
296
prompt_len,
vllm/v1/worker/gpu/sample/sampler.py
+14
-10
@@ -56,7 +56,7 @@ class Sampler:
56
def __call__(
57
self,
58
logits: torch.Tensor,
59
- idx_mapping: torch.Tensor,
59
+ expanded_idx_mapping: torch.Tensor,
60
idx_mapping_np: np.ndarray,
61
cu_num_logits_np: np.ndarray,
62
pos: torch.Tensor,
@@ -68,7 +68,7 @@ class Sampler:
68
num_nans = get_num_nans(logits) if self.compute_nans else None
69
sampled, processed_logits = self.sample(
70
logits,
71
- idx_mapping,
71
+ expanded_idx_mapping,
72
idx_mapping_np,
73
pos,
74
input_ids,
@@ -101,7 +101,7 @@ class Sampler:
101
def sample(
102
self,
103
logits: torch.Tensor,
104
- idx_mapping: torch.Tensor,
104
+ expanded_idx_mapping: torch.Tensor,
105
idx_mapping_np: np.ndarray,
106
pos: torch.Tensor,
107
input_ids: torch.Tensor,
@@ -111,12 +111,14 @@ class Sampler:
111
logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
112
113
# Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
114
- self.logit_bias_state.apply_logit_bias(logits, idx_mapping, idx_mapping_np, pos)
114
+ self.logit_bias_state.apply_logit_bias(
115
+ logits, expanded_idx_mapping, idx_mapping_np, pos
116
+ )
117
118
# Apply penalties in place.
119
self.penalties_state.apply_penalties(
120
logits,
119
- idx_mapping,
121
+ expanded_idx_mapping,
122
idx_mapping_np,
123
input_ids,
124
expanded_local_pos,
@@ -126,27 +128,29 @@ class Sampler:
128
# Apply bad words masking in place.
129
self.bad_words_state.apply_bad_words(
130
logits,
129
- idx_mapping,
131
+ expanded_idx_mapping,
132
idx_mapping_np,
133
input_ids,
134
expanded_local_pos,
135
)
136
137
# Apply temperature in place.
136
- self.sampling_states.apply_temperature(logits, idx_mapping, idx_mapping_np)
138
+ self.sampling_states.apply_temperature(
139
+ logits, expanded_idx_mapping, idx_mapping_np
140
+ )
141
142
# Apply min_p in place.
139
- self.sampling_states.apply_min_p(logits, idx_mapping, idx_mapping_np)
143
+ self.sampling_states.apply_min_p(logits, expanded_idx_mapping, idx_mapping_np)
144
145
# Apply top_k and/or top_p. This might or might not return a new tensor.
146
logits = self.sampling_states.apply_top_k_top_p(
143
- logits, idx_mapping, idx_mapping_np
147
+ logits, expanded_idx_mapping, idx_mapping_np
148
)
149
150
# Sample the next token.
151
sampled = gumbel_sample(
152
logits,
149
- idx_mapping,
153
+ expanded_idx_mapping,
154
self.sampling_states.temperature.gpu,
155
self.sampling_states.seeds.gpu,
156
pos,
vllm/v1/worker/gpu/sample/states.py
+7
-7
@@ -64,7 +64,7 @@ class SamplingStates:
64
def apply_temperature(
65
self,
66
logits: torch.Tensor,
67
- idx_mapping: torch.Tensor,
67
+ expanded_idx_mapping: torch.Tensor,
68
idx_mapping_np: np.ndarray,
69
) -> None:
70
temp_np = self.temperature.np[idx_mapping_np]
@@ -72,23 +72,23 @@ class SamplingStates:
72
# No request requires temperature. Skip the kernel launch.
73
return
74
75
- apply_temperature(logits, idx_mapping, self.temperature.gpu)
75
+ apply_temperature(logits, expanded_idx_mapping, self.temperature.gpu)
76
77
def apply_min_p(
78
self,
79
logits: torch.Tensor,
80
- idx_mapping: torch.Tensor,
80
+ expanded_idx_mapping: torch.Tensor,
81
idx_mapping_np: np.ndarray,
82
) -> None:
83
if np.all(self.min_p.np[idx_mapping_np] == 0.0):
84
# No request uses min_p. Skip the kernel launch.
85
return
86
- apply_min_p(logits, idx_mapping, self.min_p.gpu)
86
+ apply_min_p(logits, expanded_idx_mapping, self.min_p.gpu)
87
88
def apply_top_k_top_p(
89
self,
90
logits: torch.Tensor,
91
- idx_mapping: torch.Tensor,
91
+ expanded_idx_mapping: torch.Tensor,
92
idx_mapping_np: np.ndarray,
93
) -> torch.Tensor:
94
do_top_k = np.any(self.top_k.np[idx_mapping_np] != self.vocab_size)
@@ -96,8 +96,8 @@ class SamplingStates:
96
if not (do_top_k or do_top_p):
97
return logits
98
99
- top_k = self.top_k.gpu[idx_mapping] if do_top_k else None
100
- top_p = self.top_p.gpu[idx_mapping] if do_top_p else None
99
+ top_k = self.top_k.gpu[expanded_idx_mapping] if do_top_k else None
100
+ top_p = self.top_p.gpu[expanded_idx_mapping] if do_top_p else None
101
return apply_top_k_top_p(logits, top_k, top_p)
102
103
def max_num_logprobs(self, idx_mapping_np: np.ndarray) -> int: