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 10, 01:30 UTC
[TPU][Core] Enable Pipeline Parallelism on TPU backend (#28506)
Signed-off-by: Chenyaaang <chenyangli@google.com>
Chenyaaang committed
Jan 16, 2026 at 15:29 UTC
484e22bc1880bcf87c51164ba2ae4e25b7a432e9
2 files changed
+34
-14
vllm/v1/executor/multiproc_executor.py
+30
-13
@@ -104,16 +104,7 @@ class MultiprocExecutor(Executor):
104
self.shutdown_event = threading.Event()
105
self.failure_callback: FailureCallback | None = None
106
107
- self.world_size = self.parallel_config.world_size
108
- assert self.world_size % self.parallel_config.nnodes_within_dp == 0, (
109
- f"global world_size ({self.parallel_config.world_size}) must be "
110
- f"divisible by nnodes_within_dp "
111
- f"({self.parallel_config.nnodes_within_dp}). "
112
- )
113
- self.local_world_size = self.parallel_config.local_world_size
114
- tp_size = self.parallel_config.tensor_parallel_size
115
- pp_size = self.parallel_config.pipeline_parallel_size
116
- pcp_size = self.parallel_config.prefill_context_parallel_size
107
+ tp_size, pp_size, pcp_size = self._get_parallel_sizes()
108
assert self.world_size == tp_size * pp_size * pcp_size, (
109
f"world_size ({self.world_size}) must be equal to the "
110
f"tensor_parallel_size ({tp_size}) x pipeline"
@@ -154,6 +145,7 @@ class MultiprocExecutor(Executor):
145
)
146
for local_rank in range(self.local_world_size):
147
global_rank = global_start_rank + local_rank
148
+ is_driver_worker = self._is_driver_worker(global_rank)
149
unready_workers.append(
150
WorkerProc.make_worker_process(
151
vllm_config=self.vllm_config,
@@ -162,6 +154,7 @@ class MultiprocExecutor(Executor):
154
distributed_init_method=distributed_init_method,
155
input_shm_handle=scheduler_output_handle,
156
shared_worker_lock=shared_worker_lock,
157
+ is_driver_worker=is_driver_worker,
158
)
159
)
160
@@ -199,6 +192,11 @@ class MultiprocExecutor(Executor):
192
# Wait for all remote response mqs to be ready.
193
for response_mq in self.response_mqs:
194
response_mq.wait_until_ready()
195
+
196
+ self.futures_queue = deque[tuple[FutureWrapper, Callable]]()
197
+
198
+ self._post_init_executor()
199
+
200
success = True
201
finally:
202
if not success:
@@ -209,10 +207,27 @@ class MultiprocExecutor(Executor):
207
uw.death_writer.close()
208
self._ensure_worker_termination([uw.proc for uw in unready_workers])
209
212
- self.futures_queue = deque[tuple[FutureWrapper, Callable]]()
213
-
210
self.output_rank = self._get_output_rank()
211
212
+ def _get_parallel_sizes(self) -> tuple[int, int, int]:
213
+ self.world_size = self.parallel_config.world_size
214
+ assert self.world_size % self.parallel_config.nnodes_within_dp == 0, (
215
+ f"global world_size ({self.parallel_config.world_size}) must be "
216
+ f"divisible by nnodes_within_dp "
217
+ f"({self.parallel_config.nnodes_within_dp}). "
218
+ )
219
+ self.local_world_size = self.parallel_config.local_world_size
220
+ tp_size = self.parallel_config.tensor_parallel_size
221
+ pp_size = self.parallel_config.pipeline_parallel_size
222
+ pcp_size = self.parallel_config.prefill_context_parallel_size
223
+ return tp_size, pp_size, pcp_size
224
+
225
+ def _post_init_executor(self) -> None:
226
+ pass
227
+
228
+ def _is_driver_worker(self, rank: int) -> bool:
229
+ return rank % self.parallel_config.tensor_parallel_size == 0
230
+
231
def start_worker_monitor(self, inline=False) -> None:
232
workers = self.workers
233
self_ref = weakref.ref(self)
@@ -517,6 +532,7 @@ class WorkerProc:
532
distributed_init_method: str,
533
input_shm_handle: Handle,
534
shared_worker_lock: LockType,
535
+ is_driver_worker: bool,
536
):
537
self.rank = rank
538
wrapper = WorkerWrapperBase(rpc_rank=local_rank, global_rank=rank)
@@ -524,7 +540,6 @@ class WorkerProc:
540
all_kwargs: list[dict] = [
541
{} for _ in range(vllm_config.parallel_config.world_size)
542
]
527
- is_driver_worker = rank % vllm_config.parallel_config.tensor_parallel_size == 0
543
all_kwargs[local_rank] = {
544
"vllm_config": vllm_config,
545
"local_rank": local_rank,
@@ -571,6 +586,7 @@ class WorkerProc:
586
distributed_init_method: str,
587
input_shm_handle, # Receive SchedulerOutput
588
shared_worker_lock: LockType,
589
+ is_driver_worker: bool,
590
) -> UnreadyWorkerProcHandle:
591
context = get_mp_context()
592
# (reader, writer)
@@ -588,6 +604,7 @@ class WorkerProc:
604
"ready_pipe": (reader, writer),
605
"death_pipe": death_reader,
606
"shared_worker_lock": shared_worker_lock,
607
+ "is_driver_worker": is_driver_worker,
608
}
609
# Run EngineCore busy loop in background process.
610
proc = context.Process(
vllm/v1/executor/ray_utils.py
+4
-1
@@ -103,7 +103,7 @@ try:
103
output = self.worker.model_runner.execute_model(
104
scheduler_output, intermediate_tensors
105
)
106
- if isinstance(output, IntermediateTensors):
106
+ if self._is_intermediate_tensors(output):
107
return scheduler_output, grammar_output, output
108
109
if isinstance(output, AsyncModelRunnerOutput):
@@ -125,6 +125,9 @@ try:
125
def override_env_vars(self, vars: dict[str, str]):
126
os.environ.update(vars)
127
128
+ def _is_intermediate_tensors(self, output) -> bool:
129
+ return isinstance(output, IntermediateTensors)
130
+
131
ray_import_err = None
132
133
except ImportError as e: