[Bugfix][MiniCPM-o] Retire CFM DiT CUDA graphs a generation at a time (#6587)

Signed-off-by: BruceLoveDecimal <liuqihao970610@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/model_executor/models/minicpmo_4_5/test_cfm_graph_capture_gating.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""What is allowed to keep the CFM wrapper from capturing a graph.
+
+These run without a GPU: ``_capture`` is mocked, so nothing here touches a
+CUDA context. The point is the decision made *before* capture, not the capture.
+"""
+
+from typing import NamedTuple
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import (
+    CFMGraphWrapper,
+    _format_memory_delta,
+    _memory_snapshot,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _InputSignature(NamedTuple):
+    """The three properties ``replay`` reads from an input before capturing.
+
+    Everything ``replay`` does ahead of the capture decision -- the device
+    check and the cache key -- comes from these, and ``_capture`` is mocked
+    out, so the tests need no CUDA tensors and no GPU.
+    """
+
+    shape: tuple[int, ...]
+    dtype: torch.dtype
+    device: torch.device
+
+
+def _cfm_input_signatures(chunk_size: int) -> tuple[_InputSignature, ...]:
+    device = torch.device("cuda:0")
+    shapes = ((2, 16, chunk_size), (2, 1, 8), (2, 2, 8, 2), (2, 2, 1, 4, 8), (2, 2, 8, 2), (2, 2, 1, 12, 8))
+    return tuple(_InputSignature(shape, torch.float32, device) for shape in shapes)
+
+
+def _wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 4) -> CFMGraphWrapper:
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    wrapper = object.__new__(CFMGraphWrapper)
+    wrapper.max_graphs = max_graphs
+    wrapper.enabled = True
+    wrapper.graph_fn = Mock(return_value=torch.tensor([42.0]))
+    wrapper.device = torch.device("cuda:0")
+    wrapper._cache = {}
+    wrapper._unsupported = set()
+    wrapper._stats = {"calls": 0, "hits": 0, "captures": 0, "flushes": 0, "eager": 0}
+    wrapper._capture = Mock(return_value=None)
+    return wrapper
+
+
+def test_capture_is_not_gated_on_free_device_memory(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Whether to capture must not be decided from free device memory.
+
+    A capture is served by the caching allocator out of its private graph
+    pool, and the allocator falls back on memory it has already reserved. Free
+    device memory is therefore neither an upper nor a lower bound on what a
+    capture can draw, and a threshold on it decides the wrong question -- one
+    whose answer also silently varies with card size and with whatever else
+    shares the device. Anything consulting it here would fail this test.
+    """
+    import vllm_omni.platforms as platforms_module
+
+    def _saturated(device: torch.device | None = None) -> int:
+        raise AssertionError("capture must not depend on device-level free memory")
+
+    # Setting the attribute on the module shadows the lazy ``__getattr__``, so
+    # this neither builds nor requires a real platform singleton.
+    monkeypatch.setattr(
+        platforms_module,
+        "current_omni_platform",
+        Mock(get_free_memory=_saturated),
+        raising=False,
+    )
+
+    wrapper = _wrapper(monkeypatch)
+    wrapper.replay(*_cfm_input_signatures(10))
+
+    wrapper._capture.assert_called_once()
+
+
+def test_memory_reporting_is_inert_off_cuda() -> None:
+    """Capture logging must not break the platforms that have no such counters."""
+    assert _memory_snapshot(torch.device("cpu")) is None
+    assert _format_memory_delta(None, (1, 2)) == ""
+    assert _format_memory_delta((1, 2), None) == ""
+    assert "allocated" in _format_memory_delta((0, 0), (1 << 20, 2 << 20))

tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-from functools import lru_cache
+import sys
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -10,6 +10,7 @@ import torch
 import torch.nn as nn
 from vllm.platforms import current_platform
 
+import vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper as wrapper_module
 from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
     HiFTGenerator,
 )
@@ -248,50 +249,203 @@ def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes(
             torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5)
             torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5)
 
-    wrapper._capture_graph.cache_clear()
+    wrapper._flush()
+
+
+def test_importing_wrapper_does_not_resolve_platform() -> None:
+    """Importing this module must not build the OmniPlatform singleton.
+
+    The NPU platform's ``__init__`` patches Code2Wav, which imports
+    ``batched_token2wav`` -> ``cuda_graph_wrapper``. Resolving
+    ``current_omni_platform`` at module scope re-enters
+    ``platforms.__getattr__`` while the singleton is still under construction
+    and the import graph deadlocks. The module reaches for nothing on
+    ``vllm_omni.platforms`` today; this keeps it that way.
+    """
+    import importlib
+
+    import vllm_omni.platforms as platforms_module
+
+    name = "vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper"
+    importlib.import_module(name)
+    saved = platforms_module._current_omni_platform
+    try:
+        del sys.modules[name]
+        platforms_module._current_omni_platform = None
+        importlib.import_module(name)
+        assert platforms_module._current_omni_platform is None
+    finally:
+        platforms_module._current_omni_platform = saved
 
 
 def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper:
     monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
     wrapper = object.__new__(CFMGraphWrapper)
     wrapper.max_graphs = max_graphs
+    wrapper.enabled = True
     wrapper.graph_fn = Mock(return_value=torch.tensor([42.0]))
     wrapper.device = torch.device("cuda")
-
-    @lru_cache(maxsize=max_graphs)
-    def _capture_graph(key: tuple):
-        return wrapper._do_capture(key)
-
-    wrapper._capture_graph = _capture_graph
-    wrapper._do_capture = Mock(return_value=None)
+    wrapper._cache = {}
+    wrapper._unsupported = set()
+    wrapper._stats = {"calls": 0, "hits": 0, "captures": 0, "flushes": 0, "eager": 0}
+    wrapper._capture = Mock(return_value=None)
     return wrapper
 
 
 def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None:
     wrapper = _cfm_mock_wrapper(monkeypatch)
-    wrapper._capture_graph = Mock(return_value=None)
 
     inputs = _cfm_inputs(2, 10, 0)
     wrapper.replay(*inputs)
 
-    wrapper._capture_graph.assert_called_once()
+    wrapper._capture.assert_called_once()
     wrapper.graph_fn.assert_called_once()
 
 
-def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_cfm_returning_no_entry_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A capture that yields no entry must still serve the request eagerly.
+
+    ``_capture`` is mocked here, so this says nothing about ``_disable``; see
+    ``test_cfm_capture_failure_disables_further_capture`` for that.
+    """
     wrapper = _cfm_mock_wrapper(monkeypatch)
-    wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None)
 
     inputs = _cfm_inputs(2, 10, 0)
     result = wrapper.replay(*inputs)
 
     wrapper.graph_fn.assert_called_once()
     assert result[0] is wrapper.graph_fn.return_value
+    assert wrapper._stats["eager"] == 1
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Verify that LRU cache evicts oldest entries when full."""
+def test_cfm_capture_failure_disables_further_capture(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A real capture failure must stop the wrapper capturing for good.
+
+    A failed capture can leave the capture stream current and the allocator
+    still routing into the graph pool, so the next shape must not try again.
+    """
+    pool = torch.cuda.graph_pool_handle()
+    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
+
+    torch.manual_seed(0)
+    estimator = _MiniDiT().eval().cuda()
+    wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=4)
+
+    def _explode(*args: object, **kwargs: object) -> None:
+        raise RuntimeError("capture failed")
+
+    monkeypatch.setattr(torch.cuda, "graph", _explode)
+
+    with torch.inference_mode():
+        wrapper.replay(*_cfm_inputs(2, 10, 0))
+        assert wrapper.enabled is False
+        assert wrapper._cache == {}
+
+        captures_after_failure = wrapper._stats["captures"]
+        wrapper.replay(*_cfm_inputs(2, 12, 0))
+
+    assert wrapper._stats["captures"] == captures_after_failure
+    assert wrapper._stats["eager"] == 2
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_cfm_unsupported_dtype_eagers_only_that_shape(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A key that cannot round-trip is a property of one shape, not of the GPU.
+
+    It must not disable the wrapper or retire the generation the way a capture
+    failure does.
+    """
+    pool = torch.cuda.graph_pool_handle()
+    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
+
+    torch.manual_seed(0)
+    estimator = _MiniDiT().eval().cuda()
+    wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=4)
+
+    unbuildable_width = 10
+    real_tensors_from_key = wrapper_module._tensors_from_key
+
+    def _reject_one_shape(key: tuple) -> tuple:
+        if key[1][0][2] == unbuildable_width:
+            raise KeyError("torch.int64")
+        return real_tensors_from_key(key)
+
+    monkeypatch.setattr(wrapper_module, "_tensors_from_key", _reject_one_shape)
+
+    with torch.inference_mode():
+        wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0))
+        assert wrapper.enabled is True
+        assert wrapper._stats["captures"] == 0
+
+        # a capturable shape still gets a graph
+        wrapper.replay(*_cfm_inputs(2, 12, 0))
+        assert wrapper._stats["captures"] == 1
+
+        # and the rejected shape stays eager without retrying the capture
+        wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0))
+        assert wrapper._stats["captures"] == 1
+        assert wrapper._stats["eager"] == 2
+
+    wrapper._flush()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_hift_replay_survives_a_cfm_generation_flush(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Retiring a CFM generation must not disturb the vocoder.
+
+    Both wrappers capture into ``get_global_graph_pool()``, so the HiFT graphs
+    are exactly the live graphs a CFM flush could strand.
+    """
+    pool = torch.cuda.graph_pool_handle()
+    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
+
+    torch.manual_seed(0)
+    hift = _small_hift()
+    token2wav = SimpleNamespace(
+        hift=hift,
+        flow=SimpleNamespace(
+            encoder=SimpleNamespace(pre_lookahead_layer=SimpleNamespace(pre_lookahead_len=3)),
+            token_mel_ratio=2,
+        ),
+        mel_cache_len=2,
+        source_cache_len=960,
+    )
+    hift_wrapper = HiFTGraphWrapper(
+        token2wav,
+        connector_config={"codec_chunk_frames": 2, "codec_left_context_frames": 3},
+        capture_batch_sizes=[1],
+    )
+    hift_wrapper.capture()
+
+    estimator = _MiniDiT().eval().cuda()
+    cfm = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2)
+
+    speech_feat = torch.randn(1, 80, 4, device="cuda")
+    cache_source = torch.zeros(1, 1, 0, device="cuda")
+
+    with torch.inference_mode():
+        expected_speech, expected_source = hift.inference(speech_feat, cache_source)
+
+        for width in (10, 12, 14, 16):
+            cfm.replay(*_cfm_inputs(2, width, 0))
+        assert cfm._stats["flushes"] >= 1, "cache never overflowed; the test proves nothing"
+
+        actual_speech, actual_source = hift_wrapper.replay(speech_feat, cache_source)
+
+    torch.testing.assert_close(actual_speech, expected_speech, rtol=1e-4, atol=1e-5)
+    torch.testing.assert_close(actual_source, expected_source, rtol=1e-4, atol=1e-5)
+
+    cfm._flush()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_cfm_cache_flushes_whole_generation(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A full cache is retired all at once, never one graph at a time.
+
+    Destroying a single graph while its peers stay live strands them, so the
+    cache must never hold a graph that outlived one of its generation-mates.
+    """
     pool = torch.cuda.graph_pool_handle()
     monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
 
@@ -302,26 +456,32 @@ def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
     with torch.inference_mode():
         inputs_a = _cfm_inputs(2, 10, 0)
         wrapper.replay(*inputs_a)
-        assert wrapper._capture_graph.cache_info().currsize == 1
+        assert len(wrapper._cache) == 1
 
         inputs_b = _cfm_inputs(2, 12, 0)
         wrapper.replay(*inputs_b)
-        assert wrapper._capture_graph.cache_info().currsize == 2
+        assert len(wrapper._cache) == 2
 
+        # A hit must not grow the cache or trigger a flush.
         wrapper.replay(*inputs_a)
-        assert wrapper._capture_graph.cache_info().currsize == 2
+        assert len(wrapper._cache) == 2
+        assert wrapper._stats["flushes"] == 0
+        assert wrapper._stats["hits"] == 1
 
+        # The third distinct shape flushes the generation, then captures alone.
         inputs_c = _cfm_inputs(2, 14, 0)
         wrapper.replay(*inputs_c)
-        assert wrapper._capture_graph.cache_info().currsize == 2
+        assert wrapper._stats["flushes"] == 1
+        assert len(wrapper._cache) == 1
 
-        wrapper.replay(*inputs_b)
-        hits_before = wrapper._capture_graph.cache_info().hits
-        wrapper.replay(*inputs_b)
-        hits_after = wrapper._capture_graph.cache_info().hits
-        assert hits_after == hits_before + 1
+        # The flushed shapes are gone, so they capture again rather than hit.
+        hits_before = wrapper._stats["hits"]
+        wrapper.replay(*inputs_a)
+        assert wrapper._stats["hits"] == hits_before
+        wrapper.replay(*inputs_a)
+        assert wrapper._stats["hits"] == hits_before + 1
 
-    wrapper._capture_graph.cache_clear()
+    wrapper._flush()
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")

vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py

@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from functools import lru_cache
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import torch
 from torch.cuda import CUDAGraph
@@ -141,19 +139,44 @@ def _tensor_signature(value: torch.Tensor) -> tuple:
     return tuple(value.shape), str(value.dtype), str(value.device)
 
 
-_DTYPE_MAP = {
-    "torch.float32": torch.float32,
-    "torch.float16": torch.float16,
-    "torch.bfloat16": torch.bfloat16,
-    "torch.float64": torch.float64,
-}
+_DTYPE_MAP = {str(dtype): dtype for dtype in (torch.float32, torch.float16, torch.bfloat16, torch.float64)}
+
+
+def _memory_snapshot(device: torch.device) -> tuple[int, int] | None:
+    """(allocated, reserved) bytes, or None when the device cannot report them.
+
+    Capture draws on the caching allocator, so these are the numbers that say
+    what a capture cost. Free device memory is not: the allocator serves a
+    capture out of memory it has already reserved, which is most of the device
+    on a normally configured worker.
+    """
+    if device.type != "cuda":
+        return None
+    try:
+        return int(torch.accelerator.memory_allocated(device)), int(torch.accelerator.memory_reserved(device))
+    except Exception:
+        return None
+
+
+def _format_memory_delta(before: tuple[int, int] | None, after: tuple[int, int] | None) -> str:
+    if before is None or after is None:
+        return ""
+    mib = 1024 * 1024
+    return (
+        f" [allocated {after[0] / mib:.1f} MiB (+{(after[0] - before[0]) / mib:.1f}), "
+        f"reserved {after[1] / mib:.1f} MiB (+{(after[1] - before[1]) / mib:.1f})]"
+    )
 
 
 def _tensors_from_key(key: tuple) -> tuple[torch.Tensor, ...]:
-    """Rebuild zero tensors from a cache key (shape, dtype, device tuples)."""
+    """Rebuild zero tensors from a cache key (shape, dtype, device tuples).
+
+    Raises ``KeyError`` for a dtype the key cannot round-trip, so the caller
+    can fall back to eager instead of capturing wrong-dtype static buffers.
+    """
     tensors = []
     for shape, dtype_str, device_str in key[1:]:
-        dtype = _DTYPE_MAP.get(dtype_str, torch.float32)
+        dtype = _DTYPE_MAP[dtype_str]
         device = torch.device(device_str)
         tensors.append(torch.zeros(shape, dtype=dtype, device=device))
     return tuple(tensors)
@@ -166,56 +189,121 @@ class CFMGraphWrapper:
     as the graph target. The 10-step Euler loop stays in Python, replaying
     the graph 10 times per decode.
 
-    Uses functools.lru_cache for automatic LRU eviction. Cache misses
-    trigger capture; capture failures fall back to eager. Outputs are
-    cloned after replay to prevent streaming cache corruption.
+    Graphs are retired a whole generation at a time rather than one at a time.
+    Every capture shares one private memory pool, so a retired graph's blocks
+    return to that pool while its live peers still hold those addresses in
+    their recorded kernel arguments, and the next replay reads memory that now
+    belongs to something else. Retiring the whole generation bounds the cache
+    without ever leaving a live graph behind a freed one.
+
+    Cache misses capture. A capture failure disables the wrapper for the rest
+    of the process, while a key whose static buffers cannot be rebuilt sends
+    only that one shape eager. Outputs are cloned after replay to prevent
+    streaming cache corruption.
     """
 
-    def __init__(self, graph_fn, *, max_graphs: int = 32) -> None:
+    def __init__(
+        self,
+        graph_fn,
+        *,
+        max_graphs: int = 32,
+    ) -> None:
         self.graph_fn = graph_fn
         self.max_graphs = int(max_graphs)
         self.device = next(graph_fn.__self__.parameters()).device
+        # A non-positive budget means "no graphs", the same as
+        # `enable_cfm_graph: false`. Clamping to 1 would instead build a
+        # one-entry cache that flushes on every new shape.
+        self.enabled = self.max_graphs > 0
+        self._cache: dict[tuple, tuple] = {}
+        # Shapes whose key cannot round-trip: eager for those, keep the rest.
+        self._unsupported: set[tuple] = set()
+        self._stats = {
+            "calls": 0,
+            "hits": 0,
+            "captures": 0,
+            "flushes": 0,
+            "eager": 0,
+        }
 
-        @lru_cache(maxsize=self.max_graphs)
-        def _capture_graph(key: tuple):
-            return self._capture(key)
-
-        self._capture_graph = _capture_graph
+    def stats_snapshot(self) -> dict[str, int]:
+        """Bounded cumulative telemetry for the graph cache."""
+        return {**self._stats, "cache_size": len(self._cache)}
 
     def _call_graph_fn(self, args: tuple[torch.Tensor, ...]) -> torch.Tensor:
         return self.graph_fn(args[0], args[1], None, args[2], args[3], args[4], args[5])
 
+    def _eager(self, inputs: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        self._stats["eager"] += 1
+        with torch.no_grad():
+            result = self._call_graph_fn(inputs)
+        return result, inputs[4], inputs[5]
+
+    def _flush(self) -> None:
+        """Retire every captured graph at once.
+
+        The sync keeps a replay from being in flight when the graphs go, and
+        the explicit ``reset()`` tears each one down here rather than whenever
+        Python drops the last reference. Nothing process-wide is touched: the
+        cuBLAS workspace is shared with graphs this wrapper does not own.
+        """
+        if not self._cache:
+            return
+        torch.accelerator.synchronize(self.device)
+        for entry in self._cache.values():
+            entry[2].reset()
+        self._cache.clear()
+        self._stats["flushes"] += 1
+        logger.info("CFM graph cache flushed; stats=%s", self.stats_snapshot())
+
+    def _disable(self, reason: str, key: tuple) -> None:
+        logger.warning("Disabling CFM CUDA graphs (%s) for shape=%s; using eager", reason, key, exc_info=True)
+        self.enabled = False
+        self._flush()
+
     def _capture(self, key: tuple) -> tuple | None:
         """Capture a CUDA graph for the given key. Returns None on failure."""
-        static_inputs = _tensors_from_key(key)
-
-        current_stream = torch.cuda.current_stream(self.device)
-        warmup_stream = torch.cuda.Stream(device=self.device)
-        warmup_stream.wait_stream(current_stream)
-        with torch.cuda.stream(warmup_stream), torch.no_grad():
-            for _ in range(3):
-                self._call_graph_fn(static_inputs)
-        current_stream.wait_stream(warmup_stream)
-
         try:
+            static_inputs = _tensors_from_key(key)
+        except KeyError:
+            # An unsupported dtype is a property of this shape alone, so run it
+            # eager and keep the other shapes on graphs.
+            logger.warning("CFM graph key carries an unsupported dtype: %s; using eager", key)
+            self._unsupported.add(key)
+            return None
+
+        memory_before = _memory_snapshot(self.device)
+        try:
+            # Warmup runs the same kernels as the capture, so a fault here
+            # leaves the same dirty capture-stream and pool state, and must be
+            # handled the same way.
+            current_stream = torch.cuda.current_stream(self.device)
+            warmup_stream = torch.cuda.Stream(device=self.device)
+            warmup_stream.wait_stream(current_stream)
+            with torch.cuda.stream(warmup_stream), torch.no_grad():
+                for _ in range(3):
+                    warmup_output = self._call_graph_fn(static_inputs)
+            current_stream.wait_stream(warmup_stream)
+            del warmup_output
+
             graph = CUDAGraph()
             with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
                 static_output = self._call_graph_fn(static_inputs)
         except Exception:
-            logger.warning(
-                "CFM graph capture failed for shape=%s; using eager",
-                key,
-                exc_info=True,
-            )
+            # A failed capture can leave the capture stream current and the
+            # allocator still routing into the graph pool, so there is no safe
+            # way to keep capturing afterwards.
+            self._disable("capture failed", key)
             return None
 
+        self._stats["captures"] += 1
         logger.info(
-            "Captured CFM CUDA Graph for shape %s (cache=%d/%d, hits=%d, misses=%d)",
+            "Captured CFM CUDA Graph for shape %s (cache=%d/%d, stats=%s)%s",
             key,
-            self._capture_graph.cache_info().currsize,
+            len(self._cache) + 1,
             self.max_graphs,
-            self._capture_graph.cache_info().hits,
-            self._capture_graph.cache_info().misses,
+            self.stats_snapshot(),
+            _format_memory_delta(memory_before, _memory_snapshot(self.device)),
         )
         return (static_inputs, static_output, graph)
 
@@ -229,25 +317,25 @@ class CFMGraphWrapper:
         att_out: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out)
+        self._stats["calls"] += 1
 
-        if torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda":
-            with torch.no_grad():
-                result = self._call_graph_fn(inputs)
-            return result, inputs[4], inputs[5]
+        if not self.enabled or torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda":
+            return self._eager(inputs)
 
         key = ("estimator_step",) + tuple(_tensor_signature(v) for v in inputs)
-        entry = self._capture_graph(key)
+        if key in self._unsupported:
+            return self._eager(inputs)
+        entry = self._cache.get(key)
 
         if entry is None:
-            logger.debug(
-                "CFM graph eager fallback for shape=%s (hits=%d, misses=%d)",
-                key,
-                self._capture_graph.cache_info().hits,
-                self._capture_graph.cache_info().misses,
-            )
-            with torch.no_grad():
-                result = self._call_graph_fn(inputs)
-            return result, inputs[4], inputs[5]
+            if len(self._cache) >= self.max_graphs:
+                self._flush()
+            entry = self._capture(key)
+            if entry is None:
+                return self._eager(inputs)
+            self._cache[key] = entry
+        else:
+            self._stats["hits"] += 1
 
         static_inputs, static_output, graph = entry
         for static, current in zip(static_inputs, inputs, strict=True):