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 15, 21:00 UTC
[V1][BugFix] Fix for mixed top_k batch (#14301)
Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Ye Cao <caoye.cao@alibaba-inc.com>
Nick Hill committed
Mar 5, 2025 at 12:43 UTC
ac60dc7fe1b2a7ba02f4f4a06a0888e1b1c33629
1 file changed
+5
-2
vllm/v1/worker/gpu_input_batch.py
+5
-2
@@ -264,9 +264,12 @@ class InputBatch:
264
self.top_p_cpu[req_index] = sampling_params.top_p
265
if sampling_params.top_p < 1:
266
self.top_p_reqs.add(req_id)
267
- self.top_k_cpu[req_index] = sampling_params.top_k
268
- if sampling_params.top_k > 0:
267
+ top_k = sampling_params.top_k
268
+ if 0 < top_k < self.vocab_size:
269
self.top_k_reqs.add(req_id)
270
+ else:
271
+ top_k = self.vocab_size
272
+ self.top_k_cpu[req_index] = top_k
273
self.min_p_cpu[req_index] = sampling_params.min_p
274
self.frequency_penalties_cpu[
275
req_index] = sampling_params.frequency_penalty