[Bugfix] Reset cached Mamba align metadata on profiling teardown (#54044)

Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py

@@ -402,3 +402,39 @@ def test_profile_cudagraph_memory_frees_throwaway_pool(monkeypatch):
 
     assert memory["captured"] - memory["before"] >= 3 * allocation_bytes
     assert memory["after"] == memory["before"]
+
+
+def test_teardown_profiling_state_clears_mamba_align_metadata(monkeypatch):
+    """Profiling-cached Mamba align metadata must be invalidated at teardown.
+
+    ``MambaHybridModelState`` lazily caches ``_mamba_group_ids`` and
+    ``_mamba_spec`` from whichever KVCacheConfig it first sees. When the
+    profiling config's group layout differs from the real (e.g. PP-projected)
+    config, reusing the stale metadata mismatches the real block tables
+    ("expected 3 block tables, got 4" at
+    ``MambaSpecDecodeGPUContext.initialize_from_forward_context``).
+    """
+    runner: Any = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner)
+    runner.compilation_config = SimpleNamespace(static_forward_context={})
+    runner.model_state = SimpleNamespace(
+        supports_mm_inputs=False,
+        _mamba_ctx=object(),
+        _mamba_group_ids=[0, 1],
+        _mamba_spec=object(),
+    )
+    runner.cache_config = SimpleNamespace(num_gpu_blocks=1)
+    runner.kv_caches = []
+    runner.attn_groups = []
+    runner.kv_cache_config = SimpleNamespace()
+    runner.cudagraph_manager = object()
+    runner.lora_config = None
+    runner.maybe_remove_all_loras = lambda _: None
+
+    monkeypatch.setattr(cgu.torch.accelerator, "synchronize", lambda: None)
+    monkeypatch.setattr(cgu.torch.accelerator, "empty_cache", lambda: None)
+
+    cgu._teardown_profiling_state(runner)
+
+    assert runner.model_state._mamba_ctx is None
+    assert runner.model_state._mamba_group_ids == []
+    assert runner.model_state._mamba_spec is None

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -864,6 +864,14 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None:
     torch.accelerator.synchronize()
     if hasattr(runner.model_state, "_mamba_ctx"):
         runner.model_state._mamba_ctx = None
+    # 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
+    # re-derived from the real config.
+    if hasattr(runner.model_state, "_mamba_group_ids"):
+        runner.model_state._mamba_group_ids = []
+    if hasattr(runner.model_state, "_mamba_spec"):
+        runner.model_state._mamba_spec = None
     if hasattr(runner, "kv_caches"):
         runner.kv_caches.clear()
     if hasattr(runner, "attn_groups"):