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 4, 06:15 UTC
[Bugfix] Skip PP sampled-token receive on last rank during async scheduling (#40749)
Signed-off-by: Adam Winstanley <adam@winstanley.industries>
wi-adam committed
May 5, 2026 at 22:31 UTC
b53c507bc91f87e28b03e9b54bbff7c76e97d58b
2 files changed
+54
-1
tests/v1/worker/test_gpu_model_runner.py
+53
@@ -7,6 +7,7 @@ import numpy as np
7
import pytest
8
import torch
9
10
+import vllm.v1.worker.gpu_model_runner as gpu_model_runner_module
11
from vllm.config import (
12
AttentionConfig,
13
CacheConfig,
@@ -219,6 +220,58 @@ def test_select_common_block_size_uses_largest_shared_int():
220
assert selected_size == 64
221
222
223
+@pytest.mark.skip_global_cleanup
224
+@pytest.mark.parametrize(
225
+ ("world_size", "is_last_rank", "expected_calls"),
226
+ [(1, True, 0), (2, True, 0), (2, False, 1)],
227
+)
228
+def test_sample_tokens_receives_pp_sampled_ids_only_on_non_last_rank(
229
+ monkeypatch: pytest.MonkeyPatch,
230
+ world_size: int,
231
+ is_last_rank: bool,
232
+ expected_calls: int,
233
+):
234
+ runner = GPUModelRunner.__new__(GPUModelRunner)
235
+ runner.execute_model_state = None
236
+ runner.kv_connector_output = None
237
+ runner.use_async_scheduling = True
238
+ receive_calls = 0
239
+
240
+ def receive_prev_sampled_token_ids():
241
+ nonlocal receive_calls
242
+ receive_calls += 1
243
+
244
+ runner._pp_receive_prev_sampled_token_ids_to_input_batch = (
245
+ receive_prev_sampled_token_ids
246
+ )
247
+ monkeypatch.setattr(
248
+ gpu_model_runner_module,
249
+ "get_pp_group",
250
+ lambda: SimpleNamespace(world_size=world_size, is_last_rank=is_last_rank),
251
+ )
252
+
253
+ assert GPUModelRunner.sample_tokens(runner, None) is None
254
+ assert receive_calls == expected_calls
255
+
256
+
257
+@pytest.mark.skip_global_cleanup
258
+def test_sample_tokens_skips_pp_group_lookup_without_async_scheduling(
259
+ monkeypatch: pytest.MonkeyPatch,
260
+):
261
+ runner = GPUModelRunner.__new__(GPUModelRunner)
262
+ runner.execute_model_state = None
263
+ runner.kv_connector_output = None
264
+ runner.use_async_scheduling = False
265
+
266
+ monkeypatch.setattr(
267
+ gpu_model_runner_module,
268
+ "get_pp_group",
269
+ pytest.fail,
270
+ )
271
+
272
+ assert GPUModelRunner.sample_tokens(runner, None) is None
273
+
274
+
275
def test_select_common_block_size_no_valid_option():
276
backend_a = _make_mock_backend_for_kernel_block_size([64])
277
backend_b = _make_mock_backend_for_kernel_block_size([MultipleOf(16)])
vllm/v1/worker/gpu_model_runner.py
+1
-1
@@ -4182,7 +4182,7 @@ class GPUModelRunner(
4182
kv_connector_output = self.kv_connector_output
4183
self.kv_connector_output = None
4184
# receive sampled token ids from the last PP rank.
4185
- if self.use_async_scheduling and get_pp_group().world_size > 1:
4185
+ if self.use_async_scheduling and not get_pp_group().is_last_rank:
4186
self._pp_receive_prev_sampled_token_ids_to_input_batch()
4187
if not kv_connector_output:
4188
return None # type: ignore[return-value]