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
[MRV2] Skip hidden states allocation for PW CUDA graphs (#37818)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Mar 22, 2026 at 11:47 UTC
ce9b1d76cfbbbabbc72e12f99ce4ce9b8265fae8
1 file changed
+7
vllm/v1/worker/gpu/cudagraph_utils.py
+7
@@ -263,6 +263,7 @@ class ModelCudaGraphManager(CudaGraphManager):
263
decode_query_len: int,
264
):
265
super().__init__(vllm_config, device, cudagraph_mode, decode_query_len)
266
+ # Used for FULL CUDA graphs. PW CUDA graphs do not use these.
267
self.hidden_states: torch.Tensor | None = None
268
self.aux_hidden_states: list[torch.Tensor] = []
269
self.use_aux_hidden_state_outputs = False
@@ -326,6 +327,12 @@ class ModelCudaGraphManager(CudaGraphManager):
327
**model_state.prepare_dummy_inputs(num_reqs, num_tokens),
328
}
329
model_output = model(**model_inputs)
330
+
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
335
+
336
if self.use_aux_hidden_state_outputs:
337
hidden_states, aux_hidden_states = model_output
338
else: