diff --git a/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_capture_gating.py b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_capture_gating.py new file mode 100644 index 000000000..af71c07ca --- /dev/null +++ b/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)) diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 7ad35b86d..26daeff7b 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/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") diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py index 97015855e..d737905fd 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py +++ b/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):