[Bugfix][Core] Wait for the previous PP tensor sends before the next forward pass (#54962)
Signed-off-by: Derek Weitzel <djw8605@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Derek Weitzel committed
Sep 3, 2026 at 03:01 UTC
859dd395128adbce604039f2d725be1067fdf9a3
3 files changed
+85
-7
.buildkite/test_areas/misc.yaml
+1
-1
@@ -43,7 +43,7 @@ steps:
43
- label: ":nvidia: (H200 MIG 35GB) V1 Core + KV + Metrics"
44
device: h200_35gb
45
key: v1-core-kv-metrics
46
- timeout_in_minutes: 80
46
+ timeout_in_minutes: 100
47
env:
48
# GitHub intermittently rejects HTTP/2 upload-pack requests from the H200
49
# CI image. Force Git's smart protocol to use HTTP/1.1 for source installs.
tests/v1/worker/test_gpu_worker.py
+69
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from contextlib import nullcontext
5
from types import SimpleNamespace
6
from unittest.mock import patch
7
@@ -152,3 +153,71 @@ def test_profiling_fallback_declines_off_rocm(rocm):
153
result = _profile_result(consumed=-RELEASED_BY_OTHERS)
154
155
assert maybe_rocm_profiling_fallback(result) is None
156
+
157
+
158
+class _OrderedHandle:
159
+ """Send handle that logs when it is waited."""
160
+
161
+ def __init__(self, log: list[str], name: str):
162
+ self.log = log
163
+ self.name = name
164
+
165
+ def is_completed(self) -> bool:
166
+ return True
167
+
168
+ def wait(self) -> None:
169
+ self.log.append(f"wait:{self.name}")
170
+
171
+
172
+def test_execute_model_waits_previous_pp_send_before_forward(
173
+ monkeypatch: pytest.MonkeyPatch,
174
+):
175
+ """Previous device handles are waited before the forward pass; the
176
+ metadata handle is left to the GroupCoordinator's reaper."""
177
+ import torch
178
+
179
+ from vllm.sequence import IntermediateTensors
180
+
181
+ log: list[str] = []
182
+ previous_tensor_send = _OrderedHandle(log, "prev-tensor")
183
+ metadata_handle = _OrderedHandle(log, "meta")
184
+ tensor_handle = _OrderedHandle(log, "tensor")
185
+
186
+ def isend_tensor_dict(tensors, all_gather_group=None, all_gather_tensors=None):
187
+ log.append("isend")
188
+ return [metadata_handle, tensor_handle]
189
+
190
+ pp_group = SimpleNamespace(
191
+ is_first_rank=True,
192
+ is_last_rank=False,
193
+ isend_tensor_dict=isend_tensor_dict,
194
+ )
195
+ monkeypatch.setattr(gpu_worker, "get_pp_group", lambda: pp_group)
196
+ monkeypatch.setattr(gpu_worker, "get_tp_group", lambda: SimpleNamespace())
197
+
198
+ def run_model(scheduler_output, intermediate_tensors):
199
+ log.append("forward")
200
+ return IntermediateTensors({"hidden_states": torch.zeros(1)})
201
+
202
+ worker = SimpleNamespace(
203
+ vllm_config=SimpleNamespace(
204
+ compilation_config=SimpleNamespace(
205
+ pass_config=SimpleNamespace(enable_sp=False)
206
+ ),
207
+ parallel_config=SimpleNamespace(
208
+ pipeline_parallel_size=2, distributed_executor_backend="mp"
209
+ ),
210
+ ),
211
+ use_v2_model_runner=False,
212
+ model_runner=SimpleNamespace(execute_model=run_model),
213
+ annotate_profile=lambda scheduler_output: nullcontext(),
214
+ _pp_send_work=[previous_tensor_send],
215
+ )
216
+ scheduler_output = SimpleNamespace(
217
+ total_num_scheduled_tokens=4, num_scheduled_tokens={"r0": 4}
218
+ )
219
+
220
+ assert gpu_worker.Worker.execute_model(worker, scheduler_output) is None
221
+
222
+ assert log == ["wait:prev-tensor", "forward", "isend"]
223
+ assert worker._pp_send_work == [tensor_handle]
vllm/v1/worker/gpu_worker.py
+15
-6
@@ -219,6 +219,9 @@ class Worker(WorkerBase):
219
220
self.use_v2_model_runner = vllm_config.use_v2_model_runner
221
222
+ # Device handles of the previous step's PP intermediate-tensor send.
223
+ self._pp_send_work: list[Handle] = []
224
+
225
# Resolved lazily on first sleep/wake; persists worker-process state.
226
self._sleep_mode_backend: SleepModeBackend | None = None
227
@@ -1113,6 +1116,13 @@ class Worker(WorkerBase):
1116
def execute_model(
1117
self, scheduler_output: "SchedulerOutput"
1118
) -> ModelRunnerOutput | AsyncModelRunnerOutput | None:
1119
+ # Wait for the previous step's sends so this forward pass cannot
1120
+ # overwrite buffers they are still reading.
1121
+ if self._pp_send_work:
1122
+ for handle in self._pp_send_work:
1123
+ handle.wait()
1124
+ self._pp_send_work = []
1125
+
1126
intermediate_tensors = None
1127
forward_pass = scheduler_output.total_num_scheduled_tokens > 0
1128
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
@@ -1185,16 +1195,15 @@ class Worker(WorkerBase):
1195
and not get_pp_group().is_last_rank
1196
)
1197
1188
- # launch non-blocking send of intermediate tensors; the
1189
- # GroupCoordinator self-retains the source tensors and the
1190
- # metadata handle in ``_pending_isends`` and reaps them lazily on
1191
- # the next ``isend_tensor_dict(is_async=True)`` call, so this
1192
- # step returns without waiting for the receiver to drain.
1193
- get_pp_group().isend_tensor_dict(
1198
+ # Non-blocking send of the intermediate tensors. The metadata handle
1199
+ # is reaped lazily by the GroupCoordinator; the device handles are
1200
+ # waited at the top of the next step.
1201
+ handles = get_pp_group().isend_tensor_dict(
1202
output.tensors,
1203
all_gather_group=get_tp_group(),
1204
all_gather_tensors=all_gather_tensors,
1205
)
1206
+ self._pp_send_work = handles[1:]
1207
1208
return None
1209