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] Convert np array to torch tensor to index into block table for attn chunking (#24474)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
Sep 9, 2025 at 20:01 UTC
dc625ea6b8f8072afd3e0923f7ef5a39be87fc96
1 file changed
+8
-1
vllm/v1/attention/backends/utils.py
+8
-1
@@ -542,7 +542,14 @@ def make_local_attention_virtual_batches(
542
1)
543
batch_indices = np.repeat(np.arange(actual_batch_size, dtype=np.int32),
544
local_blocks * pages_per_local_batch)
545
- block_table_local = block_table[batch_indices, block_indices]\
545
+
546
+ # NOTE: https://github.com/pytorch/pytorch/pull/160256 causes performance
547
+ # regression when using numpy arrays (batch and block indices) to index into
548
+ # torch tensor (block_table). As a workaround, convert numpy arrays to torch
549
+ # tensor first, which recovers perf.
550
+ batch_indices_torch = torch.from_numpy(batch_indices)
551
+ block_indices_torch = torch.from_numpy(block_indices)
552
+ block_table_local = block_table[batch_indices_torch, block_indices_torch]\
553
.view(virtual_batches, -1)
554
555
query_start_loc_cpu = torch.from_numpy(cu_seqlens_q_local)