[Bugfix][Offloader] Preserve prefetch static-buffer slot ownership (#54975)
Signed-off-by: Big2Wheel <1404207590@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
yulun committed
Sep 7, 2026 at 10:58 UTC
1f778486fc313f3599a6b59e67f65b1c186477fd
2 files changed
+57
-1
tests/model_executor/offloader/test_prefetch.py
new
+37
@@ -0,0 +1,37 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.model_executor.offloader.prefetch import _get_next_prefetch_index
7
+
8
+pytestmark = pytest.mark.skip_global_cleanup
9
+
10
+
11
+@pytest.mark.parametrize(
12
+ ("module_count", "prefetch_step", "expected_targets"),
13
+ [
14
+ (3, 1, [1, 2, 0]),
15
+ (3, 2, [2, 1, 0]),
16
+ (4, 2, [2, 3, 0, 1]),
17
+ (5, 2, [2, 3, 4, 1, 0]),
18
+ (3, 3, [0, 1, 2]),
19
+ (3, 4, [1, 2, 0]),
20
+ ],
21
+)
22
+def test_next_prefetch_index_preserves_slot_ownership(
23
+ module_count: int,
24
+ prefetch_step: int,
25
+ expected_targets: list[int],
26
+) -> None:
27
+ targets = [
28
+ _get_next_prefetch_index(index, prefetch_step, module_count)
29
+ for index in range(module_count)
30
+ ]
31
+
32
+ assert targets == expected_targets
33
+ if prefetch_step < module_count:
34
+ assert all(
35
+ target % prefetch_step == index % prefetch_step
36
+ for index, target in enumerate(targets)
37
+ )
vllm/model_executor/offloader/prefetch.py
+20
-1
@@ -124,6 +124,21 @@ class StaticBufferPool:
124
return self._buffers[key][slot_idx % self.slot_capacity]
125
126
127
+def _get_next_prefetch_index(
128
+ index: int,
129
+ prefetch_step: int,
130
+ module_count: int,
131
+) -> int:
132
+ """Return a refill target that preserves static-buffer slot ownership."""
133
+ next_index = (index + prefetch_step) % module_count
134
+ if (
135
+ prefetch_step < module_count
136
+ and next_index % prefetch_step != index % prefetch_step
137
+ ):
138
+ next_index = index % prefetch_step
139
+ return next_index
140
+
141
+
142
class PrefetchOffloader(BaseOffloader):
143
"""Prefetching-based offloader with group-based layer selection.
144
@@ -231,7 +246,11 @@ class PrefetchOffloader(BaseOffloader):
246
247
# Start prefetch for next layer (circular)
248
# mutates_args on output_tensor creates ordering dependency
234
- next_index = (index + self.prefetch_step) % len(self.module_offloaders)
249
+ next_index = _get_next_prefetch_index(
250
+ index,
251
+ self.prefetch_step,
252
+ len(self.module_offloaders),
253
+ )
254
# Handle tuple output (e.g., (hidden_states, residual))
255
if isinstance(output, tuple):
256
torch.ops.vllm.start_prefetch(output[0], next_index)