recoverssm: release the RecoverSSM step after CUDA graph memory profiling (port of upstream #58863, fdcc43cf3e)

Port of the upstream vLLM change #58863 (Apache-2.0). The step recorded during graph-capture dummy runs held the profiling KV cache
alive after profiling; it is now released at teardown.

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -1133,6 +1133,11 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None:
     torch.accelerator.synchronize()
     if hasattr(runner.model_state, "_mamba_ctx"):
         runner.model_state._mamba_ctx = None
+    # A PIECEWISE profiling capture records its RecoverSSM step (only FULL
+    # captures skip it), and that step's commit contexts reference the
+    # profiling KV cache.
+    if (recoverssm := getattr(runner.model_state, "recoverssm", None)) is not None:
+        recoverssm.reset()
     # Invalidate the align-mode Mamba group metadata cached from the
     # profiling KVCacheConfig: the real (e.g. PP-projected) config may
     # place Mamba layers into a different group layout, so it must be

vllm/v1/worker/gpu/model_states/recoverssm.py

@@ -16,6 +16,11 @@ class RecoverSSMState:
     def __init__(self) -> None:
         self._step: tuple[RecoverSSMMetadata, ...] | None = None
 
+    def reset(self) -> None:
+        """Drop the recorded step. Its metadata holds the commit contexts, which
+        reference the KV cache tensors they were built for."""
+        self._step = None
+
     def record_step(
         self,
         attn_metadata: dict[str, Any],