@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 63f49b8bd4

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 6, 05:45 UTC

[Model Runner V2] Enable piecewise CUDA graphs for pipeline parallelism (#35162)

Signed-off-by: Zhanqiu Hu <zh338@cornell.edu> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>

zhanqiuhu committed Mar 22, 2026 at 16:48 UTC 63f49b8bd46b477da384c2cdd6613bf45ed3d515

2 files changed +102 -46

vllm/v1/worker/gpu/cudagraph_utils.py

+54 -15

@@ -11,11 +11,16 @@ from tqdm import tqdm

12 from vllm.config import VllmConfig

13 from vllm.config.compilation import CUDAGraphMode

14 -from vllm.distributed.parallel_state import graph_capture, is_global_first_rank

14 +from vllm.distributed.parallel_state import (

15 + get_pp_group,

16 + graph_capture,

17 + is_global_first_rank,

18 +)

19 from vllm.forward_context import BatchDescriptor, set_forward_context

20 from vllm.logger import init_logger

21 from vllm.model_executor.offloader.base import get_offloader

22 from vllm.platforms import current_platform

23 +from vllm.sequence import IntermediateTensors

24 from vllm.v1.kv_cache_interface import KVCacheConfig

25 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer

26 from vllm.v1.worker.gpu.block_table import BlockTables

@@ -87,7 +92,15 @@ class CudaGraphManager:

92 assert self.compilation_config is not None

93 self.cudagraph_mode = cudagraph_mode

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

104

105 self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {}

106 self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None

@@ -267,12 +280,14 @@ class ModelCudaGraphManager(CudaGraphManager):

280 self.hidden_states: torch.Tensor | None = None

281 self.aux_hidden_states: list[torch.Tensor] = []

282 self.use_aux_hidden_state_outputs = False

283 + self.intermediate_tensors: IntermediateTensors | None = None

284

285 def capture(

286 self,

287 model: nn.Module,

288 model_state: ModelState,

289 input_buffers: InputBuffers,

290 + intermediate_tensors: IntermediateTensors | None,

291 block_tables: BlockTables,

292 attn_groups: list[list[AttentionGroup]],

293 kv_cache_config: KVCacheConfig,

@@ -293,6 +308,19 @@ class ModelCudaGraphManager(CudaGraphManager):

308 if self.dp_size > 1

309 else None

310 )

311 +

312 + model_inputs = {

313 + "input_ids": input_buffers.input_ids[:num_tokens],

314 + "positions": input_buffers.positions[:num_tokens],

315 + **model_state.prepare_dummy_inputs(num_reqs, num_tokens),

316 + }

317 + if not self.is_first_pp_rank:

318 + # Update for non-first PP ranks.

319 + model_inputs["input_ids"] = None

320 + model_inputs["inputs_embeds"] = None

321 + assert intermediate_tensors is not None

322 + model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens]

323 +

324 attn_metadata, slot_mappings = prepare_inputs_to_capture(

325 num_reqs,

326 num_tokens,

@@ -318,21 +346,15 @@ class ModelCudaGraphManager(CudaGraphManager):

346 slot_mapping=slot_mappings,

347 batch_descriptor=batch_descriptor,

348 ):

321 - model_inputs = {

322 - "input_ids": input_buffers.input_ids[:num_tokens],

323 - "positions": input_buffers.positions[:num_tokens],

324 - # TODO: Pass intermediate_tensors for PP CUDA graph

325 - # support (https://github.com/vllm-project/vllm/pull/35162).

326 - "intermediate_tensors": None,

327 - **model_state.prepare_dummy_inputs(num_reqs, num_tokens),

328 - }

349 model_output = model(**model_inputs)

350

331 - if cg_mode == CUDAGraphMode.PIECEWISE:

332 - # PW CUDA graph internally handles the model outputs.

333 - # No need to keep track of the hidden states.

334 - return None

351 + if cg_mode == CUDAGraphMode.PIECEWISE:

352 + # PW CUDA graph internally handles the model outputs.

353 + # No need to keep track of the hidden states.

354 + return None

355

356 + if self.is_last_pp_rank:

357 + # Last PP rank (common case).

358 if self.use_aux_hidden_state_outputs:

359 hidden_states, aux_hidden_states = model_output

360 else:

@@ -340,13 +362,26 @@ class ModelCudaGraphManager(CudaGraphManager):

362 aux_hidden_states = []

363 if self.hidden_states is None:

364 self.hidden_states = torch.empty_like(hidden_states)

365 + self.hidden_states[:num_tokens] = hidden_states

366 if self.use_aux_hidden_state_outputs and not self.aux_hidden_states:

367 self.aux_hidden_states = [

368 torch.empty_like(x) for x in aux_hidden_states

369 ]

347 - self.hidden_states[:num_tokens] = hidden_states

370 for i, aux in enumerate(aux_hidden_states):

371 self.aux_hidden_states[i][:num_tokens] = aux

372 + else:

373 + # Non-last PP rank.

374 + intermediate_tensors = model_output

375 + assert isinstance(intermediate_tensors, IntermediateTensors)

376 + 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 + }

382 + )

383 + for k, v in intermediate_tensors.tensors.items():

384 + self.intermediate_tensors[k][:num_tokens] = v

385

386 return forward_fn

387

@@ -354,9 +389,13 @@ class ModelCudaGraphManager(CudaGraphManager):

389

390 def run_fullgraph(

391 self, desc: BatchExecutionDescriptor

357 - ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:

392 + ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]] | IntermediateTensors:

393 """Replay a captured FULL cudagraph and return hidden states."""

394 super().run_fullgraph(desc)

395 + if not self.is_last_pp_rank:

396 + assert self.intermediate_tensors is not None

397 + return self.intermediate_tensors[: desc.num_tokens]

398 +

399 assert self.hidden_states is not None

400 hidden_states = self.hidden_states[: desc.num_tokens]

401 if not self.use_aux_hidden_state_outputs:

vllm/v1/worker/gpu/model_runner.py

+48 -31

@@ -140,6 +140,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

140 else:

141 self.is_first_pp_rank = True

142 self.is_last_pp_rank = True

143 + # Persistent buffer for intermediate tensors (non-first PP ranks).

144 + self.intermediate_tensors: IntermediateTensors | None = None

145

146 # Data parallelism.

147 self.dp_size = self.parallel_config.data_parallel_size

@@ -301,6 +303,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):

303 if self.is_pooling_model and self.is_last_pp_rank:

304 self.pooling_runner = PoolingRunner(self.model)

305

306 + if not self.is_first_pp_rank:

307 + # For non-first PP ranks, create intermediate tensors sized

308 + # for the max capture size so they can be sliced per batch.

309 + # Save as persistent member so runtime can copy received data

310 + # into the same addresses that the CUDA graphs captured.

311 + self.intermediate_tensors = self.model.make_empty_intermediate_tensors(

312 + batch_size=self.max_num_tokens,

313 + dtype=self.model_config.dtype,

314 + device=self.device,

315 + )

316 +

317 def get_model(self) -> nn.Module:

318 return self.model

319

@@ -396,14 +409,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

409 # Disable any use of KVConnector for dummy runs.

410 self.kv_connector.set_disabled(True)

411

399 - # For non-first PP ranks, create dummy intermediate_tensors.

412 + # Get the intermediate tensors for the dummy run.

413 intermediate_tensors = None

414 if not self.is_first_pp_rank:

402 - intermediate_tensors = self.model.make_empty_intermediate_tensors(

403 - batch_size=num_tokens,

404 - dtype=self.model_config.dtype,

405 - device=self.device,

406 - )

415 + assert self.intermediate_tensors is not None

416 + intermediate_tensors = self.intermediate_tensors[:num_tokens]

417

418 # Execute the model.

419 self.execute_model(

@@ -528,14 +538,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

538 )

539 return 0

540

531 - # TODO (zhanqiu): support CUDA graph for PP.

532 - if self.use_pp:

533 - logger.warning_once(

534 - "Skipping CUDA graph capture because pipeline parallel is "

535 - "enabled. Pipeline parallel is currently eager-only.",

536 - )

537 - return 0

538 -

541 start_time = time.perf_counter()

542 gc.collect()

543 torch.accelerator.empty_cache()

@@ -546,6 +548,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

548 self.model,

549 self.model_state,

550 self.input_buffers,

551 + self.intermediate_tensors,

552 self.block_tables,

553 self.attn_groups,

554 self.kv_cache_config,

@@ -1010,7 +1013,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1013 "input_ids": input_batch.input_ids,

1014 "positions": input_batch.positions,

1015 "inputs_embeds": inputs_embeds,

1013 - "intermediate_tensors": intermediate_tensors,

1016 # NOTE: Values returned by `prepare_inputs` will override the default

1017 # values above.

1018 **self.model_state.prepare_inputs(input_batch, self.req_states),

@@ -1019,7 +1021,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1021 # Update for non-first PP ranks.

1022 model_inputs["input_ids"] = None

1023 model_inputs["inputs_embeds"] = None

1024 +

1025 + # Prepare the intermediate tensors.

1026 assert intermediate_tensors is not None

1027 + assert self.intermediate_tensors is not None

1028 + n = input_batch.num_tokens_after_padding

1029 + intermediate_tensors = IntermediateTensors(

1030 + {

1031 + k: v[:n].copy_(intermediate_tensors.tensors[k][:n])

1032 + for k, v in self.intermediate_tensors.tensors.items()

1033 + },

1034 + intermediate_tensors.kv_connector_output,

1035 + )

1036 + model_inputs["intermediate_tensors"] = intermediate_tensors

1037

1038 # Run model.

1039 if batch_desc.cg_mode == CUDAGraphMode.FULL:

@@ -1028,11 +1042,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1042 # because they are already copied to the CUDA graph input buffers.

1043 self.kv_connector.pre_forward(scheduler_output)

1044 model_output = self.cudagraph_manager.run_fullgraph(batch_desc)

1031 - if self.use_aux_hidden_state_outputs:

1032 - hidden_states, aux_hidden_states = model_output

1033 - else:

1034 - hidden_states = model_output

1035 - aux_hidden_states = None

1045 else:

1046 # For piecewise and eager mode, just call model().

1047 batch_descriptor = BatchDescriptor(

@@ -1052,11 +1061,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1061 ):

1062 self.kv_connector.pre_forward(scheduler_output)

1063 model_output = self.model(**model_inputs)

1055 - if self.use_aux_hidden_state_outputs:

1056 - hidden_states, aux_hidden_states = model_output

1057 - else:

1058 - hidden_states = model_output

1059 - aux_hidden_states = None

1064 +

1065 + if self.is_last_pp_rank:

1066 + if self.use_aux_hidden_state_outputs:

1067 + assert isinstance(model_output, tuple)

1068 + hidden_states, aux_hidden_states = model_output

1069 + else:

1070 + assert isinstance(model_output, torch.Tensor)

1071 + hidden_states = model_output

1072 + aux_hidden_states = None

1073 + output_intermediate_tensors = None

1074 + else:

1075 + assert isinstance(model_output, IntermediateTensors)

1076 + hidden_states = None

1077 + aux_hidden_states = None

1078 + output_intermediate_tensors = model_output

1079

1080 kv_connector_output = self.kv_connector.post_forward(scheduler_output)

1081 self.execute_model_state = ExecuteModelState(

@@ -1071,11 +1090,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1090

1091 if not self.is_last_pp_rank:

1092 # Non-last PP rank: return IntermediateTensors for sending.

1074 - assert isinstance(hidden_states, IntermediateTensors)

1075 - hidden_states.kv_connector_output = kv_connector_output

1076 - return hidden_states

1077 - # Last rank (or no PP): hidden_states is a tensor for sampling.

1078 - assert isinstance(hidden_states, torch.Tensor)

1093 + assert output_intermediate_tensors is not None

1094 + output_intermediate_tensors.kv_connector_output = kv_connector_output

1095 + return output_intermediate_tensors

1096 return None

1097

1098 @torch.inference_mode()

@@ -1259,7 +1276,7 @@ class ExecuteModelState(NamedTuple):

1276 input_batch: InputBatch

1277 attn_metadata: dict[str, Any] | None

1278 slot_mappings_by_layer: dict[str, torch.Tensor] | None

1262 - hidden_states: torch.Tensor | IntermediateTensors

1279 + hidden_states: torch.Tensor | None

1280 aux_hidden_states: list[torch.Tensor] | None

1281 kv_connector_output: KVConnectorOutput | None

1282 num_tokens_across_dp: torch.Tensor | None