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 6, 02:00 UTC
[Model Runner V2][Minor] Simplify PP logic (#38031)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 24, 2026 at 13:57 UTC
4e824d1c835d9b57db621297e8d9119bfc32fb2e
3 files changed
+23
-28
vllm/sequence.py
+9
@@ -62,3 +62,12 @@ class IntermediateTensors:
62
63
def __repr__(self) -> str:
64
return f"IntermediateTensors(tensors={self.tensors})"
65
+
66
+ @staticmethod
67
+ def empty_like(
68
+ intermediate_tensors: "IntermediateTensors",
69
+ ) -> "IntermediateTensors":
70
+ tensors = {
71
+ k: torch.empty_like(v) for k, v in intermediate_tensors.tensors.items()
72
+ }
73
+ return IntermediateTensors(tensors)
vllm/v1/worker/gpu/cudagraph_utils.py
+5
-13
@@ -94,13 +94,8 @@ class CudaGraphManager:
94
self.decode_query_len = decode_query_len
95
96
self.dp_size = vllm_config.parallel_config.data_parallel_size
97
- self.pp_size = vllm_config.parallel_config.pipeline_parallel_size
98
- if self.pp_size > 1:
99
- self.is_first_pp_rank = get_pp_group().is_first_rank
100
- self.is_last_pp_rank = get_pp_group().is_last_rank
101
- else:
102
- self.is_first_pp_rank = True
103
- self.is_last_pp_rank = True
97
+ self.is_first_pp_rank = get_pp_group().is_first_rank
98
+ self.is_last_pp_rank = get_pp_group().is_last_rank
99
100
self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {}
101
self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None
@@ -371,14 +366,11 @@ class ModelCudaGraphManager(CudaGraphManager):
366
self.aux_hidden_states[i][:num_tokens] = aux
367
else:
368
# Non-last PP rank.
369
+ assert isinstance(model_output, IntermediateTensors)
370
intermediate_tensors = model_output
375
- assert isinstance(intermediate_tensors, IntermediateTensors)
371
if self.intermediate_tensors is None:
377
- self.intermediate_tensors = IntermediateTensors(
378
- {
379
- k: torch.empty_like(v)
380
- for k, v in intermediate_tensors.tensors.items()
381
- }
372
+ self.intermediate_tensors = IntermediateTensors.empty_like(
373
+ intermediate_tensors
374
)
375
for k, v in intermediate_tensors.tensors.items():
376
self.intermediate_tensors[k][:num_tokens] = v
vllm/v1/worker/gpu/model_runner.py
+9
-15
@@ -132,14 +132,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
132
self.output_copy_event = torch.cuda.Event()
133
134
# Pipeline parallelism.
135
- self.pp_size = self.parallel_config.pipeline_parallel_size
136
- self.use_pp = self.pp_size > 1
137
- if self.use_pp:
138
- self.is_first_pp_rank = get_pp_group().is_first_rank
139
- self.is_last_pp_rank = get_pp_group().is_last_rank
140
- else:
141
- self.is_first_pp_rank = True
142
- self.is_last_pp_rank = True
135
+ self.use_pp = self.parallel_config.pipeline_parallel_size > 1
136
+ self.is_first_pp_rank = get_pp_group().is_first_rank
137
+ self.is_last_pp_rank = get_pp_group().is_last_rank
138
+
139
# Persistent buffer for intermediate tensors (non-first PP ranks).
140
self.intermediate_tensors: IntermediateTensors | None = None
141
@@ -179,7 +175,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
175
if self.speculative_config.method == "eagle3":
176
# EAGLE3 may require auxiliary hidden states from target model outputs.
177
self.use_aux_hidden_state_outputs = True
182
- if self.pp_size > 1:
178
+ if self.use_pp:
179
raise ValueError("EAGLE3 with pipeline parallel is not supported.")
180
181
# Draft tokens propagation - for spec-dec + struct outputs.
@@ -270,8 +266,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
266
logger.info("Loading model from scratch...")
267
268
self.model = model_loader.load_model(
273
- vllm_config=self.vllm_config,
274
- model_config=self.vllm_config.model_config,
269
+ vllm_config=self.vllm_config, model_config=self.vllm_config.model_config
270
)
271
if self.lora_config:
272
self.model = self.load_lora_model(
@@ -1026,14 +1021,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1021
assert intermediate_tensors is not None
1022
assert self.intermediate_tensors is not None
1023
n = input_batch.num_tokens_after_padding
1029
- intermediate_tensors = IntermediateTensors(
1024
+ model_inputs["intermediate_tensors"] = IntermediateTensors(
1025
{
1026
k: v[:n].copy_(intermediate_tensors.tensors[k][:n])
1027
for k, v in self.intermediate_tensors.tensors.items()
1033
- },
1034
- intermediate_tensors.kv_connector_output,
1028
+ }
1029
)
1036
- model_inputs["intermediate_tensors"] = intermediate_tensors
1030
+ del intermediate_tensors
1031
1032
# Run model.
1033
if batch_desc.cg_mode == CUDAGraphMode.FULL: