From 3b6a2047899b79f2ccad736ea87e4fe7d2c9084f Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Thu, 14 May 2026 18:16:17 -0500 Subject: [PATCH] [Model Runner V2][Bug Fix][DSV4] Ensure lazy attention state initializations happen during cudagraph capture (#42444) Signed-off-by: Giancarlo Delfin --- vllm/v1/worker/gpu/cudagraph_utils.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 71a15c3bd9..e76fe0247f 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -216,7 +216,6 @@ class CudaGraphManager: for desc in descs: # Prepare inputs and get forward function forward_fn, attn_state = create_forward_fn(desc) - captured_attn_states[desc] = attn_state # Warmup forward_fn(CUDAGraphMode.NONE) @@ -226,8 +225,15 @@ class CudaGraphManager: "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc ) if desc.cg_mode == CUDAGraphMode.PIECEWISE: + captured_attn_states[desc] = attn_state forward_fn(CUDAGraphMode.PIECEWISE) else: + # Capture with fresh attention state. The warmup + # attention state is discarded because some backends + # (e.g. FlashMLA) perform lazy initializations that + # must be captured in the graph. + forward_fn, attn_state = create_forward_fn(desc) + captured_attn_states[desc] = attn_state assert desc not in self.graphs, ( f"Graph already captured for {desc}" )