[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")