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}" )