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
[Perf] Remove hardcoded num_warps=1 (#26183)
Signed-off-by: Corey Lowman <clowman1993@gmail.com>
Corey Lowman committed
Oct 3, 2025 at 16:38 UTC
0879736aaba58e6fca43d487e430d38c05d40756
1 file changed
+2
-5
vllm/v1/sample/rejection_sampler.py
+2
-5
@@ -164,12 +164,12 @@ def rejection_sample(
164
assert target_probs.shape == (num_tokens, vocab_size)
165
166
# Create output buffer.
167
- output_token_ids = torch.empty(
167
+ output_token_ids = torch.full(
168
(batch_size, max_spec_len + 1),
169
+ PLACEHOLDER_TOKEN_ID,
170
dtype=torch.int32, # Consistent with SamplerOutput.sampled_token_ids.
171
device=device,
172
)
172
- output_token_ids.fill_(PLACEHOLDER_TOKEN_ID)
173
174
if sampling_metadata.all_greedy:
175
is_greedy = None
@@ -186,7 +186,6 @@ def rejection_sample(
186
bonus_token_ids,
187
is_greedy,
188
max_spec_len,
189
- num_warps=1,
189
)
190
if sampling_metadata.all_greedy:
191
return output_token_ids
@@ -227,7 +226,6 @@ def rejection_sample(
226
max_spec_len,
227
vocab_size,
228
NO_DRAFT_PROBS=draft_probs is None,
230
- num_warps=1,
229
)
230
return output_token_ids
231
@@ -329,7 +327,6 @@ def expand_batch_to_tokens(
327
replace_from,
328
replace_to,
329
MAX_NUM_TOKENS=MAX_SPEC_LEN, # To avoid recompilation.
332
- num_warps=1,
330
)
331
return expanded_x
332