[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>
Canlin Guo committed
Aug 30, 2026 at 15:43 UTC
b383e163961650c663ad2062544ecbee1b56afc8
2 files changed
+44
tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py
+36
@@ -402,3 +402,39 @@ def test_profile_cudagraph_memory_frees_throwaway_pool(monkeypatch):
402
403
assert memory["captured"] - memory["before"] >= 3 * allocation_bytes
404
assert memory["after"] == memory["before"]
405
+
406
+
407
+def test_teardown_profiling_state_clears_mamba_align_metadata(monkeypatch):
408
+ """Profiling-cached Mamba align metadata must be invalidated at teardown.
409
+
410
+ ``MambaHybridModelState`` lazily caches ``_mamba_group_ids`` and
411
+ ``_mamba_spec`` from whichever KVCacheConfig it first sees. When the
412
+ profiling config's group layout differs from the real (e.g. PP-projected)
413
+ config, reusing the stale metadata mismatches the real block tables
414
+ ("expected 3 block tables, got 4" at
415
+ ``MambaSpecDecodeGPUContext.initialize_from_forward_context``).
416
+ """
417
+ runner: Any = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner)
418
+ runner.compilation_config = SimpleNamespace(static_forward_context={})
419
+ runner.model_state = SimpleNamespace(
420
+ supports_mm_inputs=False,
421
+ _mamba_ctx=object(),
422
+ _mamba_group_ids=[0, 1],
423
+ _mamba_spec=object(),
424
+ )
425
+ runner.cache_config = SimpleNamespace(num_gpu_blocks=1)
426
+ runner.kv_caches = []
427
+ runner.attn_groups = []
428
+ runner.kv_cache_config = SimpleNamespace()
429
+ runner.cudagraph_manager = object()
430
+ runner.lora_config = None
431
+ runner.maybe_remove_all_loras = lambda _: None
432
+
433
+ monkeypatch.setattr(cgu.torch.accelerator, "synchronize", lambda: None)
434
+ monkeypatch.setattr(cgu.torch.accelerator, "empty_cache", lambda: None)
435
+
436
+ cgu._teardown_profiling_state(runner)
437
+
438
+ assert runner.model_state._mamba_ctx is None
439
+ assert runner.model_state._mamba_group_ids == []
440
+ assert runner.model_state._mamba_spec is None
vllm/v1/worker/gpu/cudagraph_utils.py
+8
@@ -864,6 +864,14 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None:
864
torch.accelerator.synchronize()
865
if hasattr(runner.model_state, "_mamba_ctx"):
866
runner.model_state._mamba_ctx = None
867
+ # Invalidate the align-mode Mamba group metadata cached from the
868
+ # profiling KVCacheConfig: the real (e.g. PP-projected) config may
869
+ # place Mamba layers into a different group layout, so it must be
870
+ # re-derived from the real config.
871
+ if hasattr(runner.model_state, "_mamba_group_ids"):
872
+ runner.model_state._mamba_group_ids = []
873
+ if hasattr(runner.model_state, "_mamba_spec"):
874
+ runner.model_state._mamba_spec = None
875
if hasattr(runner, "kv_caches"):
876
runner.kv_caches.clear()
877
if hasattr(runner, "attn_groups"):