diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 55462283db..0ea45ccf0e 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py index b02a6c19aa..3cfbecea26 100644 --- a/vllm/v1/worker/gpu/model_states/recoverssm.py +++ b/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],