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 8, 03:45 UTC
[Model Runner V2] A bit more PP simplification (#34766)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Feb 17, 2026 at 21:39 UTC
a49ea5a58fc0f8170027abd79168d6f7ca3e4789
2 files changed
+8
-12
vllm/v1/worker/gpu/model_runner.py
+5
-7
@@ -1003,15 +1003,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1003
hidden_states, input_batch, kv_connector_output = self.execute_model_state
1004
self.execute_model_state = None # type: ignore
1005
1006
- # Non-last PP rank: hidden_states is None because this rank produced
1007
- # IntermediateTensors instead of final hidden states. Receive the
1008
- # sampled tokens broadcast by the last rank and update local state.
1006
if not self.is_last_pp_rank:
1010
- received = pp_receive(
1007
+ # Non-last PP rank: hidden_states is None because this rank produced
1008
+ # IntermediateTensors instead of final hidden states. Receive the
1009
+ # sampled tokens broadcast from the last rank and update local state.
1010
+ sampled, num_sampled, num_rejected = pp_receive(
1011
input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1
1012
)
1013
- assert received is not None
1014
- sampled, num_sampled, num_rejected = received
1013
self.postprocess(input_batch, sampled, num_sampled, num_rejected)
1014
return None
1015
@@ -1020,8 +1018,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1018
hidden_states, input_batch, grammar_output
1019
)
1020
1023
- # Broadcast to non-last PP ranks (handles spec decode multi-token).
1021
if self.use_pp:
1022
+ # Broadcast to non-last PP ranks (handles spec decode multi-token).
1023
pp_broadcast(sampler_output.sampled_token_ids, num_sampled, num_rejected)
1024
1025
prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
vllm/v1/worker/gpu/pp_utils.py
+3
-5
@@ -13,8 +13,7 @@ def pp_broadcast(
13
num_rejected: torch.Tensor,
14
) -> None:
15
pp = get_pp_group()
16
- if not pp.is_last_rank:
17
- return
16
+ assert pp.is_last_rank
17
18
assert sampled_token_ids.dtype == torch.int64
19
torch.distributed.broadcast(
@@ -27,10 +26,9 @@ def pp_broadcast(
26
27
def pp_receive(
28
num_reqs: int, max_sample_len: int = 1
30
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
29
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
30
pp = get_pp_group()
32
- if pp.is_last_rank:
33
- return None
31
+ assert not pp.is_last_rank
32
33
sampled_tokens = torch.empty(
34
num_reqs, max_sample_len, dtype=torch.int64, device=pp.device