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

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

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