[Bugfix] Fix race condition in KVBlockZeroer (#48085)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Benjamin Chislett committed
Jul 9, 2026 at 18:18 UTC
1cd75b3dd4b3bf90e4ef81831b6f0dd91fde2fe1
4 files changed
+83
-20
tests/v1/worker/test_kv_block_zeroer.py
new
+44
@@ -0,0 +1,44 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.v1.worker.utils import KVBlockZeroer
8
+
9
+
10
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
11
+def test_block_ids_are_not_overwritten_while_copy_is_in_flight():
12
+ device = torch.device("cuda")
13
+ num_blocks = 4
14
+ page_size_el = 4
15
+ storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device)
16
+
17
+ # Build the minimal zeroer state directly so the test can focus on ID-buffer
18
+ # lifetime without constructing model attention groups.
19
+ zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
20
+ zeroer.device = device
21
+ zeroer.pin_memory = True
22
+ zeroer.max_concurrency = 2
23
+ zeroer._id_cap = 8
24
+ zeroer._allocate_id_buffers()
25
+ zeroer._meta = (
26
+ torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
27
+ page_size_el,
28
+ page_size_el,
29
+ 1,
30
+ )
31
+
32
+ stream = torch.cuda.Stream()
33
+ with torch.cuda.stream(stream):
34
+ # Keep the first nonblocking H2D copy pending while the host submits the
35
+ # second call. A single shared pinned source would be overwritten here.
36
+ torch.cuda._sleep(10_000_000)
37
+ zeroer.zero_block_ids([1])
38
+ zeroer.zero_block_ids([2])
39
+ stream.synchronize()
40
+
41
+ assert torch.all(storage[0] == 1)
42
+ assert torch.all(storage[1] == 0)
43
+ assert torch.all(storage[2] == 0)
44
+ assert torch.all(storage[3] == 1)
vllm/v1/worker/gpu/model_runner.py
+1
@@ -498,6 +498,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
498
kernel_block_sizes=self.kernel_block_sizes,
499
cache_dtype=self.cache_config.cache_dtype,
500
static_forward_context=self.compilation_config.static_forward_context,
501
+ max_concurrency=self.vllm_config.max_concurrent_batches,
502
)
503
504
@torch.inference_mode()
vllm/v1/worker/gpu_model_runner.py
+1
@@ -1123,6 +1123,7 @@ class GPUModelRunner(
1123
cache_dtype=self.cache_config.cache_dtype,
1124
runner_only_attn_layers=self.runner_only_attn_layers,
1125
static_forward_context=self.compilation_config.static_forward_context,
1126
+ max_concurrency=self.vllm_config.max_concurrent_batches,
1127
)
1128
1129
def _zero_block_ids(self, block_ids: list[int]) -> None:
vllm/v1/worker/utils.py
+37
-20
@@ -94,6 +94,7 @@ class KVBlockZeroer:
94
cache_dtype: str,
95
static_forward_context: dict[str, Any],
96
runner_only_attn_layers: set[str] | None = None,
97
+ max_concurrency: int = 1,
98
) -> None:
99
"""Precompute the absolute-address table for the Triton zeroing kernel.
100
@@ -109,10 +110,14 @@ class KVBlockZeroer:
110
"""
111
self.device = device
112
self.pin_memory = pin_memory
113
+ if max_concurrency < 1:
114
+ raise ValueError("max_concurrency must be at least 1")
115
+ self.max_concurrency = max_concurrency
116
self._meta: tuple[torch.Tensor, int, int, int] | None = None
117
self._id_cap: int = 0
114
- self._ids_pinned: torch.Tensor | None = None
115
- self._ids_gpu: torch.Tensor | None = None
118
+ self._ids_pinned: list[torch.Tensor] = []
119
+ self._ids_gpu: list[torch.Tensor] = []
120
+ self._id_buffer_index = 0
121
122
if runner_only_attn_layers is None:
123
runner_only_attn_layers = set()
@@ -175,12 +180,7 @@ class KVBlockZeroer:
180
181
blk_size = min(largest_power_of_2_divisor(page_size_el), 1024)
182
self._id_cap = 8192
178
- self._ids_pinned = torch.empty(
179
- self._id_cap,
180
- dtype=torch.int64,
181
- pin_memory=self.pin_memory,
182
- )
183
- self._ids_gpu = torch.empty(self._id_cap, dtype=torch.int64, device=self.device)
183
+ self._allocate_id_buffers()
184
self._meta = (
185
torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device),
186
page_size_el,
@@ -188,6 +188,21 @@ class KVBlockZeroer:
188
len(seg_addrs),
189
)
190
191
+ def _allocate_id_buffers(self) -> None:
192
+ self._ids_pinned = [
193
+ torch.empty(
194
+ self._id_cap,
195
+ dtype=torch.int64,
196
+ pin_memory=self.pin_memory,
197
+ )
198
+ for _ in range(self.max_concurrency)
199
+ ]
200
+ self._ids_gpu = [
201
+ torch.empty(self._id_cap, dtype=torch.int64, device=self.device)
202
+ for _ in range(self.max_concurrency)
203
+ ]
204
+ self._id_buffer_index = 0
205
+
206
def zero_block_ids(self, block_ids: list[int]) -> None:
207
"""Zero the KV cache memory for the given block IDs."""
208
if not block_ids or self._meta is None:
@@ -195,19 +210,21 @@ class KVBlockZeroer:
210
seg_addrs, page_size_el, blk_size, n_segs = self._meta
211
n_blocks = len(block_ids)
212
if n_blocks > self._id_cap:
213
+ # The old pinned buffers may still be the source of an in-flight
214
+ # nonblocking copy. Growing is rare, so we don't mind the sync overhead
215
+ torch.accelerator.synchronize()
216
self._id_cap = n_blocks * 2
199
- self._ids_pinned = torch.empty(
200
- self._id_cap,
201
- dtype=torch.int64,
202
- pin_memory=self.pin_memory,
203
- )
204
- self._ids_gpu = torch.empty(
205
- self._id_cap, dtype=torch.int64, device=self.device
206
- )
207
- assert self._ids_pinned is not None and self._ids_gpu is not None
208
- self._ids_pinned[:n_blocks].numpy()[:] = block_ids
209
- idx = self._ids_gpu[:n_blocks]
210
- idx.copy_(self._ids_pinned[:n_blocks], non_blocking=True)
217
+ self._allocate_id_buffers()
218
+
219
+ # The H2D copy is nonblocking, so its pinned source must not be mutated
220
+ # while this batch is in flight. Rotate through as many buffers as concurrent
221
+ # in-flight batches, to avoid collisions.
222
+ buffer_index = self._id_buffer_index
223
+ self._id_buffer_index = (buffer_index + 1) % self.max_concurrency
224
+ ids_pinned = self._ids_pinned[buffer_index]
225
+ ids_pinned[:n_blocks].numpy()[:] = block_ids
226
+ idx = self._ids_gpu[buffer_index][:n_blocks]
227
+ idx.copy_(ids_pinned[:n_blocks], non_blocking=True)
228
grid = (n_blocks * n_segs * (page_size_el // blk_size),)
229
_zero_kv_blocks_kernel[grid](
230
seg_addrs,