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

docs/design/minicpm_o45_mrv2_performance.md

@@ -12,8 +12,11 @@ require new end-to-end performance measurements after the mainline integration.
 Previously reported numbers do not describe this revised codec backend.
 
 The generic MRv2 profile and the opt-in native V1 duplex H200 profile enable
-TF32 matmuls within Code2Wav forward/capture only; the previous process policy
-is restored afterwards. cuDNN's TF32 policy is independent. TF32 changes rounding.
+ordinary TF32 for CFM DiT GEMMs within Code2Wav forward/capture only
+(`torch.backends.cuda.matmul.allow_tf32` on dense QKV/MLP, Triton
+`input_precision="tf32"` on tiled attention). This is not compensated TF32x3.
+The previous process matmul policy is restored afterwards. cuDNN's TF32 policy
+is independent. HiFT stays IEEE FP32. TF32 changes rounding.
 Duplex keeps the mainline V1 session path. Turn results do not establish duplex
 performance or interruption correctness.
 

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 32,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 95,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 96,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
         EnvironmentVariableCategory.INTERNAL: 7,
     }
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
     )
     assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
-        ModelEnvironmentVariableDisposition.PROMOTE: 66,
+        ModelEnvironmentVariableDisposition.PROMOTE: 67,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 17,

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

vllm_omni/config/environment_variable_inventory.py

@@ -172,6 +172,7 @@ _MODEL_PROMOTE = (
     "MIMO_AUDIO_TOKENIZER_DEVICE",
     "MIMO_AUDIO_TOKENIZER_PATH",
     "MING_CFM_CUDAGRAPH",
+    "MINICPMO_CODE2WAV_TF32",
     "MINICPMO_TOKEN2WAV_TRT",
     "MINICPMO_TOKEN2WAV_TRT_CACHE",
     "MINICPMO_TOKEN2WAV_TRT_DTYPE",

vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml

@@ -1,7 +1,8 @@
 # Opt-in single-H200 native full-duplex codec profile (V1).
 # Inherits the native session/stage capacities, sampling and chunk sizes.
 # Uses the mainline Whole-Euler Flow and shared prompt-state implementation.
-# TF32 changes FP32 GEMM rounding; validate audio quality for your workload.
+# Ordinary TF32 for CFM DiT GEMMs (cuBLAS allow_tf32 + Triton tf32 dots).
+# HiFT stays IEEE FP32. Validate audio quality for your workload.
 base_config: minicpmo_4_5.yaml
 
 connectors:

vllm_omni/deploy/minicpmo_4_5_turn_mrv2.yaml

@@ -14,9 +14,11 @@ connectors:
   connector_of_shared_memory:
     extra:
       hift_graph_capture_batch_sizes: [1, 2, 4, 8]
-      # TF32 tensor cores for the vocoder's fp32 flow GEMMs: ~25% less GPU
-      # time per chunk (batch 1-8). A numerics change: the HiFT input mel moves
-      # by ~3e-4 relative L2; dev-short WER/UTMOS unchanged (64 prompts).
+      # Ordinary TF32 for CFM DiT GEMMs (~25% less GPU time per chunk at
+      # batch 1-8): cuBLAS allow_tf32 on dense QKV/MLP and Triton
+      # input_precision=tf32 on tiled attention. HiFT stays IEEE FP32.
+      # Numerics: HiFT input mel moves by ~3e-4 relative L2; dev-short
+      # WER/UTMOS unchanged.
       token2wav_allow_tf32: true
 
 stages:

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+import functools
+import math
 from collections import OrderedDict
 from collections.abc import Mapping
 from contextlib import nullcontext
@@ -15,16 +17,32 @@ import torch.nn as nn
 import torch.nn.functional as F
 from vllm.logger import init_logger
 
+from vllm_omni.utils.device_copy import index_to_device
+
 from .cuda_graph_wrapper import (
     CFMGraphWrapper,
     HiFTGraphWrapper,
+    ResidentAttCache,
     WholeEulerCFMGraphWrapper,
     _att_keep_ranges,
     _copy_frame_segments,
     _euler_step,
     _euler_timeline,
+    _format_memory_delta,
+    _materialize_att_rows,
+    _memory_snapshot,
     _zero_padded_cnn_cache,
+    empty_hift_outputs,
 )
+from .dit_fused import (
+    blocks_forward_chunk_fused,
+    causal_cache_index,
+    dit_modulation,
+    gather_causal_cache,
+    supports_fused_body,
+    tiled_attention_supported,
+)
+from .flow_encoder_graph import FlowEncoderGraphs
 
 logger = init_logger(__name__)
 
@@ -37,10 +55,20 @@ _MAX_ENCODE_TOKEN_CAP = 1024
 # Streaming cache trim (stepaudio2 ``Token2wav.stream``): past ``prompt_len +
 # 100`` frames, keep the first ``prompt_len`` frames and the last 100.
 _CACHE_TRIM_SUFFIX = 100
+_TIMESTEP_MAX_PERIOD = 10000
+
+
+def _timestep_frequencies(dim: int) -> torch.Tensor:
+    """Host frequency table of ``TimestepEmbedder.timestep_embedding``."""
+    half = dim // 2
+    return torch.exp(-math.log(_TIMESTEP_MAX_PERIOD) * torch.arange(start=0, end=half) / half)
 
 
 def _trim_streaming_cache(cache: torch.Tensor, prompt_len: int) -> torch.Tensor:
     """The streaming trim of a cache whose frame axis is ``-2``."""
+    if isinstance(cache, ResidentAttCache):
+        # Its slot-pool solve already trimmed it (``WholeEulerCFMGraphWrapper._replay_slots``).
+        return cache
     ranges = _att_keep_ranges(int(cache.shape[-2]), (prompt_len, _CACHE_TRIM_SUFFIX))
     if len(ranges) == 1:
         return cache
@@ -207,6 +235,21 @@ def state_shape_signature(state: BatchedToken2WavState) -> tuple[Any, ...]:
     return flow, hift
 
 
+# Cache frame axes a row-offset merge keeps per row: the encoder runs once per conformer cache
+# length (``decode_ragged_batch``), and a slot-pool solve attends each estimator cache in its slot.
+_ROW_OFFSET_FRAME_AXES = {"conformer_att_cache": 3, "estimator_att_cache": 4}
+
+
+def row_offset_signature(state: BatchedToken2WavState) -> tuple[Any, ...]:
+    """``state_shape_signature`` without the cache lengths a row-offset merge keeps per row."""
+    flow, hift = state_shape_signature(state)
+    masked = []
+    for name, (shape, dtype, device) in flow:
+        axis = _ROW_OFFSET_FRAME_AXES.get(name)
+        masked.append((name, (shape if axis is None else (*shape[:axis], -1, *shape[axis + 1 :]), dtype, device)))
+    return tuple(masked), hift
+
+
 @dataclass(frozen=True)
 class PromptFeatures:
     cache_key: tuple[str, str]
@@ -240,6 +283,23 @@ def _undecorate_dynamo(module: nn.Module, method: str) -> None:
     logger.info("Bypassed TorchDynamo wrapper on %s.%s", type(module).__name__, method)
 
 
+def _cfm_matmul_tf32(fn):
+    """Run the flow CFM (and capture its graphs) with ordinary TF32 GEMMs when ``cfm_tf32``; HiFT stays fp32."""
+
+    @functools.wraps(fn)
+    def wrapper(self, *args, **kwargs):
+        if not getattr(self, "_cfm_tf32", False):
+            return fn(self, *args, **kwargs)
+        previous = torch.backends.cuda.matmul.allow_tf32
+        torch.backends.cuda.matmul.allow_tf32 = True
+        try:
+            return fn(self, *args, **kwargs)
+        finally:
+            torch.backends.cuda.matmul.allow_tf32 = previous
+
+    return wrapper
+
+
 class BatchedToken2Wav(nn.Module):
     """Drive Token2wav's modules with dynamically-sized, request-owned caches.
 
@@ -258,11 +318,18 @@ class BatchedToken2Wav(nn.Module):
         cfm_graph_config: Mapping[str, Any] | None = None,
         bfloat16_attention_cache: bool = False,
         setup_cache_size: int = 1,
+        cfm_tf32: bool = False,
+        encoder_graph_config: Mapping[str, Any] | None = None,
     ):
         super().__init__()
         if setup_cache_size < 0:
             raise ValueError("setup_cache_size must be >= 0")
         self._token2wav = token2wav
+        self._cfm_tf32 = bool(cfm_tf32)
+        cfm_graph_cfg = dict(cfm_graph_config or {})
+        # Ragged solves of rows whose caches differ in length (a stream's second chunk next to steady
+        # ones): each row keeps its own cache offset, in one slot-pool Whole-Euler replay when it can.
+        self._row_offset_merge = bool(cfm_graph_cfg.get("row_offset_merge", False))
         # Optional TrtDiTStepper (step_audio2_dit_trt): replaces only the
         # per-timestep DiT estimator call; encoder and HiFT stay on torch.
         self._trt_stepper = trt_stepper
@@ -274,6 +341,7 @@ class BatchedToken2Wav(nn.Module):
         encoder = getattr(self.flow, "encoder", None)
         if encoder is not None:
             _undecorate_dynamo(encoder, "forward_chunk")
+        self._encoder_graphs = self._build_encoder_graphs(encoder_graph_config or {}, connector_config or {})
         hift_parameter = next(self.hift.parameters(), None)
         if hift_parameter is not None and hift_parameter.device.type == "cuda":
             # Prime the CUDA state used by HiFT during backend construction.
@@ -336,13 +404,16 @@ class BatchedToken2Wav(nn.Module):
                     capture_batch_sizes=capture_batch_sizes,
                     max_serial_batch=max_serial_batch,
                 )
-                with torch.inference_mode(), _autocast_disabled(hift_parameter.device):
-                    self.hift_graph_wrapper.capture()
-                logger.info("HiFT CUDA Graph captured successfully")
+                # Captured by the first forward (``precapture_hift``): graphs
+                # captured while vLLM loads the weights hold far more memory.
         self._cfm_graph_wrapper: CFMGraphWrapper | None = None
         self._whole_euler_graph_wrapper: WholeEulerCFMGraphWrapper | None = None
-        cfm_graph_cfg = dict(cfm_graph_config or {})
-        if bool(cfm_graph_cfg.get("enabled", False)):
+        # Whether the ragged DiT body is the fused one (``dit_fused.py``).
+        self._ragged_fused_body = False
+        # On NPU the platform graph runner captures instead of the CUDA wrappers, so bucketing and
+        # padding key off the requested flag.
+        self._cfm_graph_enabled = bool(cfm_graph_cfg.get("enabled", False))
+        if self._cfm_graph_enabled:
             flow_parameter = next(self.flow.parameters(), None)
             if flow_parameter is not None and flow_parameter.device.type == "cuda":
                 estimator = self.flow.decoder.estimator
@@ -352,7 +423,26 @@ class BatchedToken2Wav(nn.Module):
                 max_graph_batch = int(max_graph_batch_cfg) if max_graph_batch_cfg is not None else None
                 micro_batch_size_cfg = cfm_graph_cfg.get("micro_batch_size")
                 micro_batch_size = int(micro_batch_size_cfg) if micro_batch_size_cfg is not None else 4
-                if bool(cfm_graph_cfg.get("enable_whole_euler", True)) and self._trt_stepper is None:
+                enable_whole_euler = bool(cfm_graph_cfg.get("enable_whole_euler", True))
+                if bool(cfm_graph_cfg.get("fused_body", False)):
+                    # A reduced-precision attention cache keeps the ragged body,
+                    # which attends in the activation dtype through upstream.
+                    self._ragged_fused_body = (
+                        supports_fused_body(estimator) and self._estimator_att_cache_dtype == torch.float32
+                    )
+                    if not self._ragged_fused_body:
+                        logger.warning(
+                            "cfm_fused_body requested but unsupported here (DiT blocks or a non-fp32 attention cache)"
+                        )
+                # Resident request caches need the fused body and its tiled
+                # attention, which reads keys/values out of the pool in place.
+                slot_pool = bool(cfm_graph_cfg.get("slot_pool", False))
+                if slot_pool and not (
+                    self._ragged_fused_body and tiled_attention_supported(flow_parameter.get_device())
+                ):
+                    logger.warning("cfm_slot_pool requested but needs cfm_fused_body on NVIDIA SM80+; using the arena")
+                    slot_pool = False
+                if enable_whole_euler and self._trt_stepper is None:
                     capture_frames = cfm_graph_cfg.get("capture_frames")
                     if capture_frames is None and connector_config is not None:
                         # The steady chunk's mel width, derived like HiFTGraphWrapper's
@@ -361,6 +451,11 @@ class BatchedToken2Wav(nn.Module):
                         capture_frames = int(connector_config["codec_chunk_frames"]) * int(
                             getattr(self.flow, "token_mel_ratio", 2)
                         )
+                    if isinstance(capture_frames, (list, tuple)):
+                        # Several widths (e.g. a 25-token duplex unit and a
+                        # 75-token turn chunk) each get their own graphs.
+                        query_bucket_frames: int | tuple[int, ...] = tuple(int(w) for w in capture_frames)
+                    else:
                         query_bucket_frames = int(capture_frames if capture_frames is not None else 64)
                     self._whole_euler_graph_wrapper = WholeEulerCFMGraphWrapper(
                         estimator=estimator,
@@ -373,24 +468,30 @@ class BatchedToken2Wav(nn.Module):
                         micro_batch_size=micro_batch_size,
                         query_bucket_frames=query_bucket_frames,
                         pad_max_rows=cfm_graph_cfg.get("pad_max_rows"),
-                        # Looked up at capture time, so a replaced ragged kernel is the one captured.
-                        ragged_body=(
-                            (lambda *args: self._blocks_forward_chunk_ragged(*args))
-                            if _supports_ragged_kernel(estimator)
-                            else None
-                        ),
+                        offset_bucket_frames=int(cfm_graph_cfg.get("offset_bucket_frames", 0)),
+                        # Resident request caches, a slot per stream plus a few for streams
+                        # that end while new ones start (``AttSlotPool``).
+                        att_slots=(micro_batch_size + max(2, micro_batch_size // 8)) if slot_pool else 0,
+                        # ``_ragged_body`` picks the configured body on every call, so a
+                        # replaced body is the one a capture records.
+                        ragged_body=self._ragged_body if _supports_ragged_kernel(estimator) else None,
+                        # The fused body takes each timestep's modulation, computed once outside the graphs.
+                        modulation_fn=functools.partial(dit_modulation, estimator) if self._ragged_fused_body else None,
+                        row_offsets=self._row_offset_merge,
                     )
                     logger.info(
                         "Whole-Euler CFM CUDA Graph enabled "
                         "(max_graphs=%d, max_serial_batch=%d, max_graph_batch=%s, "
-                        "micro_batch_size=%d, query_bucket_frames=%d)",
+                        "micro_batch_size=%d, query_bucket_frames=%s, offset_bucket_frames=%d, fused_body=%s)",
                         max_graphs,
                         max_serial_batch,
                         str(max_graph_batch),
                         micro_batch_size,
                         query_bucket_frames,
+                        self._whole_euler_graph_wrapper.offset_bucket_frames,
+                        self._ragged_fused_body,
                     )
-                elif self._trt_stepper is not None and bool(cfm_graph_cfg.get("enable_whole_euler", True)):
+                elif enable_whole_euler:
                     logger.info("Whole-Euler CFM CUDA Graph disabled because TensorRT stepper is configured")
                 self._cfm_graph_wrapper = CFMGraphWrapper(
                     graph_fn=estimator.blocks_forward_chunk,
@@ -407,7 +508,12 @@ class BatchedToken2Wav(nn.Module):
         # space stays small (0 disables bucketing, e.g. when graphs are off).
         self._cfm_graph_bucket_frames = (
             int(cfm_graph_cfg.get("bucket_frames", 0))
-            if (self._cfm_graph_wrapper is not None or self._whole_euler_graph_wrapper is not None)
+            if (
+                self._cfm_graph_wrapper is not None
+                or self._whole_euler_graph_wrapper is not None
+                # NPU has no wrappers; the platform runner is keyed by this flag.
+                or self._cfm_graph_enabled
+            )
             else 0
         )
         if self._cfm_graph_bucket_frames > 1:
@@ -415,6 +521,8 @@ class BatchedToken2Wav(nn.Module):
                 "CFM CUDA Graph bucketing enabled (bucket_frames=%d)",
                 self._cfm_graph_bucket_frames,
             )
+        # Device copies of the DiT timestep-embedding frequency table.
+        self._timestep_freqs: dict[tuple[Any, ...], torch.Tensor] = {}
         self._prompt_features: dict[tuple[str, str], PromptFeatures] = {}
         self._setup_cache_size = setup_cache_size
         self._setup_cache: OrderedDict[tuple[tuple[str, str], int], BatchedToken2WavState] = OrderedDict()
@@ -424,6 +532,9 @@ class BatchedToken2Wav(nn.Module):
         mel: torch.Tensor,
         source_cache: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        if int(mel.shape[2]) == 0:
+            # No frames, no samples: HiFT itself cannot run on an empty mel.
+            return empty_hift_outputs(mel)
         if self.hift_graph_wrapper is None:
             return self.hift.inference(mel, source_cache)
         return self.hift_graph_wrapper.replay(mel, source_cache)
@@ -467,6 +578,17 @@ class BatchedToken2Wav(nn.Module):
             features.mels.expand(batch_size, -1, -1),
         )
 
+    def _project_speakers(self, speakers: torch.Tensor) -> torch.Tensor:
+        """``spk_embed_affine_layer(normalize(speakers))`` with M=1, independent of the batch size.
+
+        Rows expand one prompt speaker, so one row is projected: M=1 and M>=2
+        GEMMs round differently (~6e-8), moving a row's CFM mel by ~1e-3.
+        """
+        if speakers.shape[0] > 1 and speakers.stride(0) == 0:
+            projected = self.flow.spk_embed_affine_layer(F.normalize(speakers[:1], dim=1))
+            return projected.expand(speakers.shape[0], -1).contiguous()
+        return self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1))
+
     def _autocast(self, device: torch.device):
         if device.type != "cuda":
             return nullcontext()
@@ -538,6 +660,96 @@ class BatchedToken2Wav(nn.Module):
         )
         return self.flow.encoder_proj(hidden), new_cnn, new_att
 
+    def _flow_on_cuda(self) -> bool:
+        flow_parameter = next(self.flow.parameters(), None)
+        return flow_parameter is not None and flow_parameter.device.type == "cuda"
+
+    def _build_encoder_graphs(self, config: Mapping[str, Any], connector: Mapping) -> FlowEncoderGraphs | None:
+        """``cfm_encoder_cuda_graph`` (default off): exact-shape graphs of the flow encoder's continuation chunk."""
+        if not config.get("enabled", False):
+            return None
+        if not self._flow_on_cuda():
+            logger.info("Flow encoder CUDA Graph is disabled off CUDA")
+            return None
+        encoder = self.flow.encoder
+        embed = getattr(encoder, "embed", None)
+        lookahead = self._pre_lookahead_len()
+        if (
+            lookahead is None
+            or getattr(encoder, "up_layer", None) is None
+            or not isinstance(getattr(getattr(embed, "pos_enc", None), "pe", None), torch.Tensor)
+            # ``_ensure_relpos_pe`` would grow such a table between calls.
+            or callable(getattr(embed, "extend_pe", None))
+        ):
+            logger.warning("Flow encoder CUDA Graph needs the CosyVoice2 upsample-conformer encoder; staying eager")
+            return None
+        # Every ``position_encoding`` slice a graph reads lives in these tables.
+        tables = [
+            module
+            for module in (embed, getattr(encoder, "up_embed", None))
+            if isinstance(getattr(getattr(module, "pos_enc", None), "pe", None), torch.Tensor)
+        ]
+        # Default width: the steady duplex chunk, the left context plus one unit of new frames.
+        unit = int(connector.get("initial_codec_chunk_frames", 0) or 0) or int(connector.get("codec_chunk_frames", 25))
+        graphs = FlowEncoderGraphs(
+            # With no ``extend_pe`` (checked above) this runs exactly the eager continuation's ops.
+            functools.partial(self._encode_chunk, last_chunk=False),
+            rows=config.get("rows") or [1],
+            token_widths=config.get("token_widths") or [int(connector.get("codec_left_context_frames", 3)) + unit],
+            lookahead=lookahead,
+            upsample=self._upsample_stride(),
+            held_tensors=lambda: tuple(module.pos_enc.pe for module in tables),
+        )
+        logger.info(
+            "Flow encoder CUDA Graph enabled for continuation chunks (rows %s, tokens %s)",
+            list(graphs.rows),
+            list(graphs.token_widths),
+        )
+        return graphs
+
+    def _graph_encode(
+        self, tokens: torch.Tensor, states: list[BatchedToken2WavState], copy: bool = False
+    ) -> tuple[torch.Tensor, ...] | None:
+        """A continuation chunk's encoder outputs from its graph: views rewritten by the next replay, or copies."""
+        if self._encoder_graphs is None:
+            return None
+        flows = [state.flow_cache for state in states]
+        outputs = self._encoder_graphs.run(
+            tokens, [flow["conformer_cnn_cache"] for flow in flows], [flow["conformer_att_cache"] for flow in flows]
+        )
+        return tuple(value.clone() for value in outputs) if outputs and copy else outputs
+
+    @torch.inference_mode()
+    def precapture_flow_encoder(self, features: PromptFeatures) -> int:
+        """Capture the encoder graphs of every continuation chunk this prompt's streams run; returns the count."""
+        graphs = self._encoder_graphs
+        if graphs is None or graphs.graphs:
+            return 0
+        (state,) = self.setup_batch(features, 1)
+        cnn, att = state.flow_cache["conformer_cnn_cache"], state.flow_cache["conformer_att_cache"]
+        prompt_len = int(features.mels.shape[1])
+        keys = graphs.keys_for(start=int(att.shape[3]), prompt_len=prompt_len, suffix=_CACHE_TRIM_SUFFIX)
+        torch.accelerator.synchronize(att.device)
+        memory_before = _memory_snapshot(att.device)
+        with self._autocast(att.device):
+            captured = graphs.capture(
+                keys,
+                cnn_shape=(int(cnn.shape[1]), int(cnn.shape[2])),
+                att_layout=(int(att.shape[0]), int(att.shape[2]), int(att.shape[4])),
+                hidden_dim=int(self.flow.encoder_proj.out_features),
+                dtype=att.dtype,
+                device=att.device,
+            )
+        torch.accelerator.synchronize(att.device)
+        logger.info(
+            "Captured %d flow encoder CUDA Graphs (cache frames %s): shared buffers %.1f MiB%s",
+            captured,
+            sorted({key[2] for key in keys}),
+            graphs.storage_bytes() / 2**20,
+            _format_memory_delta(memory_before, _memory_snapshot(att.device)),
+        )
+        return captured
+
     @staticmethod
     def _estimator_buffers(
         estimator: nn.Module,
@@ -558,6 +770,28 @@ class BatchedToken2Wav(nn.Module):
         att = x.new_empty((depth, batch_size, heads, old_att_len + chunk_size, att_width))
         return cnn, att
 
+    def _time_embedding(self, estimator: nn.Module, time: torch.Tensor) -> torch.Tensor:
+        """``estimator.t_embedder(time)`` without a host sync per call (bitwise identical).
+
+        Upstream rebuilds its frequency table on the host and moves it with a
+        pageable ``.to(t)`` once per CFM step, blocking the host until the GPU
+        drains; the table is built the same way once and kept on the device.
+        """
+        embedder = estimator.t_embedder
+        dim, scale, mlp = (getattr(embedder, name, None) for name in ("frequency_embedding_size", "scale", "mlp"))
+        if not isinstance(dim, int) or scale is None or mlp is None:
+            return embedder(time)
+        # The host table takes the default dtype before ``.to(t)``, as upstream.
+        key = (time.device, time.dtype, torch.get_default_dtype(), dim)
+        freqs = self._timestep_freqs.get(key)
+        if freqs is None:
+            freqs = self._timestep_freqs[key] = _timestep_frequencies(dim).to(time)
+        args = (time * scale)[:, None] * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return mlp(embedding)
+
     def _estimator_step(
         self,
         estimator: nn.Module,
@@ -585,8 +819,7 @@ class BatchedToken2Wav(nn.Module):
                 att_cache=att_cache,
             )
             return out.to(mu.dtype), new_cnn, new_att
-        if time_embedding is None:
-            time_embedding = estimator.t_embedder(time).unsqueeze(1)
+        time_embedding = self._time_embedding(estimator, time).unsqueeze(1)
         width = int(x.shape[-1])
         speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width)
         if valid_frames is not None and valid_frames < width:
@@ -621,7 +854,7 @@ class BatchedToken2Wav(nn.Module):
         if valid_lengths is not None:
             if not hasattr(estimator, "in_proj"):
                 raise RuntimeError('MiniCPMO45Code2WavBatchError {"reason":"ragged_kernel_unavailable"}')
-            result = self._blocks_forward_chunk_ragged(
+            result = self._ragged_body(
                 estimator,
                 estimator_input,
                 time_embedding,
@@ -644,15 +877,21 @@ class BatchedToken2Wav(nn.Module):
             )
         return result, cnn_out, att_out
 
-    @staticmethod
-    def _gather_causal_cache(
-        history: torch.Tensor,
-        valid_lengths: torch.Tensor,
-        width: int,
+    def _ragged_body(
+        self,
+        *args: Any,
+        modulation: torch.Tensor | None = None,
+        slots: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> torch.Tensor:
-        """The ``width`` frames of ``history`` (b, width + t, c) ending each row's valid frames, as (b, c, width)."""
-        indices = valid_lengths[:, None] + torch.arange(width, device=history.device)[None, :]
-        return history.gather(1, indices[:, :, None].expand(-1, -1, int(history.shape[2]))).transpose(1, 2)
+        """The configured ragged DiT body, looked up per call so a capture records a replaced one.
+
+        ``modulation`` (fused body only) comes from the Whole-Euler arena.
+        """
+        estimator_input = args[1]
+        if slots is not None or (self._ragged_fused_body and estimator_input.dtype == torch.float32):
+            # Slot-pool graphs are only built over the fused body (``slot_pool`` above).
+            return blocks_forward_chunk_fused(*args, modulation=modulation, slots=slots)
+        return self._blocks_forward_chunk_ragged(*args)
 
     @staticmethod
     def _causal_conv_frames_major(conv: nn.Conv1d, history: torch.Tensor) -> torch.Tensor:
@@ -765,6 +1004,8 @@ class BatchedToken2Wav(nn.Module):
                 device=estimator_input.device,
                 dtype=torch.long,
             )
+        # Causal cache gather index per conv width, shared by the blocks.
+        cache_index: dict[int, torch.Tensor] = {}
         x = estimator.in_proj(estimator_input.transpose(1, 2))
         for block_index, block in enumerate(estimator.blocks):
             (
@@ -803,11 +1044,13 @@ class BatchedToken2Wav(nn.Module):
                     )
                 )
             width = int(old_cnn.shape[2])
+            if width not in cache_index:
+                cache_index[width] = causal_cache_index(lengths, width)
             history = torch.cat((old_cnn[:, : conv.in_channels].transpose(1, 2), conv_input), dim=1)
-            new_cnn1 = BatchedToken2Wav._gather_causal_cache(history, lengths, width)
+            new_cnn1 = gather_causal_cache(history, cache_index[width])
             hidden = conv.block[4](conv.block[3](BatchedToken2Wav._causal_conv_frames_major(conv.block[1], history)))
             history = torch.cat((old_cnn[:, conv.in_channels :].transpose(1, 2), hidden), dim=1)
-            new_cnn2 = BatchedToken2Wav._gather_causal_cache(history, lengths, width)
+            new_cnn2 = gather_causal_cache(history, cache_index[width])
             x = x + gate_conv * BatchedToken2Wav._causal_conv_frames_major(conv.block[6], history)
             x = x + gate_mlp * block.mlp(block.norm2(x) * (1 + scale_mlp) + shift_mlp)
 
@@ -829,6 +1072,7 @@ class BatchedToken2Wav(nn.Module):
             self._timeline_cache[key] = cached
         return cached
 
+    @_cfm_matmul_tf32
     def _decode_cfm(
         self,
         mu: torch.Tensor,
@@ -859,7 +1103,18 @@ class BatchedToken2Wav(nn.Module):
             offset = int(att_rows[0].shape[4])
         else:
             offset = int(att_cache.shape[4]) if att_cache is not None else 0
+        # ``row_offset_merge``: ragged rows whose caches differ in length each
+        # start from the noise at their own offset and attend their own cache.
+        row_offsets = None
+        if att_rows is not None and valid_lengths is not None:
+            lengths = [int(row.shape[4]) for row in att_rows]
+            if len(set(lengths)) > 1 or self._row_offset_merge:
+                row_offsets = lengths
+                if len(set(row_offsets)) == 1 and not self._row_offset_merge:
+                    row_offsets = None
         mel_frames = int(mu.shape[2])
+        # Padding only pays off while replay is active: Whole-Euler, the CFM wrapper (still enabled,
+        # no TRT stepper), or on NPU the platform runner keyed by `_cfm_graph_enabled`.
         graphs_active = valid_lengths is None and (
             self._whole_euler_active()
             or (
@@ -867,6 +1122,7 @@ class BatchedToken2Wav(nn.Module):
                 and getattr(self._cfm_graph_wrapper, "enabled", True)
                 and self._trt_stepper is None
             )
+            or (self._cfm_graph_wrapper is None and self._cfm_graph_enabled)
         )
         pad_frames = _cfm_pad_frames(
             mel_frames=mel_frames,
@@ -881,14 +1137,18 @@ class BatchedToken2Wav(nn.Module):
             # carry less of a step change into the attention and CNN caches.
             mu = torch.nn.functional.pad(mu, (0, pad_frames), mode="replicate")
             cond = torch.nn.functional.pad(cond, (0, pad_frames), mode="replicate")
-        end = offset + int(mu.shape[2])
+        end = (offset if row_offsets is None else max(row_offsets)) + int(mu.shape[2])
         if end > int(decoder.rand_noise.shape[2]):
             raise RuntimeError(
                 "MiniCPMO45Code2WavBatchError "
                 f'{{"reason":"noise_capacity","required":{end},'
                 f'"available":{int(decoder.rand_noise.shape[2])}}}'
             )
+        if row_offsets is None:
             x = decoder.rand_noise[:, :, offset:end].expand(batch_size, -1, -1).clone()
+        else:
+            width = int(mu.shape[2])
+            x = torch.cat([decoder.rand_noise[:, :, start : start + width] for start in row_offsets], dim=0)
         if pad_frames:
             # The padded columns would otherwise carry real noise values. They
             # are excluded from this chunk's attention by ``attn_mask`` below
@@ -904,18 +1164,21 @@ class BatchedToken2Wav(nn.Module):
         if valid_lengths is not None:
             if len(valid_lengths) != batch_size:
                 raise ValueError(f"valid length count {len(valid_lengths)} != batch {batch_size}")
-            cfg_lengths = torch.tensor(
-                (*valid_lengths, *valid_lengths),
-                device=mu.device,
-            )
+            cfg_lengths = index_to_device((*valid_lengths, *valid_lengths), mu.device)
             positions = torch.arange(int(mu.shape[2]), device=mu.device)
             valid_queries = positions.unsqueeze(0) < cfg_lengths.unsqueeze(1)
             current_keys = valid_queries.unsqueeze(1).expand(-1, int(mu.shape[2]), -1)
+            if row_offsets is None:
                 old_keys = torch.ones(
                     (2 * batch_size, int(mu.shape[2]), offset),
                     dtype=torch.bool,
                     device=mu.device,
                 )
+            else:
+                # Cache columns span the longest cache; a shorter one's tail stays masked.
+                cfg_offsets = index_to_device((*row_offsets, *row_offsets), mu.device)
+                cached = torch.arange(max(row_offsets), device=mu.device).unsqueeze(0) < cfg_offsets.unsqueeze(1)
+                old_keys = cached.unsqueeze(1).expand(-1, int(mu.shape[2]), -1)
             attn_mask = valid_queries.unsqueeze(2) & torch.cat((current_keys, old_keys), dim=2)
         elif pad_frames:
             # Mask the padded keys instead of only zeroing their content: a
@@ -948,15 +1211,26 @@ class BatchedToken2Wav(nn.Module):
             )
             if whole_euler_result is not None:
                 return whole_euler_result
+        if row_offsets is not None and len(set(row_offsets)) > 1:
+            # Mixed cache lengths only. A same-length group (including every
+            # ``_decode_cfm_per_offset`` subgroup, and equal-length rows under
+            # ``row_offset_merge``) must fall through to eager: otherwise a
+            # declined Whole-Euler replay re-enters this fallback forever.
+            return self._decode_cfm_per_offset(
+                mu,
+                speakers,
+                cond,
+                cnn_cache=cnn_cache,
+                att_rows=att_rows,
+                valid_lengths=valid_lengths,
+                att_keep=att_keep,
+            )
         if att_rows is not None:
             att_cache = self._stack_estimator_att(att_rows)
 
         next_cnn: list[torch.Tensor] = []
         next_att_cache: torch.Tensor | None = None
         ragged_att_cache: list[torch.Tensor] | None = None
-        time_embeddings = [
-            estimator.t_embedder(timeline[s].expand(2 * batch_size)).unsqueeze(1) for s in range(self.n_timesteps)
-        ]
         with _token2wav_sdpa_context(mu.device):
             for step in range(self.n_timesteps):
                 old_cnn = cnn_cache[step] if cnn_cache is not None else None
@@ -974,7 +1248,6 @@ class BatchedToken2Wav(nn.Module):
                     attn_mask=attn_mask,
                     valid_lengths=valid_lengths,
                     valid_frames=mel_frames if pad_frames else None,
-                    time_embedding=time_embeddings[step],
                 )
                 if pad_frames:
                     _zero_padded_cnn_cache(step_cnn, estimator, pad_frames)
@@ -1032,6 +1305,48 @@ class BatchedToken2Wav(nn.Module):
         assert next_att_cache is not None
         return x, torch.stack(next_cnn), next_att_cache
 
+    def _decode_cfm_per_offset(
+        self,
+        mu: torch.Tensor,
+        speakers: torch.Tensor,
+        cond: torch.Tensor,
+        *,
+        cnn_cache: torch.Tensor | None,
+        att_rows: list,
+        valid_lengths: list[int],
+        **kwargs: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor, list]:
+        """``_decode_cfm`` once per cache length, in row order: a row-offset merge no single solve took."""
+        batch_size = len(att_rows)
+        groups: dict[int, list[int]] = {}
+        for row, cache in enumerate(att_rows):
+            groups.setdefault(int(cache.shape[4]), []).append(row)
+        x: torch.Tensor | None = None
+        cnn: torch.Tensor | None = None
+        att: list = [None] * batch_size
+        for rows in groups.values():
+            index = index_to_device(rows, mu.device)
+            # CFG rows are ``[cond x B | uncond x B]``.
+            cfg_index = index_to_device([*rows, *(batch_size + row for row in rows)], mu.device)
+            group_x, group_cnn, group_att = self._decode_cfm(
+                mu.index_select(0, index),
+                speakers.index_select(0, index),
+                cond.index_select(0, index),
+                cnn_cache=None if cnn_cache is None else cnn_cache.index_select(2, cfg_index),
+                att_cache=[att_rows[row] for row in rows],
+                valid_lengths=[valid_lengths[row] for row in rows],
+                **kwargs,
+            )
+            if x is None or cnn is None:
+                x = group_x.new_empty((batch_size, *group_x.shape[1:]))
+                cnn = group_cnn.new_empty((*group_cnn.shape[:2], 2 * batch_size, *group_cnn.shape[3:]))
+            x.index_copy_(0, index, group_x)
+            cnn.index_copy_(2, cfg_index, group_cnn)
+            for row, cache in zip(rows, group_att, strict=True):
+                att[row] = cache
+        assert x is not None and cnn is not None
+        return x, cnn, att
+
     def _split_flow_cache(
         self,
         cache: dict[str, Any],
@@ -1085,6 +1400,7 @@ class BatchedToken2Wav(nn.Module):
     @staticmethod
     def _stack_estimator_att(rows: list[torch.Tensor]) -> torch.Tensor:
         """Stack per-request ``(n_t, depth, 2, ...)`` caches into CFG order ``[cond x B | uncond x B]``."""
+        rows = _materialize_att_rows(rows)
         return torch.cat((*[row[:, :, 0:1] for row in rows], *[row[:, :, 1:2] for row in rows]), dim=2)
 
     def _stack_flow_cache(
@@ -1092,19 +1408,26 @@ class BatchedToken2Wav(nn.Module):
         states: list[BatchedToken2WavState],
         *,
         include_estimator_att: bool = True,
+        include_conformer: bool = True,
     ) -> dict[str, torch.Tensor | None]:
         """Stack per-request flow caches for one batched call.
 
         The estimator attention cache is by far the largest (~0.5 GiB per
         request in fp32). The Whole-Euler path reads it per request, so it
-        skips stacking it (``include_estimator_att=False``).
+        skips stacking it (``include_estimator_att=False``). A caller that
+        already ran the encoder skips the conformer caches
+        (``include_conformer=False``).
         """
         flows = [state.flow_cache for state in states]
         conditional_cnn = [flow["estimator_cnn_cache"][:, :, 0:1] for flow in flows]
         unconditional_cnn = [flow["estimator_cnn_cache"][:, :, 1:2] for flow in flows]
         return {
-            "conformer_cnn_cache": torch.cat([flow["conformer_cnn_cache"] for flow in flows], dim=0),
-            "conformer_att_cache": torch.cat([flow["conformer_att_cache"] for flow in flows], dim=1),
+            "conformer_cnn_cache": (
+                torch.cat([flow["conformer_cnn_cache"] for flow in flows], dim=0) if include_conformer else None
+            ),
+            "conformer_att_cache": (
+                torch.cat([flow["conformer_att_cache"] for flow in flows], dim=1) if include_conformer else None
+            ),
             "estimator_cnn_cache": torch.cat((*conditional_cnn, *unconditional_cnn), dim=2),
             "estimator_att_cache": (
                 self._stack_estimator_att([flow["estimator_att_cache"] for flow in flows])
@@ -1121,6 +1444,69 @@ class BatchedToken2Wav(nn.Module):
         """Whether Whole-Euler graphs also solve ragged (per-row length) batches."""
         return self._whole_euler_active() and self._whole_euler_graph_wrapper.ragged_body is not None
 
+    def _ragged_merge_fits(self, states: list[BatchedToken2WavState], token_counts: list[int]) -> bool:
+        """Ragged Whole-Euler is on, the rows share a CNN cache shape, and each row's tokens encode at once."""
+        return (
+            2 <= len(states) == len(token_counts)
+            and self._whole_euler_ragged_active()
+            and len({tuple(state.flow_cache["estimator_cnn_cache"].shape) for state in states}) == 1
+            and all(count <= self._max_encode_token_frames([s]) for count, s in zip(token_counts, states, strict=True))
+        )
+
+    def can_merge_state_shapes(self, states: list[BatchedToken2WavState], token_counts: list[int]) -> bool:
+        """Whether mixed onset/continuation rows fit ragged Whole-Euler."""
+        caches = [state.flow_cache.get("estimator_att_cache") for state in states]
+        return (
+            all(isinstance(cache, torch.Tensor) for cache in caches)
+            and len({(tuple(cache.shape[:4]), cache.dtype, cache.device) for cache in caches}) == 1
+            and self._ragged_merge_fits(states, token_counts)
+        )
+
+    def can_merge_row_offsets(self, states: list[BatchedToken2WavState], token_counts: list[int]) -> bool:
+        """Whether continuation rows whose caches differ only in length fit one slot-pool ragged solve.
+
+        ``row_offset_merge``: every estimator cache must be resident in the
+        Whole-Euler slot pool, where each row attends its own frames;
+        ``decode_ragged_batch`` runs the encoder once per conformer cache length.
+        """
+        if not (self._row_offset_merge and self._ragged_merge_fits(states, token_counts)):
+            return False
+        pool = self._whole_euler_graph_wrapper.slot_pool
+        caches = [state.flow_cache.get("estimator_att_cache") for state in states]
+        return pool is not None and all(isinstance(cache, ResidentAttCache) and cache.pool is pool for cache in caches)
+
+    @torch.inference_mode()
+    def precapture_hift(self) -> int:
+        """Capture every configured HiFT graph (prompt-independent) once, before serving; returns the count."""
+        wrapper = self.hift_graph_wrapper
+        if wrapper is None or wrapper.graph:
+            return 0
+        with _autocast_disabled(wrapper.device):
+            wrapper.capture()
+        return len(wrapper.graph)
+
+    @torch.inference_mode()
+    @_cfm_matmul_tf32
+    def precapture_whole_euler(self, features: PromptFeatures) -> int:
+        """Capture the Whole-Euler graphs of every cache length this prompt's streams pass through.
+
+        Runs in ``_decode_cfm``'s precision (``_cfm_matmul_tf32``, ``_autocast``):
+        a graph replays the kernels it was captured with. Returns the count.
+        """
+        if not self._whole_euler_active():
+            return 0
+        wrapper = self._whole_euler_graph_wrapper
+        prompt_len = int(features.mels.shape[1])
+        steady = prompt_len + _CACHE_TRIM_SUFFIX
+        with self._autocast(wrapper.device):
+            return wrapper.precapture(
+                offsets=range(prompt_len, steady + 1),
+                steady=steady,
+                channels=int(features.mels.shape[2]),
+                spk_dim=int(self.flow.spk_embed_affine_layer.out_features),
+                keep=(prompt_len, _CACHE_TRIM_SUFFIX),
+            )
+
     def _create_initial_states(
         self,
         features: PromptFeatures,
@@ -1139,7 +1525,7 @@ class BatchedToken2Wav(nn.Module):
                 cnn_cache=None,
                 att_cache=None,
             )
-            projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1))
+            projected_speakers = self._project_speakers(speakers)
             # The prompt cache never reaches the trim length; ``att_keep`` only
             # tells the Whole-Euler arena the steady cache length to reserve.
             _, estimator_cnn, estimator_att = self._decode_cfm(
@@ -1301,18 +1687,23 @@ class BatchedToken2Wav(nn.Module):
         # Whole-Euler reads and writes each request's estimator cache directly,
         # so the batch never holds a stacked copy of it.
         per_request_att = self._whole_euler_active()
-        flow_cache = self._stack_flow_cache(states, include_estimator_att=not per_request_att)
+        # A captured encoder graph stacks the conformer caches into its own inputs.
+        graphed = None if last_chunk or flush_encoder else self._graph_encode(tokens, states)
+        flow_cache = self._stack_flow_cache(
+            states, include_estimator_att=not per_request_att, include_conformer=graphed is None
+        )
         prompt_len = int(features.mels.shape[1])
         att_keep = (prompt_len, _CACHE_TRIM_SUFFIX)
         speakers = features.speaker_embedding.expand(batch_size, -1)
         with self._autocast(tokens.device):
-            hidden, conformer_cnn, conformer_att = self._encode_chunk(
+            # A graph's shared results are read below and copied per row by _split_flow_cache.
+            hidden, conformer_cnn, conformer_att = graphed or self._encode_chunk(
                 tokens,
                 last_chunk=last_chunk or flush_encoder,
                 cnn_cache=flow_cache["conformer_cnn_cache"],
                 att_cache=flow_cache["conformer_att_cache"],
             )
-            projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1))
+            projected_speakers = self._project_speakers(speakers)
             cond = torch.zeros_like(hidden).transpose(1, 2).contiguous()
             chunk_mel, estimator_cnn, estimator_att = self._decode_cfm(
                 hidden.transpose(1, 2).contiguous(),
@@ -1432,20 +1823,30 @@ class BatchedToken2Wav(nn.Module):
                     audios[row] = group_audio[group_row]
                     next_states[row] = group_states[group_row]
             return self._require_complete_ragged_outputs(audios, next_states)
-        encoder_groups: dict[tuple[int, bool], list[int]] = {}
-        for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)):
-            encoder_groups.setdefault((int(row_tokens.numel()), last_chunk), []).append(row)
+        encoder_groups: dict[tuple[Any, ...], list[int]] = {}
+        for row, (row_tokens, state, last_chunk) in enumerate(zip(tokens, states, last_chunks, strict=True)):
+            # Onset and continuation rows may differ in conformer cache length: encode them apart;
+            # the estimator caches stay per-request rows in the ragged Whole-Euler solve.
+            key = (
+                tuple(state.flow_cache["conformer_cnn_cache"].shape),
+                tuple(state.flow_cache["conformer_att_cache"].shape),
+                int(row_tokens.numel()),
+                last_chunk,
+            )
+            encoder_groups.setdefault(key, []).append(row)
 
         hidden_rows: list[torch.Tensor | None] = [None] * batch_size
         conformer_cnn_rows: list[torch.Tensor | None] = [None] * batch_size
         conformer_att_rows: list[torch.Tensor | None] = [None] * batch_size
-        for (_, last_chunk), rows in encoder_groups.items():
+        for (*_, last_chunk), rows in encoder_groups.items():
             group_states = [states[row] for row in rows]
-            # The encoder reads only the conformer caches.
-            group_cache = self._stack_flow_cache(group_states, include_estimator_att=False)
             group_tokens = torch.stack([tokens[row] for row in rows], dim=0)
+            # With several groups the next replay rewrites the shared graph results, so they are copied.
+            graphed = None if last_chunk else self._graph_encode(group_tokens, group_states, len(encoder_groups) > 1)
+            # The encoder reads only the conformer caches.
+            group_cache = {} if graphed else self._stack_flow_cache(group_states, include_estimator_att=False)
             with self._autocast(group_tokens.device):
-                hidden, conformer_cnn, conformer_att = self._encode_chunk(
+                hidden, conformer_cnn, conformer_att = graphed or self._encode_chunk(
                     group_tokens,
                     last_chunk=last_chunk,
                     cnn_cache=group_cache["conformer_cnn_cache"],
@@ -1470,11 +1871,16 @@ class BatchedToken2Wav(nn.Module):
 
         # A graphed ragged solve reads and writes each request's estimator
         # cache directly (trimmed on the way out), like _decode_batch_once.
-        flow_cache = self._stack_flow_cache(states, include_estimator_att=not per_request_att)
+        # Row-offset merges hold caches of different lengths, which never stack: the solve takes them
+        # per row, and the encoder groups above already read the conformer caches.
+        per_row_att = per_request_att or (
+            len({int(state.flow_cache["estimator_att_cache"].shape[4]) for state in states}) > 1
+        )
+        flow_cache = self._stack_flow_cache(states, include_estimator_att=not per_row_att, include_conformer=False)
         prompt_len = int(features.mels.shape[1])
         speakers = features.speaker_embedding.expand(batch_size, -1)
         with self._autocast(padded_hidden.device):
-            projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1))
+            projected_speakers = self._project_speakers(speakers)
             cond = torch.zeros_like(padded_hidden).transpose(1, 2).contiguous()
             chunk_mel, estimator_cnn, estimator_att = self._decode_cfm(
                 padded_hidden.transpose(1, 2).contiguous(),
@@ -1483,7 +1889,7 @@ class BatchedToken2Wav(nn.Module):
                 cnn_cache=flow_cache["estimator_cnn_cache"],
                 att_cache=(
                     [state.flow_cache["estimator_att_cache"] for state in states]
-                    if per_request_att
+                    if per_row_att
                     else flow_cache["estimator_att_cache"]
                 ),
                 valid_lengths=hidden_lengths,
@@ -1516,10 +1922,17 @@ class BatchedToken2Wav(nn.Module):
 
         audios: list[torch.Tensor | None] = [None] * batch_size
         next_states: list[BatchedToken2WavState | None] = [None] * batch_size
-        vocoder_groups: dict[tuple[int, bool], list[int]] = {}
+        vocoder_groups: dict[tuple[Any, ...], list[int]] = {}
         for row, last_chunk in enumerate(last_chunks):
-            vocoder_groups.setdefault((hidden_lengths[row], last_chunk), []).append(row)
-        for (hidden_length, last_chunk), rows in vocoder_groups.items():
+            key = (
+                hidden_lengths[row],
+                last_chunk,
+                tuple(states[row].hift_cache["mel"].shape),
+                tuple(states[row].hift_cache["source"].shape),
+                tuple(states[row].hift_cache["speech"].shape),
+            )
+            vocoder_groups.setdefault(key, []).append(row)
+        for (hidden_length, last_chunk, *_), rows in vocoder_groups.items():
             old_mel = torch.cat([states[row].hift_cache["mel"] for row in rows], dim=0)
             old_source = torch.cat([states[row].hift_cache["source"] for row in rows], dim=0)
             old_speech = torch.cat([states[row].hift_cache["speech"] for row in rows], dim=0)

vllm_omni/model_executor/models/minicpmo_4_5/cfm_attention.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Tiled CFM attention: float32 tensors, ordinary TF32 Tensor Core dots."""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _attention(
+    q,
+    k,
+    v,
+    mask,
+    kv_rows,
+    out,
+    qb: tl.constexpr,
+    qh: tl.constexpr,
+    qt: tl.constexpr,
+    qd: tl.constexpr,
+    kb: tl.constexpr,
+    kh: tl.constexpr,
+    kt: tl.constexpr,
+    kd: tl.constexpr,
+    vb: tl.constexpr,
+    vh: tl.constexpr,
+    vt: tl.constexpr,
+    vd: tl.constexpr,
+    mb: tl.constexpr,
+    mt: tl.constexpr,
+    mk: tl.constexpr,
+    heads: tl.constexpr,
+    nq: tl.constexpr,
+    nk: tl.constexpr,
+    dim: tl.constexpr,
+    has_mask: tl.constexpr,
+    has_kv_rows: tl.constexpr,
+    scale: tl.constexpr,
+    block_m: tl.constexpr,
+    block_n: tl.constexpr,
+    block_d: tl.constexpr,
+):
+    b, h, tile = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    # Keys and values of a pooled cache live in row ``kv_rows[b]`` of the pool.
+    kv_b = b
+    if has_kv_rows:
+        kv_b = tl.load(kv_rows + b).to(tl.int64)
+    rows = tile * block_m + tl.arange(0, block_m)
+    ds = tl.arange(0, block_d)
+    cols = tl.arange(0, block_n)
+    queries = tl.load(
+        q + b * qb + h * qh + rows[:, None] * qt + ds[None, :] * qd, (rows[:, None] < nq) & (ds[None, :] < dim), 0
+    )
+    maximum = tl.full((block_m,), -float("inf"), tl.float32)
+    denominator = tl.full((block_m,), 0, tl.float32)
+    accumulator = tl.full((block_m, block_d), 0, tl.float32)
+    for start in range(tl.cdiv(nk, block_n)):
+        keys = start * block_n + cols
+        kval = tl.load(
+            k + kv_b * kb + h * kh + keys[None, :] * kt + ds[:, None] * kd,
+            (keys[None, :] < nk) & (ds[:, None] < dim),
+            0,
+        )
+        score = tl.dot(queries, kval, input_precision="tf32") * scale
+        valid = (rows[:, None] < nq) & (keys[None, :] < nk)
+        if has_mask:
+            keep = tl.load(mask + b * mb + rows[:, None] * mt + keys[None, :] * mk, valid, 0)
+            valid = valid & (keep != 0)
+        score = tl.where(valid, score, -float("inf"))
+        new_max = tl.maximum(maximum, tl.max(score, 1))
+        # Fully masked rows produce zero, matching PyTorch SDPA.
+        # Keep the running maximum at -inf until a valid tile arrives. A
+        # permanent zero would underflow a later tile with very negative logits.
+        safe_max = tl.where(new_max == -float("inf"), 0.0, new_max)
+        correction = tl.exp(maximum - safe_max)
+        prob = tl.exp(score - safe_max[:, None])
+        denominator = denominator * correction + tl.sum(prob, 1)
+        values = tl.load(
+            v + kv_b * vb + h * vh + keys[:, None] * vt + ds[None, :] * vd,
+            (keys[:, None] < nk) & (ds[None, :] < dim),
+            0,
+        )
+        accumulator = accumulator * correction[:, None] + tl.dot(prob, values, input_precision="tf32")
+        maximum = new_max
+    result = accumulator / tl.where(denominator > 0, denominator, 1.0)[:, None]
+    tl.store(
+        out + ((b * nq + rows[:, None]) * heads + h) * dim + ds[None, :],
+        result,
+        (rows[:, None] < nq) & (ds[None, :] < dim),
+    )
+
+
+@torch.library.custom_op("vllm_omni::cfm_tiled_attention", mutates_args=())
+def cfm_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    mask: torch.Tensor | None = None,
+    kv_rows: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Tiled attention with TF32 dots; optional ``kv_rows`` indexes a pooled K/V cache."""
+    b, heads, nq, dim = q.shape
+    nk = k.shape[2]
+    if q.dtype != torch.float32 or k.dtype != torch.float32 or v.dtype != torch.float32:
+        raise ValueError("CFM tiled attention requires float32 Q/K/V")
+    num_stages = 2
+    out = torch.empty((b, nq, heads, dim), device=q.device, dtype=q.dtype)
+    _attention[(b, heads, triton.cdiv(nq, 32))](
+        q,
+        k,
+        v,
+        mask,
+        kv_rows,
+        out,
+        *q.stride(),
+        *k.stride(),
+        *v.stride(),
+        *(mask.stride() if mask is not None else (0, 0, 0)),
+        heads,
+        nq,
+        nk,
+        dim,
+        mask is not None,
+        kv_rows is not None,
+        dim**-0.5,
+        32,
+        64,
+        triton.next_power_of_2(dim),
+        num_warps=4,
+        num_stages=num_stages,
+    )
+    return out.transpose(1, 2)
+
+
+@cfm_attention.register_fake
+def _cfm_attention_fake(q, k, v, mask=None, kv_rows=None):
+    b, heads, nq, dim = q.shape
+    return q.new_empty((b, nq, heads, dim)).transpose(1, 2)

vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py

@@ -0,0 +1,284 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused Code2Wav CFM DiT block forward (~15 kernels vs ~47 eager ops per block)."""
+
+from __future__ import annotations
+
+import functools
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .ops import qkv_head_layer_norm, residual_layer_norm
+
+# adaLN chunk order of ``DiTBlock.adaLN_modulation``.
+_SHIFT_MSA, _SCALE_MSA, _GATE_MSA, _SHIFT_MLP, _SCALE_MLP, _GATE_MLP, _SHIFT_CONV, _SCALE_CONV, _GATE_CONV = range(9)
+
+
+def supports_fused_body(estimator: nn.Module) -> bool:
+    """Whether ``blocks_forward_chunk_fused`` computes this estimator's blocks."""
+    blocks = getattr(estimator, "blocks", None)
+    final_layer = getattr(estimator, "final_layer", None)
+    if not blocks or not hasattr(estimator, "in_proj") or final_layer is None:
+        return False
+    if not isinstance(getattr(final_layer, "norm_final", None), nn.LayerNorm) or not hasattr(final_layer, "linear"):
+        return False
+    for block in blocks:
+        attn = getattr(block, "attn", None)
+        conv_block = getattr(getattr(block, "conv", None), "block", None)
+        mlp = getattr(block, "mlp", None)
+        if attn is None or conv_block is None or mlp is None:
+            return False
+        if not all(hasattr(attn, name) for name in ("to_q", "to_k", "to_v", "proj", "num_heads", "head_dim")):
+            return False
+        if not (isinstance(attn.q_norm, nn.LayerNorm) and isinstance(attn.k_norm, nn.LayerNorm)):
+            return False
+        if not all(isinstance(getattr(block, name, None), nn.LayerNorm) for name in ("norm1", "norm2", "norm3")):
+            return False
+        if len(conv_block) < 7 or not isinstance(conv_block[3], nn.LayerNorm) or not isinstance(conv_block[4], nn.Mish):
+            return False
+        for conv in (conv_block[1], conv_block[6]):
+            if not isinstance(conv, nn.Conv1d) or conv.stride != (1,) or conv.dilation != (1,) or conv.groups != 1:
+                return False
+            if conv.padding != (0,) or conv.kernel_size[0] < 2:
+                return False
+        if not (hasattr(mlp, "fc1") and hasattr(mlp, "fc2") and isinstance(mlp.act, nn.GELU)):
+            return False
+        if not isinstance(getattr(mlp, "norm", nn.Identity()), nn.Identity):
+            return False
+    return True
+
+
+def dit_modulation(estimator: nn.Module, time_embedding: torch.Tensor) -> torch.Tensor:
+    """adaLN table ``(depth + 1, 9, C)`` for one timestep; scales stored as ``1 + scale``."""
+    embedding = time_embedding.reshape(1, 1, -1)
+    table = torch.stack([block.adaLN_modulation(embedding).reshape(9, -1) for block in estimator.blocks], dim=0)
+    table[:, (_SCALE_MSA, _SCALE_MLP, _SCALE_CONV)] += 1.0
+    final_shift, final_scale = estimator.final_layer.adaLN_modulation(embedding).reshape(2, -1)
+    final = table.new_zeros((1, 9, int(table.shape[2])))
+    # Where a block keeps its MSA shift/scale, so the last block's closing LayerNorm reads it like a next block.
+    final[0, _SHIFT_MSA] = final_shift
+    final[0, _SCALE_MSA] = final_scale + 1.0
+    return torch.cat((table, final), dim=0).contiguous()
+
+
+def _cached(module: nn.Module, key: str, sources: tuple[torch.Tensor, ...], build) -> Any:
+    version = tuple((t.data_ptr(), t._version, t.dtype) for t in sources)
+    cached = module.__dict__.get(key)
+    if cached is None or cached[0] != version:
+        with torch.no_grad():
+            cached = (version, build())
+        module.__dict__[key] = cached
+    return cached[1]
+
+
+def packed_qkv(attn: nn.Module) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """``[to_q; to_k; to_v]`` weight and bias as one projection, rebuilt only when a source changes."""
+    sources = (attn.to_q.weight, attn.to_k.weight, attn.to_v.weight)
+
+    def build():
+        weight = torch.cat(sources, dim=0).contiguous()
+        bias = None
+        if attn.to_q.bias is not None:
+            bias = torch.cat((attn.to_q.bias, attn.to_k.bias, attn.to_v.bias), dim=0).contiguous()
+        return weight, bias
+
+    return _cached(attn, "_packed_qkv", sources, build)
+
+
+def _conv_taps_weight(conv: nn.Conv1d) -> torch.Tensor:
+    """``(C_out, C_in, K)`` -> ``(K * C_out, C_in)``: tap ``k`` is output rows ``[k*C_out, (k+1)*C_out)``."""
+    return _cached(
+        conv,
+        "_fused_taps_weight",
+        (conv.weight,),
+        lambda: conv.weight.detach().permute(2, 0, 1).reshape(-1, conv.in_channels).contiguous(),
+    )
+
+
+def _conv_history(
+    old_cache: torch.Tensor | None, rows: int, width: int, frames: int, channels: int, like: torch.Tensor
+) -> torch.Tensor:
+    """``(N, width + frames, C)``: the cached frames, then room for the current ones."""
+    history = like.new_empty((rows, width + frames, channels))
+    if old_cache is None:
+        history[:, :width].zero_()
+    else:
+        history[:, :width].copy_(old_cache.transpose(1, 2))
+    return history
+
+
+def causal_cache_index(lengths: torch.Tensor, width: int) -> torch.Tensor:
+    """``(N, width)`` history frames that end each row's valid frames, for ``gather_causal_cache``."""
+    return lengths[:, None] + torch.arange(width, device=lengths.device)[None, :]
+
+
+def gather_causal_cache(history: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
+    """The ``index`` frames of ``history`` (N, width + T, C), as the (N, C, width) causal cache."""
+    return history.gather(1, index[:, :, None].expand(-1, -1, int(history.shape[2]))).transpose(1, 2)
+
+
+def _mlp(mlp: nn.Module, hidden: torch.Tensor) -> torch.Tensor:
+    rows = hidden.reshape(-1, int(hidden.shape[-1]))
+    if mlp.act.approximate == "tanh" and rows.is_cuda and mlp.fc1.bias is not None:
+        # cuBLASLt applies bias + tanh-GELU in the GEMM epilogue.
+        inner = torch._addmm_activation(mlp.fc1.bias, rows, mlp.fc1.weight.t(), use_gelu=True)
+    else:
+        inner = mlp.act(F.linear(rows, mlp.fc1.weight, mlp.fc1.bias))
+    return F.linear(inner, mlp.fc2.weight, mlp.fc2.bias).view(*hidden.shape[:-1], -1)
+
+
+@functools.cache
+def tiled_attention_supported(device_index: int) -> bool:
+    """``cfm_attention`` TF32 dots need NVIDIA tensor cores (SM80+)."""
+    return torch.version.hip is None and torch.cuda.get_device_capability(device_index) >= (8, 0)
+
+
+def _attention(
+    attn: nn.Module,
+    hidden: torch.Tensor,
+    att_cache: torch.Tensor | None,
+    attn_mask: torch.Tensor | None,
+    kv: torch.Tensor,
+    slots: tuple[torch.Tensor, torch.Tensor] | None = None,
+) -> torch.Tensor:
+    """``attn.forward_chunk`` over ``hidden``, writing new keys/values into ``kv``."""
+    batch, frames, _ = hidden.shape
+    weight, bias = packed_qkv(attn)
+    qkv = F.linear(hidden.reshape(batch * frames, -1), weight, bias).view(batch, frames, -1)
+    rows, positions = slots if slots is not None else (None, None)
+    if slots is None:
+        # ``att_cache`` behind the current frames: left where it is when it
+        # already sits there (the Whole-Euler arena's aliasing), else copied once.
+        cached = 0 if att_cache is None else int(att_cache.shape[2])
+        kv = kv[:, :, : frames + cached]
+        behind = kv[:, :, frames:]
+        if cached and (att_cache.data_ptr() != behind.data_ptr() or att_cache.stride() != behind.stride()):
+            behind.copy_(att_cache)
+    q = qkv_head_layer_norm(
+        qkv,
+        kv,
+        num_heads=int(attn.num_heads),
+        head_dim=int(attn.head_dim),
+        q_norm=attn.q_norm,
+        k_norm=attn.k_norm,
+        rows=rows,
+        positions=positions,
+    )
+    head_dim = int(q.shape[-1])
+    keys, values = kv[..., :head_dim], kv[..., head_dim:]
+    if rows is not None or (
+        kv.is_cuda and kv.dtype == q.dtype == torch.float32 and tiled_attention_supported(kv.get_device())
+    ):
+        from .cfm_attention import cfm_attention
+
+        out = cfm_attention(q, keys, values, attn_mask, rows)
+    else:
+        out = F.scaled_dot_product_attention(
+            q, keys, values, attn_mask=None if attn_mask is None else attn_mask.unsqueeze(1)
+        )
+    out = out.transpose(1, 2).reshape(batch * frames, -1)
+    return F.linear(out, attn.proj.weight, attn.proj.bias).view(batch, frames, -1)
+
+
+def blocks_forward_chunk_fused(
+    estimator: nn.Module,
+    estimator_input: torch.Tensor,
+    time_embedding: torch.Tensor,
+    attn_mask: torch.Tensor | None,
+    cnn_cache: Any,
+    att_cache: Any,
+    cnn_cache_buffer: torch.Tensor,
+    att_cache_buffer: torch.Tensor,
+    valid_lengths: list[int] | torch.Tensor,
+    modulation: torch.Tensor | None = None,
+    slots: tuple[torch.Tensor, torch.Tensor] | None = None,
+) -> torch.Tensor:
+    """``BatchedToken2Wav._blocks_forward_chunk_ragged`` with fused norms/GEMM/attention."""
+    if isinstance(valid_lengths, torch.Tensor):
+        lengths = valid_lengths
+    else:
+        lengths = torch.tensor((*valid_lengths, *valid_lengths), device=estimator_input.device, dtype=torch.long)
+    if modulation is None:
+        modulation = dit_modulation(estimator, time_embedding[:1])
+    blocks = estimator.blocks
+    depth = len(blocks)
+    # Causal cache gather index per conv width, shared by the blocks.
+    cache_index: dict[int, torch.Tensor] = {}
+
+    x = estimator.in_proj(estimator_input.transpose(1, 2)).contiguous()
+    rows, frames, _ = x.shape
+    first = modulation[0]
+    hidden = residual_layer_norm(x, weight=first[_SCALE_MSA], bias=first[_SHIFT_MSA], eps=blocks[0].norm1.eps)
+    for block_index, block in enumerate(blocks):
+        mod = modulation[block_index]
+        x_att = _attention(block.attn, hidden, att_cache[block_index], attn_mask, att_cache_buffer[block_index], slots)
+
+        # CausalConvBlock: [conv1, LayerNorm, Mish, conv2] on (N, width + T, C) histories.
+        conv_layers = block.conv.block
+        conv1, conv_norm, conv2 = conv_layers[1], conv_layers[3], conv_layers[6]
+        width = int(conv1.kernel_size[0]) - 1
+        old_cnn = cnn_cache[block_index]
+        split = conv1.in_channels
+        old_cnn1, old_cnn2 = (None, None) if old_cnn is None else (old_cnn[:, :split], old_cnn[:, split:])
+        history = _conv_history(old_cnn1, rows, width, frames, split, x)
+        # x += gate_msa * attention; the conv block's modulated LayerNorm lands in the history.
+        residual_layer_norm(
+            x,
+            x_att,
+            gate=mod[_GATE_MSA],
+            weight=mod[_SCALE_CONV],
+            bias=mod[_SHIFT_CONV],
+            eps=block.norm3.eps,
+            residual_out=x,
+            out=history[:, width:],
+        )
+        taps1 = F.linear(history.view(-1, conv1.in_channels), _conv_taps_weight(conv1))
+        history2 = _conv_history(old_cnn2, rows, width, frames, conv2.in_channels, x)
+        residual_layer_norm(
+            None,
+            taps1.view(rows, width + frames, -1),
+            taps=width + 1,
+            y_bias=conv1.bias,
+            weight=conv_norm.weight,
+            bias=conv_norm.bias,
+            eps=conv_norm.eps,
+            activation="mish",
+            out=history2[:, width:],
+        )
+        taps2 = F.linear(history2.view(-1, conv2.in_channels), _conv_taps_weight(conv2))
+        # x += gate_conv * conv2; the MLP's modulated LayerNorm.
+        mlp_in = residual_layer_norm(
+            x,
+            taps2.view(rows, width + frames, -1),
+            taps=width + 1,
+            y_bias=conv2.bias,
+            gate=mod[_GATE_CONV],
+            weight=mod[_SCALE_MLP],
+            bias=mod[_SHIFT_MLP],
+            eps=block.norm2.eps,
+            residual_out=x,
+        )
+        mlp_out = _mlp(block.mlp, mlp_in)
+        # x += gate_mlp * MLP, then the next block's (or the final layer's) modulated LayerNorm:
+        # table row ``depth`` holds the final layer's shift and 1 + scale where a block keeps its MSA pair.
+        last = block_index + 1 == depth
+        hidden = residual_layer_norm(
+            x,
+            mlp_out,
+            gate=mod[_GATE_MLP],
+            weight=modulation[block_index + 1][_SCALE_MSA],
+            bias=modulation[block_index + 1][_SHIFT_MSA],
+            eps=estimator.final_layer.norm_final.eps if last else blocks[block_index + 1].norm1.eps,
+            residual_out=None if last else x,
+        )
+
+        if width not in cache_index:
+            cache_index[width] = causal_cache_index(lengths, width)
+        cnn_out = cnn_cache_buffer[block_index]
+        cnn_out[:, :split].copy_(gather_causal_cache(history, cache_index[width]))
+        cnn_out[:, split:].copy_(gather_causal_cache(history2, cache_index[width]))
+
+    return estimator.final_layer.linear(hidden).transpose(1, 2)

vllm_omni/model_executor/models/minicpmo_4_5/flow_encoder_graph.py

@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exact-shape CUDA graphs of the Code2Wav flow encoder's continuation chunk.
+
+The CosyVoice2 upsample-conformer encoder (embedding, ``forward_chunk``,
+``encoder_proj``) is ~320 small kernels whose launches, not their GPU time, set
+a serving step's latency; a graph replays them with one launch. One graph per
+exact ``(rows, tokens, conformer cache frames)``, replayed only for calls of
+exactly that shape, so it runs the eager kernels and changes no output. Nothing
+is padded: rounding 3 rows up to a 4-row graph changes the M of every GEMM, and
+with it cuBLAS's kernel and FP32 rounding, which the TF32 CFM and HiFT amplified
+to 0.81 dB. The streaming trim caps the cache at ``prompt + 100`` frames, so one
+prompt's streams only pass through the lengths ``reachable_conformer_cache_frames``
+lists. All graphs take prefix views of one storage per static tensor and share
+one private pool; a replay's results stay valid until the next replay.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def reachable_conformer_cache_frames(*, start: int, prompt_len: int, suffix: int, growths: Iterable[int]) -> list[int]:
+    """Conformer attention-cache lengths a stream passes through from ``start``, for chunks of any mix of ``growths``.
+
+    Past ``prompt_len + suffix`` frames the streaming trim keeps that many (``_trim_streaming_cache``).
+    """
+    steady = int(prompt_len) + int(suffix)
+    steps = sorted({int(g) for g in growths if int(g) > 0})
+    seen, frontier = {int(start)}, [int(start)]
+    while frontier:
+        frames = frontier.pop()
+        for following in (min(frames + growth, steady) for growth in steps):
+            if following not in seen:
+                seen.add(following)
+                frontier.append(following)
+    return sorted(seen)
+
+
+def _precision_state() -> tuple[Any, ...]:
+    """The matmul/conv precision flags a captured graph's kernels were chosen under."""
+    return (
+        bool(torch.backends.cuda.matmul.allow_tf32),
+        bool(torch.backends.cudnn.allow_tf32),
+        torch.get_float32_matmul_precision(),
+    )
+
+
+class FlowEncoderGraphs:
+    """Exact-shape graphs of ``encode_fn(tokens, cnn_cache=, att_cache=)``, the eager continuation chunk.
+
+    Caches use the encoder layouts: ``cnn_cache`` ``(rows, channels, width)``,
+    ``att_cache`` ``(depth, rows, heads, frames, 2 * head_dim)``; a chunk of
+    ``tokens`` tokens emits ``upsample * (tokens - lookahead)`` frames.
+    ``held_tensors`` returns tensors the graphs read in place (the RelPos
+    tables): kept alive, and a replaced one disables the graphs.
+    """
+
+    def __init__(
+        self,
+        encode_fn: Callable[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
+        *,
+        rows: Iterable[int],
+        token_widths: Iterable[int],
+        lookahead: int,
+        upsample: int,
+        held_tensors: Callable[[], Sequence[torch.Tensor]] = tuple,
+    ) -> None:
+        self.encode_fn = encode_fn
+        self.rows = tuple(sorted({int(r) for r in rows if int(r) > 0}))
+        self.lookahead, self.upsample = int(lookahead), int(upsample)
+        self.token_widths = tuple(sorted({int(w) for w in token_widths if int(w) > self.lookahead}))
+        self._held_tensors = held_tensors
+        self._held: tuple[torch.Tensor, ...] = ()
+        self._precision: tuple[Any, ...] | None = None
+        # (rows, tokens, cache frames) -> (graph, static views by name)
+        self.graphs: dict[tuple[int, int, int], tuple[Any, dict[str, torch.Tensor]]] = {}
+        self.replays = 0
+        self.enabled = True
+        self._storage: dict[str, torch.Tensor] = {}
+        self._pool: Any = None
+
+    def output_frames(self, tokens: int) -> int:
+        return self.upsample * (int(tokens) - self.lookahead)
+
+    def keys_for(self, *, start: int, prompt_len: int, suffix: int) -> list[tuple[int, int, int]]:
+        """``(rows, tokens, cache frames)`` of every continuation chunk a prompt's streams can produce."""
+        growths = [self.output_frames(width) for width in self.token_widths]
+        frames = reachable_conformer_cache_frames(start=start, prompt_len=prompt_len, suffix=suffix, growths=growths)
+        return [(rows, width, length) for rows in self.rows for width in self.token_widths for length in frames]
+
+    def _body(self, views: dict[str, torch.Tensor]) -> None:
+        outputs = self.encode_fn(views["tokens"], cnn_cache=views["cnn"], att_cache=views["att"])
+        for name, value in zip(("hidden", "new_cnn", "new_att"), outputs, strict=True):
+            target = views[name]
+            if value.shape != target.shape or value.dtype != target.dtype:
+                raise RuntimeError(
+                    f"flow encoder {name} output {tuple(value.shape)}/{value.dtype} does not match "
+                    f"the captured layout {tuple(target.shape)}/{target.dtype}"
+                )
+            target.copy_(value)
+
+    def _record(self, views: dict[str, torch.Tensor]) -> Any:
+        """Warm up on a side stream, then capture ``_body`` into the shared private pool."""
+        device = views["att"].device
+        current = torch.cuda.current_stream(device)
+        side = torch.cuda.Stream(device=device)
+        side.wait_stream(current)
+        with torch.cuda.stream(side):
+            for _ in range(2):
+                self._body(views)
+        current.wait_stream(side)
+        torch.accelerator.synchronize(device)
+        if self._pool is None:
+            self._pool = torch.cuda.graph_pool_handle()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, pool=self._pool):
+            self._body(views)
+        return graph
+
+    def capture(
+        self,
+        keys: Iterable[tuple[int, int, int]],
+        *,
+        cnn_shape: tuple[int, int],
+        att_layout: tuple[int, int, int],
+        hidden_dim: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> int:
+        """Capture ``keys``, all or none; ``cnn_shape`` ``(channels, width)``, ``att_layout`` ``(depth, heads, 2D)``."""
+        if self.graphs:
+            raise RuntimeError("flow encoder graphs are already captured")
+        depth, heads, att_width = att_layout
+
+        def shapes(key: tuple[int, int, int]) -> dict[str, tuple[int, ...]]:
+            rows, width, frames = key
+            out = self.output_frames(width)
+            return {
+                "tokens": (rows, width),
+                "cnn": (rows, *cnn_shape),
+                "att": (depth, rows, heads, frames, att_width),
+                "hidden": (rows, out, hidden_dim),
+                "new_cnn": (rows, *cnn_shape),
+                "new_att": (depth, rows, heads, frames + out, att_width),
+            }
+
+        # Largest first, so the private pool grows once and later graphs reuse it.
+        keys = sorted(
+            {(int(r), int(w), int(f)) for r, w, f in keys},
+            key=lambda key: (key[0] * (key[2] + self.output_frames(key[1])), key[0] * key[1]),
+            reverse=True,
+        )
+        if not keys:
+            return 0
+        layouts = [shapes(key) for key in keys]
+        # Ordinary tensors, not inference tensors: replays write them in and outside inference mode.
+        with torch.inference_mode(False):
+            for name in layouts[0]:
+                numel = max(math.prod(layout[name]) for layout in layouts)
+                self._storage[name] = torch.zeros(numel, dtype=torch.long if name == "tokens" else dtype, device=device)
+        self._precision = _precision_state()
+        try:
+            for key, layout in zip(keys, layouts, strict=True):
+                with torch.inference_mode(False):
+                    views = {name: self._storage[name][: math.prod(s)].view(s) for name, s in layout.items()}
+                self.graphs[key] = (self._record(views), views)
+        except Exception:
+            # A partial set would replay some shapes and not others.
+            self.graphs.clear()
+            self._storage.clear()
+            self._pool = None
+            self.enabled = False
+            raise
+        self._held = tuple(self._held_tensors())
+        return len(self.graphs)
+
+    def storage_bytes(self) -> int:
+        """Bytes of the shared static inputs and results (not the private pool)."""
+        return sum(t.numel() * t.element_size() for t in self._storage.values())
+
+    def run(
+        self,
+        tokens: torch.Tensor,
+        cnn_rows: Sequence[torch.Tensor],
+        att_rows: Sequence[torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
+        """The continuation chunk of these streams (one cache row each) from its graph, or ``None`` (run eager).
+
+        Returns the shared result views: copy what must outlive the next replay.
+        """
+        rows = len(att_rows)
+        if not (self.enabled and rows and tokens.dim() == 2 and len(tokens) == len(cnn_rows) == rows):
+            return None
+        key = (rows, int(tokens.shape[1]), int(att_rows[0].shape[3]) if att_rows[0].dim() == 5 else -1)
+        if (entry := self.graphs.get(key)) is None:
+            return None
+        graph, views = entry
+        static_att, static_cnn = views["att"], views["cnn"]
+        att_shape, cnn_shape = (static_att.shape[0], 1, *static_att.shape[2:]), (1, *static_cnn.shape[1:])
+        if (tokens.dtype, tokens.device) != (views["tokens"].dtype, views["tokens"].device) or not all(
+            (att.shape, att.dtype, att.device, cnn.shape, cnn.dtype, cnn.device)
+            == (att_shape, static_att.dtype, static_att.device, cnn_shape, static_cnn.dtype, static_cnn.device)
+            for att, cnn in zip(att_rows, cnn_rows, strict=True)
+        ):
+            return None
+        if _precision_state() != self._precision or torch.cuda.is_current_stream_capturing():
+            return None
+        if [id(t) for t in self._held_tensors()] != [id(t) for t in self._held]:  # ``_held`` keeps them alive
+            self.enabled = False
+            logger.warning("Flow encoder RelPos tables were replaced after capture; encoder graphs disabled")
+            return None
+        with torch.no_grad():  # a replay is not differentiable; serving runs in inference mode anyway
+            torch.cat(list(cnn_rows), dim=0, out=static_cnn)
+            torch.cat(list(att_rows), dim=1, out=static_att)
+            views["tokens"].copy_(tokens)
+            graph.replay()
+        self.replays += 1
+        return views["hidden"], views["new_cnn"], views["new_att"]

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py

@@ -7,10 +7,11 @@ from __future__ import annotations
 import json
 import os
 import tempfile
+import time
 from collections import OrderedDict
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from dataclasses import dataclass
-from functools import lru_cache
+from functools import lru_cache, partial
 from hashlib import sha256
 from pathlib import Path
 from typing import Any
@@ -27,8 +28,10 @@ from vllm_omni.utils.device_copy import to_device_nonblocking
 from .batched_token2wav import (
     BatchedToken2Wav,
     BatchedToken2WavState,
+    row_offset_signature,
     state_shape_signature,
 )
+from .cuda_graph_wrapper import ResidentAttCache, _format_memory_delta, _memory_snapshot
 
 logger = init_logger(__name__)
 
@@ -47,6 +50,29 @@ def _resolve_model_dir(model_ref: str, revision: str | None = None) -> str:
     return hf_api().snapshot_download(model_ref, revision=revision, allow_patterns=["assets/*"])
 
 
+def _tf32_mode(extra: Mapping[str, Any]) -> str:
+    """Ordinary TF32 for CFM DiT GEMMs (``"off"`` / ``"tf32"``).
+
+    Enables ``allow_tf32`` on dense QKV/MLP and Triton ``input_precision="tf32"``
+    dots. Not compensated TF32x3. ``token2wav_allow_tf32`` is the shipped YAML
+    key; ``code2wav_allow_tf32`` and legacy env ``tf32x3`` / ``3xtf32`` alias it.
+    Defaults to ``"tf32"`` unless explicitly disabled.
+    """
+    raw = os.environ.get("MINICPMO_CODE2WAV_TF32")
+    value = raw if raw not in (None, "") else extra.get("token2wav_allow_tf32", extra.get("code2wav_allow_tf32", True))
+    if isinstance(value, str):
+        return {
+            "tf32": "tf32",
+            "tf32x3": "tf32",
+            "3xtf32": "tf32",
+            "flow": "tf32",
+            "1": "tf32",
+            "true": "tf32",
+            "yes": "tf32",
+        }.get(value.strip().lower(), "off")
+    return "tf32" if value else "off"
+
+
 def _batch_error(reason: str, **details: Any) -> RuntimeError:
     payload = {"reason": reason, **details}
     return RuntimeError(f"MiniCPMO45Code2WavBatchError {json.dumps(payload, sort_keys=True)}")
@@ -225,6 +251,12 @@ class _WorkItem:
     has_payload: bool = True
 
 
+def _slot_resident(item: _WorkItem) -> bool:
+    if item.previous is None:
+        return False
+    return isinstance(item.previous.token2wav.flow_cache.get("estimator_att_cache"), ResidentAttCache)
+
+
 class MiniCPMO45Code2Wav(nn.Module):
     """LLM_GENERATION model with request-owned state and compatible batching."""
 
@@ -265,12 +297,22 @@ class MiniCPMO45Code2Wav(nn.Module):
             raise ValueError("MiniCPM-o Code2Wav setup cache capacity must be >= 0")
         self._connector_config = {
             "codec_chunk_frames": int(extra.get("codec_chunk_frames", 25)),
+            "initial_codec_chunk_frames": int(extra.get("initial_codec_chunk_frames", 0)),
             "codec_left_context_frames": int(extra.get("codec_left_context_frames", 3)),
+            "hift_max_lazy_graphs": int(extra.get("hift_max_lazy_graphs", 8)),
+            "hift_graph_codec_chunk_frames": list(extra.get("hift_graph_codec_chunk_frames") or ()),
+            "hift_graph_first_chunk_frames": list(extra.get("hift_graph_first_chunk_frames") or ()),
+            "hift_graph_continuation_frames": list(extra.get("hift_graph_continuation_frames") or ()),
+            "hift_graph_exact_batch_sizes": list(extra.get("hift_graph_exact_batch_sizes") or (1,)),
         }
         if self._connector_config["codec_chunk_frames"] <= 0 or self._connector_config["codec_left_context_frames"] < 0:
             raise ValueError(f"Invalid MiniCPM-o connector chunk configuration: {self._connector_config}")
+        max_num_seqs = getattr(getattr(vllm_config, "scheduler_config", None), "max_num_seqs", None)
         raw_capture_batch_sizes = extra.get("hift_graph_capture_batch_sizes")
-        capture_batch_sizes = [1] if raw_capture_batch_sizes is None else raw_capture_batch_sizes
+        # Default: powers of two below the scheduler's batch (capped at 32) and the cap itself.
+        batch_cap = min(int(max_num_seqs), 32) if max_num_seqs else 32
+        pow2 = [1 << i for i in range(batch_cap.bit_length()) if 1 << i < batch_cap] + [batch_cap]
+        capture_batch_sizes = pow2 if raw_capture_batch_sizes is None else raw_capture_batch_sizes
         max_serial_batch = extra.get("max_serial_batch")
         max_serial_batch = 4 if max_serial_batch is None else int(max_serial_batch)
         self._hift_graph_config = {
@@ -294,11 +336,26 @@ class MiniCPMO45Code2Wav(nn.Module):
             "max_graphs": int(extra.get("cfm_max_graphs", 32)),
             "bucket_frames": int(extra.get("cfm_graph_bucket_frames", 0)),
             "capture_frames": extra.get("cfm_graph_capture_frames"),
+            "offset_bucket_frames": int(extra.get("cfm_graph_offset_bucket_frames", 50)),
             "enable_whole_euler": enable_whole_euler is None or bool(enable_whole_euler),
             "max_serial_batch": max_serial_batch,
             "max_graph_batch": max_graph_batch,
             "micro_batch_size": micro_batch_size,
             "pad_max_rows": extra.get("whole_euler_pad_max_rows"),
+            "fused_body": bool(extra.get("cfm_fused_body", False)),
+            "slot_pool": bool(extra.get("cfm_slot_pool", False)),
+            "row_offset_merge": extra.get("cfm_row_offset_merge", False) is True,
+        }
+        # Exact-shape flow-encoder graphs (``FlowEncoderGraphs``, default off).
+        # ``cfm_encoder_graph_rows`` is the largest row count or a list of counts.
+        rows = extra.get("cfm_encoder_graph_rows", 8)
+        rows = sorted({int(r) for r in rows}) if isinstance(rows, (list, tuple)) else list(range(1, int(rows) + 1))
+        if any(r < 1 for r in rows):
+            raise ValueError("MiniCPM-o cfm_encoder_graph_rows must be positive")
+        self._encoder_graph_config = {
+            "enabled": extra.get("cfm_encoder_cuda_graph", False) is True,
+            "rows": [r for r in rows if not max_num_seqs or r <= int(max_num_seqs)],
+            "token_widths": [int(w) for w in extra.get("cfm_encoder_graph_token_widths") or ()],
         }
         self._ref_max_seconds = float(extra.get("ref_audio_max_seconds", _REF_MAX_SECONDS))
         if self._ref_max_seconds <= 0:
@@ -314,6 +371,10 @@ class MiniCPMO45Code2Wav(nn.Module):
         self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio"))
         self._prompt_wav_override = extra.get("prompt_wav")
         self._default_prompt_normalized: tuple[str, str] | None = None
+        self._onset_merge = bool(extra.get("cfm_onset_merge", False))
+        self._cross_turn_buckets = extra.get("cfm_cross_turn_buckets", False) is True
+        self._row_offset_merge = self._cfm_graph_config["row_offset_merge"]
+        self._precapture_empty_cache = extra.get("cfm_precapture_empty_cache", False) is True
 
     @property
     def _default_prompt_wav(self) -> str:
@@ -658,12 +719,17 @@ class MiniCPMO45Code2Wav(nn.Module):
         )
 
     @staticmethod
-    def _bucket_key(item: _WorkItem) -> tuple[Any, ...]:
+    def _bucket_key(item: _WorkItem, *, cross_turn: bool = False, row_offsets: bool = False) -> tuple[Any, ...]:
         cache_signature: Any
         if item.previous is None:
             cache_signature = ("uninitialized",)
+        elif row_offsets and _slot_resident(item):
+            # A slot-pool row keeps its cache lengths per row (``cfm_row_offset_merge``).
+            cache_signature = ("row_offsets", row_offset_signature(item.previous.token2wav))
         else:
             cache_signature = state_shape_signature(item.previous.token2wav)
+        if cross_turn:
+            return (item.prompt_cache_id, item.prompt_wav, cache_signature)
         return (
             item.prompt_cache_id,
             item.prompt_wav,
@@ -671,6 +737,52 @@ class MiniCPMO45Code2Wav(nn.Module):
             item.cache_epoch,
         )
 
+    def _signature_groups(
+        self, bucket: list[_WorkItem], states: list[BatchedToken2WavState], *, row_offsets: bool
+    ) -> list[list[int]]:
+        """A bucket's rows grouped by cache signature, for a bucket that cannot decode as one.
+
+        With ``row_offsets``, slot-pool groups differing only in cache length stay one if the backend can solve it.
+        """
+        groups: dict[tuple[Any, ...], list[int]] = {}
+        for row, state in enumerate(states):
+            groups.setdefault(state_shape_signature(state), []).append(row)
+        if not row_offsets:
+            return list(groups.values())
+        clusters: dict[tuple[Any, ...], list[list[int]]] = {}
+        for signature, rows in groups.items():
+            if _slot_resident(bucket[rows[0]]):
+                signature = ("row_offsets", row_offset_signature(states[rows[0]]))
+            clusters.setdefault(signature, []).append(rows)
+        result: list[list[int]] = []
+        for parts in clusters.values():
+            rows = sorted(row for part in parts for row in part)
+            if len(parts) > 1 and self.backend.can_merge_row_offsets(
+                [states[row] for row in rows], [int(bucket[row].tokens.numel()) for row in rows]
+            ):
+                result.append(rows)
+            else:
+                result.extend(parts)
+        return result
+
+    def _decode_rows(
+        self, items: list[_WorkItem], states: list[BatchedToken2WavState], features: Any, *, ragged: bool
+    ) -> tuple[list[Any], list[Any]]:
+        """Decode ``items`` together: ragged when asked to, or when their token counts or last-chunk flags differ."""
+        tokens = [item.tokens for item in items]
+        last_chunks = [item.last_chunk for item in items]
+        if ragged or len({int(t.numel()) for t in tokens}) > 1 or len(set(last_chunks)) > 1:
+            return self.backend.decode_ragged_batch(tokens, features, states, last_chunks=last_chunks)
+        return self.backend.decode_batch(torch.stack(tokens, dim=0), features, states, last_chunk=last_chunks[0])
+
+    @staticmethod
+    def _onset_group_key(key: tuple[Any, ...], *, cross_turn: bool = False) -> tuple[Any, ...]:
+        """The onset-merge group of a bucket key: its prompt, plus its cache epoch unless rows merge across turns."""
+        if cross_turn:
+            return key[:2]
+        prompt_cache_id, prompt_wav, _, cache_epoch = key
+        return (prompt_cache_id, prompt_wav, cache_epoch)
+
     def _iter_limited_batches(
         self,
         buckets: Iterable[list[_WorkItem]],
@@ -738,7 +850,8 @@ class MiniCPMO45Code2Wav(nn.Module):
         # policy after eager execution/capture; cuDNN's policy is independent.
         previous_tf32 = torch.backends.cuda.matmul.allow_tf32
         try:
-            if self._extra_config().get("token2wav_allow_tf32", False):
+            extra = self._extra_config()
+            if _tf32_mode(extra) != "off":
                 torch.backends.cuda.matmul.allow_tf32 = True
             return self._forward_impl(
                 input_ids,
@@ -762,6 +875,12 @@ class MiniCPMO45Code2Wav(nn.Module):
         **kwargs: Any,
     ) -> OmniOutput:
         del positions, intermediate_tensors, inputs_embeds
+        if getattr(self, "_precapture_pending", False) and self.backend is not None:
+            try:
+                self._precapture_default_prompt()
+            finally:
+                if getattr(self, "_precapture_empty_cache", False):
+                    self._release_precapture_cache()
         ids = input_ids if isinstance(input_ids, torch.Tensor) else torch.empty(0, dtype=torch.long)
         segments = self._split_segments(ids, kwargs.get("seq_token_counts"))
         empty = torch.empty(0, dtype=torch.float32, device=ids.device)
@@ -845,9 +964,11 @@ class MiniCPMO45Code2Wav(nn.Module):
             self._trim_runtime_prompts()
             raise _batch_error("empty_nonfinal_chunk", request_ids=invalid_empty)
 
+        cross_turn = getattr(self, "_cross_turn_buckets", False) is True
+        row_offsets = getattr(self, "_row_offset_merge", False) is True
         buckets: dict[tuple[Any, ...], list[_WorkItem]] = {}
         for item in compute_items:
-            buckets.setdefault(self._bucket_key(item), []).append(item)
+            buckets.setdefault(self._bucket_key(item, cross_turn=cross_turn, row_offsets=row_offsets), []).append(item)
         undersized = [
             {
                 "size": len(bucket),
@@ -921,34 +1042,57 @@ class MiniCPMO45Code2Wav(nn.Module):
                     prompt_wav=item.prompt_wav,
                     token2wav=state,
                 )
-        for bucket in self._iter_decode_batches(buckets.values()):
+        decode_buckets: Iterable[list[_WorkItem]]
+        if self._onset_merge and self.backend is not None and self.backend._whole_euler_ragged_active():
+            # Offer a fresh onset and a continuation of one prompt/epoch, ready together, to ragged
+            # Whole-Euler; other exact-shape buckets stay as they are.
+            grouped: dict[tuple[Any, ...], list[list[_WorkItem]]] = {}
+            for key, bucket in buckets.items():
+                grouped.setdefault(self._onset_group_key(key, cross_turn=cross_turn), []).append(bucket)
+            merged: list[list[_WorkItem]] = []
+            for entries in grouped.values():
+                combined = [item for bucket in entries for item in bucket]
+                mixed = len({item.previous is None for item in combined}) > 1
+                merged.extend([combined] if mixed else entries)
+            decode_buckets = merged
+        else:
+            decode_buckets = buckets.values()
+        for bucket in self._iter_decode_batches(decode_buckets):
             batch_size = len(bucket)
             try:
                 features = self.backend.prepare_prompt(
                     bucket[0].prompt_cache_id,
                     bucket[0].prompt_wav,
                 )
-                if bucket[0].previous is None:
-                    states = self.backend.setup_batch(features, batch_size)
-                else:
-                    states = [item.previous.token2wav for item in bucket if item.previous is not None]
-                token_lengths = {int(item.tokens.numel()) for item in bucket}
-                last_chunk_values = {item.last_chunk for item in bucket}
-                if len(token_lengths) > 1 or len(last_chunk_values) > 1:
-                    audios, next_states = self.backend.decode_ragged_batch(
-                        [item.tokens for item in bucket],
-                        features,
-                        states,
-                        last_chunks=[item.last_chunk for item in bucket],
+                fresh = [item for item in bucket if item.previous is None]
+                setup = iter(self.backend.setup_batch(features, len(fresh)) if fresh else [])
+                states = [next(setup) if item.previous is None else item.previous.token2wav for item in bucket]
+                onset_mix = self._onset_merge and 0 < len(fresh) < batch_size
+                onset_ready = onset_mix and self.backend.can_merge_state_shapes(
+                    states, [int(item.tokens.numel()) for item in bucket]
                 )
+                row_groups: list[list[int]] | None = None
+                offset_mix = False
+                if onset_mix and not onset_ready:
+                    # Cache layouts ragged Whole-Euler cannot take split back into signature buckets.
+                    row_groups = self._signature_groups(bucket, states, row_offsets=row_offsets)
+                elif row_offsets and not onset_ready and batch_size > 1:
+                    groups = self._signature_groups(bucket, states, row_offsets=True)
+                    if len(groups) > 1:
+                        row_groups = groups
                     else:
-                    tokens = torch.stack([item.tokens for item in bucket], dim=0)
-                    audios, next_states = self.backend.decode_batch(
-                        tokens,
-                        features,
-                        states,
-                        last_chunk=bucket[0].last_chunk,
-                    )
+                        # One decode, ragged when it holds rows of different cache lengths.
+                        offset_mix = len({state_shape_signature(state) for state in states}) > 1
+                if row_groups is None:
+                    audios, next_states = self._decode_rows(bucket, states, features, ragged=onset_ready or offset_mix)
+                else:
+                    audios, next_states = [None] * batch_size, [None] * batch_size
+                    for rows in row_groups:
+                        group_states = [states[row] for row in rows]
+                        mixed = row_offsets and len({state_shape_signature(state) for state in group_states}) > 1
+                        decoded = self._decode_rows([bucket[row] for row in rows], group_states, features, ragged=mixed)
+                        for row, audio, next_state in zip(rows, *decoded, strict=True):
+                            audios[row], next_states[row] = audio, next_state
             except Exception as exc:
                 self._trim_runtime_prompts()
                 if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "):
@@ -1055,6 +1199,12 @@ class MiniCPMO45Code2Wav(nn.Module):
         token2wav_path = Path(self.model_path) / "assets" / "token2wav"
         if not token2wav_path.is_dir():
             raise FileNotFoundError(f"MiniCPM-o Code2Wav assets not found: {token2wav_path}")
+        # Token2wav runs in fp32, so without TF32 every flow-DiT GEMM runs on
+        # SIMT cores. Ordinary TF32 uses SM80+ Tensor Cores for dense GEMMs and
+        # Triton CFM attention dots. HiFT stays IEEE FP32.
+        tf32_mode = _tf32_mode(extra) if current_omni_platform.is_cuda() else "off"
+        if tf32_mode != "off":
+            logger.info("MiniCPM-o Code2Wav: ordinary TF32 enabled for CFM DiT GEMMs")
         use_float16 = bool(extra.get("token2wav_float16", False))
         previous_dtype = torch.get_default_dtype()
         try:
@@ -1066,6 +1216,9 @@ class MiniCPMO45Code2Wav(nn.Module):
                 str(token2wav_path),
                 float16=use_float16,
                 n_timesteps=int(extra.get("token2wav_n_timesteps", 10)),
+                # The batched backend builds its estimator caches itself and
+                # never runs the upstream chunk path that owns these buffers.
+                drop_upstream_chunk_att_buffers=extra.get("cfm_drop_upstream_att_buffer", False) is True,
             )
         finally:
             torch.set_default_dtype(previous_dtype)
@@ -1099,4 +1252,53 @@ class MiniCPMO45Code2Wav(nn.Module):
             cfm_graph_config=self._cfm_graph_config,
             bfloat16_attention_cache=bool(extra.get("code2wav_bfloat16_attention_cache", False)),
             setup_cache_size=self._setup_cache_size,
+            cfm_tf32=tf32_mode != "off",
+            encoder_graph_config=self._encoder_graph_config,
         )
+        # Captured by the first forward (the warmup run): under vLLM's weight load a graph held GiBs.
+        self._precapture_pending = bool(extra.get("cfm_graph_precapture", True))
+
+    def _release_precapture_cache(self) -> None:
+        """Return startup's unused cached allocator blocks to the driver (``cfm_precapture_empty_cache``).
+
+        Live tensors, the slot pool and graph pools stay, so no later solve changes.
+        """
+        if torch.cuda.is_available() and torch.cuda.is_initialized():
+            device = torch.device("cuda", torch.accelerator.current_device_index())
+            torch.accelerator.synchronize(device)
+            before = _memory_snapshot(device)
+            torch.accelerator.empty_cache()
+            delta = _format_memory_delta(before, _memory_snapshot(device))
+            logger.info("MiniCPM-o Code2Wav: emptied the allocator cache after precapture%s", delta)
+
+    def _precapture_default_prompt(self) -> None:
+        """Capture the HiFT, default-voice Whole-Euler and flow-encoder graphs before serving.
+
+        Otherwise each graph is captured by the first chunk that needs it,
+        stalling every stream for 1-2 s per graph.
+        """
+        self._precapture_pending = False
+        backend = self.backend
+        assert backend is not None
+        started = time.perf_counter()
+
+        def precapture(name: str, capture: Callable[[], int]) -> bool:
+            try:
+                captured = capture()
+            except Exception:
+                logger.warning("MiniCPM-o Code2Wav: %s precapture failed; graphs capture lazily", name, exc_info=True)
+                return False
+            if captured:
+                elapsed = time.perf_counter() - started
+                logger.info("MiniCPM-o Code2Wav: precaptured %d %s CUDA Graphs in %.1f s", captured, name, elapsed)
+            return True
+
+        precapture("HiFT", backend.precapture_hift)
+        prompt_wav, prompt_cache_id = self._normalized_default_prompt()
+        try:
+            features = backend.prepare_prompt(prompt_cache_id, prompt_wav)
+        except Exception:
+            logger.warning("MiniCPM-o Code2Wav: default prompt failed; graphs capture lazily", exc_info=True)
+            return
+        if precapture("Whole-Euler", partial(backend.precapture_whole_euler, features)):
+            precapture("flow encoder", partial(backend.precapture_flow_encoder, features))

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """MiniCPM-o 4.5 Token2wav adapter over in-tree ``StepAudio2Token2WavCore``.
 
 ``minicpmo_4_5_omni_tts`` historically depended on the external
@@ -64,6 +64,7 @@ class MiniCPMO45Token2wav:
         float16: bool = False,
         n_timesteps: int = 10,
         device: str | torch.device | None = None,
+        drop_upstream_chunk_att_buffers: bool = False,
     ):
         self.float16 = float16
         self.n_timesteps = n_timesteps
@@ -73,6 +74,7 @@ class MiniCPMO45Token2wav:
             float16=float16,
             device=self.device,
             n_timesteps=n_timesteps,
+            drop_upstream_chunk_att_buffers=drop_upstream_chunk_att_buffers,
         )
         # Eager-load so construction failures surface at init time (same as
         # the external Token2wav package), not on the first request.

vllm_omni/model_executor/models/minicpmo_4_5/ops/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from .fused_qkv_layer_norm import qkv_head_layer_norm
+from .fused_residual_layer_norm import residual_layer_norm
+
+__all__ = [
+    "qkv_head_layer_norm",
+    "residual_layer_norm",
+]

vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_qkv_layer_norm.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+# ruff: noqa: N803
+
+"""Fused packed-QKV split + per-head LayerNorm of q/k into the CFM KV cache."""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from .fused_residual_layer_norm import _use_triton
+
+
+@triton.jit
+def _head_layer_norm(x, mask, dims, weight_ptr, bias_ptr, eps, HEAD_DIM: tl.constexpr, HAS_AFFINE: tl.constexpr):
+    """LayerNorm over the last axis of ``(heads, dims)`` in fp32, with an optional per-dim affine."""
+    x = x.to(tl.float32)
+    mean = tl.sum(x, axis=1) / HEAD_DIM
+    centered = tl.where(mask, x - mean[:, None], 0.0)
+    variance = tl.sum(centered * centered, axis=1) / HEAD_DIM
+    x = centered * tl.rsqrt(variance + eps)[:, None]
+    if HAS_AFFINE:
+        x = x * tl.load(weight_ptr + dims, mask=dims < HEAD_DIM, other=1.0).to(tl.float32)[None, :]
+        x = x + tl.load(bias_ptr + dims, mask=dims < HEAD_DIM, other=0.0).to(tl.float32)[None, :]
+    return x
+
+
+@triton.jit
+def _qkv_head_layer_norm_kernel(
+    q_out_ptr,
+    kv_ptr,
+    qkv_ptr,
+    q_weight_ptr,
+    q_bias_ptr,
+    k_weight_ptr,
+    k_bias_ptr,
+    rows_ptr,
+    positions_ptr,
+    frames,
+    stride_qkv_n,
+    stride_qkv_t,
+    stride_qn,
+    stride_qh,
+    stride_qt,
+    stride_kvn,
+    stride_kvh,
+    stride_kvt,
+    q_eps,
+    k_eps,
+    HEADS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    HAS_Q_AFFINE: tl.constexpr,
+    HAS_K_AFFINE: tl.constexpr,
+    HAS_ROWS: tl.constexpr,
+    HAS_POSITIONS: tl.constexpr,
+):
+    row = tl.program_id(0)
+    n = row // frames
+    t = row % frames
+    heads = tl.arange(0, BLOCK_H)
+    dims = tl.arange(0, BLOCK_D)
+    mask = (heads[:, None] < HEADS) & (dims[None, :] < HEAD_DIM)
+    inner = HEADS * HEAD_DIM
+    src = qkv_ptr + n * stride_qkv_n + t * stride_qkv_t + heads[:, None] * HEAD_DIM + dims[None, :]
+
+    q = tl.load(src, mask=mask, other=0.0)
+    q = _head_layer_norm(q, mask, dims, q_weight_ptr, q_bias_ptr, q_eps, HEAD_DIM, HAS_Q_AFFINE)
+    tl.store(q_out_ptr + n * stride_qn + heads[:, None] * stride_qh + t * stride_qt + dims[None, :], q, mask=mask)
+
+    k = tl.load(src + inner, mask=mask, other=0.0)
+    k = _head_layer_norm(k, mask, dims, k_weight_ptr, k_bias_ptr, k_eps, HEAD_DIM, HAS_K_AFFINE)
+    kv_row = n
+    if HAS_ROWS:
+        kv_row = tl.load(rows_ptr + n).to(tl.int64)
+    frame = t
+    if HAS_POSITIONS:
+        frame = tl.load(positions_ptr + row).to(tl.int64)
+        mask = mask & (frame >= 0)
+    dst = kv_ptr + kv_row * stride_kvn + heads[:, None] * stride_kvh + frame * stride_kvt + dims[None, :]
+    tl.store(dst, k, mask=mask)
+
+    v = tl.load(src + 2 * inner, mask=mask, other=0.0)
+    tl.store(dst + HEAD_DIM, v, mask=mask)
+
+
+def qkv_head_layer_norm(
+    qkv: torch.Tensor,
+    kv: torch.Tensor,
+    *,
+    num_heads: int,
+    head_dim: int,
+    q_norm: torch.nn.LayerNorm,
+    k_norm: torch.nn.LayerNorm,
+    rows: torch.Tensor | None = None,
+    positions: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Split packed ``qkv``, LayerNorm q/k per head, write q and interleaved ``kv``."""
+    batch, frames, width = (int(dim) for dim in qkv.shape)
+    inner = num_heads * head_dim
+    if width != 3 * inner or qkv.stride(2) != 1:
+        raise ValueError(f"qkv_head_layer_norm: qkv must be (N, T, 3*{inner}) with contiguous channels")
+    if (
+        kv.dim() != 4
+        or (rows is None and int(kv.shape[0]) != batch)
+        or int(kv.shape[1]) != num_heads
+        or (positions is None and int(kv.shape[2]) < frames)
+        or int(kv.shape[3]) != 2 * head_dim
+        or kv.stride(3) != 1
+    ):
+        raise ValueError(f"qkv_head_layer_norm: kv must be (N, H, >= T, 2*D), got {tuple(kv.shape)}")
+    if rows is not None and tuple(rows.shape) != (batch,):
+        raise ValueError(f"qkv_head_layer_norm: rows must be ({batch},), got {tuple(rows.shape)}")
+    if positions is not None and (tuple(positions.shape) != (batch, frames) or not positions.is_contiguous()):
+        raise ValueError(f"qkv_head_layer_norm: positions must be contiguous ({batch}, {frames})")
+    q_out = torch.empty((batch, num_heads, frames, head_dim), device=qkv.device, dtype=qkv.dtype)
+
+    if not _use_triton(qkv):
+        heads = qkv.view(batch, frames, 3, num_heads, head_dim)
+        q_out.copy_(q_norm(heads[:, :, 0].transpose(1, 2)))
+        keys = k_norm(heads[:, :, 1].transpose(1, 2))
+        values = heads[:, :, 2].transpose(1, 2)
+        if rows is None and positions is None:
+            kv[:, :, :frames, :head_dim].copy_(keys)
+            kv[:, :, :frames, head_dim:].copy_(values)
+            return q_out
+        pool_rows = rows.tolist() if rows is not None else range(batch)
+        for n, pool_row in enumerate(pool_rows):
+            frame_index = positions[n] if positions is not None else torch.arange(frames, device=kv.device)
+            written = frame_index >= 0
+            destination = frame_index[written].to(kv.device)
+            kv[pool_row, :, destination, :head_dim] = keys[n][:, written]
+            kv[pool_row, :, destination, head_dim:] = values[n][:, written]
+        return q_out
+    if batch * frames == 0:
+        return q_out
+    q_weight, q_bias, q_eps = q_norm.weight, q_norm.bias, float(q_norm.eps)
+    k_weight, k_bias, k_eps = k_norm.weight, k_norm.bias, float(k_norm.eps)
+    if (q_weight is None) != (q_bias is None) or (k_weight is None) != (k_bias is None):
+        raise ValueError("qkv_head_layer_norm: LayerNorm affine needs both weight and bias")
+    _qkv_head_layer_norm_kernel[(batch * frames,)](
+        q_out,
+        kv,
+        qkv,
+        q_weight if q_weight is not None else qkv,
+        q_bias if q_bias is not None else qkv,
+        k_weight if k_weight is not None else qkv,
+        k_bias if k_bias is not None else qkv,
+        rows if rows is not None else qkv,
+        positions if positions is not None else qkv,
+        frames,
+        qkv.stride(0),
+        qkv.stride(1),
+        q_out.stride(0),
+        q_out.stride(1),
+        q_out.stride(2),
+        kv.stride(0),
+        kv.stride(1),
+        kv.stride(2),
+        q_eps,
+        k_eps,
+        HEADS=num_heads,
+        HEAD_DIM=head_dim,
+        BLOCK_H=triton.next_power_of_2(num_heads),
+        BLOCK_D=triton.next_power_of_2(head_dim),
+        HAS_Q_AFFINE=q_weight is not None,
+        HAS_K_AFFINE=k_weight is not None,
+        HAS_ROWS=rows is not None,
+        HAS_POSITIONS=positions is not None,
+        num_warps=4,
+    )
+    return q_out

vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_residual_layer_norm.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+# ruff: noqa: N803
+
+"""Fused ``act(layer_norm(residual + gate * y) * weight + bias)`` (optional conv taps)."""
+
+import torch
+import torch.nn.functional as F
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+_ACTIVATIONS = {None: 0, "mish": 1}
+
+
+def _use_triton(tensor: torch.Tensor) -> bool:
+    """Triton on CUDA tensors; everything else runs the native path."""
+    return HAS_TRITON and tensor.is_cuda
+
+
+@triton.jit
+def _residual_layer_norm_kernel(
+    out_ptr,
+    residual_out_ptr,
+    x_ptr,
+    y_ptr,
+    gate_ptr,
+    y_bias_ptr,
+    weight_ptr,
+    bias_ptr,
+    frames,
+    channels,
+    stride_on,
+    stride_out_t,
+    stride_rn,
+    stride_rt,
+    stride_xn,
+    stride_xt,
+    stride_yn,
+    stride_yt,
+    eps,
+    HAS_X: tl.constexpr,
+    HAS_Y: tl.constexpr,
+    TAPS: tl.constexpr,
+    HAS_GATE: tl.constexpr,
+    HAS_Y_BIAS: tl.constexpr,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    STORE_RESIDUAL: tl.constexpr,
+    ACTIVATION: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+):
+    row = tl.program_id(0)
+    n = row // frames
+    t = row % frames
+    cols = tl.arange(0, BLOCK_C)
+    mask = cols < channels
+
+    value = tl.zeros((BLOCK_C,), dtype=tl.float32)
+    if HAS_Y:
+        y = tl.zeros((BLOCK_C,), dtype=tl.float32)
+        for k in tl.static_range(TAPS):
+            y += tl.load(
+                y_ptr + n * stride_yn + (t + k) * stride_yt + k * channels + cols,
+                mask=mask,
+                other=0.0,
+            ).to(tl.float32)
+        if HAS_Y_BIAS:
+            y += tl.load(y_bias_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+        if HAS_GATE:
+            y = tl.load(gate_ptr + cols, mask=mask, other=0.0).to(tl.float32) * y
+        value = y
+    if HAS_X:
+        value += tl.load(x_ptr + n * stride_xn + t * stride_xt + cols, mask=mask, other=0.0).to(tl.float32)
+    if STORE_RESIDUAL:
+        tl.store(residual_out_ptr + n * stride_rn + t * stride_rt + cols, value, mask=mask)
+
+    mean = tl.sum(value, axis=0) / channels
+    centered = tl.where(mask, value - mean, 0.0)
+    variance = tl.sum(centered * centered, axis=0) / channels
+    normalized = centered * tl.rsqrt(variance + eps)
+    if HAS_WEIGHT:
+        normalized = normalized * tl.load(weight_ptr + cols, mask=mask, other=1.0).to(tl.float32)
+    if HAS_BIAS:
+        normalized = normalized + tl.load(bias_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    if ACTIVATION == 1:
+        e = tl.exp(tl.minimum(normalized, 20.0))
+        numerator = e * (e + 2.0)
+        factor = tl.where(normalized > 20.0, 1.0, numerator / (numerator + 2.0))
+        normalized = normalized * factor
+    tl.store(out_ptr + n * stride_on + t * stride_out_t + cols, normalized, mask=mask)
+
+
+def _check_rows(name: str, tensor: torch.Tensor, rows: int, frames: int, width: int) -> None:
+    if tensor.dim() != 3 or int(tensor.shape[0]) != rows or int(tensor.shape[1]) < frames:
+        raise ValueError(f"residual_layer_norm: {name} must be (N, >= T, C), got {tuple(tensor.shape)}")
+    if int(tensor.shape[2]) != width or (width > 1 and tensor.stride(2) != 1):
+        raise ValueError(f"residual_layer_norm: {name} needs {width} contiguous channels, got {tuple(tensor.shape)}")
+
+
+def _native(residual, y, gate, weight, bias, eps, taps, y_bias, activation, residual_out, out, frames, channels):
+    value = None
+    if y is not None:
+        value = sum(y[:, k : k + frames, k * channels : (k + 1) * channels] for k in range(taps))
+        if y_bias is not None:
+            value = value + y_bias
+        if gate is not None:
+            value = gate * value
+    if residual is not None:
+        value = residual[:, :frames] if value is None else residual[:, :frames] + value
+    if residual_out is not None:
+        residual_out.copy_(value)
+    result = F.layer_norm(value, (channels,), None, None, eps)
+    if weight is not None:
+        result = result * weight
+    if bias is not None:
+        result = result + bias
+    if activation == "mish":
+        result = F.mish(result)
+    out.copy_(result)
+    return out
+
+
+def residual_layer_norm(
+    residual: torch.Tensor | None,
+    y: torch.Tensor | None = None,
+    *,
+    gate: torch.Tensor | None = None,
+    weight: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    taps: int = 1,
+    y_bias: torch.Tensor | None = None,
+    activation: str | None = None,
+    residual_out: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """``act(layer_norm(residual + gate * y) * weight + bias)`` in one pass per row."""
+    if activation not in _ACTIVATIONS:
+        raise ValueError(f"residual_layer_norm: unsupported activation {activation!r}")
+    if residual is None and y is None:
+        raise ValueError("residual_layer_norm needs a residual or a y")
+    reference = residual if residual is not None else y
+    rows = int(reference.shape[0])
+    if residual is not None:
+        frames, channels = int(residual.shape[1]), int(residual.shape[2])
+        _check_rows("residual", residual, rows, frames, channels)
+    else:
+        channels = int(y.shape[2]) // taps
+        frames = int(y.shape[1]) - (taps - 1)
+    if y is not None:
+        _check_rows("y", y, rows, frames + taps - 1, taps * channels)
+    if residual_out is not None:
+        _check_rows("residual_out", residual_out, rows, frames, channels)
+    if out is None:
+        out = torch.empty((rows, frames, channels), device=reference.device, dtype=reference.dtype)
+    _check_rows("out", out, rows, frames, channels)
+
+    if not _use_triton(reference):
+        return _native(
+            residual, y, gate, weight, bias, eps, taps, y_bias, activation, residual_out, out, frames, channels
+        )
+    if rows * frames == 0:
+        return out
+
+    def strides(tensor: torch.Tensor | None) -> tuple[int, int]:
+        return (tensor.stride(0), tensor.stride(1)) if tensor is not None else (0, 0)
+
+    _residual_layer_norm_kernel[(rows * frames,)](
+        out,
+        *(t if t is not None else reference for t in (residual_out, residual, y, gate, y_bias, weight, bias)),
+        frames,
+        channels,
+        *strides(out),
+        *strides(residual_out),
+        *strides(residual),
+        *strides(y),
+        float(eps),
+        HAS_X=residual is not None,
+        HAS_Y=y is not None,
+        TAPS=int(taps),
+        HAS_GATE=gate is not None,
+        HAS_Y_BIAS=y_bias is not None,
+        HAS_WEIGHT=weight is not None,
+        HAS_BIAS=bias is not None,
+        STORE_RESIDUAL=residual_out is not None,
+        ACTIVATION=_ACTIVATIONS[activation],
+        BLOCK_C=triton.next_power_of_2(channels),
+        num_warps=4 if channels <= 1024 else 8,
+    )
+    return out

vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py

@@ -95,6 +95,22 @@ class _ConvRNNF0Predictor(nn.Module):
         return self.classifier(self.condnet(x).transpose(1, 2)).squeeze(-1).abs()
 
 
+def drop_upstream_chunk_att_buffers(flow: nn.Module) -> int:
+    """Drop upstream cosyvoice2's unused chunk-streaming ``att_cache_buffer`` tensors (~2.1 GiB); returns bytes freed.
+
+    Only ``flow.setup_cache`` / ``flow.inference_chunk`` read them, and then
+    fail instead of computing on a stand-in.
+    """
+    released = 0
+    decoder = getattr(flow, "decoder", None)
+    for module in (decoder, getattr(decoder, "estimator", None)):
+        buffer = module._buffers.get("att_cache_buffer") if isinstance(module, nn.Module) else None
+        if isinstance(buffer, torch.Tensor):
+            released += buffer.numel() * buffer.element_size()
+            del module._buffers["att_cache_buffer"]
+    return released
+
+
 def _build_hift() -> HiFTGenerator:
     return HiFTGenerator(
         sampling_rate=24000,
@@ -115,9 +131,13 @@ class StepAudio2Token2WavCore(nn.Module):
         float16: bool = False,
         device: str = "cuda",
         n_timesteps: int = DEFAULT_TOKEN2WAV_CONFIG.n_timesteps,
+        drop_upstream_chunk_att_buffers: bool = False,
     ):
         super().__init__()
         self.model_path = model_path
+        # A backend that never runs ``flow.setup_cache`` / ``flow.inference_chunk``
+        # (MiniCPM-o's ``BatchedToken2Wav``) can drop their 2.1 GiB of buffers.
+        self._drop_chunk_att_buffers = bool(drop_upstream_chunk_att_buffers)
         self.float16 = float16
         self.device = torch.device(device)
         self.n_timesteps = n_timesteps
@@ -176,6 +196,10 @@ class StepAudio2Token2WavCore(nn.Module):
         with open(f"{self.model_path}/flow.yaml") as f:
             configs = load_hyperpyyaml(f)
         self._flow = configs["flow"]
+        if self._drop_chunk_att_buffers:
+            # Before ``.to(self.device)``: the buffers never reach the device.
+            released = drop_upstream_chunk_att_buffers(self._flow)
+            logger.info("Dropped the upstream flow chunk attention buffers (%.1f MiB)", released / 2**20)
         if self.float16:
             self._flow.half()
         self._flow.load_state_dict(