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
[BugFix][FlashInfer] Fix potential race condition for paged_kv_indptr_cpu (#23737)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 28, 2025 at 14:22 UTC
7ffbf27239c3ff68d773e7d2e2cd284f1375349f
1 file changed
+10
-2
vllm/v1/attention/backends/flashinfer.py
+10
-2
@@ -237,6 +237,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
237
device="cpu",
238
pin_memory=pin_memory)
239
self.paged_kv_indptr_np = self.paged_kv_indptr_cpu.numpy()
240
+ self.paged_kv_indptr_buffer = torch.zeros_like(
241
+ self.paged_kv_indptr_cpu, pin_memory=pin_memory)
242
self.paged_kv_indices_cpu = torch.zeros(max_num_pages,
243
dtype=torch.int32,
244
device="cpu",
@@ -361,12 +363,18 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
363
dtype=np.int32,
364
out=self.paged_kv_indptr_np[1:num_reqs + 1],
365
)
366
+ # NOTE(woosuk): Because self.paged_kv_indptr_cpu can be modified
367
+ # after this line (e.g., for cuda graphs), we need to copy the data to
368
+ # self.paged_kv_indptr_buffer to avoid race condition.
369
+ self.paged_kv_indptr_buffer[:num_reqs +
370
+ 1] = (self.paged_kv_indptr_cpu[:num_reqs +
371
+ 1])
372
paged_kv_indptr = self.paged_kv_indptr[:num_reqs + 1]
365
- paged_kv_indptr.copy_(self.paged_kv_indptr_cpu[:num_reqs + 1],
373
+ paged_kv_indptr.copy_(self.paged_kv_indptr_buffer[:num_reqs + 1],
374
non_blocking=True)
375
376
# write self.paged_kv_indices inplace
369
- num_actual_pages = num_blocks_np.sum().item()
377
+ num_actual_pages = self.paged_kv_indptr_np[num_reqs]
378
paged_kv_indices = self.paged_kv_indices[:num_actual_pages]
379
_copy_page_indices_kernel[(num_reqs, )](
380
paged_kv_indices,