[Perf][MiniCPM-o] Code2Wav resident Euler slot pool, fused DiT body, and exact HiFT graphs (#8443)

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

tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py

@@ -173,7 +173,13 @@ def _config(minimum: int = 1, initial: int = 0, *, runtime_prompt_cache_size: in
 
 
 @pytest.mark.parametrize(
-    ("extra", "max_num_seqs", "micro"), [({}, 6, 6), ({}, 64, 16), ({"micro_batch_size": 2}, 6, 2)]
+    ("extra", "max_num_seqs", "micro"),
+    [
+        ({}, 6, 6),
+        ({}, 64, 16),
+        ({"micro_batch_size": 2}, 6, 2),
+        ({"max_graph_batch": 8}, 64, 8),
+    ],
 )
 def test_whole_euler_micro_batch_defaults_to_stage_concurrency(extra, max_num_seqs, micro):
     config = _config()
@@ -183,6 +189,8 @@ def test_whole_euler_micro_batch_defaults_to_stage_concurrency(extra, max_num_se
     model = MiniCPMO45Code2Wav(vllm_config=config)
 
     assert model._cfm_graph_config["micro_batch_size"] == micro
+    if "max_graph_batch" in extra:
+        assert model._cfm_graph_config["max_graph_batch"] == extra["max_graph_batch"]
 
 
 def _model(initial: int = 0, minimum: int = 1, *, runtime_prompt_cache_size: int = 4, setup_cache_size: int = 1):
@@ -1701,7 +1709,17 @@ def _padded_decode_adapter(bucket_frames: int) -> tuple[BatchedToken2Wav, int, i
     Returns the adapter, the valid width, and the padding width. Every call
     seeds the same way, so two adapters built here agree on weights and noise.
     """
-    from cosyvoice2.flow.decoder_dit import DiT
+    for name in ("cosyvoice2.flow.decoder_dit", "stepaudio2.cosyvoice2.flow.decoder_dit"):
+        try:
+            import importlib
+
+            decoder_dit = importlib.import_module(name)
+            break
+        except ImportError:
+            pass
+    else:
+        decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit")
+    DiT = decoder_dit.DiT
 
     torch.manual_seed(23)
     estimator = DiT(
@@ -1809,3 +1827,92 @@ def test_padding_does_not_change_the_valid_frames():
 
     assert int(padded_x.shape[2]) == int(exact_x.shape[2]) == mel_frames
     assert torch.allclose(padded_x, exact_x, atol=1e-5)
+
+
+def test_decode_ragged_batch_mixed_onset_and_continuation():
+    """Ragged batch merging fresh onset and continuation states does not crash stacking conformer caches."""
+    adapter = BatchedToken2Wav(_FakeToken2Wav())
+    _enable_fake_ragged_kernel(adapter)
+    prompt = adapter.prepare_prompt("shared_p", "/fake/prompt.wav")
+    onset_states = adapter.setup_batch(prompt, 1)
+    state_onset = onset_states[0]
+
+    # Run one decode step on another state so its conformer/estimator caches advance.
+    cont_states_init = adapter.setup_batch(prompt, 1)
+    _, cont_states = adapter.decode_batch(
+        torch.tensor([[10, 11]]),
+        prompt,
+        cont_states_init,
+        last_chunk=False,
+    )
+    state_continuation = cont_states[0]
+
+    # Verify conformer caches have mismatched shapes between onset and continuation states.
+    assert (
+        state_onset.flow_cache["conformer_att_cache"].shape
+        != state_continuation.flow_cache["conformer_att_cache"].shape
+    )
+
+    # Ragged decode with both onset and continuation states in a single batch.
+    audios, next_states = adapter.decode_ragged_batch(
+        [torch.tensor([20, 21]), torch.tensor([30, 31])],
+        prompt,
+        [state_onset, state_continuation],
+        last_chunks=[False, False],
+    )
+    assert len(audios) == 2
+    assert len(next_states) == 2
+
+
+def test_row_offset_fallback_reaches_eager_when_replay_declines():
+    """Declined Whole-Euler replay must eager-solve, not recurse through per-offset grouping."""
+    adapter = BatchedToken2Wav(_FakeToken2Wav(), cfm_graph_config={"row_offset_merge": True})
+    _enable_fake_ragged_kernel(adapter)
+    mu, speakers, cond = torch.ones((2, 1, 2)), torch.ones((2, 1)), torch.zeros((2, 1, 2))
+    _, _, stacked = adapter._decode_cfm(mu, speakers, cond, cnn_cache=None, att_cache=None)
+    equal_rows = [
+        torch.cat((stacked[:, :, row : row + 1], stacked[:, :, 2 + row : 3 + row]), dim=2) for row in range(2)
+    ]
+    mixed_rows = [equal_rows[0], torch.nn.functional.pad(equal_rows[1], (0, 0, 0, 3))]
+    adapter._whole_euler_graph_wrapper = SimpleNamespace(enabled=True, ragged_body=object(), replay=lambda **_: None)
+    calls = {"n": 0}
+    original = adapter._decode_cfm
+
+    def counting(*args, **kwargs):
+        calls["n"] += 1
+        assert calls["n"] < 8, "row-offset fallback recursed instead of eager-solving"
+        return original(*args, **kwargs)
+
+    adapter._decode_cfm = counting
+    previous_limit = sys.getrecursionlimit()
+    sys.setrecursionlimit(64)
+    try:
+        equal_x, _, equal_att = adapter._decode_cfm(
+            mu, speakers, cond, cnn_cache=None, att_cache=equal_rows, valid_lengths=[2, 2]
+        )
+        mixed_x, _, mixed_att = adapter._decode_cfm(
+            mu, speakers, cond, cnn_cache=None, att_cache=mixed_rows, valid_lengths=[2, 2]
+        )
+    finally:
+        sys.setrecursionlimit(previous_limit)
+    assert equal_x.shape[0] == mixed_x.shape[0] == 2
+    assert len(equal_att) == len(mixed_att) == 2
+
+
+def test_tf32_mode_reads_shipped_yaml_key(monkeypatch):
+    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import _tf32_mode
+
+    monkeypatch.delenv("MINICPMO_CODE2WAV_TF32", raising=False)
+    assert _tf32_mode({}) == "tf32"
+    assert _tf32_mode({"token2wav_allow_tf32": False}) == "off"
+    assert _tf32_mode({"code2wav_allow_tf32": False}) == "off"
+    assert _tf32_mode({"token2wav_allow_tf32": True}) == "tf32"
+    assert _tf32_mode({"code2wav_allow_tf32": True}) == "tf32"
+    monkeypatch.setenv("MINICPMO_CODE2WAV_TF32", "0")
+    assert _tf32_mode({}) == "off"
+    monkeypatch.setenv("MINICPMO_CODE2WAV_TF32", "off")
+    assert _tf32_mode({}) == "off"
+    monkeypatch.setenv("MINICPMO_CODE2WAV_TF32", "tf32")
+    assert _tf32_mode({}) == "tf32"
+    monkeypatch.setenv("MINICPMO_CODE2WAV_TF32", "tf32x3")
+    assert _tf32_mode({}) == "tf32"

tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py

@@ -103,6 +103,7 @@ def _fake_wrapper(monkeypatch: pytest.MonkeyPatch) -> HiFTGraphWrapper:
     monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
     wrapper = object.__new__(HiFTGraphWrapper)
     wrapper.capture_batch_sizes = [1]
+    wrapper._legit_shapes = {(7, 0), (9, 0)}
     wrapper.graph = {}
     wrapper.static_speech_inputs = {}
     wrapper.static_cache_source_inputs = {}
@@ -880,7 +881,7 @@ def test_whole_euler_graph_with_padding_matches_eager(
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_whole_euler_cache_flushes_whole_generation(
+def test_whole_euler_graph_boundary_enforces_budget_and_falls_back_to_eager(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: torch.cuda.graph_pool_handle())
@@ -903,7 +904,8 @@ def test_whole_euler_cache_flushes_whole_generation(
             att_cache=None,
         )
 
-    _call(10)
+    res10 = _call(10)
+    assert res10 is not None
     assert len(wrapper._cache) == 1
     assert wrapper._stats["captures"] == 1
 
@@ -912,46 +914,110 @@ def test_whole_euler_cache_flushes_whole_generation(
     assert wrapper._stats["hits"] == 1
 
     # Second distinct shape fills cache to max_graphs=2
-    _call(12)
+    res12 = _call(12)
+    assert res12 is not None
     assert len(wrapper._cache) == 2
     assert wrapper._stats["captures"] == 2
     assert wrapper._stats["flushes"] == 0
 
-    # Third distinct shape exceeds max_graphs and triggers whole-generation flush
-    _call(14)
-    assert wrapper._stats["flushes"] == 1
-    assert len(wrapper._cache) == 1
-    assert wrapper._stats["captures"] == 3
+    # A third distinct shape exceeds max_graphs=2: falls back to eager (replay returns None) without expanding budget
+    res14 = _call(14)
+    assert res14 is None
+    assert wrapper._stats["flushes"] == 0
+    assert len(wrapper._cache) == 2
+    assert wrapper.max_graphs == 2
+    assert wrapper._stats["captures"] == 2
 
-    # Replay evicted shape 10 and assert numerical match with eager
-    torch.manual_seed(42)
-    x10 = torch.randn(1, 4, 10, device="cuda")
-    mu10 = torch.randn(2, 4, 10, device="cuda")
-    spk10 = torch.randn(2, 4, device="cuda")
-    cond10 = torch.randn(2, 4, 10, device="cuda")
-    res10 = wrapper.replay(
-        x=x10.clone(),
-        mu_cfg=mu10.clone(),
-        speakers_cfg=spk10.clone(),
-        cond_cfg=cond10.clone(),
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_whole_euler_precapture_enforces_budget_boundary(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: torch.cuda.graph_pool_handle())
+
+    torch.manual_seed(0)
+    estimator = _WholeEulerDiT().eval().cuda()
+    # Configure max_graphs=1 with query_bucket_frames=16
+    wrapper = WholeEulerCFMGraphWrapper(
+        estimator=estimator,
+        n_timesteps=10,
+        max_graphs=1,
+        max_graph_batch=2,
+        micro_batch_size=2,
+        query_bucket_frames=16,
+    )
+
+    count = wrapper.precapture(
+        offsets=[0, 16, 32],
+        steady=32,
+        channels=4,
+        spk_dim=4,
+    )
+    # Must capture at most 1 graph and strictly respect max_graphs=1
+    assert count <= 1
+    assert wrapper.max_graphs == 1
+    assert len(wrapper._cache) <= 1
+
+    # An uncached shape falls back to eager (None) because max_graphs=1 is already exhausted
+    x_uncached = torch.randn(1, 4, 10, device="cuda")
+    mu_uncached = torch.randn(2, 4, 10, device="cuda")
+    spk_uncached = torch.randn(2, 4, device="cuda")
+    cond_uncached = torch.randn(2, 4, 10, device="cuda")
+    res_uncached = wrapper.replay(
+        x=x_uncached,
+        mu_cfg=mu_uncached,
+        speakers_cfg=spk_uncached,
+        cond_cfg=cond_uncached,
         cnn_cache=None,
         att_cache=None,
     )
-    assert res10 is not None
-    eager_x10, _, _ = _eager_solve_euler(
-        estimator,
-        x10.clone(),
-        mu10.clone(),
-        spk10.clone(),
-        cond10.clone(),
-        None,
-        None,
-        None,
-        wrapper.timeline,
-        mel_frames=10,
-        pad_frames=0,
+    assert res_uncached is None  # Eager fallback
+    assert wrapper.max_graphs == 1
+    assert len(wrapper._cache) == 1
+    wrapper._flush()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_whole_euler_slot_entry_enforces_budget_boundary(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: torch.cuda.graph_pool_handle())
+    torch.manual_seed(0)
+    estimator = _WholeEulerDiT().eval().cuda()
+    wrapper = WholeEulerCFMGraphWrapper(
+        estimator=estimator,
+        n_timesteps=10,
+        max_graphs=2,
+        att_slots=2,
+        ragged_body=Mock(),
+        modulation_fn=lambda t: t,
+        query_bucket_frames=16,
     )
-    torch.testing.assert_close(res10[0], eager_x10, rtol=1e-4, atol=1e-5)
+    # `_record` returns (statics, x, cnn, att, graph); `_retire` calls entry[4].reset().
+    monkeypatch.setattr(wrapper, "_record", lambda key, *args, **kwargs: (Mock(), Mock(), Mock(), Mock(), Mock()))
+    pool = wrapper._ensure_slot_pool((0, 4))
+    assert pool is not None
+
+    fill = wrapper._precapture_fill
+    # Capture 1 slot graph
+    entry1 = wrapper._slot_entry(graph_batch=1, query_cap=10, channels=4, spk_dim=4, fill=fill)
+    assert entry1 is not None
+    assert len(wrapper._slot_graphs) == 1
+    assert wrapper.stats_snapshot()["cache_size"] == 1
+
+    # Capture 1 arena graph
+    x = torch.empty((1, 4, 1), device="cuda", dtype=torch.float32)
+    entry2 = wrapper._entry(graph_batch=1, query_cap=12, offset=0, x=x, spk_dim=4, fill=fill)
+    assert entry2 is not None
+    assert len(wrapper._cache) == 1
+    assert wrapper.stats_snapshot()["cache_size"] == 2
+
+    # A 3rd graph (slot or arena) exceeds max_graphs=2: falls back to eager (returns None)
+    entry3_slot = wrapper._slot_entry(graph_batch=1, query_cap=14, channels=4, spk_dim=4, fill=fill)
+    assert entry3_slot is None
+
+    entry3_arena = wrapper._entry(graph_batch=1, query_cap=16, offset=0, x=x, spk_dim=4, fill=fill)
+    assert entry3_arena is None
+
+    assert wrapper.stats_snapshot()["cache_size"] == 2
     wrapper._flush()
 
 
@@ -2128,7 +2194,16 @@ def test_whole_euler_disabled_via_serving_config() -> None:
 
 def _tiny_upstream_dit() -> nn.Module:
     """The shipped DiT architecture at toy width, so ``_blocks_forward_chunk_ragged`` runs as in serving."""
-    decoder_dit = pytest.importorskip("stepaudio2.cosyvoice2.flow.decoder_dit")
+    for name in ("cosyvoice2.flow.decoder_dit", "stepaudio2.cosyvoice2.flow.decoder_dit"):
+        try:
+            import importlib
+
+            decoder_dit = importlib.import_module(name)
+            break
+        except ImportError:
+            pass
+    else:
+        decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit")
     torch.manual_seed(0)
     estimator = decoder_dit.DiT(in_channels=16, out_channels=4, depth=2, num_heads=2, head_dim=8, hidden_size=16)
     with torch.no_grad():
@@ -2315,7 +2390,7 @@ def test_whole_euler_request_caches_grow_and_update_in_place(monkeypatch: pytest
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def test_whole_euler_groups_are_acquired_again_after_a_flush(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_whole_euler_groups_are_acquired_again_after_a_flush() -> None:
     """Capturing a later group can retire an earlier one; no group replays until all are held at once."""
     wrapper = WholeEulerCFMGraphWrapper(estimator=_tiny_upstream_dit(), n_timesteps=10, max_graphs=8)
     calls: list[int] = []
@@ -2326,13 +2401,12 @@ def test_whole_euler_groups_are_acquired_again_after_a_flush(monkeypatch: pytest
             wrapper._stats["flushes"] += 1
         return ("entry", graph_batch, len(calls))
 
-    monkeypatch.setattr(wrapper, "_entry", entry)
     flush_always = False
     groups, fills = [(16, 16), (1, 1)], [None, None]
-    assert wrapper._group_entries(groups, fills) == [("entry", 16, 3), ("entry", 1, 4)]
+    assert wrapper._group_entries(entry, groups, fills) == [("entry", 16, 3), ("entry", 1, 4)]
     assert calls == [16, 1, 16, 1]
     flush_always = True
-    assert wrapper._group_entries(groups, fills) is None
+    assert wrapper._group_entries(entry, groups, fills) is None
 
 
 @pytest.mark.parametrize("fail", [False, True])

tests/model_executor/models/minicpmo_4_5/test_dit_fused.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused Code2Wav DiT body matches the ragged reference."""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import BatchedToken2Wav
+from vllm_omni.model_executor.models.minicpmo_4_5.dit_fused import blocks_forward_chunk_fused, supports_fused_body
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])
+
+
+def _upstream_dit(device: str) -> nn.Module:
+    for name in ("cosyvoice2.flow.decoder_dit", "stepaudio2.cosyvoice2.flow.decoder_dit"):
+        try:
+            import importlib
+
+            decoder_dit = importlib.import_module(name)
+            break
+        except ImportError:
+            pass
+    else:
+        decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit")
+    torch.manual_seed(0)
+    estimator = decoder_dit.DiT(in_channels=320, out_channels=80, depth=2, num_heads=4, head_dim=16, hidden_size=64)
+    with torch.no_grad():
+        for parameter in estimator.parameters():
+            parameter.normal_(0.0, 0.08)
+    return estimator.eval().to(device)
+
+
+def _run(body, estimator, *, rows: int, frames: int, cached: int, lengths):
+    torch.manual_seed(1)
+    depth = len(estimator.blocks)
+    channels = estimator.blocks[0].conv.block[1].in_channels
+    attn = estimator.blocks[0].attn
+    device = next(estimator.parameters()).device
+    estimator_input = torch.randn(rows, 320, frames, device=device)
+    time_embedding = estimator.t_embedder(torch.rand((), device=device).expand(rows)).unsqueeze(1)
+    mask = torch.rand(rows, frames, cached + frames, device=device) > 0.3
+    mask[..., 0] = True
+    cnn = [torch.randn(rows, 2 * channels, 2, device=device) for _ in range(depth)]
+    att = [torch.randn(rows, attn.num_heads, cached, 2 * attn.head_dim, device=device) for _ in range(depth)]
+    cnn_out = torch.zeros((depth, rows, 2 * channels, 2), device=device)
+    att_out = torch.zeros((depth, rows, attn.num_heads, cached + frames, 2 * attn.head_dim), device=device)
+    with torch.no_grad():
+        out = body(estimator, estimator_input, time_embedding, mask, cnn, att, cnn_out, att_out, lengths)
+    return out, cnn_out, att_out
+
+
+@pytest.mark.parametrize("device", _DEVICES)
+def test_fused_body_matches_ragged_body(device: str) -> None:
+    estimator = _upstream_dit(device)
+    assert supports_fused_body(estimator)
+    expected = _run(
+        BatchedToken2Wav._blocks_forward_chunk_ragged, estimator, rows=4, frames=6, cached=5, lengths=[6, 4]
+    )
+    actual = _run(blocks_forward_chunk_fused, estimator, rows=4, frames=6, cached=5, lengths=[6, 4])
+    for want, got in zip(expected, actual, strict=True):
+        torch.testing.assert_close(got, want, rtol=1e-4, atol=2e-4)
+
+
+def test_unsupported_estimator_is_refused() -> None:
+    estimator = _upstream_dit("cpu")
+    estimator.blocks[0].attn.q_norm = nn.Identity()
+    assert not supports_fused_body(estimator)

tests/model_executor/models/minicpmo_4_5/test_flow_encoder_graph.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exact-shape flow-encoder graphs (``cfm_encoder_cuda_graph``): keys, eager fallbacks, bit-exact replays."""
+
+import functools
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.model_executor.models.minicpmo_4_5.flow_encoder_graph as graph_module
+from vllm_omni.model_executor.models.minicpmo_4_5.flow_encoder_graph import (
+    FlowEncoderGraphs,
+    reachable_conformer_cache_frames,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_DEPTH, _HEADS, _WIDTH, _HIDDEN = 2, 2, 4, 8
+_CNN = (3, 2)
+_LOOKAHEAD = 1
+
+
+def _encode(tokens, *, cnn_cache, att_cache, weight):
+    """Row-independent stand-in for embedding + ``forward_chunk`` + projection: 1-token lookahead, 2x upsample."""
+    rows = int(tokens.shape[0])
+    x = (tokens.to(weight.dtype).unsqueeze(-1) @ weight[:1])[:, :-_LOOKAHEAD]
+    x = x.unsqueeze(2).expand(-1, -1, 2, -1).reshape(rows, -1, _HIDDEN)
+    x = x + cnn_cache.sum(dim=(1, 2))[:, None, None] + att_cache.mean(dim=(0, 2, 3, 4))[:, None, None]
+    hidden = x @ weight
+    new_cnn = hidden[:, -_CNN[1] :, : _CNN[0]].transpose(1, 2).contiguous()
+    added = hidden[:, :, :_WIDTH][None, :, None].expand(_DEPTH, rows, _HEADS, -1, _WIDTH)
+    return hidden, new_cnn, torch.cat((att_cache, added), dim=3)
+
+
+@pytest.fixture
+def fake_capture(monkeypatch: pytest.MonkeyPatch) -> None:
+    # A "graph" whose replay runs the recorded body on the static views.
+    record = lambda self, views: SimpleNamespace(replay=functools.partial(self._body, views))  # noqa: E731
+    monkeypatch.setattr(FlowEncoderGraphs, "_record", record)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+
+
+def _graphs(device: str, rows=(1, 2, 3), frames=(8, 18)):
+    weight = torch.randn(_HIDDEN, _HIDDEN, generator=torch.Generator().manual_seed(0)).to(device)
+    encode = functools.partial(_encode, weight=weight)
+    graphs = FlowEncoderGraphs(encode, rows=rows, token_widths=[6], lookahead=_LOOKAHEAD, upsample=2)
+    keys = [(r, 6, f) for r in rows for f in frames]
+    layout = {"cnn_shape": _CNN, "att_layout": (_DEPTH, _HEADS, _WIDTH), "hidden_dim": _HIDDEN}
+    assert graphs.capture(keys, **layout, dtype=torch.float32, device=torch.device(device)) == len(keys)
+    return graphs, encode
+
+
+def _inputs(rows: int, width: int, frames: int, device: str, seed: int = 1):
+    generator = torch.Generator().manual_seed(seed)
+    tokens = torch.randint(0, 50, (rows, width), generator=generator)
+    cnn = [torch.randn(1, *_CNN, generator=generator) for _ in range(rows)]
+    att = [torch.randn(_DEPTH, 1, _HEADS, frames, _WIDTH, generator=generator) for _ in range(rows)]
+    return tokens.to(device), [c.to(device) for c in cnn], [a.to(device) for a in att]
+
+
+@pytest.mark.parametrize(
+    ("growths", "expected"),
+    [([50], [300, 350, 400]), ([150], [300, 400]), ([30], [300, 330, 360, 390, 400]), ([], [300])],
+)
+def test_reachable_cache_frames_follow_the_streaming_trim(growths, expected) -> None:
+    # 6 s default voice (300 frames); 28-token duplex units grow the cache by 50 frames.
+    assert reachable_conformer_cache_frames(start=300, prompt_len=300, suffix=100, growths=growths) == expected
+
+
+def test_keys_cover_every_row_count_width_and_cache_length() -> None:
+    graphs = FlowEncoderGraphs(lambda *a, **k: a, rows=[2, 1, 2], token_widths=[3, 28], lookahead=3, upsample=2)
+    assert graphs.token_widths == (28,)  # a width with no new frames is dropped
+    assert graphs.keys_for(start=300, prompt_len=300, suffix=100) == [
+        (rows, 28, frames) for rows in (1, 2) for frames in (300, 350, 400)
+    ]
+
+
+def _check_replay(graphs, encode, rows: int, frames: int, device: str) -> None:
+    tokens, cnn, att = _inputs(rows, 6, frames, device, seed=rows * 100 + frames)
+    graphed = graphs.run(tokens, cnn, att)
+    expected = encode(tokens, cnn_cache=torch.cat(cnn, dim=0), att_cache=torch.cat(att, dim=1))
+    assert graphed is not None
+    for actual, reference in zip(graphed, expected, strict=True):
+        assert torch.equal(actual, reference)
+
+
+@pytest.mark.parametrize(("rows", "frames"), [(1, 8), (3, 18), (3, 8)])
+def test_replay_matches_the_eager_encoder_exactly(fake_capture, rows, frames) -> None:
+    graphs, encode = _graphs("cpu")
+    _check_replay(graphs, encode, rows, frames, "cpu")
+    assert graphs.replays == 1
+
+
+def test_other_shapes_and_layouts_run_eager(fake_capture, monkeypatch: pytest.MonkeyPatch) -> None:
+    graphs, _ = _graphs("cpu")
+    for rows, width, frames in ((4, 6, 8), (1, 5, 8), (1, 6, 10)):  # never padded, never another shape
+        assert graphs.run(*_inputs(rows, width, frames, "cpu")) is None
+    tokens, cnn, att = _inputs(2, 6, 8, "cpu")
+    assert graphs.run(tokens, cnn, [att[0], att[1].double()]) is None
+    assert graphs.run(tokens.int(), cnn, att) is None
+    assert graphs.run(tokens[:1], cnn, att) is None
+    with monkeypatch.context() as patch:
+        patch.setattr(graph_module, "_precision_state", lambda: ("other precision flags",))
+        assert graphs.run(tokens, cnn, att) is None
+    assert graphs.run(tokens, cnn, att) is not None
+    assert graphs.replays == 1
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture")
+def test_captured_graphs_match_eager_bit_exactly_on_cuda() -> None:
+    graphs, encode = _graphs("cuda")
+    for rows, frames in ((1, 8), (3, 18), (1, 18), (3, 8)):
+        _check_replay(graphs, encode, rows, frames, "cuda")
+    assert graphs.replays == 4