[Model Runner V2][Bug Fix][DSV4] Ensure lazy attention state initializations happen during cudagraph capture (#42444)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>

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