From 03fe9c65f462a9dd3cbe87b685fa3137921124ff Mon Sep 17 00:00:00 2001
From: natureofnature <wzliu@connect.hku.hk>
Date: Wed, 21 Oct 2026 06:45:16 -0400
Subject: [PATCH 1/4] [Model] Batch MiniCPM-o 4.5 cold reference CFM setup

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
---
 docs/design/minicpm_o45_mrv2_performance.md   |  12 +-
 .../minicpmo_4_5/test_code2wav_batching.py    | 254 +++++++++++++++++-
 .../models/minicpmo_4_5/batched_token2wav.py  |  79 +++++-
 .../minicpmo_4_5/minicpmo_4_5_code2wav.py     |   7 +-
 4 files changed, 338 insertions(+), 14 deletions(-)

diff --git a/docs/design/minicpm_o45_mrv2_performance.md b/docs/design/minicpm_o45_mrv2_performance.md
index c7b19698417..3703fd40a0f 100644
--- a/docs/design/minicpm_o45_mrv2_performance.md
+++ b/docs/design/minicpm_o45_mrv2_performance.md
@@ -64,8 +64,12 @@ preserving CFM step-graph and TRT dispatch. Mixed-history arena replay is
 restricted to CUDA; the existing opt-in slot path is unchanged on other
 platforms.
 
-Reference extraction and initial cache setup remain separate for each
-reference; speaker conditioning, attention lengths and cache trimming stay
+Reference extraction, encoder initialization and speaker projection remain
+per reference. With CUDA ragged Whole-Euler and zero CFM frame padding,
+compatible cold references share an initial CFM solve, bounded by the configured
+microbatch and graph batch limits. Setup-cache hits are reused, and each
+reference retains its own cached state. Other configurations keep per-reference
+initialization. Speaker conditioning, attention lengths and cache trimming stay
 per request. Cache epochs remain separate unless cross-turn batching is enabled.
 
 The MiniCPM Code2Wav stage runs CAM++ speaker embedding on CUDA and NPU using
@@ -77,8 +81,8 @@ strict FP32 arithmetic.
 The CUDA path does not require TensorRT. Set `token2wav_device_spk: false`
 in connector `extra` to use CPU ONNX Runtime for accuracy triage; explicitly
 enabling `token2wav_trt` retains its TensorRT speaker embedding path.
-This moves speaker embedding off the CPU but does not batch reference
-extraction or initial cache setup. Device kernels can change floating-point
+This moves speaker embedding off the CPU; reference extraction stays separate
+from the initial CFM batching described above. Device kernels can change floating-point
 rounding, so this does not promise bit-identical embeddings or waveforms.
 
 The ragged Whole-Euler path masks padded history and returns each request's
diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
index 099bf26b12f..aaefa419329 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
@@ -16,12 +16,16 @@
 import vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav as batched_token2wav_module
 from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
     BatchedToken2Wav,
+    BatchedToken2WavState,
+    PromptFeatures,
     _token2wav_sdpa_context,
     plan_token2wav_encode_slices,
     relpos_encode_token_budget,
 )
 from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import (
     MiniCPMO45Code2Wav,
+    _RequestState,
+    _WorkItem,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -427,6 +431,241 @@ def test_evict_prompt_clears_features_and_setup_state():
     assert all(key[0] != first.cache_key for key in adapter._setup_cache)
 
 
+def _reference_feature(voice, length=2, dtype=torch.float32):
+    return PromptFeatures(
+        cache_key=(f"voice-{voice}", f"/fake/voice-{voice}.wav"),
+        speech_tokens=torch.arange(5 + voice, 5 + voice + length).unsqueeze(0),
+        speaker_embedding=torch.tensor([[1.0 if voice % 2 == 0 else -1.0]], dtype=dtype),
+        mels=torch.full((1, length + 2, 1), float(voice + 1), dtype=dtype),
+    )
+
+
+def _reference_setup_backend(*, cache_size=4, micro_batch_size=2, max_graph_batch=8):
+    adapter = BatchedToken2Wav(_FakeToken2Wav(), setup_cache_size=cache_size)
+    # CUDA dispatch with CPU tensors; eager CFM supplies deterministic results.
+    adapter._whole_euler_graph_wrapper = SimpleNamespace(
+        enabled=True,
+        device=torch.device("cuda"),
+        ragged_body=object(),
+        replay=lambda **kwargs: None,
+        micro_batch_size=micro_batch_size,
+        max_graph_batch=max_graph_batch,
+    )
+    return adapter
+
+
+def test_setup_references_batches_distinct_conditioning_with_single_row_encoders(monkeypatch):
+    adapter = _reference_setup_backend()
+    features = [_reference_feature(voice) for voice in range(2)]
+    solves, projections = [], []
+    decode, project = adapter._decode_cfm, adapter._project_speakers
+
+    def record_decode(mu, speakers, cond, **kwargs):
+        solves.append((mu.clone(), speakers.clone(), cond.clone()))
+        return decode(mu, speakers, cond, **kwargs)
+
+    def record_project(speakers):
+        projections.append(speakers.shape[0])
+        return project(speakers)
+
+    monkeypatch.setattr(adapter, "_decode_cfm", record_decode)
+    monkeypatch.setattr(adapter, "_project_speakers", record_project)
+    states = adapter.setup_references(features)
+
+    assert adapter.flow.encoder.calls == projections == [1, 1]
+    assert len(solves) == 1
+    mu, speakers, cond = solves[0]
+    torch.testing.assert_close(mu[:, 0, 0], torch.tensor([5.0, 6.0]))
+    torch.testing.assert_close(speakers, torch.tensor([[1.0], [-1.0]]))
+    torch.testing.assert_close(cond, torch.cat([f.mels.transpose(1, 2) for f in features]))
+    independent = BatchedToken2Wav(_FakeToken2Wav())
+    for feature, state in zip(features, states, strict=True):
+        expected = independent.setup_batch(feature, 1)[0]
+        for name, value in state.flow_cache.items():
+            torch.testing.assert_close(value, expected.flow_cache[name])
+
+
+@pytest.mark.parametrize(
+    ("micro_batch_size", "max_graph_batch", "cfg_batches"),
+    [(2, 8, [2] * 6 + [4] * 2), (8, 2, [2] * 6 + [4] * 2), (2, 0, [2] * 10), (2, -1, [2] * 10)],
+)
+def test_setup_references_groups_compatible_rows_and_restores_caller_order(
+    micro_batch_size, max_graph_batch, cfg_batches
+):
+    adapter = _reference_setup_backend(micro_batch_size=micro_batch_size, max_graph_batch=max_graph_batch)
+    features = [_reference_feature(i, length=3 if i == 1 else 2) for i in range(5)]
+    features[3] = _reference_feature(3, dtype=torch.float64)
+
+    states = adapter.setup_references(features)
+
+    assert adapter.flow.encoder.calls == [1] * 5
+    assert sorted(adapter.flow.decoder.estimator.cfg_batches) == cfg_batches
+    assert [state.flow_cache["conformer_cnn_cache"].flatten()[0].item() for state in states] == [5, 6, 7, 8, 9]
+    assert list(adapter._setup_cache) == [(feature.cache_key, 0) for feature in features[1:]]
+
+
+@pytest.mark.parametrize("cache_size", [0, 1])
+def test_setup_references_retains_results_and_hits_across_cache_eviction(cache_size):
+    adapter = _reference_setup_backend(cache_size=cache_size)
+    first, second, third = [_reference_feature(i) for i in range(3)]
+    previous = adapter.setup_batch(second, 1)[0]
+    adapter.flow.encoder.calls.clear()
+
+    states = adapter.setup_references([first, second, third])
+
+    assert adapter.flow.encoder.calls == [1] * (2 if cache_size else 3)
+    assert (states[1] is previous) == bool(cache_size)
+    assert [state.flow_cache["conformer_cnn_cache"].flatten()[0].item() for state in states] == [5, 6, 7]
+    assert list(adapter._setup_cache) == ([(third.cache_key, 0)] if cache_size else [])
+
+
+@pytest.mark.parametrize("gate", ["absent", "disabled", "cpu", "npu", "exact", "padding", "trt"])
+def test_setup_references_preserves_single_reference_fallback(monkeypatch, gate):
+    adapter = _reference_setup_backend()
+    wrapper = adapter._whole_euler_graph_wrapper
+    if gate == "absent":
+        adapter._whole_euler_graph_wrapper = None
+    elif gate == "disabled":
+        wrapper.enabled = False
+    elif gate in ("cpu", "npu"):
+        wrapper.device = SimpleNamespace(type=gate)
+    elif gate == "exact":
+        wrapper.ragged_body = None
+    elif gate == "padding":
+        adapter._cfm_graph_bucket_frames = 8
+    else:
+        adapter._trt_stepper = object()
+    features = [_reference_feature(i) for i in range(2)]
+    calls = []
+
+    def setup(feature, batch_size):
+        calls.append((feature.cache_key, batch_size))
+        return [feature.cache_key]
+
+    monkeypatch.setattr(adapter, "setup_batch", setup)
+    assert adapter.setup_references(features) == [feature.cache_key for feature in features]
+    assert calls == [(feature.cache_key, 1) for feature in features]
+
+
+def test_setup_references_rechecks_disabled_graph_between_groups():
+    adapter = _reference_setup_backend()
+    wrapper = adapter._whole_euler_graph_wrapper
+
+    def disable(**kwargs):
+        wrapper.enabled = False
+        return None
+
+    wrapper.replay = disable
+    states = adapter.setup_references([_reference_feature(i) for i in range(4)])
+
+    assert len(states) == 4
+    assert adapter.flow.decoder.estimator.cfg_batches == [4, 4, 2, 2, 2, 2]
+
+
+def test_setup_references_failed_later_group_does_not_publish_or_touch_cache(monkeypatch):
+    adapter = _reference_setup_backend(cache_size=2)
+    features = [_reference_feature(i) for i in range(5)]
+    cached = [adapter.setup_batch(feature, 1)[0] for feature in features[:2]]
+    decode = adapter._decode_cfm
+    calls = 0
+
+    def fail_second_group(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            raise RuntimeError("setup failed")
+        return decode(*args, **kwargs)
+
+    monkeypatch.setattr(adapter, "_decode_cfm", fail_second_group)
+    with pytest.raises(RuntimeError, match="setup failed"):
+        adapter.setup_references([features[1], *features[2:], features[0]])
+    assert calls == 2
+    assert list(adapter._setup_cache) == [(feature.cache_key, 0) for feature in features[:2]]
+    assert list(adapter._setup_cache.values()) == cached
+
+
+def test_setup_references_owns_each_cache_after_reused_cfm_scratch(monkeypatch):
+    adapter = _reference_setup_backend()
+    decode = adapter._decode_cfm
+    scratch: list[torch.Tensor] = []
+
+    def reuse_scratch(*args, **kwargs):
+        result = decode(*args, **kwargs)
+        if not scratch:
+            scratch.extend(value.clone() for value in result)
+        else:
+            for target, value in zip(scratch, result, strict=True):
+                target.copy_(value)
+        return tuple(scratch)
+
+    monkeypatch.setattr(adapter, "_decode_cfm", reuse_scratch)
+    features = [_reference_feature(i) for i in range(4)]
+    states = adapter.setup_references(features)
+    for value in scratch:
+        value.fill_(-100)
+    independent = BatchedToken2Wav(_FakeToken2Wav())
+    for name in states[0].flow_cache:
+        assert len({state.flow_cache[name].untyped_storage().data_ptr() for state in states}) == 4
+    for feature, state in zip(features, states, strict=True):
+        expected = independent.setup_batch(feature, 1)[0]
+        for name, value in state.flow_cache.items():
+            torch.testing.assert_close(value, expected.flow_cache[name])
+
+
+@pytest.mark.parametrize("same_reference", [False, True])
+def test_prepare_decode_batch_maps_unique_fresh_references_and_continuations(monkeypatch, same_reference):
+    model, _ = _model()
+    features = [_reference_feature(i) for i in range(2)]
+    by_key = {feature.cache_key: feature for feature in features}
+    keys = [features[0].cache_key, features[0 if same_reference else 1].cache_key, features[0].cache_key]
+    continuation = BatchedToken2WavState(flow_cache={}, hift_cache={})
+    items = [
+        _WorkItem(
+            output_index=row,
+            state_id=f"request-{row}",
+            request_id=f"request-{row}",
+            cache_epoch=0,
+            chunk_seq=int(row == 1),
+            prompt_cache_id=key[0],
+            prompt_wav=key[1],
+            last_chunk=False,
+            tokens=torch.tensor([10, 11]),
+            previous=_RequestState(0, 0, *key, continuation) if row == 1 else None,
+            runtime_prompt_key=None,
+            duplex_epoch=0,
+            duplex_turn_id=0,
+            segment_text_utf8=torch.empty(0, dtype=torch.uint8),
+            tts_is_last_chunk=False,
+            segment_end=False,
+            turn_end=False,
+        )
+        for row, key in enumerate([keys[0], keys[0], keys[1], keys[2]])
+    ]
+    prepared, setups = [], []
+    fresh_states = [BatchedToken2WavState(flow_cache={}, hift_cache={}) for _ in range(1 if same_reference else 2)]
+
+    def prepare(*key):
+        prepared.append(key)
+        return by_key[key]
+
+    def setup(fresh):
+        setups.append([feature.cache_key for feature in fresh])
+        return fresh_states
+
+    monkeypatch.setattr(model.backend, "prepare_prompt", prepare)
+    monkeypatch.setattr(model.backend, "setup_references", setup)
+    states, conditioning = model._prepare_decode_batch(items)
+    expected_keys = list(dict.fromkeys(keys))
+    assert prepared == expected_keys
+    assert setups == [expected_keys]
+    expected_states = [fresh_states[0], continuation, fresh_states[-1], fresh_states[0]]
+    assert all(actual is expected for actual, expected in zip(states, expected_states, strict=True))
+    if same_reference:
+        assert conditioning is features[0]
+    else:
+        assert [feature.cache_key for feature in conditioning] == [keys[0], keys[0], keys[1], keys[2]]
+
+
 def _enable_fake_ragged_kernel(adapter: BatchedToken2Wav) -> None:
     """Keep fake-model tests focused on ragged grouping and row mapping."""
     estimator = adapter.flow.decoder.estimator
@@ -687,7 +926,7 @@ def test_adapter_runs_true_batch_cfg_and_splits_request_caches():
 
 
 @pytest.mark.parametrize("different_lengths", [False, True])
-@pytest.mark.parametrize("ragged_device", [None, "cuda", "npu"])
+@pytest.mark.parametrize("ragged_device", [None, "cuda", "cuda_disable", "npu"])
 def test_cross_reference_batch_matches_independent_streams(different_lengths, ragged_device):
     """Different voices share a solve without sharing conditioning or history."""
     batch = 8
@@ -716,13 +955,20 @@ def make_model():
         # Exercise dispatch on CPU; real CUDA replay/parity is tested separately.
         def replay(**kwargs):
             replays.append(kwargs["prompt_lengths"])
+            if ragged_device == "cuda_disable":
+                model.backend._whole_euler_graph_wrapper.enabled = False
             return None
 
         model.backend._whole_euler_graph_wrapper = SimpleNamespace(
-            enabled=True, device=SimpleNamespace(type=ragged_device), ragged_body=object(), replay=replay
+            enabled=True,
+            device=SimpleNamespace(type="cuda" if ragged_device == "cuda_disable" else ragged_device),
+            ragged_body=object(),
+            replay=replay,
+            micro_batch_size=batch,
+            max_graph_batch=batch,
         )
-    merged = ragged_device == "cuda"
     for chunk in range(8):
+        merged = ragged_device == "cuda" or (ragged_device == "cuda_disable" and chunk == 0)
         infos = []
         for row in range(batch):
             info = _info(f"voice-{row}", chunk, list(range(10 + row, 30 + row)), last_chunk=chunk == 7)
@@ -737,7 +983,7 @@ def replay(**kwargs):
         )
         assert token2wav.flow.decoder.estimator.speaker_order[-1] == expected_speakers
         assert token2wav.hift.calls[-1] == expected_batch
-        if merged:
+        if merged and ragged_device != "cuda_disable":
             assert replays[-1] == [4 + row % 3 if different_lengths else 4 for row in range(batch)]
         for row, (single, info) in enumerate(zip(independent, infos, strict=True)):
             expected = _forward(single, [info])
diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
index 3e9c8496709..58e5d11cafd 100644
--- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
+++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
@@ -1744,12 +1744,12 @@ def precapture_whole_euler(self, features: PromptFeatures) -> int:
                 tail_frames=(self._pre_lookahead_len() or 0) * int(getattr(self.flow, "token_mel_ratio", 2)),
             )
 
-    def _create_initial_states(
+    def _encode_initial_prompt(
         self,
         features: PromptFeatures,
         batch_size: int,
-    ) -> list[BatchedToken2WavState]:
-        prompt_tokens, speakers, prompt_mels = self._repeat_prompt(features, batch_size)
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        prompt_tokens, speakers, _ = self._repeat_prompt(features, batch_size)
         lookahead_width = self._pre_lookahead_len()
         lookahead = prompt_tokens.new_full(
             (batch_size, 3 if lookahead_width is None else lookahead_width),
@@ -1763,6 +1763,26 @@ def _create_initial_states(
                 att_cache=None,
             )
             projected_speakers = self._project_speakers(speakers)
+        return hidden, conformer_cnn, conformer_att, projected_speakers
+
+    def _create_initial_states(
+        self,
+        features: PromptFeatures | list[PromptFeatures],
+        batch_size: int,
+    ) -> list[BatchedToken2WavState]:
+        if isinstance(features, list):
+            # Preserve each reference's M=1 encoder and speaker projection.
+            # _encode_chunk owns its outputs, including graph replay outputs.
+            encoded = [self._encode_initial_prompt(feature, 1) for feature in features]
+            hidden = torch.cat([row[0] for row in encoded])
+            conformer_cnn = torch.cat([row[1] for row in encoded])
+            conformer_att = torch.cat([row[2] for row in encoded], dim=1)
+            projected_speakers = torch.cat([row[3] for row in encoded])
+            prompt_mels = torch.cat([feature.mels for feature in features])
+        else:
+            hidden, conformer_cnn, conformer_att, projected_speakers = self._encode_initial_prompt(features, batch_size)
+            prompt_mels = features.mels.expand(batch_size, -1, -1)
+        with self._autocast(prompt_mels.device):
             # 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(
@@ -1780,6 +1800,12 @@ def _create_initial_states(
             "estimator_att_cache": estimator_att,
         }
         split = self._split_flow_cache(flow_cache, batch_size)
+        if isinstance(features, list):
+            # Cached references have independent lifetimes. Do not let a row
+            # retain the other references' buffers after their LRU eviction.
+            for row in split:
+                for name in ("conformer_cnn_cache", "conformer_att_cache", "estimator_cnn_cache"):
+                    row[name] = row[name].clone()
         mel_channels = int(prompt_mels.shape[2])
         return [
             BatchedToken2WavState(
@@ -1831,6 +1857,53 @@ def setup_batch(self, features: PromptFeatures, batch_size: int) -> list[Batched
                     self._setup_cache.popitem(last=False)
         return [state] * batch_size
 
+    def setup_references(self, features: list[PromptFeatures]) -> list[BatchedToken2WavState]:
+        """Initialize unique references in order, batching compatible cold CFM solves."""
+        wrapper = self._whole_euler_graph_wrapper
+        if (
+            len(features) < 2
+            or not self._whole_euler_ragged_active()
+            or wrapper.device.type != "cuda"
+            or self._cfm_graph_bucket_frames != 0
+        ):
+            return [self.setup_batch(feature, 1)[0] for feature in features]
+
+        keys = [self._setup_cache_key(feature) for feature in features]
+        # Hold every hit before publishing misses can evict it from the LRU.
+        states = [self._setup_cache.get(key) for key in keys]
+        groups: dict[tuple[tuple[torch.Size, torch.dtype, torch.device], ...], list[int]] = {}
+        for row, (feature, state) in enumerate(zip(features, states, strict=True)):
+            if state is None:
+                signature = tuple(
+                    (tensor.shape, tensor.dtype, tensor.device)
+                    for tensor in (feature.speech_tokens, feature.speaker_embedding, feature.mels)
+                )
+                groups.setdefault(signature, []).append(row)
+        # A disabled graph batch budget keeps the existing singleton fallback.
+        maximum = max(1, min(wrapper.micro_batch_size, wrapper.max_graph_batch))
+        for rows in groups.values():
+            for start in range(0, len(rows), maximum):
+                indices = rows[start : start + maximum]
+                pending = [features[row] for row in indices]
+                # Capture can disable Whole-Euler during an earlier solve.
+                if len(pending) > 1 and self._whole_euler_ragged_active():
+                    created = self._create_initial_states(pending, len(pending))
+                else:
+                    created = [self._create_initial_states(feature, 1)[0] for feature in pending]
+                for row, state in zip(indices, created, strict=True):
+                    states[row] = state
+
+        result = cast(list[BatchedToken2WavState], states)
+        # Publish only after every solve succeeded, in the caller's reference
+        # order rather than shape-group order. Local references survive eviction.
+        if self._setup_cache_size > 0:
+            for key, state in zip(keys, result, strict=True):
+                self._setup_cache[key] = state
+                self._setup_cache.move_to_end(key)
+                while len(self._setup_cache) > self._setup_cache_size:
+                    self._setup_cache.popitem(last=False)
+        return result
+
     @staticmethod
     def _fade_in_out(
         speech: torch.Tensor,
diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py
index dfc23a5a0ea..8d203898369 100644
--- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py
+++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py
@@ -878,7 +878,7 @@ def _bucket_key(
     def _prepare_decode_batch(
         self, items: list[_WorkItem]
     ) -> tuple[list[BatchedToken2WavState], PromptFeatures | list[PromptFeatures]]:
-        """Prepare each voice separately, then keep conditioning in request order."""
+        """Prepare references and keep their conditioning and states in request order."""
         prompts: dict[tuple[str, str], PromptFeatures] = {}
         fresh: dict[tuple[str, str], list[int]] = {}
         states = [item.previous.token2wav if item.previous is not None else None for item in items]
@@ -890,8 +890,9 @@ def _prepare_decode_batch(
             features.append(prompts[key])
             if states[row] is None:
                 fresh.setdefault(key, []).append(row)
-        for key, rows in fresh.items():
-            for row, state in zip(rows, self.backend.setup_batch(prompts[key], len(rows)), strict=True):
+        initial = self.backend.setup_references([prompts[key] for key in fresh])
+        for rows, state in zip(fresh.values(), initial, strict=True):
+            for row in rows:
                 states[row] = state
         if any(state is None for state in states):
             raise _batch_error("incomplete_setup")

From da8df8a98bebb1f589e211f49aa0c75c1a5219f1 Mon Sep 17 00:00:00 2001
From: natureofnature <wzliu@connect.hku.hk>
Date: Wed, 21 Oct 2026 06:58:18 -0400
Subject: [PATCH 2/4] [Model] Preserve reference setup for other precision
 policies

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
---
 docs/design/minicpm_o45_mrv2_performance.md              | 7 ++++---
 .../models/minicpmo_4_5/test_code2wav_batching.py        | 9 +++++++--
 .../models/minicpmo_4_5/batched_token2wav.py             | 3 +++
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/docs/design/minicpm_o45_mrv2_performance.md b/docs/design/minicpm_o45_mrv2_performance.md
index 3703fd40a0f..d47330befd0 100644
--- a/docs/design/minicpm_o45_mrv2_performance.md
+++ b/docs/design/minicpm_o45_mrv2_performance.md
@@ -65,11 +65,12 @@ restricted to CUDA; the existing opt-in slot path is unchanged on other
 platforms.
 
 Reference extraction, encoder initialization and speaker projection remain
-per reference. With CUDA ragged Whole-Euler and zero CFM frame padding,
+per reference. With FP32/TF32 CUDA ragged Whole-Euler and zero CFM frame padding,
 compatible cold references share an initial CFM solve, bounded by the configured
 microbatch and graph batch limits. Setup-cache hits are reused, and each
-reference retains its own cached state. Other configurations keep per-reference
-initialization. Speaker conditioning, attention lengths and cache trimming stay
+reference retains its own cached state. FP16, TF32 opt-out and other execution
+configurations keep per-reference initialization. Speaker conditioning,
+attention lengths and cache trimming stay
 per request. Cache epochs remain separate unless cross-turn batching is enabled.
 
 The MiniCPM Code2Wav stage runs CAM++ speaker embedding on CUDA and NPU using
diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
index aaefa419329..0e69726d852 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
@@ -441,7 +441,7 @@ def _reference_feature(voice, length=2, dtype=torch.float32):
 
 
 def _reference_setup_backend(*, cache_size=4, micro_batch_size=2, max_graph_batch=8):
-    adapter = BatchedToken2Wav(_FakeToken2Wav(), setup_cache_size=cache_size)
+    adapter = BatchedToken2Wav(_FakeToken2Wav(), setup_cache_size=cache_size, cfm_tf32=True)
     # CUDA dispatch with CPU tensors; eager CFM supplies deterministic results.
     adapter._whole_euler_graph_wrapper = SimpleNamespace(
         enabled=True,
@@ -519,7 +519,7 @@ def test_setup_references_retains_results_and_hits_across_cache_eviction(cache_s
     assert list(adapter._setup_cache) == ([(third.cache_key, 0)] if cache_size else [])
 
 
-@pytest.mark.parametrize("gate", ["absent", "disabled", "cpu", "npu", "exact", "padding", "trt"])
+@pytest.mark.parametrize("gate", ["absent", "disabled", "cpu", "npu", "exact", "padding", "trt", "tf32_off", "float16"])
 def test_setup_references_preserves_single_reference_fallback(monkeypatch, gate):
     adapter = _reference_setup_backend()
     wrapper = adapter._whole_euler_graph_wrapper
@@ -533,6 +533,10 @@ def test_setup_references_preserves_single_reference_fallback(monkeypatch, gate)
         wrapper.ragged_body = None
     elif gate == "padding":
         adapter._cfm_graph_bucket_frames = 8
+    elif gate == "tf32_off":
+        adapter._cfm_tf32 = False
+    elif gate == "float16":
+        adapter.float16 = True
     else:
         adapter._trt_stepper = object()
     features = [_reference_feature(i) for i in range(2)]
@@ -944,6 +948,7 @@ def prepare_prompt(path):
 
     def make_model():
         model, token2wav = _model()
+        model.backend._cfm_tf32 = True
         token2wav._prepare_prompt = prepare_prompt
         token2wav.flow.decoder.rand_noise = torch.zeros(1, 1, 1024)
         return model, token2wav
diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
index 58e5d11cafd..0b320bdafaa 100644
--- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
+++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
@@ -1860,8 +1860,11 @@ def setup_batch(self, features: PromptFeatures, batch_size: int) -> list[Batched
     def setup_references(self, features: list[PromptFeatures]) -> list[BatchedToken2WavState]:
         """Initialize unique references in order, batching compatible cold CFM solves."""
         wrapper = self._whole_euler_graph_wrapper
+        # Keep other precision policies on their existing M=1 setup path.
         if (
             len(features) < 2
+            or self.float16
+            or not self._cfm_tf32
             or not self._whole_euler_ragged_active()
             or wrapper.device.type != "cuda"
             or self._cfm_graph_bucket_frames != 0

From 41cb868615e8480e5d75f6b9dfc9732ccac3eddc Mon Sep 17 00:00:00 2001
From: natureofnature <wzliu@connect.hku.hk>
Date: Wed, 21 Oct 2026 10:09:00 -0400
Subject: [PATCH 3/4] [Model] Preserve MiniCPM-o scalar setup on unfused or
 refused graphs

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
---
 docs/design/minicpm_o45_mrv2_performance.md   | 10 +-
 .../minicpmo_4_5/test_code2wav_batching.py    | 99 +++++++++++++++++--
 .../models/minicpmo_4_5/batched_token2wav.py  | 55 +++++++++--
 3 files changed, 143 insertions(+), 21 deletions(-)

diff --git a/docs/design/minicpm_o45_mrv2_performance.md b/docs/design/minicpm_o45_mrv2_performance.md
index d47330befd0..593ecc6bcfa 100644
--- a/docs/design/minicpm_o45_mrv2_performance.md
+++ b/docs/design/minicpm_o45_mrv2_performance.md
@@ -65,12 +65,14 @@ restricted to CUDA; the existing opt-in slot path is unchanged on other
 platforms.
 
 Reference extraction, encoder initialization and speaker projection remain
-per reference. With FP32/TF32 CUDA ragged Whole-Euler and zero CFM frame padding,
+per reference. With FP32/TF32 CUDA fused ragged Whole-Euler and zero CFM frame padding,
 compatible cold references share an initial CFM solve, bounded by the configured
 microbatch and graph batch limits. Setup-cache hits are reused, and each
-reference retains its own cached state. FP16, TF32 opt-out and other execution
-configurations keep per-reference initialization. Speaker conditioning,
-attention lengths and cache trimming stay
+reference retains its own cached state. A refused graph replay retries setup
+per reference, preserving the existing solver selection. FP16, TF32 opt-out,
+unfused bodies (including BF16 attention caches) and other execution
+configurations keep per-reference initialization. Speaker
+conditioning, attention lengths and cache trimming stay
 per request. Cache epochs remain separate unless cross-turn batching is enabled.
 
 The MiniCPM Code2Wav stage runs CAM++ speaker embedding on CUDA and NPU using
diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
index 0e69726d852..0d389be5f42 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
@@ -442,15 +442,34 @@ def _reference_feature(voice, length=2, dtype=torch.float32):
 
 def _reference_setup_backend(*, cache_size=4, micro_batch_size=2, max_graph_batch=8):
     adapter = BatchedToken2Wav(_FakeToken2Wav(), setup_cache_size=cache_size, cfm_tf32=True)
-    # CUDA dispatch with CPU tensors; eager CFM supplies deterministic results.
-    adapter._whole_euler_graph_wrapper = SimpleNamespace(
+    adapter._ragged_fused_body = True
+    # CUDA dispatch with CPU tensors; model a successful graph with fake math.
+    decode = adapter._decode_cfm
+    wrapper = SimpleNamespace(
         enabled=True,
         device=torch.device("cuda"),
         ragged_body=object(),
-        replay=lambda **kwargs: None,
         micro_batch_size=micro_batch_size,
         max_graph_batch=max_graph_batch,
     )
+
+    def replay(**kwargs):
+        batch = kwargs["x"].shape[0]
+        wrapper.enabled = False
+        try:
+            return decode(
+                kwargs["mu_cfg"][:batch],
+                kwargs["speakers_cfg"][:batch],
+                kwargs["cond_cfg"][:batch],
+                cnn_cache=kwargs["cnn_cache"],
+                att_cache=kwargs["att_cache"],
+                att_keep=kwargs["att_keep"],
+            )
+        finally:
+            wrapper.enabled = True
+
+    wrapper.replay = replay
+    adapter._whole_euler_graph_wrapper = wrapper
     return adapter
 
 
@@ -519,7 +538,10 @@ def test_setup_references_retains_results_and_hits_across_cache_eviction(cache_s
     assert list(adapter._setup_cache) == ([(third.cache_key, 0)] if cache_size else [])
 
 
-@pytest.mark.parametrize("gate", ["absent", "disabled", "cpu", "npu", "exact", "padding", "trt", "tf32_off", "float16"])
+@pytest.mark.parametrize(
+    "gate",
+    ["absent", "disabled", "cpu", "npu", "exact", "padding", "trt", "tf32_off", "float16", "unfused", "bf16_cache"],
+)
 def test_setup_references_preserves_single_reference_fallback(monkeypatch, gate):
     adapter = _reference_setup_backend()
     wrapper = adapter._whole_euler_graph_wrapper
@@ -537,6 +559,10 @@ def test_setup_references_preserves_single_reference_fallback(monkeypatch, gate)
         adapter._cfm_tf32 = False
     elif gate == "float16":
         adapter.float16 = True
+    elif gate in ("unfused", "bf16_cache"):
+        adapter._ragged_fused_body = False
+        if gate == "bf16_cache":
+            adapter._estimator_att_cache_dtype = torch.bfloat16
     else:
         adapter._trt_stepper = object()
     features = [_reference_feature(i) for i in range(2)]
@@ -551,19 +577,69 @@ def setup(feature, batch_size):
     assert calls == [(feature.cache_key, 1) for feature in features]
 
 
-def test_setup_references_rechecks_disabled_graph_between_groups():
+@pytest.mark.parametrize("refusal", ["enabled", "disabled", "full_budget"])
+def test_setup_references_retries_refused_graphs_with_scalar_setup(refusal):
     adapter = _reference_setup_backend()
     wrapper = adapter._whole_euler_graph_wrapper
+    successful_replay = wrapper.replay
+    features = [_reference_feature(i) for i in range(4)]
+    original_features = [
+        tuple(value.clone() for value in (f.speech_tokens, f.speaker_embedding, f.mels)) for f in features
+    ]
+    original_noise = adapter.flow.decoder.rand_noise.clone()
+    single_key = ("whole_euler", 1, 4, 0, "torch.float32", "cpu")
+    graph_type = batched_token2wav_module.WholeEulerCFMGraphWrapper
+    graph_cache = graph_type.__new__(graph_type)
+    graph_cache._cache = {single_key: (None,) * 5}
+    graph_cache._slot_graphs = {}
+    graph_cache._unsupported = set()
+    graph_cache._stats = {"hits": 0}
+    graph_cache.max_graphs = 1
+    replay_batches = []
 
-    def disable(**kwargs):
-        wrapper.enabled = False
+    def replay(**kwargs):
+        batch = int(kwargs["x"].shape[0])
+        replay_batches.append(batch)
+        if refusal == "full_budget":
+            # Exercise the real admission check: the M=1 key is still usable
+            # but a new batched key cannot fit, without disabling the wrapper.
+            entry = batched_token2wav_module.WholeEulerCFMGraphWrapper._entry(
+                graph_cache,
+                graph_batch=batch,
+                query_cap=int(kwargs["mu_cfg"].shape[2]),
+                offset=0,
+                x=kwargs["x"],
+                spk_dim=int(kwargs["speakers_cfg"].shape[1]),
+                fill=None,
+            )
+            if entry is not None:
+                return successful_replay(**kwargs)
+        elif batch == 1:
+            return successful_replay(**kwargs)
+        if refusal == "disabled":
+            wrapper.enabled = False
         return None
 
-    wrapper.replay = disable
-    states = adapter.setup_references([_reference_feature(i) for i in range(4)])
+    wrapper.replay = replay
+    states = adapter.setup_references(features)
 
     assert len(states) == 4
-    assert adapter.flow.decoder.estimator.cfg_batches == [4, 4, 2, 2, 2, 2]
+    assert adapter.flow.decoder.estimator.cfg_batches == [2] * 8
+    assert replay_batches == ([2] if refusal == "disabled" else [2, 1, 1, 2, 1, 1])
+    if refusal == "full_budget":
+        assert graph_cache._stats["hits"] == 4
+        assert graph_cache._unsupported == {("whole_euler", 2, 4, 0, "torch.float32", "cpu")}
+        assert wrapper.enabled
+    independent = BatchedToken2Wav(_FakeToken2Wav())
+    for feature, state, original in zip(features, states, original_features, strict=True):
+        expected = independent.setup_batch(feature, 1)[0]
+        for name, value in state.flow_cache.items():
+            torch.testing.assert_close(value, expected.flow_cache[name])
+        for value, prior in zip(
+            (feature.speech_tokens, feature.speaker_embedding, feature.mels), original, strict=True
+        ):
+            torch.testing.assert_close(value, prior, rtol=0, atol=0)
+    torch.testing.assert_close(adapter.flow.decoder.rand_noise, original_noise, rtol=0, atol=0)
 
 
 def test_setup_references_failed_later_group_does_not_publish_or_touch_cache(monkeypatch):
@@ -949,6 +1025,9 @@ def prepare_prompt(path):
     def make_model():
         model, token2wav = _model()
         model.backend._cfm_tf32 = True
+        model.backend._ragged_fused_body = True
+        # The fake kernel covers row mapping for either production body.
+        model.backend._ragged_body = model.backend._blocks_forward_chunk_ragged
         token2wav._prepare_prompt = prepare_prompt
         token2wav.flow.decoder.rand_noise = torch.zeros(1, 1, 1024)
         return model, token2wav
diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
index 0b320bdafaa..5cdc6619742 100644
--- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
+++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
@@ -10,7 +10,7 @@
 from collections.abc import Mapping
 from contextlib import nullcontext
 from dataclasses import dataclass
-from typing import Any, cast
+from typing import Any, Literal, cast, overload
 
 import torch
 import torch.nn as nn
@@ -60,6 +60,8 @@
 # strided in-place attention once eliminating history copies pays for it.
 _NPU_IN_PLACE_ATTENTION_MIN_CACHE_FRAMES = 512
 
+_CFMOutput = tuple[torch.Tensor, torch.Tensor, torch.Tensor | list[torch.Tensor]]
+
 
 def _scale_shift(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor:
     """Apply DiT modulation with one Ascend elementwise kernel."""
@@ -1315,6 +1317,36 @@ def _get_timeline(self, device: torch.device, dtype: torch.dtype) -> tuple[torch
             self._timeline_cache[key] = cached
         return cached
 
+    @overload
+    def _decode_cfm(
+        self,
+        mu: torch.Tensor,
+        speakers: torch.Tensor,
+        cond: torch.Tensor,
+        *,
+        cnn_cache: torch.Tensor | None,
+        att_cache: torch.Tensor | list[torch.Tensor] | None,
+        valid_lengths: list[int] | None = None,
+        att_keep: tuple[int, int] | None = None,
+        prompt_lengths: list[int] | None = None,
+        whole_euler_only: Literal[False] = False,
+    ) -> _CFMOutput: ...
+
+    @overload
+    def _decode_cfm(
+        self,
+        mu: torch.Tensor,
+        speakers: torch.Tensor,
+        cond: torch.Tensor,
+        *,
+        cnn_cache: torch.Tensor | None,
+        att_cache: torch.Tensor | list[torch.Tensor] | None,
+        valid_lengths: list[int] | None = None,
+        att_keep: tuple[int, int] | None = None,
+        prompt_lengths: list[int] | None = None,
+        whole_euler_only: bool,
+    ) -> _CFMOutput | None: ...
+
     @_cfm_matmul_tf32
     def _decode_cfm(
         self,
@@ -1327,17 +1359,16 @@ def _decode_cfm(
         valid_lengths: list[int] | None = None,
         att_keep: tuple[int, int] | None = None,
         prompt_lengths: list[int] | None = None,
-    ) -> tuple[
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor | list[torch.Tensor],
-    ]:
+        whole_euler_only: bool = False,
+    ) -> _CFMOutput | None:
         """Run the CFM Euler solve for one chunk.
 
         ``att_cache`` may be the stacked estimator cache or one tensor per
         request (the Whole-Euler path). A Whole-Euler solve returns the new
         cache in the same form, already trimmed by ``att_keep``; the eager
         path returns it stacked and untrimmed (ragged: per request, untrimmed).
+        ``whole_euler_only`` returns ``None`` on refused replay, allowing cold
+        reference setup to retain its scalar step/eager fallback policy.
         """
         decoder = self.flow.decoder
         estimator = decoder.estimator
@@ -1445,6 +1476,8 @@ def _decode_cfm(
             )
             if whole_euler_result is not None:
                 return whole_euler_result
+        if whole_euler_only:
+            return None
         # Whole-Euler builds its capture mask directly from mel_frames.
         # Only a refused replay needs the separate step/eager padding mask.
         if pad_frames:
@@ -1785,14 +1818,21 @@ def _create_initial_states(
         with self._autocast(prompt_mels.device):
             # 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(
+            initial = self._decode_cfm(
                 hidden.transpose(1, 2).contiguous(),
                 projected_speakers,
                 prompt_mels.transpose(1, 2).contiguous(),
                 cnn_cache=None,
                 att_cache=None,
                 att_keep=(int(prompt_mels.shape[1]), _CACHE_TRIM_SUFFIX),
+                whole_euler_only=isinstance(features, list),
             )
+        if initial is None:
+            # A full graph budget or failed capture can refuse this batch even
+            # while M=1 graphs remain usable. Preserve each reference's solver.
+            assert isinstance(features, list)
+            return [self._create_initial_states(feature, 1)[0] for feature in features]
+        _, estimator_cnn, estimator_att = initial
         flow_cache = {
             "conformer_cnn_cache": conformer_cnn,
             "conformer_att_cache": conformer_att,
@@ -1865,6 +1905,7 @@ def setup_references(self, features: list[PromptFeatures]) -> list[BatchedToken2
             len(features) < 2
             or self.float16
             or not self._cfm_tf32
+            or not self._ragged_fused_body
             or not self._whole_euler_ragged_active()
             or wrapper.device.type != "cuda"
             or self._cfm_graph_bucket_frames != 0

From 532e9ca1145cef8b7cdd3f5334cff20c6ad9c0ac Mon Sep 17 00:00:00 2001
From: natureofnature <wzliu@connect.hku.hk>
Date: Wed, 21 Oct 2026 12:11:45 -0400
Subject: [PATCH 4/4] [Model] Reuse MiniCPM-o setup encoding and precapture
 setup graphs

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
---
 docs/design/minicpm_o45_mrv2_performance.md   | 12 ++++++-
 .../minicpmo_4_5/test_cfm_graph_bucketing.py  | 23 +++++++++----
 .../minicpmo_4_5/test_code2wav_batching.py    | 29 ++++++++++++++--
 .../models/minicpmo_4_5/batched_token2wav.py  | 34 +++++++++++++------
 .../models/minicpmo_4_5/cuda_graph_wrapper.py |  8 ++++-
 5 files changed, 85 insertions(+), 21 deletions(-)

diff --git a/docs/design/minicpm_o45_mrv2_performance.md b/docs/design/minicpm_o45_mrv2_performance.md
index 593ecc6bcfa..572a7917467 100644
--- a/docs/design/minicpm_o45_mrv2_performance.md
+++ b/docs/design/minicpm_o45_mrv2_performance.md
@@ -69,12 +69,22 @@ per reference. With FP32/TF32 CUDA fused ragged Whole-Euler and zero CFM frame p
 compatible cold references share an initial CFM solve, bounded by the configured
 microbatch and graph batch limits. Setup-cache hits are reused, and each
 reference retains its own cached state. A refused graph replay retries setup
-per reference, preserving the existing solver selection. FP16, TF32 opt-out,
+per reference using the already-computed encoder and speaker-projection outputs,
+preserving the existing solver selection. FP16, TF32 opt-out,
 unfused bodies (including BF16 attention caches) and other execution
 configurations keep per-reference initialization. Speaker
 conditioning, attention lengths and cache trimming stay
 per request. Cache epochs remain separate unless cross-turn batching is enabled.
 
+Cold setup shares the `cfm_max_graphs` budget with decode graphs. For eligible
+batched setup, non-slot precapture warms the default prompt's offset-zero graphs
+after decode shapes, using remaining budget with M=1 first. Other reference
+widths capture lazily or fall back to per-reference solves. During batched setup,
+the stacked CFG attention output and independent per-reference cache copies
+coexist, increasing transient memory with the group size. The microbatch bound
+limits this peak; independent copies keep one reference from retaining another
+reference's buffers after cache eviction.
+
 The MiniCPM Code2Wav stage runs CAM++ speaker embedding on CUDA and NPU using
 the in-tree torch model with FP32 weights and inputs, loaded from the
 checkpoint's `campplus.onnx`. On CUDA, reference preparation follows the
diff --git a/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py
index 6751330ebd6..bd6621777f7 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py
@@ -305,8 +305,9 @@ def test_zero_padded_cnn_cache_uniform_blocks_use_one_write():
     assert torch.count_nonzero(cache[..., 3:]) == 0
 
 
-@pytest.mark.parametrize("max_graphs", [15, 32])
-def test_precapture_covers_final_chunk_without_displacing_steady_graphs(monkeypatch, max_graphs):
+@pytest.mark.parametrize("max_graphs", [15, 30, 31, 35])
+@pytest.mark.parametrize("setup_frames", [0, 273])
+def test_precapture_covers_final_chunk_without_displacing_steady_graphs(monkeypatch, max_graphs, setup_frames):
     from vllm_omni.model_executor.models.minicpmo_4_5 import cuda_graph_wrapper as graph_module
 
     captured = []
@@ -336,14 +337,24 @@ def entry(*, graph_batch, query_cap, offset, **kwargs):
     )
     monkeypatch.setattr(graph_module, "_memory_snapshot", lambda device: None)
     count = WholeEulerCFMGraphWrapper.precapture(
-        wrapper, offsets=range(300, 401), steady=400, channels=4, spk_dim=4, tail_frames=6
+        wrapper,
+        offsets=range(300, 401),
+        steady=400,
+        channels=4,
+        spk_dim=4,
+        tail_frames=6,
+        setup_frames=setup_frames,
     )
     assert captured[:15] == [(b, 50, o) for b in [16, 8, 4, 2, 1] for o in [400, 350, 300]]
-    assert count == min(max_graphs, 30)
+    expected = [(b, 50, o) for b in [16, 8, 4, 2, 1] for o in [400, 350, 300]]
+    expected += [(b, 100, o) for b in [16, 8, 4, 2, 1] for o in [400, 350, 300]]
+    if setup_frames:
+        expected += [(b, 300, 0) for b in [1, 2, 4, 8, 16]]
+    assert captured == expected[:max_graphs]
+    assert count == len(captured)
     assert wrapper.query_widths == (50,)
-    if max_graphs == 32:
+    if max_graphs >= 30:
         assert (1, 100, 350) in captured
-        assert captured[15:] == [(b, 100, o) for b in [16, 8, 4, 2, 1] for o in [400, 350, 300]]
 
 
 @pytest.mark.parametrize("width,valid", [(0, 0), (3, 2), (8, 8)])
diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
index 0d389be5f42..7a476d15d08 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
@@ -4,7 +4,7 @@
 import gc
 import subprocess
 import sys
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -578,7 +578,7 @@ def setup(feature, batch_size):
 
 
 @pytest.mark.parametrize("refusal", ["enabled", "disabled", "full_budget"])
-def test_setup_references_retries_refused_graphs_with_scalar_setup(refusal):
+def test_setup_references_retries_refused_graphs_with_scalar_setup(monkeypatch, refusal):
     adapter = _reference_setup_backend()
     wrapper = adapter._whole_euler_graph_wrapper
     successful_replay = wrapper.replay
@@ -596,6 +596,14 @@ def test_setup_references_retries_refused_graphs_with_scalar_setup(refusal):
     graph_cache._stats = {"hits": 0}
     graph_cache.max_graphs = 1
     replay_batches = []
+    encode = adapter._encode_initial_prompt
+    encoded_rows = []
+
+    def encode_once(feature, batch_size):
+        encoded_rows.append((feature.cache_key, batch_size))
+        return encode(feature, batch_size)
+
+    monkeypatch.setattr(adapter, "_encode_initial_prompt", encode_once)
 
     def replay(**kwargs):
         batch = int(kwargs["x"].shape[0])
@@ -624,6 +632,9 @@ def replay(**kwargs):
     states = adapter.setup_references(features)
 
     assert len(states) == 4
+    # Both the first refusal and later groups with the memoized unsupported
+    # shape reuse their M=1 encoder and speaker-projection outputs.
+    assert encoded_rows == [(feature.cache_key, 1) for feature in features]
     assert adapter.flow.decoder.estimator.cfg_batches == [2] * 8
     assert replay_batches == ([2] if refusal == "disabled" else [2, 1, 1, 2, 1, 1])
     if refusal == "full_budget":
@@ -642,6 +653,20 @@ def replay(**kwargs):
     torch.testing.assert_close(adapter.flow.decoder.rand_noise, original_noise, rtol=0, atol=0)
 
 
+@pytest.mark.parametrize("batched_setup", [True, False])
+def test_precapture_requests_setup_shapes_only_for_eligible_solver(monkeypatch, mocker, batched_setup):
+    adapter = _reference_setup_backend()
+    adapter._ragged_fused_body = batched_setup
+    adapter.flow.spk_embed_affine_layer = nn.Linear(1, 1)
+    wrapper = adapter._whole_euler_graph_wrapper
+    wrapper.precapture = mocker.Mock(return_value=0)
+    monkeypatch.setattr(adapter, "_autocast", lambda device: nullcontext())
+    feature = _reference_feature(0)
+
+    assert adapter.precapture_whole_euler(feature) == 0
+    assert wrapper.precapture.call_args.kwargs["setup_frames"] == (feature.mels.shape[1] if batched_setup else 0)
+
+
 def test_setup_references_failed_later_group_does_not_publish_or_touch_cache(monkeypatch):
     adapter = _reference_setup_backend(cache_size=2)
     features = [_reference_feature(i) for i in range(5)]
diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
index 5cdc6619742..b97b6997ac1 100644
--- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
+++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
@@ -1713,6 +1713,17 @@ def _whole_euler_ragged_active(self) -> bool:
         """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 _batch_reference_setup_enabled(self) -> bool:
+        """Whether reference setup can use the validated batched CFM solver."""
+        return (
+            not self.float16
+            and self._cfm_tf32
+            and self._ragged_fused_body
+            and self._whole_euler_ragged_active()
+            and self._whole_euler_graph_wrapper.device.type == "cuda"
+            and self._cfm_graph_bucket_frames == 0
+        )
+
     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 (
@@ -1775,6 +1786,7 @@ def precapture_whole_euler(self, features: PromptFeatures) -> int:
                 spk_dim=int(self.flow.spk_embed_affine_layer.out_features),
                 keep=(prompt_len, _CACHE_TRIM_SUFFIX),
                 tail_frames=(self._pre_lookahead_len() or 0) * int(getattr(self.flow, "token_mel_ratio", 2)),
+                setup_frames=prompt_len if self._batch_reference_setup_enabled() else 0,
             )
 
     def _encode_initial_prompt(
@@ -1802,7 +1814,10 @@ def _create_initial_states(
         self,
         features: PromptFeatures | list[PromptFeatures],
         batch_size: int,
+        *,
+        encoded_prompt: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
     ) -> list[BatchedToken2WavState]:
+        """Solve prompt states, reusing an M=1 encoding after a batched graph refusal."""
         if isinstance(features, list):
             # Preserve each reference's M=1 encoder and speaker projection.
             # _encode_chunk owns its outputs, including graph replay outputs.
@@ -1813,7 +1828,9 @@ def _create_initial_states(
             projected_speakers = torch.cat([row[3] for row in encoded])
             prompt_mels = torch.cat([feature.mels for feature in features])
         else:
-            hidden, conformer_cnn, conformer_att, projected_speakers = self._encode_initial_prompt(features, batch_size)
+            if encoded_prompt is None:
+                encoded_prompt = self._encode_initial_prompt(features, batch_size)
+            hidden, conformer_cnn, conformer_att, projected_speakers = encoded_prompt
             prompt_mels = features.mels.expand(batch_size, -1, -1)
         with self._autocast(prompt_mels.device):
             # The prompt cache never reaches the trim length; ``att_keep`` only
@@ -1831,7 +1848,10 @@ def _create_initial_states(
             # A full graph budget or failed capture can refuse this batch even
             # while M=1 graphs remain usable. Preserve each reference's solver.
             assert isinstance(features, list)
-            return [self._create_initial_states(feature, 1)[0] for feature in features]
+            return [
+                self._create_initial_states(feature, 1, encoded_prompt=row)[0]
+                for feature, row in zip(features, encoded, strict=True)
+            ]
         _, estimator_cnn, estimator_att = initial
         flow_cache = {
             "conformer_cnn_cache": conformer_cnn,
@@ -1901,15 +1921,7 @@ def setup_references(self, features: list[PromptFeatures]) -> list[BatchedToken2
         """Initialize unique references in order, batching compatible cold CFM solves."""
         wrapper = self._whole_euler_graph_wrapper
         # Keep other precision policies on their existing M=1 setup path.
-        if (
-            len(features) < 2
-            or self.float16
-            or not self._cfm_tf32
-            or not self._ragged_fused_body
-            or not self._whole_euler_ragged_active()
-            or wrapper.device.type != "cuda"
-            or self._cfm_graph_bucket_frames != 0
-        ):
+        if len(features) < 2 or not self._batch_reference_setup_enabled():
             return [self.setup_batch(feature, 1)[0] for feature in features]
 
         keys = [self._setup_cache_key(feature) for feature in features]
diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py
index 544cb927175..c6ae9058c0b 100644
--- a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py
+++ b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py
@@ -1885,8 +1885,9 @@ def precapture(
         max_graph_bytes: int = 1 << 30,
         keep: tuple[int, int] | None = None,
         tail_frames: int = 0,
+        setup_frames: int = 0,
     ) -> int:
-        """Capture Whole-Euler graphs for one prompt before serving. Returns the capture count."""
+        """Capture decode graphs, then optional non-slot setup shapes within the same budget."""
         if not self.enabled:
             return 0
         if self.att_slots and keep is not None and self._ensure_slot_pool(keep) is not None:
@@ -1905,6 +1906,11 @@ def precapture(
             keys.extend(
                 (b, w, o) for b in sorted(self._graph_batches(), reverse=True) for w in tail_widths for o in grid
             )
+        if setup_frames > 0:
+            # Decode keeps priority. Reserve M=1 setup before larger batches so
+            # a tight budget still permits the existing per-reference solver.
+            setup_width = _capture_query_width(setup_frames, widths)
+            keys.extend((b, setup_width, 0) for b in sorted(self._graph_batches()))
         room = self.max_graphs - (len(self._cache) + len(self._slot_graphs))
         if len(keys) > room:
             logger.warning(