@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e15f20258b
Fork 0

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 2, 03:30 UTC

[ModelRunnerV2] Avoid pipeline parallel bubbles (#42187)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Jun 2, 2026 at 14:02 UTC e15f20258b4448999eee5c133577bf86bf1693a2
19 files changed +615 -144
.buildkite/test_areas/model_runner_v2.yaml
+2
@@ -94,11 +94,13 @@ steps:
94 - vllm/v1/worker/gpu_worker.py
95 - tests/distributed/test_pipeline_parallel.py
96 - tests/distributed/test_pp_cudagraph.py
97 + - tests/v1/distributed/test_pp_dp_v2.py
98 commands:
99 - set -x
100 - export VLLM_USE_V2_MODEL_RUNNER=1
101 - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
102 - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
103 + - pytest -v -s v1/distributed/test_pp_dp_v2.py
104
105 - label: Model Runner V2 Spec Decode
106 device: h200_35gb
tests/v1/distributed/test_pp_dp_v2.py new
+178
@@ -0,0 +1,178 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""V2 ModelRunner + pipeline parallel + data parallel integration tests.
4 +
5 +Covers the interaction between the V2 model runner's PP sampled-token
6 +broadcast and the DP per-step all-reduce across a few concurrency
7 +regimes. Requires 4 GPUs (DP=2, PP=2, TP=1) on CUDA.
8 +"""
9 +
10 +import asyncio
11 +import contextlib
12 +import os
13 +from contextlib import ExitStack
14 +
15 +import pytest
16 +
17 +from vllm import SamplingParams
18 +from vllm.engine.arg_utils import AsyncEngineArgs
19 +from vllm.platforms import current_platform
20 +from vllm.sampling_params import RequestOutputKind
21 +from vllm.v1.engine.async_llm import AsyncLLM
22 +
23 +PP_DP_MODEL = "ibm-research/PowerMoE-3b" # smallest cached MoE that supports PP
24 +PROMPT = "This is a test of data parallel and pipeline parallel together"
25 +
26 +
27 +def _gpu_skip_reason() -> str | None:
28 + if not current_platform.is_cuda():
29 + return "requires CUDA"
30 + n = current_platform.device_count()
31 + if n < 4:
32 + return f"requires 4 GPUs, got {n}"
33 + return None
34 +
35 +
36 +_GPU_SKIP = _gpu_skip_reason()
37 +
38 +pytestmark = [
39 + pytest.mark.skipif(
40 + os.environ.get("VLLM_USE_V2_MODEL_RUNNER", "0") != "1",
41 + reason="VLLM_USE_V2_MODEL_RUNNER=1 required",
42 + ),
43 + pytest.mark.skipif(_GPU_SKIP is not None, reason=_GPU_SKIP or ""),
44 +]
45 +
46 +
47 +def _engine_args(async_scheduling: bool) -> AsyncEngineArgs:
48 + return AsyncEngineArgs(
49 + model=PP_DP_MODEL,
50 + pipeline_parallel_size=2,
51 + data_parallel_size=2,
52 + data_parallel_backend="mp",
53 + tensor_parallel_size=1,
54 + max_model_len=4096,
55 + max_num_batched_tokens=2048,
56 + max_num_seqs=256,
57 + async_scheduling=async_scheduling,
58 + enable_prefix_caching=False,
59 + enforce_eager=False,
60 + enable_expert_parallel=False,
61 + )
62 +
63 +
64 +async def _generate(engine: AsyncLLM, prompt: str, max_tokens: int) -> int:
65 + """Run one streaming completion and return the number of tokens it yielded."""
66 + sampling_params = SamplingParams(
67 + max_tokens=max_tokens,
68 + ignore_eos=True,
69 + output_kind=RequestOutputKind.DELTA,
70 + temperature=0.0,
71 + )
72 + request_id = f"req-{id(prompt):x}-{max_tokens}"
73 + total = 0
74 + async for out in engine.generate(
75 + request_id=request_id, prompt=prompt, sampling_params=sampling_params
76 + ):
77 + total += len(out.outputs[0].token_ids)
78 + return total
79 +
80 +
81 +@pytest.mark.asyncio
82 +@pytest.mark.parametrize("async_scheduling", [True, False])
83 +async def test_pp_dp_v2_low_concurrency(async_scheduling: bool):
84 + """A single in-flight request at a time, repeated, to exercise the
85 + PP slot ring under empty batches between decodes."""
86 + with ExitStack() as after:
87 + engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling))
88 + after.callback(engine.shutdown)
89 +
90 + for _ in range(4):
91 + n = await _generate(engine, PROMPT, max_tokens=16)
92 + assert n == 16
93 +
94 +
95 +@pytest.mark.asyncio
96 +@pytest.mark.parametrize("async_scheduling", [True, False])
97 +async def test_pp_dp_v2_mid_concurrency(async_scheduling: bool):
98 + """64 concurrent requests, staggered, to exercise the steady-state
99 + DP all-reduce + PP slot-ring path."""
100 + with ExitStack() as after:
101 + engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling))
102 + after.callback(engine.shutdown)
103 +
104 + async def _one(i: int) -> int:
105 + await asyncio.sleep(0.01 * i) # stagger so DP load-balances
106 + return await _generate(engine, f"{PROMPT} {i}", max_tokens=64)
107 +
108 + results = await asyncio.gather(*[_one(i) for i in range(64)])
109 + assert all(n == 64 for n in results), results
110 +
111 +
112 +@pytest.mark.asyncio
113 +async def test_pp_dp_v2_abort_mid_decode():
114 + """Cancel half the in-flight requests mid-stream and confirm the
115 + engine survives the abort storm."""
116 +
117 + with ExitStack() as after:
118 + engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True))
119 + after.callback(engine.shutdown)
120 +
121 + async def _maybe_cancel(i: int):
122 + sampling_params = SamplingParams(
123 + max_tokens=64,
124 + ignore_eos=True,
125 + output_kind=RequestOutputKind.DELTA,
126 + temperature=0.0,
127 + )
128 + request_id = f"abort-req-{i}"
129 + count = 0
130 + cancel_at = 4 if i % 2 == 0 else 64
131 + async for out in engine.generate(
132 + request_id=request_id,
133 + prompt=f"{PROMPT} {i}",
134 + sampling_params=sampling_params,
135 + ):
136 + count += len(out.outputs[0].token_ids)
137 + if count >= cancel_at:
138 + break
139 + return count, i
140 +
141 + results = await asyncio.gather(*[_maybe_cancel(i) for i in range(32)])
142 + for count, i in results:
143 + if i % 2 == 0:
144 + assert count >= 4
145 + else:
146 + assert count == 64
147 +
148 + # Engine must still serve after the abort storm.
149 + final = await _generate(engine, "post-abort warmup", max_tokens=8)
150 + assert final == 8
151 +
152 +
153 +@pytest.mark.asyncio
154 +async def test_pp_dp_v2_pause_resume():
155 + """Pause an engine with a request in flight, then resume and confirm
156 + new requests still work."""
157 +
158 + with ExitStack() as after:
159 + engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True))
160 + after.callback(engine.shutdown)
161 +
162 + # Start a long-running generation, let some decoding happen, then
163 + # pause (abort mode) and confirm the in-flight task terminates.
164 + inflight = asyncio.create_task(_generate(engine, PROMPT, max_tokens=128))
165 + await asyncio.sleep(0.5)
166 +
167 + assert not await engine.is_paused()
168 + await engine.pause_generation(mode="abort")
169 + assert await engine.is_paused()
170 +
171 + with contextlib.suppress(Exception):
172 + await inflight
173 +
174 + await engine.resume_generation()
175 + assert not await engine.is_paused()
176 +
177 + n = await _generate(engine, PROMPT, max_tokens=8)
178 + assert n == 8
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+15 -12
@@ -167,7 +167,9 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
167 events = []
168 runner = _make_runner(is_last_pp_rank=False, num_speculative_steps=0)
169 runner.execute_model_state = SimpleNamespace(
170 - input_batch=SimpleNamespace(num_reqs=2),
170 + input_batch=SimpleNamespace(
171 + num_reqs=2, idx_mapping=torch.zeros(2, dtype=torch.int32)
172 + ),
173 attn_metadata=None,
174 slot_mappings_by_layer=None,
175 hidden_states=None,
@@ -175,18 +177,19 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
177 finished_req_ids=set(),
178 num_tokens_across_dp=None,
179 )
178 - runner.postprocess = lambda *args, **kwargs: events.append("postprocess")
179 - runner.eplb.step = lambda *args, **kwargs: events.append("eplb")
180 - monkeypatch.setattr(
181 - mrv2,
182 - "pp_receive",
183 - lambda *args, **kwargs: (
184 - torch.zeros((2, 1), dtype=torch.long),
185 - torch.ones(2, dtype=torch.int32),
186 - torch.zeros(2, dtype=torch.int32),
187 - ),
180 + runner.req_states = SimpleNamespace()
181 +
182 + def fake_receive(*args, **kwargs):
183 + events.append("receive")
184 + # all_decode_next=True, so model_state.postprocess_state is skipped.
185 + return True
186 +
187 + runner.pp_handler = SimpleNamespace(receive=fake_receive)
188 + runner.postprocess_num_computed_tokens = lambda *args, **kwargs: events.append(
189 + "postprocess_num_computed_tokens"
190 )
191 + runner.eplb.step = lambda *args, **kwargs: events.append("eplb")
192
193 output = mrv2.GPUModelRunner.sample_tokens(runner, None)
194 assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)
192 - assert events == ["postprocess", "eplb"]
195 + assert events == ["receive", "postprocess_num_computed_tokens", "eplb"]
vllm/config/vllm.py
+7 -3
@@ -498,9 +498,13 @@ class VllmConfig:
498 # PP requires PP-size concurrent batches to fill the pipeline.
499 # Async scheduling requires 2 concurrent batches to overlap.
500 pp_size = self.parallel_config.pipeline_parallel_size
501 - if pp_size > 1:
502 - return pp_size
503 - return 2 if self.scheduler_config.async_scheduling else 1
501 + if self.scheduler_config.async_scheduling:
502 + if self.use_v2_model_runner:
503 + return pp_size + 1
504 + # V1 Model Runner does not fully support async scheduling with PP.
505 + if pp_size <= 1:
506 + return 2
507 + return pp_size
508
509 @property
510 def num_speculative_tokens(self) -> int:
vllm/distributed/parallel_state.py
+19
@@ -359,6 +359,9 @@ class GroupCoordinator:
359 assert self_cpu_group is not None
360 assert self_device_group is not None
361
362 + self.group_ranks = group_ranks
363 + self.torch_distributed_backend = torch_distributed_backend
364 +
365 self.cpu_group = self_cpu_group
366 self.device_group = self_device_group
367
@@ -406,6 +409,22 @@ class GroupCoordinator:
409 and getattr(self.device_communicator, "supports_tensor_dict", False)
410 )
411
412 + def make_sibling_device_group(self, group_desc: str | None = None) -> ProcessGroup:
413 + """Create a new device-side ProcessGroup with the same per-rank membership
414 + as this coordinator's `device_group`, but backed by a distinct communicator.
415 + This is a collective call: every world rank must invoke it. Used where we
416 + want to issue ops that can run concurrently with ops on `device_group`.
417 + """
418 + sibling: ProcessGroup | None = None
419 + for ranks in self.group_ranks:
420 + pg = torch.distributed.new_group(
421 + ranks, backend=self.torch_distributed_backend, group_desc=group_desc
422 + )
423 + if self.rank in ranks:
424 + sibling = pg
425 + assert sibling is not None
426 + return sibling
427 +
428 def create_mq_broadcaster(
429 self, writer_rank=0, external_writer_handle=None, blocking=True
430 ):
vllm/v1/core/sched/async_scheduler.py
+6
@@ -14,6 +14,7 @@ class AsyncScheduler(Scheduler):
14 super().__init__(*args, **kwargs)
15 # reusable read-only placeholder list for speculative decoding.
16 self._spec_token_placeholders: list[int] = [-1] * self.num_spec_tokens
17 + self.pp_size = self.parallel_config.pipeline_parallel_size
18
19 def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
20 super()._update_after_schedule(scheduler_output)
@@ -34,6 +35,11 @@ class AsyncScheduler(Scheduler):
35 # We will update the actual spec token ids in the worker process.
36 request.spec_token_ids = self._spec_token_placeholders
37
38 + if self.use_v2_model_runner:
39 + # Set the next step index in which this request is eligible to be
40 + # scheduled for decode (for PP microbatching).
41 + request.next_decode_eligible_step = self.current_step + self.pp_size
42 +
43 def _update_request_with_output(
44 self, request: Request, new_token_ids: list[int]
45 ) -> tuple[list[int], bool]:
vllm/v1/core/sched/scheduler.py
+10
@@ -248,6 +248,9 @@ class Scheduler(SchedulerInterface):
248
249 self.use_pp = self.parallel_config.pipeline_parallel_size > 1
250 self.use_v2_model_runner = vllm_config.use_v2_model_runner
251 + # Scheduler iteration counter. Drives the V2+PP+async decode-throttle
252 + # cadence (`next_decode_eligible_step`).
253 + self.current_step = 0
254 self.scheduler_reserve_full_isl = (
255 self.scheduler_config.scheduler_reserve_full_isl
256 )
@@ -333,6 +336,7 @@ class Scheduler(SchedulerInterface):
336 return num_new_tokens
337
338 def schedule(self) -> SchedulerOutput:
339 + self.current_step += 1
340 # NOTE(woosuk) on the scheduling algorithm:
341 # There's no "decoding phase" nor "prefill phase" in the scheduler.
342 # Each request just has the num_computed_tokens and
@@ -388,6 +392,12 @@ class Scheduler(SchedulerInterface):
392 req_index += 1
393 continue
394
395 + if self.current_step < request.next_decode_eligible_step:
396 + # V2+PP+async: enforce `pp_size` steps between same-req decodes
397 + # to match worker-side sampled-tokens broadcast slot ring cadence.
398 + req_index += 1
399 + continue
400 +
401 num_new_tokens = (
402 request.num_tokens_with_spec
403 + request.num_output_placeholders
vllm/v1/engine/core.py
+3 -5
@@ -534,14 +534,12 @@ class EngineCore:
534 if not deferred_scheduler_output:
535 # Add this step's future to the queue.
536 batch_queue.appendleft((future, scheduler_output, exec_future))
537 - if (
538 - model_executed
539 - and len(batch_queue) < self.batch_queue_size
540 - and not batch_queue[-1][0].done()
537 + if len(batch_queue) < self.batch_queue_size and (
538 + model_executed or self.scheduler.has_requests()
539 ):
540 # Don't block on next worker response unless the queue is full
541 # or there are no more requests to schedule.
544 - return None, True
542 + return None, model_executed
543
544 elif not batch_queue:
545 # Queue is empty. We should not reach here since this method should
vllm/v1/request.py
+4
@@ -141,6 +141,10 @@ class Request:
141 self.num_output_placeholders = 0
142 self.async_tokens_to_discard = 0
143
144 + # V2+PP+async: Enforces `pp_size` cadence between same-request decode steps
145 + # so the worker's broadcast slot ring stays consistent.
146 + self.next_decode_eligible_step = 0
147 +
148 self.spec_token_ids: list[int] = []
149 self.num_computed_tokens = 0
150 self.cache_salt: str | None = cache_salt
vllm/v1/worker/gpu/buffer_utils.py
+19 -3
@@ -13,6 +13,15 @@ from vllm.utils.torch_utils import (
13 get_accelerator_view_from_cpu_tensor,
14 )
15
16 +# Default round-robin depth for the UVA buffer pools. Must be >= the number of
17 +# concurrent in-flight steps (engine batch_queue_size).
18 +_DEFAULT_MAX_CONCURRENCY = 2
19 +
20 +
21 +def set_default_max_concurrency(n: int) -> None:
22 + global _DEFAULT_MAX_CONCURRENCY
23 + _DEFAULT_MAX_CONCURRENCY = max(2, n)
24 +
25
26 def async_copy_to_gpu(
27 x: torch.Tensor | np.ndarray,
@@ -47,8 +56,10 @@ class UvaBufferPool:
56 self,
57 size: int | Sequence[int],
58 dtype: torch.dtype,
50 - max_concurrency: int = 2,
59 + max_concurrency: int | None = None,
60 ):
61 + if max_concurrency is None:
62 + max_concurrency = _DEFAULT_MAX_CONCURRENCY
63 self.size = size
64 self.dtype = dtype
65 self.max_concurrency = max_concurrency
@@ -80,7 +91,10 @@ class UvaBufferPool:
91
92 class UvaBackedTensor:
93 def __init__(
83 - self, size: int | Sequence[int], dtype: torch.dtype, max_concurrency: int = 2
94 + self,
95 + size: int | Sequence[int],
96 + dtype: torch.dtype,
97 + max_concurrency: int | None = None,
98 ):
99 self.dtype = dtype
100
@@ -104,9 +118,11 @@ class StagedWriteTensor:
118 size: int | Sequence[int],
119 dtype: torch.dtype,
120 device: torch.device,
107 - max_concurrency: int = 2,
121 + max_concurrency: int | None = None,
122 uva_instead_of_gpu: bool = False,
123 ):
124 + if max_concurrency is None:
125 + max_concurrency = _DEFAULT_MAX_CONCURRENCY
126 supported_dtypes = [torch.int32, torch.int64, torch.float32]
127 if dtype not in supported_dtypes:
128 raise ValueError(
vllm/v1/worker/gpu/input_batch.py
+32 -12
@@ -67,9 +67,18 @@ class InputBatch:
67 seq_lens_cpu_upper_bound: torch.Tensor
68 # [num_reqs]
69 dcp_local_seq_lens: torch.Tensor | None
70 - # [num_reqs] CPU bool array.
70 + # [num_reqs]
71 + num_computed_tokens_np: np.ndarray
72 + # [num_reqs]
73 + prefill_len_np: np.ndarray
74 + # [num_reqs]
75 + num_computed_prefill_tokens_np: np.ndarray
76 + # [num_reqs] CPU bool array == (num_computed_prefill_tokens_np < prefill_len_np).
77 is_prefilling_np: np.ndarray
78
79 + # [num_reqs] only populated when pipeline parallelism is enabled.
80 + max_seq_len_np: np.ndarray | None
81 +
82 # [num_tokens_after_padding]
83 input_ids: torch.Tensor
84 # [num_tokens_after_padding]
@@ -148,7 +157,11 @@ class InputBatch:
157 seq_lens=seq_lens,
158 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
159 dcp_local_seq_lens=None,
160 + num_computed_tokens_np=np.zeros(num_reqs, dtype=np.int32),
161 + prefill_len_np=np.zeros(num_reqs, dtype=np.int32),
162 + num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),
163 is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
164 + max_seq_len_np=None,
165 input_ids=input_ids,
166 positions=positions,
167 logits_indices=logits_indices,
@@ -438,6 +451,9 @@ def _post_update_kernel(
451 ):
452 req_id = tl.program_id(0)
453 req_state_idx = tl.load(idx_mapping_ptr + req_id)
454 + if req_state_idx < 0:
455 + # Filter rows with negative index entries.
456 + return
457
458 total_len = tl.load(total_len_ptr + req_state_idx)
459 num_sampled = tl.load(num_sampled_ptr + req_id)
@@ -464,18 +480,22 @@ def _post_update_kernel(
480 count = tl.load(token_ptr)
481 tl.store(token_ptr, count + 1)
482
467 - query_start = tl.load(query_start_loc_ptr + req_id)
468 - query_end = tl.load(query_start_loc_ptr + req_id + 1)
469 - query_len = query_end - query_start
483 + if query_start_loc_ptr is None:
484 + query_len = 0
485 + else:
486 + query_start = tl.load(query_start_loc_ptr + req_id)
487 + query_end = tl.load(query_start_loc_ptr + req_id + 1)
488 + query_len = query_end - query_start
489 num_rejected = tl.load(num_rejected_ptr + req_id)
490
472 - num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
473 - num_computed += query_len - num_rejected
474 - tl.store(num_computed_tokens_ptr + req_state_idx, num_computed)
491 + computed_delta = query_len - num_rejected
492 + if computed_delta != 0:
493 + num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
494 + tl.store(num_computed_tokens_ptr + req_state_idx, num_computed + computed_delta)
495
496
497 def post_update(
478 - # [num_reqs]
498 + # [num_reqs] batch_idx -> req_state_idx; negative index means skip.
499 idx_mapping: torch.Tensor,
500 # [max_num_reqs]
501 num_computed_tokens: torch.Tensor,
@@ -490,7 +510,7 @@ def post_update(
510 # [num_reqs]
511 num_rejected: torch.Tensor,
512 # [num_reqs + 1]
493 - query_start_loc: torch.Tensor,
513 + query_start_loc: torch.Tensor | None,
514 # [max_num_reqs, max_model_len]
515 all_token_ids: torch.Tensor,
516 # [max_num_reqs]
@@ -516,7 +536,7 @@ def post_update(
536
537
538 @triton.jit
519 -def _post_update_pool_kernel(
539 +def _post_update_num_computed_tokens_kernel(
540 idx_mapping_ptr,
541 num_computed_tokens_ptr,
542 query_start_loc_ptr,
@@ -531,7 +551,7 @@ def _post_update_pool_kernel(
551 tl.store(num_computed_tokens_ptr + req_state_idx, num_computed + query_len)
552
553
534 -def post_update_pool(
554 +def post_update_num_computed_tokens(
555 # [num_reqs]
556 idx_mapping: torch.Tensor,
557 # [max_num_reqs]
@@ -540,7 +560,7 @@ def post_update_pool(
560 query_start_loc: torch.Tensor,
561 ) -> None:
562 num_reqs = idx_mapping.shape[0]
543 - _post_update_pool_kernel[(num_reqs,)](
563 + _post_update_num_computed_tokens_kernel[(num_reqs,)](
564 idx_mapping,
565 num_computed_tokens,
566 query_start_loc,
vllm/v1/worker/gpu/model_runner.py
+94 -41
@@ -61,7 +61,10 @@ from vllm.v1.worker.gpu.attn_utils import (
61 init_kv_cache,
62 )
63 from vllm.v1.worker.gpu.block_table import BlockTables
64 -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
64 +from vllm.v1.worker.gpu.buffer_utils import (
65 + async_copy_to_gpu,
66 + set_default_max_concurrency,
67 +)
68 from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
69 from vllm.v1.worker.gpu.cudagraph_utils import (
70 BatchExecutionDescriptor,
@@ -77,7 +80,7 @@ from vllm.v1.worker.gpu.input_batch import (
80 expand_idx_mapping,
81 get_num_sampled_and_rejected,
82 post_update,
80 - post_update_pool,
83 + post_update_num_computed_tokens,
84 prepare_pos_seq_lens,
85 prepare_prefill_inputs,
86 )
@@ -90,7 +93,7 @@ from vllm.v1.worker.gpu.lora_utils import LoraState
93 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
94 from vllm.v1.worker.gpu.model_states import init_model_state
95 from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
93 -from vllm.v1.worker.gpu.pp_utils import pp_broadcast, pp_receive
96 +from vllm.v1.worker.gpu.pp_utils import PPHandler
97 from vllm.v1.worker.gpu.sample.output import SamplerOutput
98 from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
99 from vllm.v1.worker.gpu.sample.sampler import Sampler
@@ -149,6 +152,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
152 self.is_first_pp_rank = get_pp_group().is_first_rank
153 self.is_last_pp_rank = get_pp_group().is_last_rank
154
155 + # Size the UVA buffer pools to the max number of concurrent in-flight
156 + # steps. Must run before any pooled buffer is constructed
157 + set_default_max_concurrency(vllm_config.max_concurrent_batches)
158 +
159 + # PP broadcast/recv helper. Runs the collective on a side stream.
160 + self.pp_handler: PPHandler | None = None
161 +
162 # Persistent buffer for intermediate tensors (non-first PP ranks).
163 self.intermediate_tensors: IntermediateTensors | None = None
164
@@ -210,6 +220,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
220 device=self.device,
221 )
222
223 + if self.use_pp:
224 + self.pp_handler = PPHandler(
225 + max_num_reqs=self.max_num_reqs,
226 + num_speculative_steps=self.num_speculative_steps,
227 + device=self.device,
228 + )
229 +
230 self.sampler: Sampler | None = None
231 self.rejection_sampler: RejectionSampler | None = None
232 self.prompt_logprobs_worker: PromptLogprobsWorker | None = None
@@ -688,8 +705,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
705 return cuda_graph_size
706
707 def _remove_request(self, req_id: str) -> bool:
691 - if not self.req_states.remove_request(req_id):
708 + req_idx = self.req_states.remove_request(req_id)
709 + if req_idx is None:
710 return False
711 + if self.pp_handler is not None:
712 + self.pp_handler.on_req_idx_freed(req_idx)
713 if self.encoder_cache is not None:
714 self.encoder_cache.remove_request(req_id)
715 if self.prompt_logprobs_worker is not None:
@@ -710,6 +730,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
730 for mm_hash in scheduler_output.free_encoder_mm_hashes:
731 self.encoder_cache.free_encoder_cache(mm_hash)
732
733 + def update_pp_decode_requests(self):
734 + # For non-last PP ranks, update decode requests with sampler output from
735 + # the prior step in which they were scheduled (pp_size steps ago).
736 + if self.pp_handler is not None:
737 + outputs = self.pp_handler.get_prev_sampled_outputs()
738 + if outputs is not None:
739 + self.postprocess_sampled(**outputs)
740 +
741 def add_requests(self, scheduler_output: SchedulerOutput) -> None:
742 for new_req_data in scheduler_output.scheduled_new_reqs:
743 assert new_req_data.prompt_token_ids is not None
@@ -722,11 +750,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
750 self._remove_request(req_id)
751
752 prompt_len = len(new_req_data.prompt_token_ids)
753 + sampling_params = new_req_data.sampling_params
754 self.req_states.add_request(
755 req_id=req_id,
756 prompt_len=prompt_len,
757 all_token_ids=new_req_data.prefill_token_ids,
758 num_computed_tokens=new_req_data.num_computed_tokens,
759 + max_tokens=sampling_params.max_tokens if sampling_params else 1, # type: ignore[arg-type]
760 )
761 req_index = self.req_states.req_id_to_index[req_id]
762
@@ -769,7 +799,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
799 req_index, req_new_block_ids, overwrite=False
800 )
801
772 - # Update num_computed_prefill_tokens.
802 + # Update CPU num_computed_prefill_tokens.
803 np.minimum(
804 self.req_states.num_computed_tokens_np,
805 self.req_states.prefill_len.np,
@@ -848,7 +878,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
878 async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)
879 query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]
880 query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1]
851 - is_prefilling_np = self.req_states.is_prefilling(idx_mapping_np)
881 + prefill_len_np = self.req_states.prefill_len.np[idx_mapping_np]
882 + computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens
883 + num_computed_prefill_tokens_np = computed_prefill_tokens_np[idx_mapping_np]
884 + is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np
885
886 # Get prefill tokens if any.
887 if np.any(is_prefilling_np):
@@ -900,13 +933,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
933 )
934
935 # CPU upper bound on seq_lens; padded entries left at zero.
936 + num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np]
937 seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)
938 np.add(
905 - self.req_states.num_computed_tokens_np[idx_mapping_np],
939 + num_computed_tokens_np,
940 num_scheduled_tokens,
941 out=seq_lens_cpu_upper_bound_np[:num_reqs],
942 )
943 seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)
944 +
945 + max_seq_len_np = None
946 + if self.use_pp:
947 + # max_seq_len is only consumed by the PP `compute_need_sampled_mask`
948 + max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np]
949 return InputBatch(
950 req_ids=req_ids,
951 num_reqs=num_reqs,
@@ -925,7 +964,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
964 seq_lens=seq_lens,
965 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
966 dcp_local_seq_lens=dcp_local_seq_lens,
967 + num_computed_tokens_np=num_computed_tokens_np,
968 + prefill_len_np=prefill_len_np,
969 + num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,
970 is_prefilling_np=is_prefilling_np,
971 + max_seq_len_np=max_seq_len_np,
972 input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
973 positions=self.input_buffers.positions[:num_tokens_after_padding],
974 logits_indices=logits_indices,
@@ -1005,12 +1048,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1048 )
1049 return sampler_output, num_sampled, num_rejected
1050
1008 - def postprocess(
1051 + def postprocess_sampled(
1052 self,
1010 - input_batch: InputBatch,
1053 + idx_mapping: torch.Tensor, # May include -1 for masked entries
1054 sampled_tokens: torch.Tensor,
1055 num_sampled: torch.Tensor,
1056 num_rejected: torch.Tensor,
1057 + query_start_loc: torch.Tensor | None = None,
1058 ) -> None:
1059 # Update the number of computed tokens.
1060 if self.is_last_pp_rank:
@@ -1019,19 +1063,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1063 else:
1064 output_bin_counts = None
1065 post_update(
1022 - input_batch.idx_mapping,
1066 + idx_mapping,
1067 self.req_states.num_computed_tokens.gpu,
1068 self.req_states.last_sampled_tokens,
1069 output_bin_counts,
1070 sampled_tokens,
1071 num_sampled,
1072 num_rejected,
1029 - input_batch.query_start_loc,
1073 + query_start_loc,
1074 self.req_states.all_token_ids.gpu,
1075 self.req_states.total_len.gpu,
1076 )
1077
1034 - self.model_state.postprocess_state(input_batch, num_sampled)
1078 + self.model_state.postprocess_state(idx_mapping, num_sampled)
1079
1080 @torch.inference_mode()
1081 def execute_model(
@@ -1044,6 +1088,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1088 ) -> ModelRunnerOutput | IntermediateTensors | None:
1089 if not dummy_run:
1090 # Update the request states.
1091 + self.update_pp_decode_requests()
1092 self.finish_requests(scheduler_output)
1093 self.free_states(scheduler_output)
1094 self.add_requests(scheduler_output)
@@ -1162,9 +1207,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1207 # NOTE(woosuk): We must call get_mm_embeddings even during dummy runs
1208 # to obtain inputs_embeds, because the compiled model expects this input.
1209 inputs_embeds = self.model_state.get_mm_embeddings(
1165 - scheduler_output.scheduled_encoder_inputs,
1166 - input_batch,
1167 - self.req_states,
1210 + scheduler_output.scheduled_encoder_inputs, input_batch
1211 )
1212
1213 model_inputs = {
@@ -1184,12 +1227,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1227 assert intermediate_tensors is not None
1228 assert self.intermediate_tensors is not None
1229 n = input_batch.num_tokens_after_padding
1187 - model_inputs["intermediate_tensors"] = IntermediateTensors(
1188 - {
1189 - k: v[:n].copy_(intermediate_tensors.tensors[k][:n])
1190 - for k, v in self.intermediate_tensors.tensors.items()
1191 - }
1192 - )
1230 + new_tensors = {
1231 + k: v[:n]
1232 + if dummy_run
1233 + else v[:n].copy_(intermediate_tensors.tensors[k][:n])
1234 + for k, v in self.intermediate_tensors.tensors.items()
1235 + }
1236 + model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors)
1237 del intermediate_tensors
1238
1239 # Run model.
@@ -1281,10 +1325,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1325 # Non-last PP rank: hidden_states is None because this rank produced
1326 # IntermediateTensors instead of final hidden states. Receive the
1327 # sampled tokens broadcast from the last rank and update local state.
1284 - sampled, num_sampled, num_rejected = pp_receive(
1285 - input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1
1286 - )
1287 - self.postprocess(input_batch, sampled, num_sampled, num_rejected)
1328 + assert self.pp_handler is not None
1329 + all_decode_next = self.pp_handler.receive(input_batch)
1330 + # Optimistically update num_computed_tokens for entire batch here.
1331 + # Will be adjusted for rejections if necessary in update_requests.
1332 + self.postprocess_num_computed_tokens(input_batch)
1333 + if not all_decode_next:
1334 + # Might contain non-final prefill chunks, which will be scheduled
1335 + # in the immediate next step (rather than in pp_size steps).
1336 + self.model_state.postprocess_state(input_batch.idx_mapping, 0)
1337
1338 # Post-step KV connector related operations.
1339 kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
@@ -1295,9 +1344,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1344 hidden_states, input_batch, grammar_output
1345 )
1346
1298 - if self.use_pp:
1347 + if self.pp_handler is not None:
1348 # Broadcast to non-last PP ranks (handles spec decode multi-token).
1300 - pp_broadcast(sampler_output.sampled_token_ids, num_sampled, num_rejected)
1349 + self.pp_handler.broadcast(
1350 + sampler_output.sampled_token_ids,
1351 + num_sampled,
1352 + num_rejected,
1353 + input_batch,
1354 + )
1355
1356 assert self.prompt_logprobs_worker is not None
1357 prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
@@ -1307,8 +1361,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1361 self.req_states.all_token_ids.gpu,
1362 self.req_states.num_computed_tokens.gpu,
1363 self.req_states.prompt_len.np,
1310 - self.req_states.prefill_len.np,
1311 - self.req_states.num_computed_prefill_tokens,
1364 )
1365
1366 # Prepare the model runner output.
@@ -1334,17 +1386,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1386 # Get cached multimodal embeddings for draft forward.
1387 # NOTE: This is done here because postprocess updates
1388 # num_computed_prefill_tokens.
1337 - prefill_lens = self.req_states.prefill_len.np[input_batch.idx_mapping_np]
1338 - computed_prefill_lens = self.req_states.num_computed_prefill_tokens[
1339 - input_batch.idx_mapping_np
1340 - ]
1389 mm_inputs = self.model_state.encoder_runner.gather_mm_embeddings(
1390 input_batch.req_ids,
1391 input_batch.num_tokens,
1392 input_batch.num_scheduled_tokens,
1393 input_batch.query_start_loc_np,
1346 - prefill_lens,
1347 - computed_prefill_lens + 1, # +1 to consider the skew in eagle
1394 + input_batch.prefill_len_np,
1395 + # +1 to consider the skew in eagle
1396 + input_batch.num_computed_prefill_tokens_np + 1,
1397 )
1398
1399 # Postprocess results and update request states.
@@ -1352,8 +1401,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1401 # ensuring that `copy_event` is recorded before calling postprocess.
1402 # This sequencing may slightly reduce latency as async D2H copy does not
1403 # need to wait for the postprocess to finish.
1355 - self.postprocess(
1356 - input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
1404 + self.postprocess_sampled(
1405 + input_batch.idx_mapping,
1406 + sampler_output.sampled_token_ids,
1407 + num_sampled,
1408 + num_rejected,
1409 + input_batch.query_start_loc,
1410 )
1411
1412 if self.speculator is not None:
@@ -1410,7 +1463,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1463 kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
1464
1465 if not self.is_last_pp_rank:
1413 - self.postprocess_pool(input_batch)
1466 + self.postprocess_num_computed_tokens(input_batch)
1467 return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
1468
1469 assert self.pooling_runner is not None
@@ -1432,14 +1485,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1485 copy_stream=self.output_copy_stream,
1486 )
1487
1435 - self.postprocess_pool(input_batch)
1488 + self.postprocess_num_computed_tokens(input_batch)
1489 if self.use_async_scheduling:
1490 return async_output
1491 return async_output.get_output()
1492
1440 - def postprocess_pool(self, input_batch: InputBatch) -> None:
1493 + def postprocess_num_computed_tokens(self, input_batch: InputBatch) -> None:
1494 # Update the number of computed tokens.
1442 - post_update_pool(
1495 + post_update_num_computed_tokens(
1496 input_batch.idx_mapping,
1497 self.req_states.num_computed_tokens.gpu,
1498 input_batch.query_start_loc,
vllm/v1/worker/gpu/model_states/default.py
+2 -3
@@ -102,7 +102,6 @@ class DefaultModelState(ModelState):
102 self,
103 scheduled_encoder_inputs: dict[str, list[int]],
104 input_batch: InputBatch,
105 - req_states: RequestState,
105 ) -> torch.Tensor:
106 mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
107 scheduled_encoder_inputs
@@ -118,8 +117,8 @@ class DefaultModelState(ModelState):
117 input_batch.num_tokens,
118 input_batch.num_scheduled_tokens,
119 input_batch.query_start_loc_np,
121 - req_states.prefill_len.np[input_batch.idx_mapping_np],
122 - req_states.num_computed_prefill_tokens[input_batch.idx_mapping_np],
120 + input_batch.prefill_len_np,
121 + input_batch.num_computed_prefill_tokens_np,
122 )
123 # Use unpadded input_ids to match is_mm_embed size (num_tokens).
124 # input_batch.input_ids may be padded for CUDA graphs.
vllm/v1/worker/gpu/model_states/interface.py
+2 -7
@@ -57,18 +57,13 @@ class ModelState(ABC):
57 return None
58
59 def postprocess_state(
60 - self,
61 - input_batch: InputBatch,
62 - num_sampled: torch.Tensor,
60 + self, idx_mapping: torch.Tensor, num_sampled: torch.Tensor
61 ) -> None:
62 return None
63
64 @abstractmethod
65 def get_mm_embeddings(
68 - self,
69 - scheduled_encoder_inputs: dict[str, list[int]],
70 - input_batch: InputBatch,
71 - req_states: RequestState,
66 + self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch
67 ) -> torch.Tensor | None:
68 raise NotImplementedError
69
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+31 -12
@@ -9,6 +9,7 @@ import torch.nn as nn
9
10 from vllm.config import VllmConfig
11 from vllm.config.compilation import CUDAGraphMode
12 +from vllm.triton_utils import tl, triton
13 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
14 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
15 from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -112,13 +113,10 @@ class MambaHybridModelState(DefaultModelState):
113 # need the -1 sentinel rather than a raw zero draft count.
114 num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32)
115 if input_batch.num_draft_tokens_per_req is not None:
115 - spec_decode_mask = (
116 - input_batch.num_draft_tokens_per_req > 0
117 - ) & ~input_batch.is_prefilling_np
116 + has_draft_tokens = input_batch.num_draft_tokens_per_req > 0
117 + spec_decode_mask = has_draft_tokens & ~input_batch.is_prefilling_np
118 num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where(
119 - spec_decode_mask,
120 - input_batch.num_draft_tokens_per_req,
121 - -1,
119 + spec_decode_mask, input_batch.num_draft_tokens_per_req, -1
120 )
121 num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np)
122
@@ -145,12 +143,33 @@ class MambaHybridModelState(DefaultModelState):
143 )
144
145 def postprocess_state(
148 - self,
149 - input_batch: InputBatch,
150 - num_sampled: torch.Tensor,
146 + self, idx_mapping: torch.Tensor, num_sampled: torch.Tensor | int
147 ) -> None:
148 # Chunked prefill does not sample a token, so num_sampled can be 0.
149 # Mamba treats num_accepted_tokens=1 as the neutral non-spec value.
154 - self.num_accepted_tokens_gpu[input_batch.idx_mapping] = torch.clamp(
155 - num_sampled, min=1
156 - )
150 + if not isinstance(num_sampled, int):
151 + # idx_mapping may contain -1 sentinels (filtered rows) under PP; the
152 + # kernel skips them rather than scattering with a host-side gather.
153 + num_reqs = idx_mapping.shape[0]
154 + if num_reqs:
155 + _scatter_num_accepted_kernel[(num_reqs,)](
156 + idx_mapping, num_sampled, self.num_accepted_tokens_gpu
157 + )
158 + return
159 +
160 + # Fill with single value.
161 + self.num_accepted_tokens_gpu.index_fill_(0, idx_mapping, max(num_sampled, 1))
162 +
163 +
164 +@triton.jit
165 +def _scatter_num_accepted_kernel(
166 + idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip)
167 + num_sampled_ptr, # [num_reqs]
168 + num_accepted_ptr, # [max_num_reqs]
169 +):
170 + row = tl.program_id(0)
171 + req_state_idx = tl.load(idx_mapping_ptr + row)
172 + if req_state_idx < 0:
173 + return
174 + num_sampled = tl.load(num_sampled_ptr + row)
175 + tl.store(num_accepted_ptr + req_state_idx, tl.maximum(num_sampled, 1))
vllm/v1/worker/gpu/model_states/whisper.py
+1 -4
@@ -84,10 +84,7 @@ class WhisperModelState(ModelState):
84 return ("transcription",)
85
86 def get_mm_embeddings(
87 - self,
88 - scheduled_encoder_inputs: dict[str, list[int]],
89 - input_batch: InputBatch,
90 - req_states: RequestState,
87 + self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch
88 ) -> None:
89 # Ensure encoder inputs are ordered consistently with input_batch.req_ids.
90 encoder_inputs: dict[str, list[int]] = {}
vllm/v1/worker/gpu/pp_utils.py
+179 -26
@@ -2,40 +2,193 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Pipeline Parallelism utils for V2 Model Runner."""
4
5 +from collections import deque
6 +from dataclasses import dataclass
7 +
8 +import numpy as np
9 import torch
10
11 from vllm.distributed.parallel_state import get_pp_group
12 +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
13 +from vllm.v1.worker.gpu.input_batch import InputBatch
14 +
15 +
16 +@dataclass
17 +class PendingRecv:
18 + """Per-step slot data for a deferred postprocess on the main stream."""
19 +
20 + event: torch.cuda.Event
21 +
22 + sampled_tokens: torch.Tensor # [num_reqs, max_sample_len]
23 + num_sampled: torch.Tensor # [num_reqs]
24 + num_rejected: torch.Tensor # [num_reqs]
25 + idx_mapping: torch.Tensor # [num_reqs]
26 + idx_mapping_np: np.ndarray # [num_reqs]
27 + # Records which rows need a deferred postprocess (bool).
28 + need_sampled_mask: np.ndarray # [num_reqs]
29 + # Snapshot of slot generation counters at receive time, used to
30 + # detect requests aborted since then.
31 + gen_at_receive_np: np.ndarray # [num_reqs]
32 +
33 +
34 +def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None:
35 + """Return a bool array of shape `[input_batch.num_reqs]` marking requests
36 + with outputs that might be needed in a subsequent (decode) step.
37 + Returns None if no sampled outputs are needed in the requests' next step."""
38 +
39 + old_computed = input_batch.num_computed_tokens_np
40 + prefill_len = input_batch.prefill_len_np
41 + max_seq_len = input_batch.max_seq_len_np
42 + assert max_seq_len is not None # always populated under PP
43 + # Exclude non-final prefill chunks (they don't produce a sample).
44 + produces_sample = old_computed + input_batch.num_scheduled_tokens >= prefill_len
45 + # Exclude requests that we know are finished.
46 + not_finishing = np.maximum(old_computed, prefill_len) + 1 < max_seq_len
47 + need_sampled_mask = produces_sample & not_finishing
48 + return need_sampled_mask if need_sampled_mask.any() else None
49 +
50 +
51 +class PPHandler:
52 + """Runs the PP sampled-token broadcast/recv on a side stream so the
53 + default stream isn't gated by the matching peer call. Step T's recv is
54 + consumed at step T+pp_size via `get_prev_sampled_outputs`.
55 +
56 + Uses a dedicated NCCL communicator (sibling of the PP `device_group`)
57 + for the broadcast so it does not serialize on the wire with the
58 + inter-stage hidden-state p2p send/recv ops.
59 + """
60 +
61 + def __init__(
62 + self, max_num_reqs: int, num_speculative_steps: int, device: torch.device
63 + ):
64 + self.is_last_rank = get_pp_group().is_last_rank
65 + self.last_rank = get_pp_group().last_rank
66 + self.max_sample_len = num_speculative_steps + 1
67 + self.device = device
68 + self.main_stream = torch.cuda.current_stream(device)
69 + self.broadcast_stream = torch.cuda.Stream(device)
70 +
71 + # On non-last ranks, a FIFO with one entry per in-flight step: the entry
72 + # pushed by step T's `receive` is consumed pp_size steps later. Pre-seeded
73 + # with pp_size None placeholders so the first pp_size consumes are no-ops.
74 + # None means no postprocess is pending for that step (broadcast skipped).
75 + self.queue: deque[PendingRecv | None] = (
76 + deque() if self.is_last_rank else deque([None] * get_pp_group().world_size)
77 + )
78 +
79 + # Per req-index generation counter, incremented every time a request
80 + # index is freed in RequestStats. Used for invalidating freed req data
81 + # between PP decodes.
82 + self.req_idx_gen_np = np.zeros(max_num_reqs, dtype=np.int32)
83 +
84 + # Dedicated subgroup for the sampled-token broadcast.
85 + self.broadcast_group = get_pp_group().make_sibling_device_group(
86 + group_desc="pp_broadcast"
87 + )
88 +
89 + def on_req_idx_freed(self, req_idx: int) -> None:
90 + self.req_idx_gen_np[req_idx] += 1
91
92 + def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None:
93 + """Consume the entry from pp_size steps ago and wait for its recv event,
94 + then filter out entries whose request was freed since `receive`.
95 + """
96 + if not self.queue:
97 + return None
98 + slot = self.queue.popleft()
99 + # Reserve this step's slot; `receive` overwrites it if applicable.
100 + self.queue.append(None)
101 + if slot is None:
102 + return None
103
10 -def pp_broadcast(
11 - sampled_token_ids: torch.Tensor,
12 - num_sampled: torch.Tensor,
13 - num_rejected: torch.Tensor,
14 -) -> None:
15 - pp = get_pp_group()
16 - assert pp.is_last_rank
104 + # Skip requests which did not need sampled output and/or those already
105 + # finished. The post_update kernel skips the -1 entries.
106 + freed = self.req_idx_gen_np[slot.idx_mapping_np] != slot.gen_at_receive_np
107 + exclude_mask = freed | ~slot.need_sampled_mask
108 + idx_mapping = slot.idx_mapping
109 + if exclude_mask.any():
110 + if exclude_mask.all():
111 + # No states require update anymore.
112 + return None
113 + # Filter excluded request indices.
114 + idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np)
115 + idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
116
18 - assert sampled_token_ids.dtype == torch.int64
19 - torch.distributed.broadcast(
20 - sampled_token_ids.contiguous(), src=pp.last_rank, group=pp.device_group
21 - )
117 + self.main_stream.wait_event(slot.event)
118 + return dict(
119 + sampled_tokens=slot.sampled_tokens,
120 + num_sampled=slot.num_sampled,
121 + num_rejected=slot.num_rejected,
122 + idx_mapping=idx_mapping,
123 + )
124
23 - combined = torch.stack((num_sampled, num_rejected), dim=0)
24 - torch.distributed.broadcast(combined, src=pp.last_rank, group=pp.device_group)
125 + def receive(self, input_batch: InputBatch) -> bool:
126 + """Returns True iff sampled tokens need to be gathered from *all*
127 + requests in the batch."""
128 + assert not self.is_last_rank
129 + need_sampled_mask = compute_need_sampled_mask(input_batch)
130 + if need_sampled_mask is None:
131 + # Leave this step's reserved slot as None.
132 + return False
133
134 + # Snapshot the per-slot generation counter so a later free of any of
135 + # these RequestStates request indices is detectable at consume time.
136 + gen_at_receive_np = self.req_idx_gen_np[input_batch.idx_mapping_np]
137
27 -def pp_receive(
28 - num_reqs: int, max_sample_len: int = 1
29 -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
30 - pp = get_pp_group()
31 - assert not pp.is_last_rank
138 + num_reqs = input_batch.num_reqs
139 + with torch.cuda.stream(self.broadcast_stream):
140 + self.broadcast_stream.wait_stream(self.main_stream)
141 + sampled_tokens = torch.empty(
142 + num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device
143 + )
144 + combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device)
145 + torch.distributed.broadcast(
146 + sampled_tokens, src=self.last_rank, group=self.broadcast_group
147 + )
148 + torch.distributed.broadcast(
149 + combined, src=self.last_rank, group=self.broadcast_group
150 + )
151 + event = self.broadcast_stream.record_event()
152 + num_sampled, num_rejected = combined.unbind(dim=0)
153 + # Must record_stream since these were allocated on broadcast stream but
154 + # later used on the main stream.
155 + sampled_tokens.record_stream(self.main_stream)
156 + combined.record_stream(self.main_stream)
157 + self.queue[-1] = PendingRecv(
158 + event,
159 + sampled_tokens,
160 + num_sampled,
161 + num_rejected,
162 + input_batch.idx_mapping,
163 + input_batch.idx_mapping_np,
164 + need_sampled_mask,
165 + gen_at_receive_np,
166 + )
167 + return bool(need_sampled_mask.all())
168
33 - sampled_tokens = torch.empty(
34 - num_reqs, max_sample_len, dtype=torch.int64, device=pp.device
35 - )
36 - torch.distributed.broadcast(sampled_tokens, src=pp.last_rank, group=pp.device_group)
169 + def broadcast(
170 + self,
171 + sampled_token_ids: torch.Tensor,
172 + num_sampled: torch.Tensor,
173 + num_rejected: torch.Tensor,
174 + input_batch: InputBatch,
175 + ) -> None:
176 + assert self.is_last_rank
177 + if compute_need_sampled_mask(input_batch) is None:
178 + # No request needs sampled outputs for a subsequent decode step.
179 + return
180
38 - combined = torch.empty(2, num_reqs, dtype=torch.int32, device=pp.device)
39 - torch.distributed.broadcast(combined, src=pp.last_rank, group=pp.device_group)
40 - num_sampled, num_rejected = combined.unbind(dim=0)
41 - return sampled_tokens, num_sampled, num_rejected
181 + assert sampled_token_ids.dtype == torch.int64
182 + with torch.cuda.stream(self.broadcast_stream):
183 + self.broadcast_stream.wait_stream(self.main_stream)
184 + torch.distributed.broadcast(
185 + sampled_token_ids.contiguous(),
186 + src=self.last_rank,
187 + group=self.broadcast_group,
188 + )
189 + combined = torch.stack((num_sampled, num_rejected), dim=0)
190 + torch.distributed.broadcast(
191 + combined, src=self.last_rank, group=self.broadcast_group
192 + )
193 + for tensor in (sampled_token_ids, num_sampled, num_rejected):
194 + tensor.record_stream(self.broadcast_stream)
vllm/v1/worker/gpu/sample/prompt_logprob.py
+2 -6
@@ -42,10 +42,6 @@ class PromptLogprobsWorker:
42 num_computed_tokens: torch.Tensor,
43 # [max_num_reqs]
44 prompt_lens: np.ndarray,
45 - # [max_num_reqs]
46 - prefill_lens: np.ndarray,
47 - # [max_num_reqs]
48 - num_computed_prefill_tokens: np.ndarray,
45 ) -> dict[str, LogprobsTensors]:
46 idx_mapping_np = input_batch.idx_mapping_np
47 needs_prompt_logprobs = self.uses_prompt_logprobs[idx_mapping_np]
@@ -55,11 +51,11 @@ class PromptLogprobsWorker:
51
52 num_prompt_logprobs = self.num_prompt_logprobs[idx_mapping_np]
53 prompt_lens = prompt_lens[idx_mapping_np]
58 - computed_prefill = num_computed_prefill_tokens[idx_mapping_np]
54 + computed_prefill = input_batch.num_computed_prefill_tokens_np
55 includes_prompt = computed_prefill < prompt_lens
56 # NOTE(woosuk): If the request was resumed after preemption, its prompt
57 # logprobs must have been computed before preemption. Skip.
62 - resumed_after_prompt = prompt_lens < prefill_lens[idx_mapping_np]
58 + resumed_after_prompt = prompt_lens < input_batch.prefill_len_np
59 needs_prompt_logprobs &= includes_prompt & ~resumed_after_prompt
60 if not np.any(needs_prompt_logprobs):
61 return {}
vllm/v1/worker/gpu/states.py
+9 -10
@@ -65,6 +65,9 @@ class RequestState:
65 self.max_num_reqs, 1, dtype=torch.int64, device=device
66 )
67
68 + # Max total seq length (prompt_len + max_tokens).
69 + self.max_seq_len = np.zeros(self.max_num_reqs, dtype=np.int32)
70 +
71 # Draft tokens.
72 self.draft_tokens = torch.zeros(
73 self.max_num_reqs,
@@ -87,12 +90,14 @@ class RequestState:
90 prompt_len: int,
91 all_token_ids: list[int],
92 num_computed_tokens: int,
93 + max_tokens: int,
94 ) -> None:
95 assert len(self.free_indices) > 0, "No free indices"
96 req_idx = self.free_indices.pop()
97 self.req_id_to_index[req_id] = req_idx
98 self.index_to_req_id[req_idx] = req_id
99
100 + self.max_seq_len[req_idx] = prompt_len + max_tokens
101 self.prompt_len.np[req_idx] = prompt_len
102 prefill_len = len(all_token_ids)
103 assert prefill_len >= prompt_len, (
@@ -124,17 +129,11 @@ class RequestState:
129 self.all_token_ids.apply_write()
130 self.num_computed_tokens.apply_write()
131
127 - def remove_request(self, req_id: str) -> bool:
132 + def remove_request(self, req_id: str) -> int | None:
133 + """Return the freed slot index, or None if the request was not found."""
134 req_idx = self.req_id_to_index.pop(req_id, None)
135 if req_idx is None:
130 - # Request not found.
131 - return False
136 + return None
137 self.index_to_req_id.pop(req_idx, None)
138 self.free_indices.append(req_idx)
134 - return True
135 -
136 - def is_prefilling(self, idx_mapping_np: np.ndarray) -> np.ndarray:
137 - return (
138 - self.num_computed_prefill_tokens[idx_mapping_np]
139 - < self.prefill_len.np[idx_mapping_np]
140 - )
139 + return req_idx