[Feature][MiniCPM-o] Support configurable concurrent duplex sessions (#6021)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>

pyproject.toml

@@ -73,6 +73,8 @@ dev = [
     "FlagEmbedding==1.4.0", # for hunyuan_image3 tests
     "kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
     "s3tokenizer==0.3.0", # for tts test
+    # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
+    # whose top-level stepaudio2 package conflicts with this distribution.
     "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
 ]
 

tests/core/sched/test_omni_ar_scheduler_streaming.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
 
 """Unit tests for Omni AR streaming-session async placeholder handling."""
 
@@ -64,6 +64,9 @@ def _run_resumable_segment_stop(
     session: Request,
     *,
     session_finished: bool = False,
+    handle_stopped=None,
+    chunk_transfer_adapter=None,
+    inter_stage_output=None,
 ):
     sched = MagicMock()
     sched.requests = {session.request_id: session}
@@ -75,11 +78,14 @@ def _run_resumable_segment_stop(
         return [42], True
 
     sched._update_request_with_output.side_effect = stop_request
-    sched._handle_stopped_request.return_value = session_finished
     # vLLM 0.26 returns (kv_xfer_params, ec_xfer_params); an unconfigured
     # MagicMock iterates empty and fails to unpack at the call site.
     sched._free_request.return_value = (None, None)
-    sched.chunk_transfer_adapter = None
+    if handle_stopped is None:
+        sched._handle_stopped_request.return_value = session_finished
+    else:
+        sched._handle_stopped_request.side_effect = handle_stopped
+    sched.chunk_transfer_adapter = chunk_transfer_adapter
     sched.running = [session]
     sched.waiting_for_transfer_free = set()
     sched.transfer_triggered_requests = set()
@@ -105,6 +111,7 @@ def _run_resumable_segment_stop(
     model_runner_output.cudagraph_stats = None
     model_runner_output.req_id_to_index = {session.request_id: 0}
     model_runner_output.routed_experts = None
+    model_runner_output.inter_stage_outputs = [inter_stage_output] if inter_stage_output is not None else None
 
     return OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
 
@@ -160,6 +167,41 @@ def test_update_from_output_settles_in_flight_tokens() -> None:
     assert session.num_in_flight_tokens == 0
 
 
+def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> None:
+    """The old segment's transfer must not observe the next segment's reset.
+
+    A queued streaming update can replace the same mutable Request while the
+    stop output is still being handled. The connector needs the confirmed
+    token count from before that replacement to avoid dropping the boundary.
+    """
+    session = _make_request()
+    session.status = RequestStatus.RUNNING
+    session.resumable = True
+    session.num_computed_tokens = 26
+    adapter = MagicMock()
+    adapter._confirmed_num_computed_tokens.return_value = 26
+    inter_stage_output = {"codes": {"audio": [7]}}
+
+    def replace_with_next_segment(request: Request) -> bool:
+        request.num_computed_tokens = 0
+        request.status = RequestStatus.WAITING
+        return False
+
+    _run_resumable_segment_stop(
+        session,
+        handle_stopped=replace_with_next_segment,
+        chunk_transfer_adapter=adapter,
+        inter_stage_output=inter_stage_output,
+    )
+
+    adapter.save_async.assert_called_once_with(
+        inter_stage_output,
+        session,
+        True,
+        confirmed_num_computed_tokens=26,
+    )
+
+
 def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> None:
     """A decode step that neither stops nor carries an inter-stage payload.
 

tests/distributed/omni_connectors/test_chunk_transfer_adapter.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
 
 import threading
 from collections import deque
@@ -9,9 +9,10 @@ from unittest.mock import Mock, patch
 import pytest
 import torch
 from pytest_mock import MockerFixture
+from vllm.sampling_params import SamplingParams
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
-from vllm.v1.request import RequestStatus
+from vllm.v1.request import Request, RequestStatus
 
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
@@ -337,6 +338,121 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad
     assert len(adapter._pending_save_reqs) == 1
 
 
+def test_segment_boundary_starts_new_send_watermark_before_background_flush(build_adapter):
+    """A queued boundary owns the end of its deduplication generation.
+
+    The next segment can start before the save thread sends the old boundary.
+    Sending that old task later must not erase the new segment's watermark.
+    """
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
+    request.resumable = True
+    request.num_computed_tokens = 0
+    adapter.requests_num_chunks_sent["ext-stream"] = 26
+
+    adapter.save_async(
+        multimodal_output=None,
+        request=request,
+        is_segment_finished=True,
+        confirmed_num_computed_tokens=26,
+    )
+
+    assert len(adapter._pending_save_reqs) == 1
+    boundary_task = adapter._pending_save_reqs.popleft()
+    assert "ext-stream" not in adapter.requests_num_chunks_sent
+
+    request.num_computed_tokens = 3
+    adapter.save_async(
+        multimodal_output=None,
+        request=request,
+        is_segment_finished=False,
+    )
+
+    assert len(adapter._pending_save_reqs) == 1
+    assert adapter.requests_num_chunks_sent["ext-stream"] == 3
+
+    adapter._send_single_request(boundary_task)
+
+    assert adapter.requests_num_chunks_sent["ext-stream"] == 3
+
+
+def test_background_send_uses_enqueued_request_snapshot(build_adapter):
+    """A queued segment must not observe later in-place request mutations."""
+    adapter, _ = build_adapter(stage_id=1)
+    request = Request(
+        request_id="req-stream",
+        prompt_token_ids=[1, 2],
+        sampling_params=SamplingParams(max_tokens=8),
+        pooling_params=None,
+        resumable=True,
+    )
+    request.external_req_id = "ext-stream"
+    ref_audio = torch.tensor([0.1, -0.1])
+    request.additional_information = {"codes": {"ref": ref_audio}, "meta": {"segment": "old"}}
+    request.append_output_token_ids([7])
+    seen_requests = []
+
+    def recording_processor(**kwargs):
+        queued = kwargs["request"]
+        seen_requests.append(
+            (
+                queued.additional_information["meta"]["segment"],
+                list(queued.prompt_token_ids),
+                list(queued.output_token_ids),
+                list(queued.all_token_ids),
+                queued.additional_information["codes"]["ref"] is ref_audio,
+            )
+        )
+        return OmniPayloadStruct()
+
+    adapter.custom_process_next_stage_input_func = recording_processor
+    adapter.save_async(
+        multimodal_output=None,
+        request=request,
+        is_segment_finished=True,
+    )
+    request.additional_information["meta"]["segment"] = "next"
+    request.prompt_token_ids.append(3)
+    request.append_output_token_ids([8])
+    adapter.save_async(
+        multimodal_output=None,
+        request=request,
+        is_segment_finished=True,
+    )
+    request.additional_information["meta"]["segment"] = "later"
+    request.prompt_token_ids.append(4)
+    request.append_output_token_ids([9])
+
+    assert len(adapter._pending_save_reqs) == 2
+    first_task = adapter._pending_save_reqs.popleft()
+    second_task = adapter._pending_save_reqs.popleft()
+    assert first_task["request"].additional_information is not request.additional_information
+    assert first_task["request"].additional_information["meta"] is not request.additional_information["meta"]
+
+    adapter._send_single_request(first_task)
+    adapter._send_single_request(second_task)
+
+    assert seen_requests == [
+        ("old", [1, 2], [7], [1, 2, 7], True),
+        ("next", [1, 2, 3], [7, 8], [1, 2, 7, 8], True),
+    ]
+
+
+def test_save_without_custom_processor_does_not_snapshot_request(build_adapter, mocker):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-direct", RequestStatus.RUNNING, external_req_id="ext-direct")
+    snapshot = mocker.patch.object(
+        adapter,
+        "_snapshot_processor_request",
+        side_effect=AssertionError("unexpected snapshot"),
+    )
+
+    adapter.save_async(multimodal_output=None, request=request)
+
+    snapshot.assert_not_called()
+    assert adapter._pending_save_reqs.popleft()["request"] is request
+
+
 def test_send_single_request_terminal_chunk_still_flushes_processor(build_adapter, monkeypatch):
     """A terminal stop is not a segment boundary (#5383), but the producer-side
     processor must still receive the flush signal on the terminal chunk.

tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """MiniCPM-o 4.5 Daily-Omni + Seed-TTS accuracy regression coverage.
 
 Daily-Omni settings follow the MiniCPM interleaved AV recipe that reaches
@@ -21,13 +22,25 @@ from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import (
     build_acc_benchmark_cli_argv,
     find_vllm_cli,
 )
-from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import SERVER_PARAMS as DUPLEX_TEST_PARAMS
+from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
+    DEPLOY_CONFIG as _DUPLEX_DEPLOY_CONFIG,
+)
+from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
+    SERVER_PARAMS as DUPLEX_TEST_PARAMS,
+)
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 _MODEL = os.environ.get("VLLM_TEST_MINICPMO_4_5_MODEL", "openbmb/MiniCPM-o-4_5")
 _DEPLOY_CONFIG = get_deploy_config_path("minicpmo_4_5.yaml")
+_DUPLEX_BF16_DEPLOY_CONFIG = modify_stage_config(
+    _DUPLEX_DEPLOY_CONFIG,
+    updates={
+        "base_config": _DEPLOY_CONFIG,
+        "connectors.connector_of_shared_memory.extra.code2wav_bfloat16_attention_cache": True,
+    },
+)
 _RESULT_DIR = Path(
     os.environ.get(
         "ACC_BENCH_RESULT_DIR",
@@ -91,6 +104,20 @@ seed_test_params = [
 ]
 
 
+duplex_accuracy_test_params = [
+    *DUPLEX_TEST_PARAMS,
+    pytest.param(
+        OmniServerParams(
+            model=_MODEL,
+            stage_config_path=_DUPLEX_BF16_DEPLOY_CONFIG,
+            use_stage_cli=False,
+            server_args=list(_SEED_TTS_SERVER_ARGS),
+        ),
+        id="three-stage-single-gpu-bf16-attention-cache",
+    ),
+]
+
+
 def _require_vllm_cli() -> None:
     try:
         find_vllm_cli()
@@ -174,7 +201,7 @@ def test_minicpmo_4_5_seed_tts_wer_bench(omni_server) -> None:
 
 
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", DUPLEX_TEST_PARAMS, indirect=True)
+@pytest.mark.parametrize("omni_server", duplex_accuracy_test_params, indirect=True)
 def test_minicpmo_4_5_duplex_seed_tts_wer_bench(omni_server) -> None:
     """Gate Seed-TTS WER through the explicit Realtime TTS contract."""
     _require_vllm_cli()

tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Strict MiniCPM-o 4.5 Realtime duplex E2E scenario harness.
 
 This script is intentionally scenario-based instead of a generic chat client.
@@ -1351,6 +1354,9 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
         try:
             await ws.send(json.dumps(_session_update_event(args)))
             await _wait_for(state, lambda: state.count("session.created") > 0, timeout_s=20, label="session.created")
+            start_barrier = getattr(args, "start_barrier", None)
+            if start_barrier is not None:
+                await start_barrier.wait()
 
             turn_transcripts = _turn_transcripts(args.first_turn_transcript, turns=args.turns)
             turn_specs = list(zip(turn_transcripts, turn_durations, strict=True))

tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Concurrent MiniCPM-o Realtime duplex and resumable-session E2E driver."""
 
 from __future__ import annotations
@@ -14,6 +17,7 @@ from types import SimpleNamespace
 from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
 
 import websockets
+from websockets.asyncio.client import ClientConnection
 from websockets.exceptions import ConnectionClosed
 
 SCRIPT_DIR = Path(__file__).resolve().parent
@@ -42,6 +46,67 @@ async def run_demo(args):
     return await _scenario_module().run_demo(args)
 
 
+class _SynchronizedStartGate:
+    """One-shot, failure-propagating start gate for Python 3.10+."""
+
+    def __init__(self, parties: int, *, timeout_s: float):
+        if parties < 1:
+            raise ValueError("synchronized start parties must be positive")
+        if timeout_s <= 0:
+            raise ValueError("synchronized start timeout must be positive")
+        self._parties = parties
+        self._timeout_s = timeout_s
+        self._arrived = 0
+        self._released = False
+        self._failure: BaseException | None = None
+        self._event = asyncio.Event()
+        self._lock = asyncio.Lock()
+
+    async def wait(self) -> None:
+        async with self._lock:
+            if self._failure is not None:
+                raise RuntimeError(f"synchronized start aborted: {self._failure!r}") from self._failure
+            if self._released:
+                return
+            self._arrived += 1
+            if self._arrived == self._parties:
+                self._released = True
+                self._event.set()
+                return
+
+        try:
+            await asyncio.wait_for(self._event.wait(), timeout=self._timeout_s)
+        except asyncio.TimeoutError as exc:
+            failure = TimeoutError(
+                f"synchronized start timed out after {self._timeout_s:.1f}s "
+                f"({self._arrived}/{self._parties} sessions ready)"
+            )
+            await self.abort(failure)
+            if self._released and self._failure is None:
+                return
+            raise failure from exc
+
+        if self._failure is not None:
+            raise RuntimeError(f"synchronized start aborted: {self._failure!r}") from self._failure
+
+    async def abort(self, failure: BaseException) -> None:
+        async with self._lock:
+            if self._released or self._failure is not None:
+                return
+            self._failure = failure
+            self._event.set()
+
+
+async def _run_demo_with_start_gate(args: SimpleNamespace):
+    start_gate = getattr(args, "start_barrier", None)
+    try:
+        return await run_demo(args)
+    except BaseException as exc:
+        if start_gate is not None:
+            await start_gate.abort(exc)
+        raise
+
+
 def _with_resume_mode(url: str) -> str:
     parts = urlsplit(url)
     query = dict(parse_qsl(parts.query, keep_blank_values=True))
@@ -101,14 +166,24 @@ def _validate_semantic_isolation(
         return True
     if len(expected_tokens) != len(results):
         return False
-    for result, expected_token in zip(results, expected_tokens, strict=True):
+    normalized_expected_tokens = [token.strip().casefold() for token in expected_tokens]
+    if any(not token for token in normalized_expected_tokens):
+        return False
+    for result_index, (result, expected_token) in enumerate(zip(results, normalized_expected_tokens, strict=True)):
         details = result.get("transcript_integrity")
         transcripts = (
             [str(item.get("transcript", "")) for item in details if isinstance(item, dict)]
             if isinstance(details, list)
             else []
         )
-        if expected_token not in "".join(transcripts):
+        joined_transcripts = "".join(transcripts).casefold()
+        if expected_token not in joined_transcripts:
+            return False
+        if any(
+            other_token in joined_transcripts
+            for other_index, other_token in enumerate(normalized_expected_tokens)
+            if other_index != result_index
+        ):
             return False
     return True
 
@@ -130,10 +205,14 @@ async def _receive_until(ws, event_type: str, *, timeout_s: float) -> tuple[dict
     return await asyncio.wait_for(receive(), timeout=timeout_s)
 
 
+def _server_event_sequences(events: list[dict[str, object]]) -> list[int]:
+    return [sequence for event in events if isinstance(sequence := event.get("server_event_seq"), int)]
+
+
 async def _open_admission_session(
     args: argparse.Namespace,
     session_id: str,
-) -> tuple[object, dict[str, object]]:
+) -> tuple[ClientConnection, dict[str, object]]:
     url = _url_with_model(
         args.url,
         args.model,
@@ -159,7 +238,7 @@ async def _open_admission_session(
     return ws, created
 
 
-async def _close_admission_session(ws, *, timeout_s: float) -> None:
+async def _close_admission_session(ws: ClientConnection, *, timeout_s: float) -> None:
     await ws.send(json.dumps({"type": "session.close"}))
     await _receive_until(ws, "session.closed", timeout_s=timeout_s)
     await ws.close()
@@ -169,8 +248,8 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str,
     if limit < 1:
         raise ValueError("admission limit must be positive")
     prefix = f"admission-{uuid.uuid4().hex}"
-    accepted: list[tuple[object, dict[str, object]]] = []
-    replacement: tuple[object, dict[str, object]] | None = None
+    accepted: list[tuple[ClientConnection, dict[str, object]]] = []
+    replacement: tuple[ClientConnection, dict[str, object]] | None = None
     overflow_code = None
     try:
         for index in range(limit):
@@ -274,14 +353,7 @@ async def _resume_probe(
         generation = created.get("attachment_generation")
         if not isinstance(token, str) or not isinstance(incarnation, int):
             raise RuntimeError("session.created omitted resumable credentials")
-        last_seq = max(
-            (
-                event.get("server_event_seq", 0)
-                for event in first_events
-                if isinstance(event.get("server_event_seq"), int)
-            ),
-            default=0,
-        )
+        last_seq = max(_server_event_sequences(first_events), default=0)
 
     delay_s = args.expire_after_s if expect_expired else args.resume_after_ms / 1000
     if delay_s > 0:
@@ -324,7 +396,7 @@ async def _resume_probe(
         await second.send(json.dumps({"type": "session.close"}))
         closed, close_events = await _receive_until(second, "session.closed", timeout_s=args.timeout_s)
 
-    replay_sequences = [event["server_event_seq"] for event in replay if isinstance(event.get("server_event_seq"), int)]
+    replay_sequences = _server_event_sequences(replay)
     return {
         "ok": (
             resumed.get("session_id") == session_id
@@ -386,14 +458,7 @@ async def _takeover_probe(
         generation = created.get("attachment_generation")
         if not isinstance(token, str) or not isinstance(incarnation, int) or not isinstance(generation, int):
             raise RuntimeError("session.created omitted takeover credentials")
-        last_seq = max(
-            (
-                event.get("server_event_seq", 0)
-                for event in first_events
-                if isinstance(event.get("server_event_seq"), int)
-            ),
-            default=0,
-        )
+        last_seq = max(_server_event_sequences(first_events), default=0)
 
         second = await websockets.connect(_with_resume_mode(url), max_size=64 * 1024 * 1024)
         await second.send(
@@ -457,7 +522,9 @@ async def _takeover_probe(
         await first.close()
 
 
-def _demo_args(args: argparse.Namespace, index: int) -> SimpleNamespace:
+def _demo_args(
+    args: argparse.Namespace, index: int, start_barrier: _SynchronizedStartGate | None = None
+) -> SimpleNamespace:
     validation_mode = "response-required" if args.response_required else "model-policy"
     input_wav = args.session_input_wav[index] if args.session_input_wav else args.input_wav
     return SimpleNamespace(
@@ -477,14 +544,16 @@ def _demo_args(args: argparse.Namespace, index: int) -> SimpleNamespace:
         omit_transcript_hints=True,
         validation_mode=validation_mode,
         temperature=args.temperature,
-        scenario="sequential",
+        scenario=getattr(args, "scenario", "sequential"),
+        silence_ms=getattr(args, "silence_ms", 500),
         require_audio=args.response_required,
         require_distinct_inputs=False,
-        expect_empty_turn=[],
+        expect_empty_turn=list(getattr(args, "expect_empty_turn", []) or []),
         short_ack_ms=350,
         turns=args.turns,
         timeout_s=args.timeout_s,
         model_policy_settle_ms=args.model_policy_settle_ms,
+        start_barrier=start_barrier,
     )
 
 
@@ -511,8 +580,16 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]:
         )
     lifecycle_result = await run_lifecycle_probes(args)
 
+    start_barrier = (
+        _SynchronizedStartGate(
+            args.sessions,
+            timeout_s=min(float(args.timeout_s), 30.0),
+        )
+        if getattr(args, "synchronized_start", False)
+        else None
+    )
     session_results = await asyncio.gather(
-        *(run_demo(_demo_args(args, index)) for index in range(args.sessions)),
+        *(_run_demo_with_start_gate(_demo_args(args, index, start_barrier)) for index in range(args.sessions)),
         return_exceptions=True,
     )
     failures = [repr(result) for result in session_results if isinstance(result, BaseException)]
@@ -596,7 +673,14 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--turns", type=int, default=1)
     parser.add_argument("--first-turn-ms", type=int, default=1400)
     parser.add_argument("--turn-duration-ms", type=int, action="append", default=[])
+    parser.add_argument("--expect-empty-turn", type=int, action="append", default=[])
     parser.add_argument("--response-required", action="store_true")
+    parser.add_argument(
+        "--scenario",
+        choices=["sequential", "listen-only-overlap"],
+        default="sequential",
+    )
+    parser.add_argument("--silence-ms", type=int, default=500)
     parser.add_argument("--temperature", type=float, default=None)
     parser.add_argument("--disconnect-session-index", type=int)
     parser.add_argument("--takeover-session-index", type=int)
@@ -604,6 +688,11 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--expire-session-index", type=int)
     parser.add_argument("--expire-after-s", type=float, default=6.0)
     parser.add_argument("--verify-admission-limit", type=int)
+    parser.add_argument(
+        "--synchronized-start",
+        action="store_true",
+        help="Wait until every session is created before starting input.",
+    )
     parser.add_argument("--model-policy-settle-ms", type=int, default=600)
     parser.add_argument("--timeout-s", type=float, default=120.0)
     args = parser.parse_args()
@@ -615,6 +704,13 @@ def parse_args() -> argparse.Namespace:
         parser.error("provide exactly one --session-expected-token per session")
     if args.session_expected_token and not args.session_input_wav:
         parser.error("--session-expected-token requires --session-input-wav")
+    normalized_expected_tokens = [token.strip().casefold() for token in args.session_expected_token]
+    for token_index, token in enumerate(normalized_expected_tokens):
+        if any(
+            token in other_token or other_token in token
+            for other_token in normalized_expected_tokens[token_index + 1 :]
+        ):
+            parser.error("--session-expected-token values must not overlap after normalization")
     return args
 
 

tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
 import base64
 import importlib.util
@@ -758,6 +761,34 @@ def test_realtime_duplex_multi_session_rejects_cross_session_response_identity()
     )
 
 
+def test_realtime_duplex_multi_session_rejects_overlapping_expected_tokens(monkeypatch, capsys):
+    demo = _load_multi_demo_module()
+    monkeypatch.setattr(
+        demo.sys,
+        "argv",
+        [
+            "multi.py",
+            "--sessions",
+            "2",
+            "--input-wav",
+            "fallback.wav",
+            "--session-input-wav",
+            "a.wav",
+            "--session-input-wav",
+            "b.wav",
+            "--session-expected-token",
+            "cat",
+            "--session-expected-token",
+            "cater",
+        ],
+    )
+
+    with pytest.raises(SystemExit):
+        demo.parse_args()
+
+    assert "must not overlap after normalization" in capsys.readouterr().err
+
+
 def test_realtime_duplex_multi_session_reads_nested_terminal_identity():
     demo = _load_multi_demo_module()
 
@@ -768,6 +799,67 @@ def test_realtime_duplex_multi_session_reads_nested_terminal_identity():
     )
 
 
+def test_realtime_duplex_synchronized_start_gate_releases_all_participants():
+    demo = _load_multi_demo_module()
+
+    async def exercise():
+        gate = demo._SynchronizedStartGate(3, timeout_s=0.5)
+        released = []
+
+        async def participant(index):
+            await gate.wait()
+            released.append(index)
+
+        await asyncio.wait_for(
+            asyncio.gather(*(participant(index) for index in range(3))),
+            timeout=1.0,
+        )
+        return released
+
+    assert set(asyncio.run(exercise())) == {0, 1, 2}
+
+
+def test_realtime_duplex_synchronized_start_gate_aborts_waiters_on_failure(monkeypatch):
+    demo = _load_multi_demo_module()
+
+    async def fake_run_demo(args):
+        if args.fail:
+            raise RuntimeError("session creation failed")
+        await args.start_barrier.wait()
+        return {"ok": True}
+
+    monkeypatch.setattr(demo, "run_demo", fake_run_demo)
+
+    async def exercise():
+        gate = demo._SynchronizedStartGate(2, timeout_s=0.5)
+        return await asyncio.wait_for(
+            asyncio.gather(
+                demo._run_demo_with_start_gate(SimpleNamespace(start_barrier=gate, fail=True)),
+                demo._run_demo_with_start_gate(SimpleNamespace(start_barrier=gate, fail=False)),
+                return_exceptions=True,
+            ),
+            timeout=1.0,
+        )
+
+    failed, aborted = asyncio.run(exercise())
+    assert isinstance(failed, RuntimeError)
+    assert str(failed) == "session creation failed"
+    assert isinstance(aborted, RuntimeError)
+    assert "synchronized start aborted" in str(aborted)
+    assert "session creation failed" in str(aborted)
+
+
+def test_realtime_duplex_synchronized_start_gate_has_bounded_wait():
+    demo = _load_multi_demo_module()
+
+    async def exercise():
+        gate = demo._SynchronizedStartGate(2, timeout_s=0.01)
+        await gate.wait()
+
+    with pytest.raises(TimeoutError, match=r"1/2 sessions ready"):
+        asyncio.run(exercise())
+
+
 def test_realtime_duplex_demo_resolves_distinct_turn_inputs():
     demo = _load_demo_module()
     primary = Path("first.wav")

tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py

@@ -1,3 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import subprocess
+import sys
+from contextlib import contextmanager
 from pathlib import Path
 from types import SimpleNamespace
 
@@ -5,8 +11,10 @@ import pytest
 import torch
 import torch.nn as nn
 
+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,
+    _token2wav_sdpa_context,
     plan_token2wav_encode_slices,
     relpos_encode_token_budget,
 )
@@ -52,6 +60,9 @@ class _FakeEstimator(nn.Module):
         self.blocks = [_FakeBlock()]
         self.cfg_batches: list[int] = []
         self.speaker_order: list[list[float]] = []
+        self.attention_cache_dtypes: list[torch.dtype | None] = []
+        self.register_buffer("att_cache_buffer", torch.ones(1), persistent=False)
+        self.register_buffer("cnn_cache_buffer", torch.ones(1), persistent=False)
 
     def t_embedder(self, time):
         return time[:, None]
@@ -66,7 +77,8 @@ class _FakeEstimator(nn.Module):
         cnn_out,
         att_out,
     ):
-        del time, mask, cnn_cache, att_cache
+        del time, mask, cnn_cache
+        self.attention_cache_dtypes.append(att_cache.dtype if isinstance(att_cache, torch.Tensor) else None)
         self.cfg_batches.append(inputs.shape[0])
         self.speaker_order.append(inputs[:, 2, 0].tolist())
         marker = inputs[:, 1, 0]
@@ -100,16 +112,17 @@ class _FakeHiFT(nn.Module):
         super().__init__()
         self.calls: list[int] = []
 
-    def forward(self, mel, source):
+    def forward(self, batch, device):
+        mel = batch["speech_feat"].transpose(1, 2).to(device)
+        return self.inference(mel, mel.new_zeros((mel.shape[0], 1, 0)))
+
+    def inference(self, mel, source):
         del source
         self.calls.append(mel.shape[0])
         speech = mel[:, 0].repeat_interleave(3, dim=1)
         generated_source = speech[:, None]
         return speech, generated_source
 
-    def inference(self, mel, source):
-        return self.forward(mel, source)
-
 
 class _FakeToken2Wav:
     def __init__(self):
@@ -140,13 +153,14 @@ class _FakeToken2Wav:
         raise AssertionError("sequential __call__ fallback must never be called")
 
 
-def _config(minimum: int = 1):
+def _config(minimum: int = 1, initial: int = 0):
     return SimpleNamespace(
         model_config=SimpleNamespace(
             model="/fake/model",
             stage_connector_config={
                 "extra": {
                     "code2wav_min_batch_size": minimum,
+                    "code2wav_initial_batch_size": initial,
                     "prompt_cache_id": "shared",
                     "prompt_wav": "/fake/prompt.wav",
                 }
@@ -155,14 +169,45 @@ def _config(minimum: int = 1):
     )
 
 
-def _model():
+def _model(initial: int = 0, minimum: int = 1):
     token2wav = _FakeToken2Wav()
     backend = BatchedToken2Wav(token2wav)
-    model = MiniCPMO45Code2Wav(vllm_config=_config())
+    _enable_fake_ragged_kernel(backend)
+    model = MiniCPMO45Code2Wav(vllm_config=_config(minimum=minimum, initial=initial))
     model.backend = backend
     return model, token2wav
 
 
+def _enable_fake_ragged_kernel(adapter: BatchedToken2Wav) -> None:
+    """Keep fake-model tests focused on ragged grouping and row mapping."""
+    estimator = adapter.flow.decoder.estimator
+    estimator.in_proj = nn.Identity()
+
+    def fake_ragged_kernel(
+        estimator,
+        estimator_input,
+        time_embedding,
+        attn_mask,
+        cnn_cache,
+        att_cache,
+        cnn_cache_buffer,
+        att_cache_buffer,
+        valid_lengths,
+    ):
+        del valid_lengths
+        return estimator.blocks_forward_chunk(
+            estimator_input,
+            time_embedding,
+            attn_mask,
+            cnn_cache,
+            att_cache,
+            cnn_cache_buffer,
+            att_cache_buffer,
+        )
+
+    adapter._blocks_forward_chunk_ragged = fake_ragged_kernel  # type: ignore[method-assign]
+
+
 def _info(
     request_id: str,
     chunk_seq: int,
@@ -256,6 +301,400 @@ def test_decode_batch_splits_overlong_prefill_inside_relpos_budget():
     assert audios[0].numel() > 0
 
 
+def test_ragged_overlong_prefill_uses_relpos_safe_exact_slices():
+    token2wav = _FakeToken2Wav()
+    token2wav.flow.encoder.pre_lookahead_layer = SimpleNamespace(pre_lookahead_len=3)
+    adapter = BatchedToken2Wav(token2wav)
+    _enable_fake_ragged_kernel(adapter)
+    prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+    states = adapter.setup_batch(prompt, 2)
+    setup_encodes = len(token2wav.flow.encoder.calls)
+    adapter._max_encode_token_frames = lambda _states: 4
+
+    tokens = [
+        torch.arange(10, dtype=torch.long),
+        torch.arange(4, dtype=torch.long),
+    ]
+    last_chunks = [True, False]
+    audios, next_states = adapter.decode_ragged_batch(
+        tokens,
+        prompt,
+        states,
+        last_chunks=last_chunks,
+    )
+
+    assert token2wav.flow.encoder.calls[setup_encodes:] == [1, 1, 1, 1]
+    assert token2wav.flow.encoder.last_chunk_calls[setup_encodes:] == [False, False, True, False]
+    assert len(audios) == len(next_states) == 2
+    assert all(audio.numel() > 0 for audio in audios)
+
+    for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)):
+        reference_token2wav = _FakeToken2Wav()
+        reference_token2wav.flow.encoder.pre_lookahead_layer = SimpleNamespace(pre_lookahead_len=3)
+        reference_adapter = BatchedToken2Wav(reference_token2wav)
+        reference_prompt = reference_adapter.prepare_prompt("shared", "/fake/prompt.wav")
+        reference_states = reference_adapter.setup_batch(reference_prompt, 1)
+        reference_adapter._max_encode_token_frames = lambda _states: 4
+        reference_audio, reference_state = reference_adapter.decode_batch(
+            row_tokens.unsqueeze(0),
+            reference_prompt,
+            reference_states,
+            last_chunk=last_chunk,
+        )
+        torch.testing.assert_close(audios[row], reference_audio[0])
+        for cache_name, expected in reference_state[0].flow_cache.items():
+            torch.testing.assert_close(
+                next_states[row].flow_cache[cache_name],
+                expected,
+            )
+        for cache_name, expected in reference_state[0].hift_cache.items():
+            torch.testing.assert_close(
+                next_states[row].hift_cache[cache_name],
+                expected,
+            )
+
+
+def test_ragged_diffusion_batches_eight_rows_and_preserves_exact_results():
+    token2wav = _FakeToken2Wav()
+    adapter = BatchedToken2Wav(token2wav)
+    _enable_fake_ragged_kernel(adapter)
+    prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+    states = adapter.setup_batch(prompt, 8)
+    token2wav.flow.encoder.calls.clear()
+    token2wav.flow.decoder.estimator.cfg_batches.clear()
+    token2wav.hift.calls.clear()
+    tokens = [
+        *(torch.tensor([10 + row, 11 + row, 12 + row]) for row in range(6)),
+        torch.tensor([20, 21]),
+        torch.tensor([30, 31]),
+    ]
+    last_chunks = [False] * 6 + [True, True]
+
+    audios, next_states = adapter.decode_ragged_batch(
+        tokens,
+        prompt,
+        states,
+        last_chunks=last_chunks,
+    )
+
+    assert token2wav.flow.encoder.calls == [6, 2]
+    assert token2wav.flow.decoder.estimator.cfg_batches == [16, 16]
+    assert token2wav.hift.calls == [6, 2]
+    assert len(audios) == len(next_states) == 8
+
+    for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)):
+        reference_adapter = BatchedToken2Wav(_FakeToken2Wav())
+        reference_prompt = reference_adapter.prepare_prompt("shared", "/fake/prompt.wav")
+        reference_states = reference_adapter.setup_batch(reference_prompt, 1)
+        reference_audio, reference_state = reference_adapter.decode_batch(
+            row_tokens.unsqueeze(0),
+            reference_prompt,
+            reference_states,
+            last_chunk=last_chunk,
+        )
+        torch.testing.assert_close(audios[row], reference_audio[0])
+        for cache_name, expected in reference_state[0].flow_cache.items():
+            torch.testing.assert_close(
+                next_states[row].flow_cache[cache_name],
+                expected,
+            )
+        for cache_name, expected in reference_state[0].hift_cache.items():
+            torch.testing.assert_close(
+                next_states[row].hift_cache[cache_name],
+                expected,
+            )
+
+
+def test_npu_patch_preserves_ragged_estimator_contract():
+    script = r"""
+import torch
+
+from tests.model_executor.models.minicpmo_4_5.test_code2wav_batching import (
+    _FakeToken2Wav,
+    _enable_fake_ragged_kernel,
+)
+from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import BatchedToken2Wav
+import vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav as npu_patch
+
+
+class _FailGraphRunner:
+    def run(self, *args, **kwargs):
+        raise AssertionError("ragged decode must bypass the exact-shape NPU Graph")
+
+
+npu_patch.apply_minicpmo_4_5_code2wav_patch()
+adapter = BatchedToken2Wav(_FakeToken2Wav())
+_enable_fake_ragged_kernel(adapter)
+prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+states = adapter.setup_batch(prompt, 2)
+npu_patch._backend_graph_runners[adapter] = _FailGraphRunner()
+audios, next_states = adapter.decode_ragged_batch(
+    [
+        torch.tensor([10, 11, 12]),
+        torch.tensor([20, 21]),
+    ],
+    prompt,
+    states,
+    last_chunks=[False, True],
+)
+assert len(audios) == len(next_states) == 2
+"""
+    completed = subprocess.run(
+        [sys.executable, "-c", script],
+        cwd=Path(__file__).resolve().parents[4],
+        capture_output=True,
+        text=True,
+    )
+
+    assert completed.returncode == 0, completed.stderr
+
+
+def test_decode_cfm_enters_platform_sdpa_context(monkeypatch):
+    entered: list[str] = []
+
+    @contextmanager
+    def recording_context():
+        entered.append("enter")
+        yield
+        entered.append("exit")
+
+    monkeypatch.setattr(
+        batched_token2wav_module,
+        "_token2wav_sdpa_context",
+        lambda _device: recording_context(),
+    )
+    adapter = BatchedToken2Wav(_FakeToken2Wav())
+    _enable_fake_ragged_kernel(adapter)
+    adapter._decode_cfm(
+        torch.ones((2, 1, 2)),
+        torch.ones((2, 1)),
+        torch.zeros((2, 1, 2)),
+        cnn_cache=None,
+        att_cache=None,
+        valid_lengths=[2, 1],
+    )
+
+    assert entered == ["enter", "exit"]
+
+
+def test_ragged_decode_bypasses_exact_shape_accelerators():
+    adapter = BatchedToken2Wav(_FakeToken2Wav())
+    _enable_fake_ragged_kernel(adapter)
+
+    class ForbiddenAccelerator:
+        def replay(self, *args):
+            pytest.fail("ragged decode must not use the exact-shape CFM graph")
+
+        def step(self, **kwargs):
+            pytest.fail("ragged decode must not use the exact-shape TensorRT stepper")
+
+    accelerator = ForbiddenAccelerator()
+    adapter._cfm_graph_wrapper = accelerator
+    adapter._trt_stepper = accelerator
+    adapter._decode_cfm(
+        torch.ones((2, 1, 2)),
+        torch.ones((2, 1)),
+        torch.zeros((2, 1, 2)),
+        cnn_cache=None,
+        att_cache=None,
+        valid_lengths=[2, 1],
+    )
+
+
+def test_ragged_decode_fails_when_estimator_has_no_ragged_kernel():
+    adapter = BatchedToken2Wav(_FakeToken2Wav())
+
+    with pytest.raises(
+        RuntimeError,
+        match=r'"reason":"ragged_kernel_unavailable"',
+    ):
+        adapter._decode_cfm(
+            torch.ones((2, 1, 2)),
+            torch.ones((2, 1)),
+            torch.zeros((2, 1, 2)),
+            cnn_cache=None,
+            att_cache=None,
+            valid_lengths=[2, 1],
+        )
+
+
+def test_cfm_graph_receives_bfloat16_cache_in_compute_dtype():
+    adapter = BatchedToken2Wav(_FakeToken2Wav(), bfloat16_attention_cache=True)
+    seen_dtypes: list[torch.dtype] = []
+
+    class RecordingGraph:
+        def replay(self, estimator_input, time_embedding, cnn_cache, att_cache, cnn_out, att_out):
+            del time_embedding, cnn_cache
+            seen_dtypes.append(att_cache.dtype)
+            return estimator_input[:, :1], cnn_out, att_out
+
+    adapter._cfm_graph_wrapper = RecordingGraph()
+    adapter._estimator_step(
+        adapter.flow.decoder.estimator,
+        x=torch.ones((2, 1, 2)),
+        mu=torch.ones((2, 1, 2)),
+        time=torch.zeros(2),
+        speakers=torch.ones((2, 1)),
+        cond=torch.zeros((2, 1, 2)),
+        cnn_cache=torch.zeros((1, 2, 2, 1)),
+        att_cache=torch.zeros((1, 2, 1, 1, 2), dtype=torch.bfloat16),
+    )
+
+    assert seen_dtypes == [torch.float32]
+
+
+def test_npu_sdpa_context_uses_platform_patch(monkeypatch):
+    entered: list[str] = []
+
+    @contextmanager
+    def recording_context():
+        entered.append("enter")
+        yield
+        entered.append("exit")
+
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms.npu.models.step_audio2_token2wav",
+        SimpleNamespace(npu_token2wav_sdpa_context=recording_context),
+    )
+    with _token2wav_sdpa_context(SimpleNamespace(type="npu")):
+        assert entered == ["enter"]
+
+    assert entered == ["enter", "exit"]
+
+
+def test_ragged_outputs_fail_closed_on_middle_hole():
+    adapter = BatchedToken2Wav(_FakeToken2Wav())
+    prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+    states = adapter.setup_batch(prompt, 3)
+    audios: list[torch.Tensor | None] = [
+        torch.zeros(1),
+        None,
+        torch.ones(1),
+    ]
+    next_states = [states[0], None, states[2]]
+
+    with pytest.raises(
+        RuntimeError,
+        match=r'"reason":"incomplete_ragged_output","audio_rows":\[1\],"state_rows":\[1\]',
+    ):
+        adapter._require_complete_ragged_outputs(audios, next_states)
+
+
+@pytest.mark.parametrize(
+    "bfloat16_attention_cache",
+    [False, True],
+    ids=["fp32-cache", "bf16-cache"],
+)
+def test_ragged_dit_two_step_decode_matches_per_row_exact_decode(
+    bfloat16_attention_cache: bool,
+):
+    from cosyvoice2.flow.decoder_dit import DiT
+
+    torch.manual_seed(17)
+    template = DiT(
+        in_channels=4,
+        out_channels=1,
+        depth=2,
+        num_heads=2,
+        head_dim=2,
+        hidden_size=4,
+    ).eval()
+    template_state = {name: value.detach().clone() for name, value in template.state_dict().items()}
+    rand_noise = torch.randn(1, 1, 32)
+
+    def make_adapter() -> BatchedToken2Wav:
+        token2wav = _FakeToken2Wav()
+        estimator = DiT(
+            in_channels=4,
+            out_channels=1,
+            depth=2,
+            num_heads=2,
+            head_dim=2,
+            hidden_size=4,
+        ).eval()
+        estimator.load_state_dict(template_state)
+        token2wav.flow.decoder.estimator = estimator
+        token2wav.flow.decoder.rand_noise = rand_noise.clone()
+        return BatchedToken2Wav(
+            token2wav,
+            bfloat16_attention_cache=bfloat16_attention_cache,
+        )
+
+    ragged_adapter = make_adapter()
+    ragged_prompt = ragged_adapter.prepare_prompt("shared", "/fake/prompt.wav")
+    ragged_states = ragged_adapter.setup_batch(ragged_prompt, 2)
+
+    exact_adapters = [make_adapter(), make_adapter()]
+    exact_prompts = [adapter.prepare_prompt("shared", "/fake/prompt.wav") for adapter in exact_adapters]
+    exact_states = [
+        adapter.setup_batch(prompt, 1) for adapter, prompt in zip(exact_adapters, exact_prompts, strict=True)
+    ]
+    tolerance = (2e-2, 2e-3) if bfloat16_attention_cache else (1e-5, 1e-6)
+    steps = [
+        (
+            [torch.tensor([10, 11, 12]), torch.tensor([20, 21, 22])],
+            [False, True],
+        ),
+        (
+            [torch.tensor([13, 14]), torch.tensor([22, 23, 24])],
+            [True, True],
+        ),
+    ]
+    first_step_history_width: int | None = None
+
+    for step, (tokens, last_chunks) in enumerate(steps):
+        if step == 1:
+            assert first_step_history_width is not None
+            assert all(
+                int(state.flow_cache["estimator_att_cache"].shape[4]) == first_step_history_width
+                for state in ragged_states
+            )
+        ragged_audios, ragged_states = ragged_adapter.decode_ragged_batch(
+            tokens,
+            ragged_prompt,
+            ragged_states,
+            last_chunks=last_chunks,
+        )
+        for row, (adapter, prompt) in enumerate(zip(exact_adapters, exact_prompts, strict=True)):
+            exact_audios, exact_states[row] = adapter.decode_batch(
+                tokens[row].unsqueeze(0),
+                prompt,
+                exact_states[row],
+                last_chunk=last_chunks[row],
+            )
+            torch.testing.assert_close(
+                ragged_audios[row],
+                exact_audios[0],
+                rtol=tolerance[0],
+                atol=tolerance[1],
+            )
+            for cache_name, expected in exact_states[row][0].flow_cache.items():
+                actual = ragged_states[row].flow_cache[cache_name]
+                assert actual.dtype == expected.dtype
+                torch.testing.assert_close(
+                    actual,
+                    expected,
+                    rtol=tolerance[0],
+                    atol=tolerance[1],
+                )
+            for cache_name, expected in exact_states[row][0].hift_cache.items():
+                torch.testing.assert_close(
+                    ragged_states[row].hift_cache[cache_name],
+                    expected,
+                    rtol=tolerance[0],
+                    atol=tolerance[1],
+                )
+
+        estimator_att = ragged_states[0].flow_cache["estimator_att_cache"]
+        assert estimator_att.shape[4] > 0
+        if step == 0:
+            first_step_history_width = int(estimator_att.shape[4])
+        else:
+            assert first_step_history_width is not None
+            assert int(estimator_att.shape[4]) > first_step_history_width
+
+
 def test_fade_in_out_limits_overlap_to_available_previous_audio():
     speech = torch.arange(6, dtype=torch.float32).reshape(1, -1)
     previous = torch.full((1, 3), 2.0)
@@ -295,6 +734,69 @@ def test_estimator_cache_stack_split_round_trip_preserves_cfg_rows():
         )
 
 
+def test_bfloat16_estimator_attention_cache_materializes_only_current_timestep():
+    token2wav = _FakeToken2Wav()
+    adapter = BatchedToken2Wav(token2wav, bfloat16_attention_cache=True)
+    prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+    states = adapter.setup_batch(prompt, 2)
+
+    assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in states)
+    assert (
+        states[0].flow_cache["estimator_att_cache"].data_ptr() != states[1].flow_cache["estimator_att_cache"].data_ptr()
+    )
+
+    stacked = adapter._stack_flow_cache(states)
+    assert stacked["estimator_att_cache"].dtype == torch.bfloat16
+    restored = adapter._split_flow_cache(stacked, 2)
+    for original, round_tripped in zip(states, restored, strict=True):
+        assert round_tripped["estimator_att_cache"].dtype == torch.bfloat16
+        torch.testing.assert_close(
+            round_tripped["estimator_att_cache"],
+            original.flow_cache["estimator_att_cache"],
+        )
+
+    _, next_states = adapter.decode_batch(
+        torch.tensor([[10, 11], [20, 21]]),
+        prompt,
+        states,
+        last_chunk=False,
+    )
+    assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in next_states)
+    assert token2wav.flow.decoder.estimator.attention_cache_dtypes[-2:] == [torch.float32, torch.float32]
+
+
+def test_bfloat16_estimator_attention_cache_materializes_each_timestep():
+    adapter = BatchedToken2Wav(_FakeToken2Wav(), bfloat16_attention_cache=True)
+
+    _, estimator_cnn, estimator_att = adapter._decode_cfm(
+        torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]),
+        torch.ones((2, 1)),
+        torch.zeros((2, 1, 2)),
+        cnn_cache=None,
+        att_cache=None,
+    )
+
+    assert estimator_cnn.dtype == torch.float32
+    assert estimator_att.dtype == torch.bfloat16
+    assert estimator_att.shape[:3] == (2, 1, 4)
+
+
+def test_float32_estimator_attention_cache_materializes_each_timestep():
+    adapter = BatchedToken2Wav(_FakeToken2Wav())
+
+    _, estimator_cnn, estimator_att = adapter._decode_cfm(
+        torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]),
+        torch.ones((2, 1)),
+        torch.zeros((2, 1, 2)),
+        cnn_cache=None,
+        att_cache=None,
+    )
+
+    assert estimator_cnn.dtype == torch.float32
+    assert estimator_att.dtype == torch.float32
+    assert estimator_att.shape[:3] == (2, 1, 4)
+
+
 def test_model_preserves_output_slots_and_prefers_runtime_codes():
     model, token2wav = _model()
     output = _forward(
@@ -386,6 +888,28 @@ def test_initial_empty_segment_marker_initializes_stream_without_audio():
     assert "duplex" in model._states
 
 
+def test_initial_empty_segment_markers_respect_initial_batch_limit():
+    model, token2wav = _model(initial=2)
+    names = ["a", "b", "c", "d", "e"]
+    boundaries = []
+    for name in names:
+        boundary = _info(name, 0, [])
+        boundary["meta"].update(
+            {
+                "code_flat_numel": 0,
+                "tts_is_last_chunk": True,
+                "turn_end": False,
+            }
+        )
+        boundaries.append(boundary)
+
+    output = _forward(model, boundaries)
+
+    assert token2wav.flow.encoder.calls == [2, 2, 1]
+    assert [audio.numel() for audio in output.multimodal_outputs["model_outputs"]] == [0] * len(names)
+    assert set(model._states) == set(names)
+
+
 def test_shared_runtime_prompt_recreates_missing_file_before_second_owner(tmp_path, monkeypatch):
     monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
     model, _ = _model()
@@ -647,11 +1171,56 @@ def test_singleton_and_mixed_shape_buckets_use_same_batched_backend_without_fall
     assert token2wav.hift.calls[-2:] == [1, 1]
 
 
+def test_initial_batch_limit_allows_later_full_batch():
+    model, token2wav = _model(initial=4)
+    names = [f"request-{index}" for index in range(8)]
+
+    for chunk_seq in (0, 1):
+        infos = [_info(name, chunk_seq, [1, 2]) for name in names]
+        _forward(model, infos)
+        assert token2wav.hift.calls[-2:] == [4, 4]
+
+    token2wav.hift.calls.clear()
+    _forward(model, [_info(name, 2, [1, 2]) for name in names])
+    assert token2wav.hift.calls == [8]
+
+
+def test_initial_batch_partition_preserves_minimum():
+    model, token2wav = _model(initial=4, minimum=3)
+    names = [f"request-{index}" for index in range(6)]
+
+    _forward(model, [_info(name, 0, [1, 2]) for name in names])
+
+    assert token2wav.hift.calls == [3, 3]
+
+
+def test_initial_batch_partition_rejects_impossible_remainder():
+    model, _ = _model(initial=4, minimum=3)
+    names = [f"request-{index}" for index in range(5)]
+
+    with pytest.raises(RuntimeError, match="initial_batch_partition_below_minimum"):
+        _forward(model, [_info(name, 0, [1, 2]) for name in names])
+
+
+@pytest.mark.parametrize("chunk_seqs", [[0, 2, 2, 2], [0, 0, 0, 2]])
+def test_initial_steady_partition_rejects_undersized_wave(chunk_seqs):
+    model, _ = _model(initial=4, minimum=3)
+    bucket = [SimpleNamespace(chunk_seq=chunk_seq) for chunk_seq in chunk_seqs]
+
+    with pytest.raises(RuntimeError, match="decode_wave_below_minimum"):
+        list(model._iter_decode_batches([bucket]))
+
+
 def test_backend_failure_does_not_commit_any_request_state(monkeypatch):
     model, _ = _model()
     _forward(
         model,
-        [_info(name, 0, [index + 1, index + 2]) for index, name in enumerate(("a", "b", "c", "d"))],
+        [
+            _info("a", 0, [1, 2], cache_epoch=0),
+            _info("b", 0, [2, 3], cache_epoch=0),
+            _info("c", 0, [3, 4], cache_epoch=1),
+            _info("d", 0, [4, 5], cache_epoch=1),
+        ],
     )
     before = dict(model._states)
     original = model.backend.decode_batch
@@ -671,8 +1240,8 @@ def test_backend_failure_does_not_commit_any_request_state(monkeypatch):
             [
                 _info("a", 1, [5, 6]),
                 _info("b", 1, [7, 8]),
-                _info("c", 1, [9, 10, 11]),
-                _info("d", 1, [12, 13, 14]),
+                _info("c", 1, [9, 10], cache_epoch=1),
+                _info("d", 1, [12, 13], cache_epoch=1),
             ],
         )
     assert call_count == 2

tests/platforms/npu/test_graph_tools.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
 
 from contextlib import contextmanager
 from types import SimpleNamespace
@@ -196,6 +196,7 @@ def test_code2wav_patch_reads_stage_additional_config(
 
     model = SimpleNamespace(
         backend=None,
+        _extra_config=lambda: {},
         vllm_config=SimpleNamespace(
             additional_config={
                 "code2wav_enable_npu_graph": enabled,
@@ -224,6 +225,47 @@ def test_code2wav_patch_reads_stage_additional_config(
         assert model.backend not in code2wav_patch._backend_graph_runners
 
 
+@pytest.mark.parametrize("value", [True, "true", "false", "0", "off"])
+def test_code2wav_patch_rejects_bfloat16_attention_cache(monkeypatch, value):
+    built = False
+
+    model = SimpleNamespace(
+        backend=None,
+        _extra_config=lambda: {"code2wav_bfloat16_attention_cache": value},
+    )
+
+    def build_backend(_instance):
+        nonlocal built
+        built = True
+
+    monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend)
+
+    with pytest.raises(ValueError, match="currently supported only on CUDA"):
+        code2wav_patch._patched_build_backend(model)
+
+    assert not built
+
+
+@pytest.mark.parametrize("value", [False, None, 0, ""])
+def test_code2wav_patch_allows_disabled_bfloat16_attention_cache(monkeypatch, value):
+    built = False
+
+    model = SimpleNamespace(
+        backend=None,
+        _extra_config=lambda: {"code2wav_bfloat16_attention_cache": value},
+    )
+
+    def build_backend(instance):
+        nonlocal built
+        built = True
+        instance.backend = object()
+
+    monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend)
+    code2wav_patch._patched_build_backend(model)
+
+    assert built
+
+
 @pytest.mark.parametrize("with_cache", [False, True])
 def test_code2wav_estimator_graph_dispatch_is_platform_owned(monkeypatch, with_cache):
     calls = []

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 from collections import defaultdict
@@ -463,7 +466,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     request.resumable = False
                     stopped = True
 
+            confirmed_num_computed_tokens = None
             if stopped:
+                if self.chunk_transfer_adapter is not None:
+                    confirmed_num_computed_tokens = self.chunk_transfer_adapter._confirmed_num_computed_tokens(request)
                 if model_runner_output.routed_experts is not None:
                     routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
 
@@ -555,6 +561,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     inter_stage_output,
                     request,
                     is_segment_finished,
+                    confirmed_num_computed_tokens=confirmed_num_computed_tokens,
                 )
 
         self._remove_stopped_requests_from_queues(

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import time

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+import copy
 import importlib
 import time
 from collections import defaultdict, deque
@@ -10,6 +11,7 @@ from typing import Any
 import torch
 from vllm.v1.metrics.stats import PrefillStats
 from vllm.v1.request import Request, RequestStatus
+from vllm.v1.utils import ConstantList
 
 from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload
 
@@ -118,6 +120,38 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         num_placeholders = int(getattr(request, "num_output_placeholders", 0) or 0)
         return max(0, num_computed - num_placeholders)
 
+    @staticmethod
+    def _snapshot_processor_request(request: Request) -> Request:
+        """Snapshot mutable processor inputs at save-queue admission time."""
+
+        def snapshot_container(value: Any) -> Any:
+            if isinstance(value, dict):
+                return {key: snapshot_container(item) for key, item in value.items()}
+            if isinstance(value, list):
+                return [snapshot_container(item) for item in value]
+            if isinstance(value, tuple):
+                return tuple(snapshot_container(item) for item in value)
+            return value
+
+        snapshot = copy.copy(request)
+        for name in (
+            "additional_information",
+            "prompt_token_ids",
+        ):
+            if hasattr(request, name):
+                setattr(snapshot, name, snapshot_container(getattr(request, name)))
+
+        for private_name, public_name in (
+            ("_all_token_ids", "all_token_ids"),
+            ("_output_token_ids", "output_token_ids"),
+        ):
+            token_ids = getattr(request, private_name, None)
+            if isinstance(token_ids, list):
+                frozen_ids = token_ids.copy()
+                setattr(snapshot, private_name, frozen_ids)
+                setattr(snapshot, public_name, ConstantList(frozen_ids))
+        return snapshot
+
     @staticmethod
     def _refresh_generation_chunk_prefill_state(request: Request) -> None:
         request.num_prompt_tokens = len(request.prompt_token_ids)
@@ -169,6 +203,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         multimodal_output: dict[str, Any] | None = None,
         request: Request | None = None,
         is_segment_finished: bool = False,
+        confirmed_num_computed_tokens: int | None = None,
     ):
         """Build and enqueue one chunk for asynchronous sending.
 
@@ -186,10 +221,13 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             multimodal_output: Per-request multimodal output dictionary
             request: Request object
             is_segment_finished: whether the segment of request is finished
+            confirmed_num_computed_tokens: committed token count captured
+                before a streaming transition can mutate ``request``
         """
         is_finished = request.is_finished() and not request.resumable
 
-        confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request)
+        if confirmed_num_computed_tokens is None:
+            confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request)
 
         # If the request is preempted, skip the already saved chunks.
         if confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(request.external_req_id, 0):
@@ -204,11 +242,17 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.requests_num_chunks_sent[request.external_req_id] = confirmed_num_computed_tokens
         task = {
             "multimodal_output": multimodal_output,
-            "request": request,
+            "request": (
+                self._snapshot_processor_request(request) if self.custom_process_next_stage_input_func else request
+            ),
             "is_finished": is_finished,
             "is_segment_finished": is_segment_finished,
         }
         self._pending_save_reqs.append(task)
+        if is_segment_finished:
+            # The queued FIFO item now owns the old segment. Start the next
+            # segment's deduplication watermark before the worker sends it.
+            self.requests_num_chunks_sent.pop(request.external_req_id, None)
         with self._save_cond:
             self._save_cond.notify()
 
@@ -425,7 +469,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         if is_segment_finished:
             self.code_prompt_token_ids.pop(external_req_id, None)
-            self.requests_num_chunks_sent.pop(external_req_id, None)
             self.ramp_chunk_count.pop(external_req_id, None)
             cached_ic = getattr(self, "_cached_ic", None)
             if cached_ic is not None:

vllm_omni/model_executor/models/minicpmo_4_5/batched_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
 """Strict, state-explicit batching for MiniCPM-o 4.5 Token2wav."""
 
 from __future__ import annotations
@@ -7,7 +7,7 @@ from __future__ import annotations
 from collections.abc import Mapping
 from contextlib import nullcontext
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, cast
 
 import torch
 import torch.nn as nn
@@ -94,6 +94,17 @@ def _autocast_disabled(device: torch.device):
         return nullcontext()
 
 
+def _token2wav_sdpa_context(device: torch.device):
+    if device.type != "npu":
+        return nullcontext()
+
+    from vllm_omni.platforms.npu.models.step_audio2_token2wav import (
+        npu_token2wav_sdpa_context,
+    )
+
+    return npu_token2wav_sdpa_context()
+
+
 def tensor_signature(value: torch.Tensor) -> tuple[tuple[int, ...], str, str]:
     return tuple(value.shape), str(value.dtype), value.device.type
 
@@ -152,6 +163,7 @@ class BatchedToken2Wav(nn.Module):
         connector_config: Mapping[str, int] | None = None,
         hift_graph_config: Mapping[str, Any] | None = None,
         cfm_graph_config: Mapping[str, Any] | None = None,
+        bfloat16_attention_cache: bool = False,
     ):
         super().__init__()
         self._token2wav = token2wav
@@ -185,6 +197,10 @@ class BatchedToken2Wav(nn.Module):
             del speech, source
             torch.accelerator.empty_cache()
         self.float16 = bool(token2wav.float16)
+        self._estimator_att_compute_dtype = torch.float16 if self.float16 else torch.float32
+        self._estimator_att_cache_dtype = (
+            torch.bfloat16 if bfloat16_attention_cache else self._estimator_att_compute_dtype
+        )
         self.n_timesteps = int(token2wav.n_timesteps)
         self.mel_cache_len = int(token2wav.mel_cache_len)
         self.source_cache_len = int(token2wav.source_cache_len)
@@ -310,9 +326,7 @@ class BatchedToken2Wav(nn.Module):
         return max(1, int(stride)) if stride is not None else 2
 
     def _max_encode_token_frames(self, states: list[BatchedToken2WavState]) -> int:
-        att = None
-        if states:
-            att = self._stack_flow_cache(states).get("conformer_att_cache")
+        att = states[0].flow_cache.get("conformer_att_cache") if states else None
         offset1 = int(att.shape[3] // 2) if att is not None else 0
         lookahead = self._pre_lookahead_len() or 0
         return relpos_encode_token_budget(
@@ -382,8 +396,10 @@ class BatchedToken2Wav(nn.Module):
         cond: torch.Tensor,
         cnn_cache: torch.Tensor | None,
         att_cache: torch.Tensor | None,
+        attn_mask: torch.Tensor | None = None,
+        valid_lengths: list[int] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        if self._trt_stepper is not None:
+        if self._trt_stepper is not None and valid_lengths is None:
             out, new_cnn, new_att = self._trt_stepper.step(
                 x=x,
                 mu=mu,
@@ -399,32 +415,139 @@ class BatchedToken2Wav(nn.Module):
         speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width)
         estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1)
         cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache)
-        if self._cfm_graph_wrapper is not None:
-            if cnn_cache is None:
-                cnn_cache = torch.zeros_like(cnn_out)
-            if att_cache is None:
-                att_cache = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
+        old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks)
+        old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks)
+        if isinstance(old_att, torch.Tensor) and old_att.dtype != estimator_input.dtype:
+            old_att = old_att.to(dtype=estimator_input.dtype)
+        if self._cfm_graph_wrapper is not None and valid_lengths is None:
+            graph_cnn = torch.zeros_like(cnn_out) if cnn_cache is None else cnn_cache
+            graph_att = (
+                estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
+                if att_cache is None
+                else old_att
+            )
             return self._cfm_graph_wrapper.replay(
                 estimator_input,
                 time_embedding,
-                cnn_cache,
-                att_cache,
+                graph_cnn,
+                graph_att,
+                cnn_out,
+                att_out,
+            )
+        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(
+                estimator,
+                estimator_input,
+                time_embedding,
+                attn_mask,
+                old_cnn,
+                old_att,
+                cnn_out,
+                att_out,
+                valid_lengths,
+            )
+        else:
+            result = estimator.blocks_forward_chunk(
+                estimator_input,
+                time_embedding,
+                attn_mask,
+                old_cnn,
+                old_att,
                 cnn_out,
                 att_out,
             )
-        old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks)
-        old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks)
-        result = estimator.blocks_forward_chunk(
-            estimator_input,
-            time_embedding,
-            None,
-            old_cnn,
-            old_att,
-            cnn_out,
-            att_out,
-        )
         return result, cnn_out, att_out
 
+    @staticmethod
+    def _gather_causal_cache(
+        history: torch.Tensor,
+        valid_lengths: torch.Tensor,
+        width: int,
+    ) -> torch.Tensor:
+        indices = valid_lengths[:, None] + torch.arange(width, device=history.device)[None, :]
+        return history.gather(
+            2,
+            indices[:, None, :].expand(-1, int(history.shape[1]), -1),
+        )
+
+    def _blocks_forward_chunk_ragged(
+        self,
+        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:
+        """Run one padded DiT batch while capturing exact per-row CNN state."""
+        lengths = torch.tensor(
+            (*valid_lengths, *valid_lengths),
+            device=estimator_input.device,
+            dtype=torch.long,
+        )
+        x = estimator.in_proj(estimator_input.transpose(1, 2))
+        for block_index, block in enumerate(estimator.blocks):
+            (
+                shift_msa,
+                scale_msa,
+                gate_msa,
+                shift_mlp,
+                scale_mlp,
+                gate_mlp,
+                shift_conv,
+                scale_conv,
+                gate_conv,
+            ) = block.adaLN_modulation(time_embedding).chunk(9, dim=-1)
+
+            normalized = block.norm1(x) * (1 + scale_msa) + shift_msa
+            x_att, new_att = block.attn.forward_chunk(
+                normalized,
+                att_cache[block_index],
+                attn_mask,
+            )
+            x = x + gate_msa * x_att
+
+            conv_input = block.norm3(x) * (1 + scale_conv) + shift_conv
+            old_cnn = cnn_cache[block_index]
+            if old_cnn is None:
+                width = int(block.conv.block[1].causal_padding[0])
+                old_cnn = conv_input.new_zeros(
+                    (
+                        int(conv_input.shape[0]),
+                        int(block.conv.in_channels + block.conv.out_channels),
+                        width,
+                    )
+                )
+            old_cnn1, old_cnn2 = old_cnn.split((block.conv.in_channels, block.conv.out_channels), dim=1)
+
+            conv1_input = block.conv.block[0](conv_input)
+            conv1_output, _ = block.conv.block[1].forward_chunk(conv1_input, old_cnn1)
+            new_cnn1 = self._gather_causal_cache(
+                torch.cat((old_cnn1, conv1_input), dim=2),
+                lengths,
+                int(old_cnn1.shape[2]),
+            )
+
+            conv2_input = block.conv.block[2:6](conv1_output)
+            conv2_output, _ = block.conv.block[6].forward_chunk(conv2_input, old_cnn2)
+            new_cnn2 = self._gather_causal_cache(
+                torch.cat((old_cnn2, conv2_input), dim=2),
+                lengths,
+                int(old_cnn2.shape[2]),
+            )
+            x = x + gate_conv * block.conv.block[7](conv2_output)
+            x = x + gate_mlp * block.mlp(block.norm2(x) * (1 + scale_mlp) + shift_mlp)
+
+            cnn_cache_buffer[block_index].copy_(torch.cat((new_cnn1, new_cnn2), dim=1))
+            att_cache_buffer[block_index][:, :, : int(new_att.shape[2]), :].copy_(new_att)
+
+        return estimator.final_layer(x, time_embedding).transpose(1, 2)
+
     def _decode_cfm(
         self,
         mu: torch.Tensor,
@@ -433,7 +556,12 @@ class BatchedToken2Wav(nn.Module):
         *,
         cnn_cache: torch.Tensor | None,
         att_cache: torch.Tensor | None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        valid_lengths: list[int] | None = None,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor | list[torch.Tensor],
+    ]:
         decoder = self.flow.decoder
         estimator = decoder.estimator
         batch_size = int(mu.shape[0])
@@ -458,36 +586,114 @@ class BatchedToken2Wav(nn.Module):
         mu_cfg = torch.cat((mu, torch.zeros_like(mu)), dim=0)
         speakers_cfg = torch.cat((speakers, torch.zeros_like(speakers)), dim=0)
         cond_cfg = torch.cat((cond, torch.zeros_like(cond)), dim=0)
-        next_cnn: list[torch.Tensor] = []
-        next_att: list[torch.Tensor] = []
-        dt = timeline[1] - timeline[0]
-        for step in range(self.n_timesteps):
-            old_cnn = cnn_cache[step] if cnn_cache is not None else None
-            old_att = att_cache[step] if att_cache is not None else None
-            estimate, step_cnn, step_att = self._estimator_step(
-                estimator,
-                x=torch.cat((x, x), dim=0),
-                mu=mu_cfg,
-                time=torch.cat((time, time), dim=0),
-                speakers=speakers_cfg,
-                cond=cond_cfg,
-                cnn_cache=old_cnn,
-                att_cache=old_att,
+        attn_mask = None
+        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,
             )
-            conditional, unconditional = estimate.split(batch_size, dim=0)
-            velocity = (1.0 + decoder.inference_cfg_rate) * conditional - decoder.inference_cfg_rate * unconditional
-            x = x + dt * velocity
-            time = time + dt
-            if step + 1 < self.n_timesteps:
-                dt = timeline[step + 2] - time[0]
-            next_cnn.append(step_cnn)
-            next_att.append(step_att)
-        return x, torch.stack(next_cnn), torch.stack(next_att)
+            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)
+            old_keys = torch.ones(
+                (2 * batch_size, int(mu.shape[2]), offset),
+                dtype=torch.bool,
+                device=mu.device,
+            )
+            attn_mask = valid_queries.unsqueeze(2) & torch.cat((current_keys, old_keys), dim=2)
+        next_cnn: list[torch.Tensor] = []
+        next_att_cache: torch.Tensor | None = None
+        ragged_att_cache: list[torch.Tensor] | None = None
+        dt = timeline[1] - timeline[0]
+        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
+                old_att = att_cache[step] if att_cache is not None else None
+                estimate, step_cnn, step_att = self._estimator_step(
+                    estimator,
+                    x=torch.cat((x, x), dim=0),
+                    mu=mu_cfg,
+                    time=torch.cat((time, time), dim=0),
+                    speakers=speakers_cfg,
+                    cond=cond_cfg,
+                    cnn_cache=old_cnn,
+                    att_cache=old_att,
+                    attn_mask=attn_mask,
+                    valid_lengths=valid_lengths,
+                )
+                conditional, unconditional = estimate.split(batch_size, dim=0)
+                velocity = (1.0 + decoder.inference_cfg_rate) * conditional - decoder.inference_cfg_rate * unconditional
+                x = x + dt * velocity
+                time = time + dt
+                if step + 1 < self.n_timesteps:
+                    dt = timeline[step + 2] - time[0]
+                next_cnn.append(step_cnn)
+                if valid_lengths is not None:
+                    if ragged_att_cache is None:
+                        ragged_att_cache = [
+                            torch.empty(
+                                (
+                                    self.n_timesteps,
+                                    int(step_att.shape[0]),
+                                    2,
+                                    int(step_att.shape[2]),
+                                    valid_length + offset,
+                                    int(step_att.shape[4]),
+                                ),
+                                device=step_att.device,
+                                dtype=self._estimator_att_cache_dtype,
+                            )
+                            for valid_length in valid_lengths
+                        ]
+                    current_width = int(mu.shape[2])
+                    for row, (valid_length, row_cache) in enumerate(zip(valid_lengths, ragged_att_cache, strict=True)):
+                        for cfg_row, source_row in enumerate((row, batch_size + row)):
+                            row_cache[step, :, cfg_row, :, :valid_length].copy_(
+                                step_att[:, source_row, :, :valid_length]
+                            )
+                            if offset:
+                                row_cache[step, :, cfg_row, :, valid_length:].copy_(
+                                    step_att[:, source_row, :, current_width:]
+                                )
+                    del step_att
+                    continue
+                if next_att_cache is None:
+                    next_att_cache = torch.empty(
+                        (self.n_timesteps, *step_att.shape),
+                        device=step_att.device,
+                        dtype=self._estimator_att_cache_dtype,
+                    )
+                next_att_cache[step].copy_(step_att)
+                del step_att
+        if ragged_att_cache is not None:
+            return x, torch.stack(next_cnn), ragged_att_cache
+        assert next_att_cache is not None
+        return x, torch.stack(next_cnn), next_att_cache
 
-    @staticmethod
-    def _split_flow_cache(cache: dict[str, torch.Tensor], batch_size: int) -> list[dict[str, torch.Tensor]]:
+    def _split_flow_cache(self, cache: dict[str, torch.Tensor], batch_size: int) -> list[dict[str, torch.Tensor]]:
         result: list[dict[str, torch.Tensor]] = []
         for row in range(batch_size):
+            estimator_att = cache["estimator_att_cache"]
+            if estimator_att.dtype == self._estimator_att_cache_dtype:
+                request_att = torch.cat(
+                    (
+                        estimator_att[:, :, row : row + 1],
+                        estimator_att[:, :, batch_size + row : batch_size + row + 1],
+                    ),
+                    dim=2,
+                ).detach()
+            else:
+                request_shape = list(estimator_att.shape)
+                request_shape[2] = 2
+                request_att = torch.empty(
+                    request_shape,
+                    device=estimator_att.device,
+                    dtype=self._estimator_att_cache_dtype,
+                )
+                request_att[:, :, 0:1].copy_(estimator_att[:, :, row : row + 1])
+                request_att[:, :, 1:2].copy_(estimator_att[:, :, batch_size + row : batch_size + row + 1])
             result.append(
                 {
                     "conformer_cnn_cache": cache["conformer_cnn_cache"][row : row + 1].detach().clone(),
@@ -499,29 +705,23 @@ class BatchedToken2Wav(nn.Module):
                         ),
                         dim=2,
                     ).detach(),
-                    "estimator_att_cache": torch.cat(
-                        (
-                            cache["estimator_att_cache"][:, :, row : row + 1],
-                            cache["estimator_att_cache"][:, :, batch_size + row : batch_size + row + 1],
-                        ),
-                        dim=2,
-                    ).detach(),
+                    "estimator_att_cache": request_att,
                 }
             )
         return result
 
-    @staticmethod
-    def _stack_flow_cache(states: list[BatchedToken2WavState]) -> dict[str, torch.Tensor]:
+    def _stack_flow_cache(self, states: list[BatchedToken2WavState]) -> dict[str, torch.Tensor]:
         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]
         conditional_att = [flow["estimator_att_cache"][:, :, 0:1] for flow in flows]
         unconditional_att = [flow["estimator_att_cache"][:, :, 1:2] for flow in flows]
+        estimator_att = torch.cat((*conditional_att, *unconditional_att), dim=2)
         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),
             "estimator_cnn_cache": torch.cat((*conditional_cnn, *unconditional_cnn), dim=2),
-            "estimator_att_cache": torch.cat((*conditional_att, *unconditional_att), dim=2),
+            "estimator_att_cache": estimator_att,
         }
 
     def setup_batch(
@@ -729,3 +929,188 @@ class BatchedToken2Wav(nn.Module):
         ]
         audios = [emitted[row].reshape(-1).to(dtype=torch.float32) for row in range(batch_size)]
         return audios, next_states
+
+    @staticmethod
+    def _require_complete_ragged_outputs(
+        audios: list[torch.Tensor | None],
+        next_states: list[BatchedToken2WavState | None],
+    ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]:
+        missing_audio_rows = [row for row, audio in enumerate(audios) if audio is None]
+        missing_state_rows = [row for row, state in enumerate(next_states) if state is None]
+        if missing_audio_rows or missing_state_rows:
+            raise RuntimeError(
+                "MiniCPMO45Code2WavBatchError "
+                f'{{"reason":"incomplete_ragged_output","audio_rows":{missing_audio_rows},'
+                f'"state_rows":{missing_state_rows}}}'
+            )
+        return (
+            cast(list[torch.Tensor], audios),
+            cast(list[BatchedToken2WavState], next_states),
+        )
+
+    def decode_ragged_batch(
+        self,
+        tokens: list[torch.Tensor],
+        features: PromptFeatures,
+        states: list[BatchedToken2WavState],
+        *,
+        last_chunks: list[bool],
+    ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]:
+        batch_size = len(tokens)
+        if batch_size != len(states) or batch_size != len(last_chunks):
+            raise ValueError("ragged token, state, and final-flag batches must have the same size")
+        if batch_size == 0:
+            return [], []
+
+        lookahead = self._pre_lookahead_len()
+        for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)):
+            num_frames = int(row_tokens.numel())
+            if lookahead is not None and not last_chunk and num_frames <= lookahead:
+                raise RuntimeError(
+                    "MiniCPMO45Code2WavBatchError "
+                    f'{{"reason":"chunk_below_lookahead_window","row":{row},'
+                    f'"frames":{num_frames},"minimum":{lookahead + 1}}}'
+                )
+
+        max_frames = self._max_encode_token_frames(states)
+        if any(int(row_tokens.numel()) > max_frames for row_tokens in tokens):
+            # decode_batch owns the RelPos-safe encoder slicing used by
+            # non-async Talker dumps. Preserve that path for overlong rows;
+            # ordinary short duplex chunks continue through one padded DiT.
+            exact_groups: dict[tuple[int, bool], list[int]] = {}
+            for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)):
+                exact_groups.setdefault((int(row_tokens.numel()), last_chunk), []).append(row)
+            audios: list[torch.Tensor | None] = [None] * batch_size
+            next_states: list[BatchedToken2WavState | None] = [None] * batch_size
+            for (_, last_chunk), rows in exact_groups.items():
+                group_audio, group_states = self.decode_batch(
+                    torch.stack([tokens[row] for row in rows], dim=0),
+                    features,
+                    [states[row] for row in rows],
+                    last_chunk=last_chunk,
+                )
+                for group_row, row in enumerate(rows):
+                    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)
+
+        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():
+            group_states = [states[row] for row in rows]
+            group_cache = self._stack_flow_cache(group_states)
+            group_tokens = torch.stack([tokens[row] for row in rows], dim=0)
+            with self._autocast(group_tokens.device):
+                hidden, conformer_cnn, conformer_att = self._encode_chunk(
+                    group_tokens,
+                    last_chunk=last_chunk,
+                    cnn_cache=group_cache["conformer_cnn_cache"],
+                    att_cache=group_cache["conformer_att_cache"],
+                )
+            for group_row, row in enumerate(rows):
+                hidden_rows[row] = hidden[group_row : group_row + 1]
+                conformer_cnn_rows[row] = conformer_cnn[group_row : group_row + 1]
+                conformer_att_rows[row] = conformer_att[:, group_row : group_row + 1]
+
+        missing_hidden_rows = [row for row, value in enumerate(hidden_rows) if value is None]
+        if missing_hidden_rows:
+            raise RuntimeError(
+                f'MiniCPMO45Code2WavBatchError {{"reason":"incomplete_encoder_output","rows":{missing_hidden_rows}}}'
+            )
+        resolved_hidden = cast(list[torch.Tensor], hidden_rows)
+        hidden_lengths = [int(value.shape[1]) for value in resolved_hidden]
+        max_hidden_length = max(hidden_lengths)
+        padded_hidden = resolved_hidden[0].new_zeros((batch_size, max_hidden_length, int(resolved_hidden[0].shape[2])))
+        for row, hidden in enumerate(resolved_hidden):
+            padded_hidden[row, : int(hidden.shape[1])].copy_(hidden[0])
+
+        flow_cache = self._stack_flow_cache(states)
+        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))
+            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(),
+                projected_speakers,
+                cond,
+                cnn_cache=flow_cache["estimator_cnn_cache"],
+                att_cache=flow_cache["estimator_att_cache"],
+                valid_lengths=hidden_lengths,
+            )
+
+        prompt_len = int(features.mels.shape[1])
+        assert isinstance(estimator_att, list)
+        new_flow: list[dict[str, torch.Tensor]] = []
+        for row, row_estimator_att in enumerate(estimator_att):
+            conformer_cnn = conformer_cnn_rows[row]
+            conformer_att = conformer_att_rows[row]
+            assert conformer_cnn is not None and conformer_att is not None
+            if row_estimator_att.shape[4] > prompt_len + 100:
+                row_estimator_att = torch.cat(
+                    (
+                        row_estimator_att[..., :prompt_len, :],
+                        row_estimator_att[..., -100:, :],
+                    ),
+                    dim=4,
+                )
+            if conformer_att.shape[3] > prompt_len + 100:
+                conformer_att = torch.cat(
+                    (
+                        conformer_att[..., :prompt_len, :],
+                        conformer_att[..., -100:, :],
+                    ),
+                    dim=3,
+                )
+            new_flow.append(
+                {
+                    "conformer_cnn_cache": conformer_cnn.detach().clone(),
+                    "conformer_att_cache": conformer_att.detach().clone(),
+                    "estimator_cnn_cache": torch.cat(
+                        (
+                            estimator_cnn[:, :, row : row + 1],
+                            estimator_cnn[:, :, batch_size + row : batch_size + row + 1],
+                        ),
+                        dim=2,
+                    ).detach(),
+                    "estimator_att_cache": row_estimator_att,
+                }
+            )
+
+        audios: list[torch.Tensor | None] = [None] * batch_size
+        next_states: list[BatchedToken2WavState | None] = [None] * batch_size
+        vocoder_groups: dict[tuple[int, bool], 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():
+            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)
+            group_mel = torch.cat(
+                [chunk_mel[row : row + 1, :, :hidden_length] for row in rows],
+                dim=0,
+            )
+            mel = torch.cat((old_mel, group_mel), dim=2)
+            speech, source = self._hift_inference(mel, old_source)
+            if old_speech.shape[-1] > 0:
+                window = self.speech_window.to(device=speech.device, dtype=speech.dtype)
+                speech = self._fade_in_out(speech, old_speech, window)
+            next_hift = {
+                "mel": mel[..., -self.mel_cache_len :].detach(),
+                "source": source[..., -self.source_cache_len :].detach(),
+                "speech": speech[..., -self.source_cache_len :].detach(),
+            }
+            emitted = speech if last_chunk else speech[..., : -self.source_cache_len]
+            for group_row, row in enumerate(rows):
+                audios[row] = emitted[group_row].reshape(-1).to(dtype=torch.float32)
+                next_states[row] = BatchedToken2WavState(
+                    flow_cache=new_flow[row],
+                    hift_cache={
+                        name: value[group_row : group_row + 1].detach().clone() for name, value in next_hift.items()
+                    },
+                )
+
+        return self._require_complete_ragged_outputs(audios, next_states)

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.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
 """Strict batched codec-to-waveform stage for MiniCPM-o 4.5."""
 
 from __future__ import annotations
@@ -135,7 +135,7 @@ class _WorkItem:
 
 
 class MiniCPMO45Code2Wav(nn.Module):
-    """LLM_GENERATION model that admits only true exact-shape GPU batches."""
+    """LLM_GENERATION model with request-owned state and compatible batching."""
 
     input_modalities = "audio"
     have_multimodal_outputs = True
@@ -184,6 +184,11 @@ class MiniCPMO45Code2Wav(nn.Module):
         self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1))
         if self._min_batch_size < 1:
             raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1")
+        self._initial_batch_size = int(extra.get("code2wav_initial_batch_size", 0))
+        if self._initial_batch_size < 0:
+            raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be >= 0")
+        if self._initial_batch_size and self._initial_batch_size < self._min_batch_size:
+            raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be 0 or >= code2wav_min_batch_size")
         self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio"))
         self._prompt_wav_override = extra.get("prompt_wav")
 
@@ -482,13 +487,63 @@ class MiniCPMO45Code2Wav(nn.Module):
         return (
             item.prompt_cache_id,
             item.prompt_wav,
-            int(item.tokens.numel()),
             cache_signature,
-            item.last_chunk,
-            item.tts_is_last_chunk,
             item.cache_epoch,
         )
 
+    def _iter_limited_batches(
+        self,
+        buckets: Iterable[list[_WorkItem]],
+        *,
+        maximum: int,
+    ) -> Iterable[list[_WorkItem]]:
+        for bucket in buckets:
+            if not bucket:
+                continue
+            if not maximum:
+                yield bucket
+                continue
+            batch_count = (len(bucket) + maximum - 1) // maximum
+            if len(bucket) < batch_count * self._min_batch_size:
+                raise _batch_error(
+                    "initial_batch_partition_below_minimum",
+                    size=len(bucket),
+                    minimum=self._min_batch_size,
+                    maximum=maximum,
+                )
+            batch_size, larger_batches = divmod(len(bucket), batch_count)
+            start = 0
+            for batch_index in range(batch_count):
+                stop = start + batch_size + (batch_index < larger_batches)
+                yield bucket[start:stop]
+                start = stop
+
+    def _iter_decode_batches(
+        self,
+        buckets: Iterable[list[_WorkItem]],
+    ) -> Iterable[list[_WorkItem]]:
+        for bucket in buckets:
+            if not self._initial_batch_size:
+                yield bucket
+                continue
+            initial = [item for item in bucket if item.chunk_seq <= 1]
+            steady = [item for item in bucket if item.chunk_seq > 1]
+
+            undersized = [
+                {"wave": wave, "size": len(items)}
+                for wave, items in (("initial", initial), ("steady", steady))
+                if items and len(items) < self._min_batch_size
+            ]
+            if undersized:
+                raise _batch_error(
+                    "decode_wave_below_minimum",
+                    minimum=self._min_batch_size,
+                    waves=undersized,
+                )
+            yield from self._iter_limited_batches([initial], maximum=self._initial_batch_size)
+            if steady:
+                yield steady
+
     @torch.inference_mode()
     def forward(
         self,
@@ -624,7 +679,10 @@ class MiniCPMO45Code2Wav(nn.Module):
                     (item.prompt_cache_id, item.prompt_wav),
                     [],
                 ).append(item)
-        for bucket in initial_marker_buckets.values():
+        for bucket in self._iter_limited_batches(
+            initial_marker_buckets.values(),
+            maximum=self._initial_batch_size,
+        ):
             try:
                 features = self.backend.prepare_prompt(
                     bucket[0].prompt_cache_id,
@@ -656,7 +714,7 @@ class MiniCPMO45Code2Wav(nn.Module):
                     prompt_wav=item.prompt_wav,
                     token2wav=state,
                 )
-        for bucket in buckets.values():
+        for bucket in self._iter_decode_batches(buckets.values()):
             batch_size = len(bucket)
             try:
                 features = self.backend.prepare_prompt(
@@ -667,13 +725,23 @@ class MiniCPMO45Code2Wav(nn.Module):
                     states = self.backend.setup_batch(features, batch_size)
                 else:
                     states = [item.previous.token2wav for item in bucket if item.previous is not None]
-                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,
-                )
+                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],
+                    )
+                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,
+                    )
             except Exception as exc:
                 self._prune_unowned_runtime_prompts()
                 if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "):
@@ -822,4 +890,5 @@ class MiniCPMO45Code2Wav(nn.Module):
             connector_config=self._connector_config,
             hift_graph_config=self._hift_graph_config,
             cfm_graph_config=self._cfm_graph_config,
+            bfloat16_attention_cache=bool(extra.get("code2wav_bfloat16_attention_cache", False)),
         )

vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.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
 """Inject MiniCPM-o Code2Wav NPUGraph acceleration on Ascend."""
 
 from __future__ import annotations
@@ -7,6 +7,7 @@ from __future__ import annotations
 import os
 from collections.abc import Mapping
 from contextlib import nullcontext
+from typing import cast
 from weakref import WeakKeyDictionary
 
 import torch
@@ -24,6 +25,7 @@ _original_decode_batch = None
 _backend_graph_runners: WeakKeyDictionary[object, NPUExactGraphRunner] = WeakKeyDictionary()
 _ENABLE_KEY = "code2wav_enable_npu_graph"
 _MAX_GRAPHS_KEY = "code2wav_max_npu_graphs"
+_BF16_ATTENTION_CACHE_KEY = "code2wav_bfloat16_attention_cache"
 
 
 def _config_bool(value: object, default: bool) -> bool:
@@ -104,10 +106,18 @@ def _patched_estimator_step(
     cond,
     cnn_cache,
     att_cache,
+    attn_mask=None,
+    valid_lengths=None,
 ):
     assert _original_estimator_step is not None
     graph_runner = _backend_graph_runners.get(self)
-    if graph_runner is None or self._trt_stepper is not None or self._cfm_graph_wrapper is not None:
+    if (
+        graph_runner is None
+        or self._trt_stepper is not None
+        or self._cfm_graph_wrapper is not None
+        or attn_mask is not None
+        or valid_lengths is not None
+    ):
         return _original_estimator_step(
             self,
             estimator,
@@ -118,6 +128,8 @@ def _patched_estimator_step(
             cond=cond,
             cnn_cache=cnn_cache,
             att_cache=att_cache,
+            attn_mask=attn_mask,
+            valid_lengths=valid_lengths,
         )
     if (cnn_cache is None) != (att_cache is None):
         raise ValueError("estimator CNN and attention caches must both be present or absent")
@@ -198,8 +210,15 @@ def _patched_build_backend(self) -> None:
     if self.backend is not None:
         return
 
+    extra = self._extra_config()
+    if bool(extra.get(_BF16_ATTENTION_CACHE_KEY, False)):
+        raise ValueError(
+            "MiniCPM-o Code2Wav code2wav_bfloat16_attention_cache is "
+            "currently supported only on CUDA; leave it unset or false on NPU."
+        )
+
     config = _graph_config(self)
-    max_graphs = max(0, int(config.get(_MAX_GRAPHS_KEY, 32)))
+    max_graphs = max(0, int(cast(int | str, config.get(_MAX_GRAPHS_KEY, 32))))
     graph_enabled = max_graphs > 0 and _config_bool(config.get(_ENABLE_KEY), False)
     if graph_enabled:
         # NPUOmniPlatform enables internal format for quantized LLM kernels.