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 3, 15:15 UTC
[Model Runner V2][Bug Fix][DSV4] Ensure lazy attention state initializations happen during cudagraph capture (#42444)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
May 14, 2026 at 18:16 UTC
3b6a2047899b79f2ccad736ea87e4fe7d2c9084f
1 file changed
+7
-1
vllm/v1/worker/gpu/cudagraph_utils.py
+7
-1
@@ -216,7 +216,6 @@ class CudaGraphManager:
216
for desc in descs:
217
# Prepare inputs and get forward function
218
forward_fn, attn_state = create_forward_fn(desc)
219
- captured_attn_states[desc] = attn_state
219
220
# Warmup
221
forward_fn(CUDAGraphMode.NONE)
@@ -226,8 +225,15 @@ class CudaGraphManager:
225
"CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc
226
)
227
if desc.cg_mode == CUDAGraphMode.PIECEWISE:
228
+ captured_attn_states[desc] = attn_state
229
forward_fn(CUDAGraphMode.PIECEWISE)
230
else:
231
+ # Capture with fresh attention state. The warmup
232
+ # attention state is discarded because some backends
233
+ # (e.g. FlashMLA) perform lazy initializations that
234
+ # must be captured in the graph.
235
+ forward_fn, attn_state = create_forward_fn(desc)
236
+ captured_attn_states[desc] = attn_state
237
assert desc not in self.graphs, (
238
f"Graph already captured for {desc}"
239
)