[WIP][Bugfix][MiniCPM-o] Fix Full duplex bugs for minicpm-challenge (#5496)

Signed-off-by: Sy03 <1370724210@qq.com>

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -348,9 +348,16 @@ class BatchedToken2Wav(nn.Module):
         previous: torch.Tensor,
         window: torch.Tensor,
     ) -> torch.Tensor:
-        overlap = int(window.shape[0] // 2)
+        overlap = min(
+            int(window.shape[0] // 2),
+            int(speech.shape[-1]),
+            int(previous.shape[-1]),
+        )
         result = speech.clone()
-        result[..., :overlap] = result[..., :overlap] * window[:overlap] + previous[..., -overlap:] * window[overlap:]
+        if overlap > 0:
+            result[..., :overlap] = (
+                result[..., :overlap] * window[:overlap] + previous[..., -overlap:] * window[-overlap:]
+            )
         return result
 
     def decode_batch(
@@ -360,6 +367,7 @@ class BatchedToken2Wav(nn.Module):
         states: list[BatchedToken2WavState],
         *,
         last_chunk: bool,
+        flush_encoder: bool = False,
     ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]:
         batch_size = int(tokens.shape[0])
         if batch_size != len(states):
@@ -382,7 +390,7 @@ class BatchedToken2Wav(nn.Module):
         with self._autocast(tokens.device):
             hidden, conformer_cnn, conformer_att = self._encode_chunk(
                 tokens,
-                last_chunk=last_chunk,
+                last_chunk=last_chunk or flush_encoder,
                 cnn_cache=flow_cache["conformer_cnn_cache"],
                 att_cache=flow_cache["conformer_att_cache"],
             )

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py

@@ -4,12 +4,12 @@
 
 from __future__ import annotations
 
-import hashlib
 import json
 import os
 import tempfile
 from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import dataclass
+from hashlib import sha256
 from pathlib import Path
 from typing import Any
 
@@ -77,6 +77,13 @@ class _RequestState:
     token2wav: BatchedToken2WavState
 
 
+@dataclass
+class _RuntimePrompt:
+    cache_id: str
+    path: str
+    owners: set[str]
+
+
 @dataclass(frozen=True)
 class _WorkItem:
     output_index: int
@@ -89,9 +96,11 @@ class _WorkItem:
     last_chunk: bool
     tokens: torch.Tensor
     previous: _RequestState | None
-    segment_text: str
-    duplex_turn_id: int | None
-    duplex_epoch: int | None
+    runtime_prompt_key: str | None
+    duplex_epoch: int
+    duplex_turn_id: int
+    segment_text_utf8: torch.Tensor
+    tts_is_last_chunk: bool
     segment_end: bool
     turn_end: bool
     has_payload: bool = True
@@ -120,12 +129,17 @@ class MiniCPMO45Code2Wav(nn.Module):
         self.model_path = str(vllm_config.model_config.model)
         self.backend: BatchedToken2Wav | None = None
         self._states: dict[str, _RequestState] = {}
-        self._owned_prompt_wavs: dict[str, tuple[str, str]] = {}
+        self._runtime_prompts: dict[str, _RuntimePrompt] = {}
+        self._request_prompt_keys: dict[str, str] = {}
+        self._runtime_prompt_dir = tempfile.TemporaryDirectory(
+            prefix="minicpmo45-runtime-prompts-",
+        )
         extra = self._extra_config()
         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._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio"))
+        self._prompt_wav_explicit = "prompt_wav" in extra
         self._default_prompt_wav = str(
             extra.get(
                 "prompt_wav",
@@ -133,6 +147,34 @@ class MiniCPMO45Code2Wav(nn.Module):
             )
         )
 
+    def _resolve_model_root(self) -> Path:
+        model_root = Path(self.model_path)
+        if model_root.is_dir():
+            return model_root
+
+        from vllm_omni.model_executor.model_loader.weight_utils import (
+            download_weights_from_hf_specific,
+        )
+
+        model_config = self.vllm_config.model_config
+        load_config = getattr(self.vllm_config, "load_config", None)
+        model_root = Path(
+            download_weights_from_hf_specific(
+                self.model_path,
+                getattr(load_config, "download_dir", None),
+                allow_patterns=[
+                    "assets/HT_ref_audio.wav",
+                    "assets/token2wav/*",
+                ],
+                revision=getattr(model_config, "revision", None),
+                require_all=True,
+            )
+        )
+        self.model_path = str(model_root)
+        if not self._prompt_wav_explicit:
+            self._default_prompt_wav = str(model_root / "assets" / "HT_ref_audio.wav")
+        return model_root
+
     def _extra_config(self) -> dict[str, Any]:
         model_config = getattr(self.vllm_config, "model_config", None)
         connector = getattr(model_config, "stage_connector_config", None)
@@ -142,63 +184,122 @@ class MiniCPMO45Code2Wav(nn.Module):
             extra = getattr(connector, "extra", None)
         return dict(extra) if isinstance(extra, Mapping) else {}
 
-    def _release_prompt(self, state_id: str) -> None:
-        owned = self._owned_prompt_wavs.pop(state_id, None)
-        if owned is None:
-            return
-        prompt_cache_id, prompt_wav = owned
-        if self.backend is not None:
-            self.backend.evict_prompt(prompt_cache_id, prompt_wav)
-        try:
-            os.unlink(prompt_wav)
-        except FileNotFoundError:
-            pass
-
-    def _write_reference_prompt(
-        self,
-        state_id: str,
-        reference: object,
-        sample_rate: object,
-    ) -> tuple[str, str]:
-        waveform = torch.as_tensor(reference, dtype=torch.float32).reshape(-1).cpu()
-        if waveform.numel() == 0:
-            return self._default_prompt_id, self._default_prompt_wav
-        sample_rate_hz = int(_scalar(sample_rate, 24000))
-        if sample_rate_hz <= 0:
-            raise _batch_error(
-                "invalid_ref_audio_sample_rate",
-                request_id=state_id,
-                sample_rate=sample_rate_hz,
-            )
-        digest = hashlib.sha256()
-        digest.update(str(sample_rate_hz).encode("ascii"))
-        digest.update(waveform.numpy().tobytes())
-        prompt_cache_id = f"{state_id}:{digest.hexdigest()}"
-        current = self._owned_prompt_wavs.get(state_id)
-        if current is not None and current[0] == prompt_cache_id and os.path.exists(current[1]):
-            return current
-
-        self._release_prompt(state_id)
-        handle = tempfile.NamedTemporaryFile(prefix="minicpmo45-ref-", suffix=".wav", delete=False)
-        prompt_wav = handle.name
-        handle.close()
-        try:
-            sf.write(prompt_wav, waveform.numpy(), sample_rate_hz, format="WAV")
-        except Exception:
-            try:
-                os.unlink(prompt_wav)
-            except FileNotFoundError:
-                pass
-            raise
-        self._owned_prompt_wavs[state_id] = (prompt_cache_id, prompt_wav)
-        return prompt_cache_id, prompt_wav
-
     def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
         return torch.zeros((input_ids.numel(), 1), device=input_ids.device, dtype=torch.float32)
 
     def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None:
         return None
 
+    def _materialize_runtime_prompt(
+        self,
+        ref_audio: Any,
+        sample_rate: Any,
+    ) -> tuple[str, _RuntimePrompt]:
+        sample_rate_hz = int(_scalar(sample_rate, 0))
+        waveform = torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1).cpu().contiguous()
+        if sample_rate_hz <= 0:
+            raise _batch_error("invalid_ref_audio_sample_rate", sample_rate=sample_rate_hz)
+        if waveform.numel() == 0:
+            raise _batch_error("empty_ref_audio")
+        if not bool(torch.isfinite(waveform).all().item()):
+            raise _batch_error("non_finite_ref_audio")
+
+        digest = sha256()
+        digest.update(waveform.numpy().tobytes())
+        digest.update(str(sample_rate_hz).encode())
+        cache_key = digest.hexdigest()
+        cache_id = f"runtime-ref-{cache_key[:24]}-{sample_rate_hz}"
+        path = str(Path(self._runtime_prompt_dir.name) / f"minicpmo45_ref_{cache_key[:24]}_{sample_rate_hz}.wav")
+        entry = self._runtime_prompts.get(cache_key)
+        if entry is None:
+            entry = _RuntimePrompt(cache_id=cache_id, path=path, owners=set())
+            self._runtime_prompts[cache_key] = entry
+        prompt_path = Path(entry.path)
+        if not prompt_path.is_file():
+            with tempfile.NamedTemporaryFile(
+                dir=prompt_path.parent,
+                prefix=f".{prompt_path.stem}-",
+                suffix=".wav",
+                delete=False,
+            ) as handle:
+                temporary_path = Path(handle.name)
+            try:
+                sf.write(
+                    temporary_path,
+                    waveform.numpy(),
+                    sample_rate_hz,
+                    format="WAV",
+                )
+                os.replace(temporary_path, prompt_path)
+            finally:
+                temporary_path.unlink(missing_ok=True)
+        return cache_key, entry
+
+    def _resolve_prompt(
+        self,
+        state_id: str,
+        info: Mapping[str, Any],
+        meta: Mapping[str, Any],
+        previous: _RequestState | None,
+    ) -> tuple[str, str, str | None]:
+        codes = info.get("codes")
+        ref_audio = codes.get("ref") if isinstance(codes, Mapping) else None
+        if ref_audio is not None:
+            cache_key, entry = self._materialize_runtime_prompt(
+                ref_audio,
+                meta.get("ref_audio_sr"),
+            )
+            return entry.cache_id, entry.path, cache_key
+
+        if previous is not None:
+            return previous.prompt_cache_id, previous.prompt_wav, self._request_prompt_keys.get(state_id)
+
+        cache_key = self._request_prompt_keys.get(state_id)
+        entry = self._runtime_prompts.get(cache_key) if cache_key is not None else None
+        if entry is not None:
+            return entry.cache_id, entry.path, cache_key
+
+        return (
+            str(_scalar(meta.get("prompt_cache_id"), self._default_prompt_id)),
+            str(_scalar(meta.get("prompt_wav"), self._default_prompt_wav)),
+            None,
+        )
+
+    def _release_request_prompt(self, state_id: str) -> None:
+        cache_key = self._request_prompt_keys.pop(state_id, None)
+        entry = self._runtime_prompts.get(cache_key) if cache_key is not None else None
+        if entry is None:
+            return
+        entry.owners.discard(state_id)
+        if entry.owners:
+            return
+        if self.backend is not None:
+            self.backend.evict_prompt(entry.cache_id, entry.path)
+        Path(entry.path).unlink(missing_ok=True)
+        self._runtime_prompts.pop(cache_key, None)
+
+    def _commit_runtime_prompt_owners(self, items: list[_WorkItem]) -> None:
+        for item in items:
+            cache_key = item.runtime_prompt_key
+            if cache_key is None:
+                continue
+            previous_key = self._request_prompt_keys.get(item.state_id)
+            if previous_key != cache_key:
+                self._release_request_prompt(item.state_id)
+            entry = self._runtime_prompts.get(cache_key)
+            if entry is not None:
+                entry.owners.add(item.state_id)
+                self._request_prompt_keys[item.state_id] = cache_key
+
+    def _prune_unowned_runtime_prompts(self) -> None:
+        for cache_key, entry in list(self._runtime_prompts.items()):
+            if entry.owners:
+                continue
+            if self.backend is not None:
+                self.backend.evict_prompt(entry.cache_id, entry.path)
+            Path(entry.path).unlink(missing_ok=True)
+            self._runtime_prompts.pop(cache_key, None)
+
     @staticmethod
     def _split_segments(input_ids: torch.Tensor, counts: Any) -> list[torch.Tensor]:
         flat = input_ids.reshape(-1)
@@ -247,9 +348,11 @@ class MiniCPMO45Code2Wav(nn.Module):
                 last_chunk=False,
                 tokens=segment.new_empty(0, dtype=torch.long),
                 previous=None,
-                segment_text="",
-                duplex_turn_id=None,
-                duplex_epoch=None,
+                runtime_prompt_key=None,
+                duplex_epoch=-1,
+                duplex_turn_id=-1,
+                segment_text_utf8=torch.empty(0, dtype=torch.uint8),
+                tts_is_last_chunk=False,
                 segment_end=False,
                 turn_end=False,
                 has_payload=False,
@@ -266,14 +369,15 @@ class MiniCPMO45Code2Wav(nn.Module):
                 chunk_seq=chunk_seq,
             )
         last_chunk = bool(_scalar(meta.get("last_chunk"), False))
+        tts_is_last_chunk = bool(_scalar(meta.get("tts_is_last_chunk"), False))
         codes = info.get("codes")
         audio = codes.get("audio") if isinstance(codes, Mapping) else None
-        reference = codes.get("ref") if isinstance(codes, Mapping) else None
         tokens = _codec_tensor(audio, segment)
-        if last_chunk and int(_scalar(meta.get("code_flat_numel"), tokens.numel())) == 0:
+        if int(_scalar(meta.get("code_flat_numel"), tokens.numel())) == 0:
             # The generation scheduler reserves one placeholder token for an
-            # empty terminal chunk. The producer's explicit length is the
-            # authority, so do not decode that placeholder as codec data.
+            # empty terminal or segment-boundary chunk. The producer's
+            # explicit length is the authority, so do not decode that
+            # placeholder as codec data.
             tokens = segment.new_empty(0, dtype=torch.long)
         previous = self._states.get(state_id)
         if previous is None:
@@ -299,7 +403,6 @@ class MiniCPMO45Code2Wav(nn.Module):
                     cache_epoch=cache_epoch,
                     chunk_seq=chunk_seq,
                 )
-            self._release_prompt(state_id)
             previous = None
         elif chunk_seq != previous.chunk_seq + 1:
             raise _batch_error(
@@ -308,22 +411,12 @@ class MiniCPMO45Code2Wav(nn.Module):
                 expected=previous.chunk_seq + 1,
                 actual=chunk_seq,
             )
-
-        explicit_prompt_id = _scalar(meta.get("prompt_cache_id"))
-        explicit_prompt_wav = _scalar(meta.get("prompt_wav"))
-        if previous is not None and explicit_prompt_id is None and explicit_prompt_wav is None:
-            prompt_cache_id = previous.prompt_cache_id
-            prompt_wav = previous.prompt_wav
-        elif reference is not None:
-            prompt_cache_id, prompt_wav = self._write_reference_prompt(
-                state_id,
-                reference,
-                meta.get("ref_audio_sr"),
-            )
-        else:
-            prompt_cache_id = str(explicit_prompt_id or self._default_prompt_id)
-            prompt_wav = str(explicit_prompt_wav or self._default_prompt_wav)
-
+        prompt_cache_id, prompt_wav, runtime_prompt_key = self._resolve_prompt(
+            state_id,
+            info,
+            meta,
+            previous,
+        )
         if previous is not None and prompt_cache_id != previous.prompt_cache_id:
             raise _batch_error(
                 "prompt_changed_midstream",
@@ -338,8 +431,9 @@ class MiniCPMO45Code2Wav(nn.Module):
                 expected=previous.prompt_wav,
                 actual=prompt_wav,
             )
-        duplex_turn_id = _scalar(meta.get("duplex_turn_id"))
-        duplex_epoch = _scalar(meta.get("duplex_epoch"))
+        segment_text_utf8 = meta.get("llm_output_text_utf8")
+        if not isinstance(segment_text_utf8, torch.Tensor):
+            segment_text_utf8 = torch.empty(0, dtype=torch.uint8)
         return _WorkItem(
             output_index=index,
             state_id=state_id,
@@ -351,9 +445,11 @@ class MiniCPMO45Code2Wav(nn.Module):
             last_chunk=last_chunk,
             tokens=tokens,
             previous=previous,
-            segment_text=str(_scalar(meta.get("native_duplex_segment_text"), "")),
-            duplex_turn_id=int(duplex_turn_id) if duplex_turn_id is not None else None,
-            duplex_epoch=int(duplex_epoch) if duplex_epoch is not None else None,
+            runtime_prompt_key=runtime_prompt_key,
+            duplex_epoch=int(_scalar(meta.get("duplex_epoch"), -1)),
+            duplex_turn_id=int(_scalar(meta.get("duplex_turn_id"), -1)),
+            segment_text_utf8=segment_text_utf8,
+            tts_is_last_chunk=tts_is_last_chunk,
             segment_end=bool(_scalar(meta.get("segment_end"), False)),
             turn_end=bool(_scalar(meta.get("turn_end"), False)),
         )
@@ -371,6 +467,7 @@ class MiniCPMO45Code2Wav(nn.Module):
             int(item.tokens.numel()),
             cache_signature,
             item.last_chunk,
+            item.tts_is_last_chunk,
             item.cache_epoch,
         )
 
@@ -435,26 +532,37 @@ class MiniCPMO45Code2Wav(nn.Module):
                 request_ids=len(state_ids),
             )
         items: list[_WorkItem] = []
-        for index, (state_id, segment, info) in enumerate(
-            zip(state_ids, segments, runtime_additional_information, strict=True)
-        ):
-            if not isinstance(info, Mapping):
-                raise _batch_error(
-                    "invalid_runtime_info",
-                    output_index=index,
-                    value_type=type(info).__name__,
-                )
-            items.append(self._parse_item(index, str(state_id), segment, info))
+        try:
+            for index, (state_id, segment, info) in enumerate(
+                zip(state_ids, segments, runtime_additional_information, strict=True)
+            ):
+                if not isinstance(info, Mapping):
+                    raise _batch_error(
+                        "invalid_runtime_info",
+                        output_index=index,
+                        value_type=type(info).__name__,
+                    )
+                items.append(self._parse_item(index, str(state_id), segment, info))
+        except Exception:
+            self._prune_unowned_runtime_prompts()
+            raise
         state_ids = [item.state_id for item in items]
         if len(state_ids) != len(set(state_ids)):
+            self._prune_unowned_runtime_prompts()
             raise _batch_error("duplicate_request_in_forward", request_ids=state_ids)
         outputs = [empty for _ in segments]
         sentinels = [item for item in items if item.last_chunk and item.tokens.numel() == 0]
+        segment_markers = [
+            item for item in items if not item.last_chunk and item.tts_is_last_chunk and item.tokens.numel() == 0
+        ]
         compute_items = [item for item in items if item.tokens.numel() > 0]
         invalid_empty = [
-            item.request_id for item in items if item.has_payload and not item.last_chunk and item.tokens.numel() == 0
+            item.request_id
+            for item in items
+            if item.has_payload and not item.last_chunk and not item.tts_is_last_chunk and item.tokens.numel() == 0
         ]
         if invalid_empty:
+            self._prune_unowned_runtime_prompts()
             raise _batch_error("empty_nonfinal_chunk", request_ids=invalid_empty)
 
         buckets: dict[tuple[Any, ...], list[_WorkItem]] = {}
@@ -470,6 +578,7 @@ class MiniCPMO45Code2Wav(nn.Module):
             if len(bucket) < self._min_batch_size
         ]
         if undersized:
+            self._prune_unowned_runtime_prompts()
             raise _batch_error(
                 "exact_shape_bucket_below_minimum",
                 minimum=self._min_batch_size,
@@ -477,6 +586,58 @@ class MiniCPMO45Code2Wav(nn.Module):
             )
 
         pending: dict[str, _RequestState | None] = {item.state_id: None for item in sentinels}
+        pending.update(
+            {
+                item.state_id: _RequestState(
+                    cache_epoch=item.cache_epoch,
+                    chunk_seq=item.chunk_seq,
+                    prompt_cache_id=item.prompt_cache_id,
+                    prompt_wav=item.prompt_wav,
+                    token2wav=item.previous.token2wav,
+                )
+                for item in segment_markers
+                if item.previous is not None
+            }
+        )
+        initial_marker_buckets: dict[tuple[str, str], list[_WorkItem]] = {}
+        for item in segment_markers:
+            if item.previous is None:
+                initial_marker_buckets.setdefault(
+                    (item.prompt_cache_id, item.prompt_wav),
+                    [],
+                ).append(item)
+        for bucket in initial_marker_buckets.values():
+            try:
+                features = self.backend.prepare_prompt(
+                    bucket[0].prompt_cache_id,
+                    bucket[0].prompt_wav,
+                )
+                states = self.backend.setup_batch(features, len(bucket))
+            except Exception as exc:
+                self._prune_unowned_runtime_prompts()
+                if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "):
+                    raise
+                raise _batch_error(
+                    "backend_unsupported_or_failed",
+                    request_ids=[item.request_id for item in bucket],
+                    error_type=type(exc).__name__,
+                    error=str(exc),
+                ) from exc
+            if len(states) != len(bucket):
+                self._prune_unowned_runtime_prompts()
+                raise _batch_error(
+                    "backend_result_size_mismatch",
+                    expected=len(bucket),
+                    states=len(states),
+                )
+            for item, state in zip(bucket, states, strict=True):
+                pending[item.state_id] = _RequestState(
+                    cache_epoch=item.cache_epoch,
+                    chunk_seq=item.chunk_seq,
+                    prompt_cache_id=item.prompt_cache_id,
+                    prompt_wav=item.prompt_wav,
+                    token2wav=state,
+                )
         for bucket in buckets.values():
             batch_size = len(bucket)
             try:
@@ -496,6 +657,7 @@ class MiniCPMO45Code2Wav(nn.Module):
                     last_chunk=bucket[0].last_chunk,
                 )
             except Exception as exc:
+                self._prune_unowned_runtime_prompts()
                 if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "):
                     raise
                 raise _batch_error(
@@ -505,6 +667,7 @@ class MiniCPMO45Code2Wav(nn.Module):
                     error=str(exc),
                 ) from exc
             if len(audios) != batch_size or len(next_states) != batch_size:
+                self._prune_unowned_runtime_prompts()
                 raise _batch_error(
                     "backend_result_size_mismatch",
                     expected=batch_size,
@@ -525,10 +688,10 @@ class MiniCPMO45Code2Wav(nn.Module):
                     )
                 )
 
+        self._commit_runtime_prompt_owners(items)
         for request_id, state in pending.items():
             if state is None:
                 self._states.pop(request_id, None)
-                self._release_prompt(request_id)
             else:
                 self._states[request_id] = state
         sample_rate_tensor = torch.as_tensor(sample_rate, dtype=torch.int32)
@@ -537,24 +700,15 @@ class MiniCPMO45Code2Wav(nn.Module):
             multimodal_outputs={
                 "model_outputs": outputs,
                 "sr": [sample_rate_tensor.clone() for _ in outputs],
-                "meta.llm_output_text_utf8": [
-                    torch.tensor(list(item.segment_text.encode("utf-8")), dtype=torch.uint8) for item in items
-                ],
-                "meta.tts_is_last_chunk": [torch.tensor(item.last_chunk, dtype=torch.bool) for item in items],
+                # Generation runner wire payloads are flat and tensor-only.
+                # Dotted metadata keys are unflattened again by the output
+                # processor before the full-duplex data plane consumes them.
+                "meta.duplex_epoch": [torch.tensor(item.duplex_epoch, dtype=torch.int32) for item in items],
+                "meta.duplex_turn_id": [torch.tensor(item.duplex_turn_id, dtype=torch.int32) for item in items],
+                "meta.llm_output_text_utf8": [item.segment_text_utf8 for item in items],
+                "meta.tts_is_last_chunk": [torch.tensor(item.tts_is_last_chunk, dtype=torch.bool) for item in items],
                 "meta.segment_end": [torch.tensor(item.segment_end, dtype=torch.bool) for item in items],
                 "meta.turn_end": [torch.tensor(item.turn_end, dtype=torch.bool) for item in items],
-                "meta.duplex_turn_id": [
-                    torch.tensor(item.duplex_turn_id, dtype=torch.int64)
-                    if item.duplex_turn_id is not None
-                    else torch.empty(0, dtype=torch.int64)
-                    for item in items
-                ],
-                "meta.duplex_epoch": [
-                    torch.tensor(item.duplex_epoch, dtype=torch.int64)
-                    if item.duplex_epoch is not None
-                    else torch.empty(0, dtype=torch.int64)
-                    for item in items
-                ],
             },
         )
 
@@ -562,7 +716,7 @@ class MiniCPMO45Code2Wav(nn.Module):
         for request_id in finished_req_ids:
             state_id = str(request_id)
             self._states.pop(state_id, None)
-            self._release_prompt(state_id)
+            self._release_request_prompt(state_id)
 
     def make_omni_output(self, model_outputs: Any, **_: Any) -> OmniOutput:
         if isinstance(model_outputs, OmniOutput):
@@ -600,10 +754,11 @@ class MiniCPMO45Code2Wav(nn.Module):
             from stepaudio2.token2wav import Token2wav
 
         extra = self._extra_config()
+        model_root = self._resolve_model_root()
         prompt_path = Path(self._default_prompt_wav)
         if not prompt_path.is_file():
             raise FileNotFoundError(f"MiniCPM-o Code2Wav prompt audio not found: {prompt_path}")
-        token2wav_path = Path(self.model_path) / "assets" / "token2wav"
+        token2wav_path = model_root / "assets" / "token2wav"
         if not token2wav_path.is_dir():
             raise FileNotFoundError(f"MiniCPM-o Code2Wav assets not found: {token2wav_path}")
         use_float16 = bool(extra.get("token2wav_float16", False))

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py

@@ -175,7 +175,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
 
         token_ids = self._minicpmo45_native_duplex_token_ids()
         listen_id = int(token_ids.get("listen_token_id", -1))
-        tts_bos_id = int(token_ids.get("tts_bos_token_id", -1))
         turn_eos_id = int(token_ids.get("turn_eos_token_id", -1))
         if listen_id < 0 or listen_id >= logits.shape[-1]:
             return
@@ -199,7 +198,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
                 continue
             force_listen = payload.get("force_listen") is True
             is_speech = payload.get("is_speech")
-            redirect_listen = False
             segment_key = (row.request_id, row.seq if row.seq is not None else -1)
             session_key = (row.session_id, row.incarnation) if row.session_id is not None else None
             if turn_eos_id >= 0 and session_key is not None:
@@ -209,8 +207,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
                 )
                 if is_speech is True:
                     if state is not None:
-                        if not getattr(state, "current_turn_ended", True):
-                            redirect_listen = True
                         with suppress(Exception):
                             state.last_terminator_token = None
                 else:
@@ -218,7 +214,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
                     if turn_ended and not pending_speech_context:
                         force_listen = True
 
-            if not force_listen and not redirect_listen:
+            if not force_listen:
                 continue
             if force_listen and segment_key in force_listen_segments:
                 continue
@@ -226,10 +222,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
                 logits[row_idx, :] = float("-inf")
                 logits[row_idx, listen_id] = 0.0
                 force_listen_segments.add(segment_key)
-            elif 0 <= tts_bos_id < logits.shape[-1]:
-                if logits[row_idx, listen_id] > logits[row_idx, tts_bos_id]:
-                    logits[row_idx, tts_bos_id] = logits[row_idx, listen_id]
-                logits[row_idx, listen_id] = float("-inf")
 
     # -------------------- Device utilities --------------------
     @staticmethod
@@ -730,6 +722,10 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
         output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or []
         raw_recent_tokens = output_token_ids[row_idx] if row_idx < len(output_token_ids) else []
         recent_tokens = [int(token_id) for token_id in raw_recent_tokens if isinstance(token_id, int) and token_id >= 0]
+        temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7))
+        top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100))
+        top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8))
+        state = self._minicpmo45_duplex_state_for_row(row_idx)
         if chunk_eos_id >= 0 and chunk_eos_id < logits.shape[-1]:
             max_speak_tokens = int(
                 getattr(
@@ -765,23 +761,19 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
             if valid_forbidden:
                 logits[:, valid_forbidden] = float("-inf")
 
-        special_ids = self._minicpmo45_native_special_token_ids(token_ids)
-        state = self._minicpmo45_duplex_state_for_row(row_idx)
-        generated_text_tokens = getattr(state, "generated_text_tokens", None)
-        repetition_tokens = generated_text_tokens if generated_text_tokens else recent_tokens
+        generated_tokens = getattr(state, "generated_tokens", None)
+        repetition_tokens = generated_tokens if generated_tokens else recent_tokens
         repetition_penalty = 1.05
         if repetition_penalty != 1.0 and repetition_tokens:
             history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
             for token_id in set(repetition_tokens[-history_size:]):
-                if token_id in special_ids or token_id < 0 or token_id >= logits.shape[-1]:
+                if token_id < 0 or token_id >= logits.shape[-1]:
                     continue
                 logits[0, token_id] /= repetition_penalty
 
-        temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7))
-        top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100))
-        top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8))
         if getattr(sampling_metadata, "all_greedy", False) or temperature <= 0:
             sampled = int(torch.argmax(logits, dim=-1).item())
+            self._record_minicpmo45_duplex_generation_token(row_idx, sampled)
             sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk(
                 sampled,
                 recent_tokens,
@@ -793,6 +785,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
         logits = self._top_k_top_p_filter(logits, top_k=top_k, top_p=top_p)
         probs = F.softmax(logits, dim=-1)
         sampled = int(torch.multinomial(probs, num_samples=1, generator=generator).item())
+        self._record_minicpmo45_duplex_generation_token(row_idx, sampled)
         sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk(
             sampled,
             recent_tokens,
@@ -901,6 +894,28 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
             return int(tts_bos_id)
         return int(sampled)
 
+    def _record_minicpmo45_duplex_generation_token(self, row_idx: int, sampled: int) -> None:
+        """Track tokens returned by the model-policy decoder.
+
+        The released StreamDecoder is constructed without special_token_ids, so
+        every normally decoded token participates in repetition penalty. Forced
+        listen bypasses that decoder, while chunk_eos boundary decisions return
+        before this method is called.
+        """
+        state = self._minicpmo45_duplex_state_for_row(row_idx)
+        if state is None:
+            return
+        payload = self._minicpmo45_duplex_payload_for_row(row_idx)
+        if isinstance(payload, dict) and payload.get("force_listen") is True:
+            return
+        generated_tokens = getattr(state, "generated_tokens", None)
+        if not isinstance(generated_tokens, list):
+            generated_tokens = []
+            state.generated_tokens = generated_tokens
+        generated_tokens.append(int(sampled))
+        history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
+        del generated_tokens[:-history_size]
+
     def _record_minicpmo45_duplex_terminator(self, row_idx: int, sampled: int, token_ids: dict[str, int]) -> None:
         """Remember sampled unit state for the next append.
 
@@ -928,15 +943,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
                 with suppress(Exception):
                     state.pending_speech_response_open = False
             return
-        special_ids = self._minicpmo45_native_special_token_ids(token_ids)
-        if sampled not in special_ids:
-            generated_text_tokens = getattr(state, "generated_text_tokens", None)
-            if not isinstance(generated_text_tokens, list):
-                generated_text_tokens = []
-                state.generated_text_tokens = generated_text_tokens
-            generated_text_tokens.append(int(sampled))
-            history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
-            del generated_text_tokens[:-history_size]
         if (
             sampled == tts_bos_id
             and getattr(state, "current_turn_ended", True)

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py

@@ -43,6 +43,7 @@ _CODEC_TOP_K = 25
 _CODEC_TOP_P = 0.85
 _CODEC_REPETITION_PENALTY = 1.05
 _CODEC_MIN_TOKENS = 50
+_DUPLEX_CODEC_TOKENS_PER_CHUNK = 26
 
 
 def _max_audio_tokens(condition_tokens: int) -> int:
@@ -61,12 +62,7 @@ def _max_audio_tokens(condition_tokens: int) -> int:
 
 def _restore_weight_norm_weight(weight_g: torch.Tensor, weight_v: torch.Tensor) -> torch.Tensor:
     """Materialize ``weight_norm(..., dim=0)`` checkpoint parameters."""
-    row_norm = torch.linalg.vector_norm(
-        weight_v.float(),
-        dim=1,
-        keepdim=True,
-    ).clamp_min(torch.finfo(torch.float32).tiny)
-    return weight_v * (weight_g.to(dtype=weight_v.dtype) / row_norm.to(dtype=weight_v.dtype))
+    return torch._weight_norm(weight_v, weight_g, dim=0)
 
 
 def _apply_repetition_penalty(
@@ -218,6 +214,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
         self,
         tts_token_ids: torch.Tensor,
         tts_hidden_states: torch.Tensor,
+        *,
+        native_duplex: bool = False,
     ) -> torch.Tensor:
         if tts_token_ids.numel() == 0 or tts_hidden_states.numel() == 0:
             # The thinker can legally emit an empty speech segment (<|tts_bos|>
@@ -239,7 +237,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
         hidden_embeds = self.projector_semantic(hidden)
         if self._normalize:
             hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1)
-        return torch.cat([text_embeds + hidden_embeds, self._boundary_embeddings()], dim=0)
+        audio_bos = self.emb_text(torch.tensor([self._tts_bos_id], device=device, dtype=torch.long))
+        condition = text_embeds + hidden_embeds
+        if native_duplex:
+            # Match MiniCPMTTS.generate_chunk's streaming condition.
+            return torch.cat([condition, audio_bos], dim=0)
+        return torch.cat([condition, self._boundary_embeddings()], dim=0)
 
     def preprocess(
         self,
@@ -281,18 +284,22 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
                     "MiniCPM-o Talker received an empty condition (request %s); this request produces no audio.",
                     info_dict.get("request_id"),
                 )
-            full_embeds = self._build_condition_embeddings(token_ids, hidden_states)
+            native_duplex = bool(info_dict.get("native_duplex", False))
+            full_embeds = self._build_condition_embeddings(
+                token_ids,
+                hidden_states,
+                native_duplex=native_duplex,
+            )
             offset = int(info_dict.get("_omni_num_computed_tokens", 0))
-            if offset == 0 and span_len > full_embeds.shape[0]:
-                # Async-chunk prewarms Stage 1 with placeholder token IDs. Two
-                # Thinker handoffs can arrive before the first one executes,
-                # so the scheduler prompt may already include the newer
-                # placeholder prefix while additional_information still
-                # carries the preceding condition. Materialize that prefix as
-                # the embeddings of its actual zero token IDs and keep the
-                # complete Talker condition tail-aligned (audio BOS remains
-                # the final prefill token).
-                prefix_len = span_len - full_embeds.shape[0]
+            request_id = str(info_dict.get("request_id", "0"))
+            meta = info_dict.get("meta")
+            # The handoff rebuilds only the tail-aligned Talker condition.
+            # Materialize zero-token embeddings for any scheduler prompt
+            # prefix so chunked prefill can slice from a non-zero offset.
+            prompt_len = info_dict.get("_omni_prompt_len")
+            target_len = int(prompt_len) if prompt_len is not None else offset + span_len
+            prefix_len = target_len - full_embeds.shape[0]
+            if prefix_len > 0:
                 placeholder_ids = torch.zeros(
                     prefix_len,
                     dtype=torch.long,
@@ -308,13 +315,21 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
                     f"tts_ids={token_ids.shape[0]} tts_hidden={hidden_states.shape[0]} "
                     f"prompt_len={info_dict.get('_omni_prompt_len')}"
                 )
-            max_tokens = _max_audio_tokens(int(token_ids.numel()))
+            duplex_boundary = isinstance(meta, dict) and (
+                bool(meta.get("turn_start", False)) or bool(meta.get("turn_end", False))
+            )
+            if native_duplex:
+                max_tokens = _DUPLEX_CODEC_TOKENS_PER_CHUNK
+                min_tokens = 0 if duplex_boundary else _DUPLEX_CODEC_TOKENS_PER_CHUNK
+            else:
+                max_tokens = _max_audio_tokens(int(token_ids.numel()))
+                min_tokens = self._codec_min_tokens
             state = {
                 "step": 0,
                 "max_tokens": max_tokens,
+                "min_tokens": min_tokens,
                 "finished": empty_condition,
             }
-            request_id = str(info_dict.get("request_id", "0"))
             request_states = getattr(self, "_request_audio_states", None)
             if request_states is None:
                 request_states = {}
@@ -370,7 +385,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
             penalty=self._codec_repetition_penalty,
             window_size=_REPETITION_WINDOW,
         )
-        if step < self._codec_min_tokens:
+        request_states = getattr(self, "_request_audio_states", {})
+        state = request_states.get(request_id)
+        min_tokens = (
+            int(state.get("min_tokens", self._codec_min_tokens)) if isinstance(state, dict) else self._codec_min_tokens
+        )
+        if step < min_tokens:
             logits[..., eos_id] = float("-inf")
         logits = _apply_top_k_top_p(
             logits,
@@ -404,12 +424,61 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
             raise RuntimeError(
                 f"MiniCPM-o continuous Talker received {len(sample_eligible)} sampling flags for {len(infos)} requests"
             )
+        emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos)
 
         stop_rows: list[torch.Tensor] = []
         codec_deltas: list[torch.Tensor] = []
         terminal_flags: list[torch.Tensor] = []
+        native_duplex_flags: list[torch.Tensor] = []
+        duplex_epochs: list[torch.Tensor] = []
+        duplex_turn_ids: list[torch.Tensor] = []
+        segment_texts_utf8: list[torch.Tensor] = []
+        turn_end_flags: list[torch.Tensor] = []
         empty_delta = hidden.new_empty((0, 1), dtype=torch.long)
         for index, info in enumerate(infos):
+            info_dict = info if isinstance(info, dict) else {}
+            native_duplex = info_dict.get("native_duplex") is True
+            if emit_duplex_metadata:
+                duplex_info = info_dict.get("duplex")
+                if not isinstance(duplex_info, dict):
+                    duplex_info = {}
+                epoch = duplex_info.get("epoch", -1)
+                turn_id = duplex_info.get("turn_id", -1)
+                if native_duplex and not all(
+                    isinstance(value, int) and not isinstance(value, bool) and value >= 0 for value in (epoch, turn_id)
+                ):
+                    raise RuntimeError(
+                        "MiniCPM-o native duplex Talker requires non-negative integer "
+                        f"epoch and turn_id, got epoch={epoch!r}, turn_id={turn_id!r}"
+                    )
+                meta_info = info_dict.get("meta")
+                if not isinstance(meta_info, dict):
+                    meta_info = {}
+                segment_text = meta_info.get("native_duplex_segment_text", "") if native_duplex else ""
+                if not isinstance(segment_text, str):
+                    segment_text = ""
+                turn_eos_id = meta_info.get("turn_eos_token_id")
+                ids_info = info_dict.get("ids")
+                tts_ids = ids_info.get("tts") if native_duplex and isinstance(ids_info, dict) else None
+                if isinstance(tts_ids, torch.Tensor):
+                    contains_turn_eos = isinstance(turn_eos_id, int) and bool(
+                        torch.any(tts_ids.reshape(-1) == turn_eos_id).item()
+                    )
+                elif isinstance(tts_ids, (list, tuple)):
+                    contains_turn_eos = isinstance(turn_eos_id, int) and turn_eos_id in tts_ids
+                else:
+                    contains_turn_eos = False
+                native_duplex_flags.append(torch.tensor(native_duplex, dtype=torch.bool))
+                duplex_epochs.append(torch.tensor(epoch if isinstance(epoch, int) else -1, dtype=torch.long))
+                duplex_turn_ids.append(torch.tensor(turn_id if isinstance(turn_id, int) else -1, dtype=torch.long))
+                segment_texts_utf8.append(
+                    torch.tensor(
+                        list(segment_text.encode("utf-8")),
+                        dtype=torch.uint8,
+                    )
+                )
+                turn_end_flags.append(torch.tensor(native_duplex and contains_turn_eos, dtype=torch.bool))
+
             if not isinstance(info, dict):
                 stop_rows.append(hidden.new_tensor([0.0, float("-inf")]))
                 codec_deltas.append(empty_delta)
@@ -453,12 +522,15 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
                 codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1)
             step = int(state.get("step", 0))
             sampled = self._sample_audio_code(hidden[end - 1 : end], codes, request_id, step)
-            is_eos = int(sampled.item()) == self._num_audio_tokens - 1
+            sampled_id = int(sampled.item())
+            is_eos = sampled_id == self._num_audio_tokens - 1
             state["step"] = int(state.get("step", 0)) + 1
             reached_limit = int(state["step"]) >= int(state.get("max_tokens", 2048))
             finished = is_eos or reached_limit
             state["finished"] = finished
-            if not is_eos:
+            # MiniCPMTTS.generate_chunk consumes the boundary sample but
+            # returns only codes that were fed into the retained KV state.
+            if not is_eos and not reached_limit:
                 codes = torch.cat([codes[-(_REPETITION_WINDOW - 1) :], sampled.reshape(1)])
                 delta = sampled.reshape(1, 1)
             else:
@@ -476,9 +548,20 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
         self._batch_stop_logits = torch.stack(stop_rows, dim=0) if stop_rows else hidden.new_empty((0, 2))
         # Lists are deliberate: the runner routes element i to request i,
         # preserving compaction alignment while emitting only this step's code.
+        meta_outputs = {"finished": terminal_flags}
+        if emit_duplex_metadata:
+            meta_outputs.update(
+                {
+                    "native_duplex": native_duplex_flags,
+                    "duplex_epoch": duplex_epochs,
+                    "duplex_turn_id": duplex_turn_ids,
+                    "llm_output_text_utf8": segment_texts_utf8,
+                    "turn_end": turn_end_flags,
+                }
+            )
         multimodal_outputs: dict[str, Any] = {
             "codes": {"audio": codec_deltas},
-            "meta": {"finished": terminal_flags},
+            "meta": meta_outputs,
         }
         return OmniOutput(
             text_hidden_states=hidden,

vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py

@@ -21,6 +21,7 @@ logger = logging.getLogger(__name__)
 _MINICPMO45_ASYNC_STATE = "_minicpmo45_async_codec_state"
 _MINICPMO45_STREAM_RECORD = "_minicpmo45_async_stream_record"
 _MINICPMO45_SILENCE_CODE = 4218
+_MINICPMO45_MIN_STREAM_BODY_FRAMES = 5
 
 
 class _MiniCPMO45MetaStruct(MetaStruct):
@@ -28,6 +29,7 @@ class _MiniCPMO45MetaStruct(MetaStruct):
 
     ref_audio_sr: int | None = None
     native_duplex_segment_text: str | None = None
+    llm_output_text_utf8: torch.Tensor | None = None
     duplex_turn_id: int | None = None
     duplex_epoch: int | None = None
     segment_end: bool | None = None
@@ -212,6 +214,16 @@ def tts2code2wav_async_chunk(
     internal_id = getattr(request, "request_id", None)
     request_id = str(external_id if external_id is not None else internal_id)
     internal_id = str(internal_id if internal_id is not None else request_id)
+    output_meta = multimodal_output.get("meta") if isinstance(multimodal_output, Mapping) else None
+    output_meta = output_meta if isinstance(output_meta, Mapping) else {}
+    native_duplex = bool(_coerce_int(output_meta.get("native_duplex")))
+    duplex_epoch = _coerce_int(output_meta.get("duplex_epoch"))
+    duplex_turn_id = _coerce_int(output_meta.get("duplex_turn_id"))
+    segment_text_utf8 = output_meta.get("llm_output_text_utf8")
+    if not isinstance(segment_text_utf8, torch.Tensor):
+        segment_text_utf8 = None
+    turn_end = bool(_coerce_int(output_meta.get("turn_end")))
+    duplex_turn_key = (duplex_epoch, duplex_turn_id)
 
     request_payload = getattr(transfer_manager, "request_payload", None)
     if request_payload is None:
@@ -227,7 +239,9 @@ def tts2code2wav_async_chunk(
         record = {
             "internal_id": internal_id,
             "cache_epoch": 0,
+            "chunk_seq": 0,
             "retired_internal_ids": set(),
+            "last_terminal_turn": None,
         }
         container[_MINICPMO45_STREAM_RECORD] = record
     elif internal_id in record["retired_internal_ids"]:
@@ -236,6 +250,8 @@ def tts2code2wav_async_chunk(
         record["retired_internal_ids"].add(record["internal_id"])
         record["internal_id"] = internal_id
         record["cache_epoch"] = int(record["cache_epoch"]) + 1
+        record["chunk_seq"] = 0
+        record["last_terminal_turn"] = None
         _drop_codec_state(transfer_manager, request_id)
 
     if _is_aborted(request):
@@ -243,44 +259,46 @@ def tts2code2wav_async_chunk(
         _drop_codec_state(transfer_manager, request_id)
         return None
 
+    if native_duplex and turn_end and record.get("last_terminal_turn") == duplex_turn_key:
+        return None
+
     state = container.get(_MINICPMO45_ASYNC_STATE)
     if not isinstance(state, dict):
-        request_info = getattr(request, "additional_information", None)
-        if not isinstance(request_info, Mapping):
-            request_info = {}
-        codes_info = request_info.get("codes")
-        if not isinstance(codes_info, Mapping):
-            codes_info = {}
-        meta_info = request_info.get("meta")
-        if not isinstance(meta_info, Mapping):
-            meta_info = {}
-        duplex_info = request_info.get("duplex")
-        if not isinstance(duplex_info, Mapping):
-            duplex_info = {}
         state = {
             "internal_id": internal_id,
             "pending": [],
+            "pending_text_utf8": [],
+            "segment_text_recorded": False,
             "left_context": [],
             "codec_end": 0,
-            "ref_audio": codes_info.get("ref"),
-            "ref_audio_sr": meta_info.get("ref_audio_sr"),
-            "segment_text": meta_info.get("native_duplex_segment_text"),
-            "segment_end": bool(meta_info.get("segment_end", False)),
-            "turn_end": bool(meta_info.get("turn_end", False)),
-            "duplex_turn_id": duplex_info.get("model_turn_id", duplex_info.get("turn_id")),
-            "duplex_epoch": duplex_info.get("epoch"),
         }
         container[_MINICPMO45_ASYNC_STATE] = state
 
     pending = state["pending"]
     pending.extend(_extract_codec_delta(multimodal_output, request_id))
+    pending_text_utf8 = state.setdefault("pending_text_utf8", [])
+    current_text_utf8 = (
+        segment_text_utf8.detach().to(device="cpu", dtype=torch.uint8).reshape(-1).tolist()
+        if isinstance(segment_text_utf8, torch.Tensor)
+        else []
+    )
+    if native_duplex and current_text_utf8 and not state.get("segment_text_recorded", False):
+        # Talker repeats the unit text on every codec step. Queue it once for
+        # the first Code2Wav payload that can carry audio for this unit.
+        pending_text_utf8.extend(current_text_utf8)
+        state["segment_text_recorded"] = True
     request_finished = getattr(request, "is_finished", None)
     finished = bool(is_finished or (callable(request_finished) and request_finished()))
     chunk_frames, left_context_frames = _codec_config(transfer_manager)
-    if not finished and len(pending) < chunk_frames:
+    flush_pending = finished
+    last_chunk = bool(flush_pending and (not native_duplex or turn_end))
+    if not flush_pending and len(pending) < chunk_frames:
         return None
 
-    new_token_count = len(pending) if finished else chunk_frames
+    hold_short_unit = (
+        native_duplex and flush_pending and not last_chunk and 0 < len(pending) < _MINICPMO45_MIN_STREAM_BODY_FRAMES
+    )
+    new_token_count = 0 if hold_short_unit else (len(pending) if flush_pending else chunk_frames)
     new_codes = pending[:new_token_count]
     del pending[:new_token_count]
     codec_start = int(state["codec_end"])
@@ -292,53 +310,79 @@ def tts2code2wav_async_chunk(
         else:
             context = list(state["left_context"])
         output_codes = [*context, *new_codes]
-        history = [*state["left_context"], *new_codes]
+        history = output_codes
         state["left_context"] = history[-left_context_frames:] if left_context_frames else []
-    elif finished and codec_start > 0 and state["left_context"]:
+    elif last_chunk and codec_start > 0 and state["left_context"]:
         context = list(state["left_context"])
         output_codes = context
     else:
         context = []
         output_codes = []
     state["codec_end"] = codec_end
+    code_flat_numel = len(output_codes)
+    if native_duplex and code_flat_numel > 0:
+        segment_text_utf8 = torch.tensor(pending_text_utf8, dtype=torch.uint8)
+        pending_text_utf8.clear()
+    if native_duplex and finished:
+        state["segment_text_recorded"] = False
+    if flush_pending and not last_chunk and code_flat_numel == 0:
+        # Keep the generic generation connector model-agnostic: a real token
+        # makes this control-only TTS boundary schedulable, while the explicit
+        # zero length tells Code2Wav to discard the placeholder.
+        output_codes = [0]
 
-    last_chunk = bool(finished and not pending)
-    if last_chunk:
+    if last_chunk and not native_duplex:
         record["retired_internal_ids"].add(internal_id)
         _drop_codec_state(transfer_manager, request_id)
 
-    chunk_seq = int(getattr(transfer_manager, "put_req_chunk", {}).get(request_id, 0))
+    chunk_seq = int(record["chunk_seq"])
+    record["chunk_seq"] = chunk_seq + 1
+    ref_audio = None
+    ref_audio_sr = None
+    if int(record["cache_epoch"]) == 0 and chunk_seq == 0:
+        request_info = getattr(request, "additional_information", None)
+        if isinstance(request_info, Mapping):
+            codes_info = request_info.get("codes")
+            meta_info = request_info.get("meta")
+            raw_ref_audio = codes_info.get("ref") if isinstance(codes_info, Mapping) else None
+            raw_ref_audio_sr = meta_info.get("ref_audio_sr") if isinstance(meta_info, Mapping) else None
+            ref_audio_sr = _coerce_int(raw_ref_audio_sr)
+            if raw_ref_audio is not None:
+                ref_audio = torch.as_tensor(raw_ref_audio, dtype=torch.float32).reshape(-1).cpu()
     finished_tensor = torch.tensor(last_chunk, dtype=torch.bool)
-    ref_audio = state.get("ref_audio") if codec_start == 0 else None
-    return OmniPayloadStruct(
+    payload = OmniPayloadStruct(
         codes=CodesStruct(
             audio=torch.tensor(output_codes, dtype=torch.long),
-            ref=torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1) if ref_audio is not None else None,
+            ref=ref_audio,
         ),
         meta=_MiniCPMO45MetaStruct(
             request_id=request_id,
             chunk_seq=chunk_seq,
             cache_epoch=int(record["cache_epoch"]),
-            code_flat_numel=len(output_codes),
+            code_flat_numel=code_flat_numel,
             codec_chunk_frames=new_token_count,
             codec_left_context_frames=len(context),
             left_context_size=len(context),
             last_chunk=last_chunk,
             stream_finished=finished_tensor,
             finished=finished_tensor,
+            is_segment_finished=finished_tensor,
             req_id=[request_id],
-            ref_audio_sr=_coerce_int(state.get("ref_audio_sr")),
-            native_duplex_segment_text=(
-                str(state["segment_text"]) if isinstance(state.get("segment_text"), str) else None
-            ),
-            duplex_turn_id=_coerce_int(state.get("duplex_turn_id")),
-            duplex_epoch=_coerce_int(state.get("duplex_epoch")),
-            segment_end=bool(state.get("segment_end", False)),
-            turn_end=bool(state.get("turn_end", False)),
-            tts_is_last_chunk=last_chunk,
+            duplex_epoch=duplex_epoch,
+            duplex_turn_id=duplex_turn_id,
+            llm_output_text_utf8=segment_text_utf8,
+            tts_is_last_chunk=flush_pending,
+            turn_end=turn_end and last_chunk,
+            ref_audio_sr=ref_audio_sr,
         ),
         request_id=request_id,
     )
+    if last_chunk and native_duplex:
+        record["last_terminal_turn"] = duplex_turn_key
+        record["cache_epoch"] = int(record["cache_epoch"]) + 1
+        record["chunk_seq"] = 0
+        _drop_codec_state(transfer_manager, request_id)
+    return payload
 
 
 def tts2code2wav_full_payload(
@@ -496,13 +540,33 @@ def _native_tts_boundary_token_ids(special_token_ids):
     }
 
 
+def _decode_native_duplex_token_ids(
+    token_ids: list[int],
+    streaming_context,
+    *,
+    request_id: str,
+) -> str | None:
+    decode_token_ids = getattr(streaming_context, "source_token_decoder", None)
+    if not isinstance(decode_token_ids, Callable):
+        return None
+    decode_ids = [int(token_id) for token_id in token_ids]
+    try:
+        try:
+            return str(decode_token_ids(decode_ids, skip_special_tokens=True))
+        except TypeError:
+            return str(decode_token_ids(decode_ids))
+    except Exception:
+        logger.exception("Failed to decode MiniCPM-o duplex token delta for request_id=%s", request_id)
+        return ""
+
+
 def _native_duplex_segment_output_ids(
     output_ids: list[int],
     output_text: str,
     streaming_context,
     *,
     request_id: str,
-) -> tuple[list[int], str]:
+) -> tuple[list[int], str, bool]:
     """Slice the cumulative thinker output down to the current segment.
 
     Tracks how many output tokens were already handed to the talker in the
@@ -511,7 +575,7 @@ def _native_duplex_segment_output_ids(
     """
     bridge_states = getattr(streaming_context, "bridge_states", None)
     if not isinstance(bridge_states, dict):
-        return output_ids, output_text
+        return output_ids, output_text, True
     state = bridge_states.setdefault("minicpmo45_tts_handoff", {})
     duplex_state = bridge_states.get("duplex")
     turn_id = duplex_state.get("model_turn_id", duplex_state.get("turn_id")) if isinstance(duplex_state, dict) else None
@@ -521,17 +585,12 @@ def _native_duplex_segment_output_ids(
         state["request_id"] = request_id
         state["sent_output_len"] = 0
         state["sent_output_ids"] = []
-        state["acc_tts_ids"] = []
-        state["acc_tts_hidden"] = []
+    previous_turn_id = state.get("turn_id")
     sent_len = state.get("sent_output_len", 0)
     prev_output_ids = state.get("sent_output_ids", [])
-    prev_turn_id = state.get("turn_id")
     if not isinstance(sent_len, int) or sent_len < 0 or sent_len > len(output_ids):
-        # Shrunken cumulative output = epoch reset after barge-in: the talker
-        # condition history is stale too.
+        # Shrunken cumulative output = epoch reset after barge-in.
         sent_len = 0
-        state["acc_tts_ids"] = []
-        state["acc_tts_hidden"] = []
     elif sent_len and isinstance(prev_output_ids, list):
         prev_prefix = prev_output_ids[:sent_len]
         current_prefix = output_ids[:sent_len]
@@ -540,27 +599,15 @@ def _native_duplex_segment_output_ids(
             # others keep returning cumulative ids. Detect restart from the
             # token prefix instead of clearing the cursor at turn_eos.
             sent_len = 0
-            state["acc_tts_ids"] = []
-            state["acc_tts_hidden"] = []
-    if isinstance(prev_turn_id, int) and isinstance(turn_id, int) and prev_turn_id != turn_id:
-        # The thinker output list is cumulative across clean turns, so the
-        # output cursor must stay put. The talker condition is per assistant
-        # turn, though; carrying it across turn_id boundaries makes stage 1
-        # replay the previous turn after its consumed cursor has been reset.
-        state["acc_tts_ids"] = []
-        state["acc_tts_hidden"] = []
+    turn_start = sent_len == 0 or previous_turn_id != turn_id
     segment_ids = output_ids[sent_len:]
-    decode_token_ids = getattr(streaming_context, "source_token_decoder", None)
-    if isinstance(decode_token_ids, Callable):
-        decode_ids = [int(token_id) for token_id in segment_ids]
-        try:
-            try:
-                segment_text = str(decode_token_ids(decode_ids, skip_special_tokens=True))
-            except TypeError:
-                segment_text = str(decode_token_ids(decode_ids))
-        except Exception:
-            logger.exception("Failed to decode MiniCPM-o duplex token delta for request_id=%s", request_id)
-            segment_text = ""
+    decoded_segment_text = _decode_native_duplex_token_ids(
+        segment_ids,
+        streaming_context,
+        request_id=request_id,
+    )
+    if decoded_segment_text is not None:
+        segment_text = decoded_segment_text
     elif sent_len == 0:
         # Non-orchestrator unit tests and legacy callers may not install a
         # decoder. Never slice cumulative text with a stale character cursor.
@@ -575,22 +622,7 @@ def _native_duplex_segment_output_ids(
     state["sent_output_len"] = len(output_ids)
     state["sent_output_ids"] = list(output_ids)
     state["turn_id"] = turn_id
-    return segment_ids, segment_text
-
-
-def _reset_native_tts_handoff(streaming_context) -> None:
-    bridge_states = getattr(streaming_context, "bridge_states", None)
-    if not isinstance(bridge_states, dict):
-        return
-    state = bridge_states.get("minicpmo45_tts_handoff")
-    if not isinstance(state, dict):
-        return
-    # Keep the output cursor: the thinker can keep reporting cumulative
-    # output after turn_eos on the same resumable request. If a runtime really
-    # restarts token ids, _native_duplex_segment_output_ids detects the prefix
-    # mismatch and resets the slice cursor there.
-    state["acc_tts_ids"] = []
-    state["acc_tts_hidden"] = []
+    return segment_ids, segment_text, turn_start
 
 
 def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | None:
@@ -623,32 +655,6 @@ def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] |
     return metadata
 
 
-def _accumulate_native_tts_handoff(streaming_context, new_ids, new_hidden):
-    """Hand the talker the FULL accumulated condition on every handoff.
-
-    The runner's streaming buffer update is not merge-safe for a resumable
-    stage-1 request: in-place updates merge sub-keys, but a resume prefill
-    REPLACES the buffer, silently dropping every earlier segment's tts
-    tokens/hiddens (observed losing alternating reply segments — the talker
-    vocalized text it never saw between islands it did). Accumulate here so
-    the latest handoff always carries the complete history and downstream
-    replace semantics are lossless; the talker consumes by cursor.
-    """
-    bridge_states = getattr(streaming_context, "bridge_states", None)
-    if not isinstance(bridge_states, dict):
-        return new_ids, new_hidden
-    state = bridge_states.setdefault("minicpmo45_tts_handoff", {})
-    acc_ids = state.setdefault("acc_tts_ids", [])
-    acc_hidden = state.setdefault("acc_tts_hidden", [])
-    if new_ids:
-        acc_ids.extend(int(t) for t in new_ids)
-        if new_hidden:
-            acc_hidden.extend(new_hidden)
-    if not acc_ids:
-        return None, None
-    return list(acc_ids), list(acc_hidden)
-
-
 def _build_tts_scheduler_prompt_token_ids(
     tts_token_ids: torch.Tensor | None,
     llm_output_ids: list[int],
@@ -723,6 +729,7 @@ def llm2tts(
         # buffer overflows.
         llm_output_ids = list(llm_output_ids)
         thinker_text = getattr(output, "text", "") or ""
+        native_turn_start = False
         if _has_native_duplex_prompt_metadata(mm_output):
             # The thinker's resumable duplex request reports cumulative
             # output ids/text, but earlier segments are already folded into
@@ -732,7 +739,7 @@ def llm2tts(
             # alignment holds, the talker prompt grows linearly with new
             # tokens instead of quadratically, and downstream transcripts
             # carry per-unit deltas instead of re-sending the whole reply.
-            llm_output_ids, thinker_text = _native_duplex_segment_output_ids(
+            llm_output_ids, thinker_text, native_turn_start = _native_duplex_segment_output_ids(
                 llm_output_ids,
                 thinker_text,
                 _streaming_context,
@@ -875,6 +882,18 @@ def llm2tts(
                         .to(torch.float32)
                         .contiguous()
                     )
+        handoff_ids = _coerce_token_id_list(tts_token_ids_slice) if tts_token_ids_slice is not None else None
+        if is_native_duplex_handoff and handoff_ids:
+            handoff_text = _decode_native_duplex_token_ids(
+                handoff_ids,
+                _streaming_context,
+                request_id=str(llm_output.request_id),
+            )
+            if handoff_text is not None:
+                # Match the released streaming_generate contract: transcript
+                # text comes from total_ids_in_unit, the same slice that
+                # conditions the Talker.
+                thinker_text = handoff_text
         model_intermediate_buffer = build_duplex_intermediate_buffer(
             request_id=str(llm_output.request_id),
             prompt_token_ids=prompt_token_ids,
@@ -900,6 +919,7 @@ def llm2tts(
                 # The talker detects turn end from <|turn_eos|> inside the
                 # handed condition (official conditions on its embedding).
                 meta["turn_eos_token_id"] = int(turn_eos_id)
+            meta["turn_start"] = native_turn_start
             if native_segment_end:
                 meta["segment_end"] = True
         req_mm_data = multi_modal_data.get(llm_output.request_id)
@@ -911,30 +931,27 @@ def llm2tts(
         if ref_audio is not None:
             ref_waveform, ref_sr = ref_audio
             set_ref_audio(model_intermediate_buffer, _to_transport_list(ref_waveform), ref_sr)
-        handoff_ids = _coerce_token_id_list(tts_token_ids_slice) if tts_token_ids_slice is not None else None
         handoff_hidden = _to_transport_list(tts_hidden_slice) if tts_hidden_slice is not None else None
         native_turn_end_handoff = False
         if is_native_duplex_handoff:
             turn_eos_id = special_token_ids.get("turn_eos_token_id")
             native_turn_end_handoff = turn_eos_id is not None and handoff_ids is not None and turn_eos_id in handoff_ids
-            handoff_ids, handoff_hidden = _accumulate_native_tts_handoff(
-                _streaming_context,
-                handoff_ids,
-                handoff_hidden,
-            )
             if not handoff_ids:
                 continue
         set_tts_handoff(model_intermediate_buffer, handoff_ids, handoff_hidden)
         if native_turn_end_handoff:
             model_intermediate_buffer.setdefault("meta", {})["turn_end"] = True
-            _reset_native_tts_handoff(_streaming_context)
 
         if handoff_ids is not None and handoff_hidden is not None:
-            condition_length = max(len(handoff_ids), len(handoff_hidden)) + 2
+            condition_suffix_length = 1 if is_native_duplex_handoff else 2
+            condition_length = max(len(handoff_ids), len(handoff_hidden)) + condition_suffix_length
             scheduler_prompt_token_ids = [0] * condition_length
             handoff_meta = model_intermediate_buffer.setdefault("meta", {})
-            handoff_meta["replace_streaming_prompt"] = True
             handoff_meta["next_stage_prompt_len"] = condition_length
+            # Native duplex resumes one Talker request within a turn, but a new
+            # assistant turn must discard the previous turn's prompt and KV.
+            if not is_native_duplex_handoff or native_turn_start:
+                handoff_meta["replace_streaming_prompt"] = True
         else:
             scheduler_prompt_token_ids = _build_tts_scheduler_prompt_token_ids(
                 tts_token_ids_slice,

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -685,19 +685,32 @@ def talker2code2wav_async_chunk(
     """
     Multimodal output version.
     """
-    if not isinstance(multimodal_output, Mapping):
-        return None
-    talker_codes = multimodal_output.get("codes", {})
-    if not isinstance(talker_codes, dict):
-        return None
-    code_predictor_codes = talker_codes.get("audio")
-    if code_predictor_codes is None:
-        return None
+    request_id = request.external_req_id
+    code_predictor_codes = None
+    if isinstance(multimodal_output, Mapping):
+        talker_codes = multimodal_output.get("codes", {})
+        if isinstance(talker_codes, dict):
+            code_predictor_codes = talker_codes.get("audio")
 
-    if code_predictor_codes.numel() == 0:
-        return None
+    sampling_params = getattr(request, "sampling_params", None)
+    stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or [])
+    stop_token_id = getattr(sampling_params, "stop_token_id", None)
+    if stop_token_id is not None:
+        stop_token_ids.add(stop_token_id)
 
-    if not code_predictor_codes.any():
+    append_codes = (
+        isinstance(code_predictor_codes, torch.Tensor)
+        and code_predictor_codes.numel() > 0
+        and bool(code_predictor_codes.any())
+    )
+    if append_codes:
+        first_codebook = int(code_predictor_codes[0, 0].item())
+        if first_codebook in stop_token_ids:
+            logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook)
+            append_codes = False
+    if append_codes:
+        transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes)
+    elif not is_finished:
         return None
 
     connector = getattr(transfer_manager, "connector", None)
@@ -707,20 +720,10 @@ def talker2code2wav_async_chunk(
     left_context_size_config = int(cfg.get("codec_left_context_frames", 25))
     configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0)
 
-    sampling_params = getattr(request, "sampling_params", None)
-    stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or [])
-    stop_token_id = getattr(sampling_params, "stop_token_id", None)
-    if stop_token_id is not None:
-        stop_token_ids.add(stop_token_id)
-    first_codebook = int(code_predictor_codes[0, 0].item())
-    if first_codebook in stop_token_ids:
-        logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook)
-        return None
-
-    request_id = request.external_req_id
     chunk_id = transfer_manager.put_req_chunk[request_id]
-    transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes)
     length = len(transfer_manager.code_prompt_token_ids[request_id])
+    if length <= 0:
+        return None
 
     if configured_initial_chunk_size > 0:
         if chunk_id == 0:
@@ -732,6 +735,9 @@ def talker2code2wav_async_chunk(
     if chunk_length != 0 and not is_finished:
         return None
 
+    if is_finished and not append_codes and chunk_length == 0:
+        return None
+
     context_length = chunk_length if chunk_length != 0 else chunk_size_config
     # ensure left context does not exceed available length
     if configured_initial_chunk_size > 0 and chunk_id == 1: