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, 12:15 UTC
[Bugfix] fix tmp_out and exp_sums dimensions (#17438)
Signed-off-by: Hui Liu <96135754+hliuca@users.noreply.github.com>
Hui Liu committed
May 2, 2025 at 09:44 UTC
4c33d6732148fdaeb9780fa86fca1f87f2a93c19
1 file changed
+1
-1
vllm/attention/ops/chunked_prefill_paged_decode.py
+1
-1
@@ -289,7 +289,7 @@ def chunked_prefill_paged_decode(
289
max_num_partitions = ((max_seq_len + _PARTITION_SIZE_ROCM - 1) //
290
_PARTITION_SIZE_ROCM)
291
assert _PARTITION_SIZE_ROCM % block_size == 0
292
- total_num_seq = query.shape[0]
292
+ total_num_seq = block_table.shape[0]
293
tmp_output = torch.empty(
294
size=(total_num_seq, num_query_heads, max_num_partitions,
295
head_size),