From 07e1d53587b62fbd3c75a0b98b9c338af8dbf3ae Mon Sep 17 00:00:00 2001
From: LJX-xixi <lijinxi413@icloud.com>
Date: Fri, 10 Jul 2026 17:27:49 +0800
Subject: [PATCH 1/5] [Bugfix] Qwen2.5-Omni: route talker phases by runner
 metadata and handle decode spans statelessly

talker_preprocess guessed prefill vs decode from span length, predating the
per-request phase metadata from #3662. Multi-position decode spans
(speculative-decoding verify, preempt/resume replay) were misrouted to the
prefill branch; the decode branch filled only one position (later rows stayed
on uninitialized torch.empty scratch) and consumed the thinker_reply queue
statefully, drifting after any multi-token accept or replay.

Route by _omni_is_prefill (with num_computed/prompt_len and span-length
fallbacks, same layered pattern as qwen3_tts_talker), and assemble decode
spans statelessly: thinker_reply[num_computed - prompt_len + i] +
codec_embedding per position, span embeddings built once from
embed_input_ids, no queue consumption. Callers without the metadata keep the
unchanged one-step path.

FIX #5018

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: LJX-xixi <lijinxi413@icloud.com>
---
 .../test_qwen2_5_omni_talker_preprocess.py    | 168 ++++++++++++++++++
 .../models/qwen2_5_omni/qwen2_5_omni.py       |  87 +++++++--
 2 files changed, 242 insertions(+), 13 deletions(-)
 create mode 100644 tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py

diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
new file mode 100644
index 00000000000..e6e29b039c4
--- /dev/null
+++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit tests for the Qwen2.5-Omni talker preprocess phase routing and the
+stateless decode-span embedding assembly.
+
+Regression tests for the span-length prefill/decode heuristic: a
+speculative-decoding verify step schedules num_draft_tokens + 1 decode
+positions in one span, and a chunked prefill can end in a 1-token tail —
+both misroute under ``span_len > 1``. The decode path must also assemble
+one ``thinker_reply`` row per position, statelessly, so a preempt/resume
+replay span reconstructs exactly the embeddings the original steps used.
+"""
+
+import functools
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+HIDDEN = 4
+
+
+@functools.lru_cache(maxsize=1)
+def _qwen2_5_omni_mod():
+    """Import lazily so collection does not pull model_executor too early."""
+    import vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni as m
+
+    return m
+
+
+def _fake_embed_input_ids(ids: torch.Tensor) -> torch.Tensor:
+    # Deterministic, id-distinguishing stand-in for the codec embedding table.
+    return ids.reshape(-1, 1).to(torch.float32).expand(-1, HIDDEN).contiguous()
+
+
+def _make_minimal_model():
+    cls = _qwen2_5_omni_mod().Qwen2_5OmniForConditionalGeneration
+    # ``__new__`` skips ``__init__`` (no weights, no config); set only what
+    # the decode path touches.
+    model = cls.__new__(cls)
+    model.talker = SimpleNamespace(embed_input_ids=_fake_embed_input_ids)
+    return model
+
+
+def _thinker_reply(num_rows: int) -> torch.Tensor:
+    # Row k filled with 100 * (k + 1) so any offset error is visible.
+    return (torch.arange(1, num_rows + 1, dtype=torch.float32).reshape(-1, 1) * 100.0).expand(-1, HIDDEN).contiguous()
+
+
+def test_spec_verify_span_gets_one_thinker_row_per_position():
+    # A verify step with 2 draft tokens schedules a 3-position decode span at
+    # generation offset 2 (prompt_len 4, num_computed 6). Every position must
+    # get thinker_reply[offset + i] + codec_embedding(token_i).
+    model = _make_minimal_model()
+    input_ids = torch.tensor([11, 22, 33], dtype=torch.long)
+    q = _thinker_reply(8)
+
+    out_ids, out_embeds, update = model.talker_preprocess(
+        input_ids,
+        None,
+        embed={"thinker_reply": q},
+        _omni_is_prefill=False,
+        _omni_prompt_len=4,
+        _omni_num_computed_tokens=6,
+    )
+
+    assert torch.equal(out_ids, input_ids)
+    expected = _fake_embed_input_ids(input_ids) + q[2:5]
+    assert torch.equal(out_embeds, expected)
+    # Stateless: the reply buffer must not be consumed or rewritten.
+    assert update == {}
+    assert torch.equal(q, _thinker_reply(8))
+
+
+def test_decode_span_is_replayable():
+    # Preempt/resume replays already-generated tokens as one multi-token span.
+    # Because the reply buffer is never consumed, calling preprocess twice with
+    # the same metadata must produce identical embeddings (the stock consuming
+    # path would have shifted by one queue row per call).
+    model = _make_minimal_model()
+    input_ids = torch.tensor([7, 8], dtype=torch.long)
+    payload = dict(
+        embed={"thinker_reply": _thinker_reply(6)},
+        _omni_is_prefill=False,
+        _omni_prompt_len=3,
+        _omni_num_computed_tokens=4,
+    )
+
+    _, first, _ = model.talker_preprocess(input_ids, None, **payload)
+    _, second, _ = model.talker_preprocess(input_ids, None, **payload)
+
+    assert torch.equal(first, second)
+
+
+def test_positions_past_thinker_reply_keep_codec_embedding():
+    # Near the end of the thinker text the span can extend past the last
+    # reply row; those positions keep the plain codec embedding (the verifier
+    # rejects drafts there naturally).
+    model = _make_minimal_model()
+    input_ids = torch.tensor([1, 2, 3], dtype=torch.long)
+    q = _thinker_reply(6)
+
+    _, out_embeds, _ = model.talker_preprocess(
+        input_ids,
+        None,
+        embed={"thinker_reply": q},
+        _omni_is_prefill=False,
+        _omni_prompt_len=2,
+        _omni_num_computed_tokens=7,  # offset 5: one reply row left
+    )
+
+    base = _fake_embed_input_ids(input_ids)
+    assert torch.equal(out_embeds[0], base[0] + q[5])
+    assert torch.equal(out_embeds[1:], base[1:])
+
+
+def test_one_token_prefill_tail_routes_to_prefill():
+    # Chunked prefill can end in a span_len == 1 tail. With the runner
+    # metadata saying is_prefill, the row must go down the prefill path and
+    # must not touch the decode-side reply buffer (the span heuristic would
+    # have misrouted it to decode and consumed one row).
+    model = _make_minimal_model()
+    sentinel = (torch.tensor([0]), torch.zeros(1, HIDDEN), {"via": "prefill"})
+    calls = []
+
+    def fake_prefill_process(input_ids, input_embeds, payload):
+        calls.append(input_ids)
+        return sentinel
+
+    model.thinker_to_talker_process = fake_prefill_process
+
+    out = model.talker_preprocess(
+        torch.tensor([5], dtype=torch.long),
+        None,
+        embed={"thinker_reply": _thinker_reply(4)},
+        _omni_is_prefill=True,
+        _omni_prompt_len=8,
+        _omni_num_computed_tokens=7,
+    )
+
+    assert len(calls) == 1
+    assert out == sentinel
+
+
+def test_missing_runner_metadata_falls_back_to_legacy_single_step():
+    # Out-of-tree callers that predate the #3662 runner metadata must keep the
+    # stock consuming single-step behavior, bit for bit.
+    model = _make_minimal_model()
+    sentinel = (torch.tensor([9]), torch.ones(1, HIDDEN), {"via": "legacy"})
+    calls = []
+
+    def fake_one_step(input_ids, input_embeds, payload):
+        calls.append(input_ids)
+        return sentinel
+
+    model.thinker_to_talker_decode_one_step = fake_one_step
+
+    out = model.talker_preprocess(
+        torch.tensor([9], dtype=torch.long),
+        None,
+        embed={"thinker_reply": _thinker_reply(4)},
+    )
+
+    assert len(calls) == 1
+    assert out == sentinel
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
index 22dbcfac32c..44a1ceebc3a 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
@@ -692,26 +692,38 @@ def talker_preprocess(
         # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
         # Rules:
         # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
-        # - Decode segments consist of a single non-special token.
+        # - Decode segments may span one position (sequential decoding) or several
+        #   (speculative-decoding verification, preempt/resume KV replay).
         # - If additional_information is provided (can be a list split by request or a
         #   concatenated tensor plus a list of shapes), then for each request, reconstruct
         #   the thinker→talker input embeddings for the Prefill segments;
-        # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional),
-        #   add them; otherwise, keep the original embedding.
+        # - For Decode segments, add the request's ``thinker_reply`` row for each
+        #   position; past the end of the thinker reply, keep the codec embedding.
 
         payload: OmniPayload = info_dict
 
-        # Ensure we have base embeddings when only ids are provided
-        if input_embeds is None and input_ids is not None:
-            input_embeds = self.talker.embed_input_ids(input_ids)
-
-        span_len = input_ids.shape[0]
-        if span_len > 1:
-            # prefill
-            return self.thinker_to_talker_process(input_ids, input_embeds, payload)
+        span_len = int(input_ids.shape[0])
+        # Phase must come from the runner's per-request metadata (#3662), not from
+        # the span length: a speculative-decoding verify step schedules span_len ==
+        # num_draft_tokens + 1 decode positions, and a chunked prefill can end in a
+        # 1-token tail — both misroute under the span_len > 1 heuristic.
+        is_prefill_raw = payload.get("_omni_is_prefill")
+        if isinstance(is_prefill_raw, bool):
+            is_prefill = is_prefill_raw
         else:
-            # decode
-            return self.thinker_to_talker_decode_one_step(input_ids, input_embeds, payload)
+            try:
+                is_prefill = int(payload["_omni_num_computed_tokens"]) < int(payload["_omni_prompt_len"])
+            except (KeyError, TypeError, ValueError):
+                # Legacy callers without the runner metadata: fall back to the
+                # span-length heuristic (pre-#3662 behavior).
+                is_prefill = span_len > 1
+
+        if is_prefill:
+            # Ensure we have base embeddings when only ids are provided
+            if input_embeds is None and input_ids is not None:
+                input_embeds = self.talker.embed_input_ids(input_ids)
+            return self.thinker_to_talker_process(input_ids, input_embeds, payload)
+        return self.thinker_to_talker_decode(input_ids, input_embeds, payload)
 
     def thinker_to_talker_process(
         self,
@@ -802,6 +814,55 @@ def _thinker_to_talker_prefill(
             )
         return prompt_token_ids_processed, prompt_embeds
 
+    def thinker_to_talker_decode(self, input_ids, input_embeds, payload: OmniPayload):
+        """Assemble decode-side embeddings for a span of one or more positions.
+
+        Each decode position needs ``thinker_reply[k] + codec_embedding`` where
+        ``k`` is the position's generation offset. The offset is derived
+        statelessly from the runner metadata (``_omni_num_computed_tokens -
+        _omni_prompt_len``) instead of consuming one queue row per call, so the
+        same call handles the sequential single-token step, the multi-position
+        span of a speculative-decoding verify step, and the multi-token replay
+        span submitted when a preempted request resumes — a replay reconstructs
+        exactly the embeddings the original steps used. Positions past the end
+        of the thinker reply keep the plain codec embedding.
+
+        When the runner metadata is absent (legacy out-of-tree callers), the
+        single-position stateful path :meth:`thinker_to_talker_decode_one_step`
+        is used unchanged.
+        """
+        try:
+            offset = max(
+                0,
+                int(payload["_omni_num_computed_tokens"]) - int(payload["_omni_prompt_len"]),
+            )
+        except (KeyError, TypeError, ValueError):
+            # Pre-#3662 callers: keep the consuming single-step behavior.
+            if input_embeds is None and input_ids is not None:
+                input_embeds = self.talker.embed_input_ids(input_ids)
+            return self.thinker_to_talker_decode_one_step(input_ids, input_embeds, payload)
+
+        span_len = int(input_ids.shape[0])
+        # Build the span embeddings from the codec ids rather than trusting the
+        # incoming buffer: in a mixed batch the runner hands decode requests a
+        # slice of a ``torch.empty`` scratch tensor (uninitialized memory), and
+        # this is the only embedding build for the span (the entry-point one is
+        # prefill-only).
+        out_dtype = input_embeds.dtype if input_embeds is not None else None
+        input_embeds = self.talker.embed_input_ids(input_ids)
+        if out_dtype is not None and input_embeds.dtype != out_dtype:
+            input_embeds = input_embeds.to(out_dtype)
+
+        embed = payload.get("embed", {})
+        q = embed.get("thinker_reply", None)
+        if isinstance(q, torch.Tensor) and q.numel() > 0:
+            rows = q[offset : offset + span_len]
+            n = int(rows.shape[0])
+            if n > 0:
+                input_embeds[:n] = input_embeds[:n] + rows.to(dtype=input_embeds.dtype, device=input_embeds.device)
+        # Stateless: ``thinker_reply`` is not consumed, so no update to merge.
+        return input_ids, input_embeds, {}
+
     def thinker_to_talker_decode_one_step(self, input_ids, input_embeds, payload: OmniPayload):
         embed = payload.get("embed", {})
 

From fd95fedfc50baef5fe5c855291847e8d3e35c99e Mon Sep 17 00:00:00 2001
From: LJX-xixi <lijinxi413@icloud.com>
Date: Sun, 18 Oct 2026 21:24:10 +0800
Subject: [PATCH 2/5] [Bugfix] Qwen2.5-Omni: slice prefill spans at the runner
 offset and condition replayed decode suffixes

The prefill branch returned the complete reconstructed Talker prompt, and the
runner copies the returned rows into the positions it scheduled. A prefill
chunk at a non-zero offset (e.g. prompt_len=8, num_computed_tokens=7,
span_len=1) therefore received prompt row 0 instead of row 7.

After preemption the scheduler resets num_computed_tokens to 0 while keeping
the generated tokens, and may schedule the prompt and those tokens in one
span. The prefill branch returned only prompt rows, leaving the generated
suffix without thinker-conditioned decode embeddings.

Slice the prompt rows (ids and embeddings) at _omni_num_computed_tokens, and
assemble any generated-token suffix with the stateless decode helper, using
the thinker_reply reconstructed in the same call. Without runner metadata the
complete prompt is still returned.

The tests now run the real prefill helper with distinguishable rows: offset
chunks, the 1-token tail, replays crossing the prompt boundary from offsets 0
and 5 (compared against the original step-by-step decode embeddings), a
repeated replay, and the legacy path.

Signed-off-by: LJX-xixi <lijinxi413@icloud.com>
---
 .../test_qwen2_5_omni_talker_preprocess.py    | 200 ++++++++++++++++--
 .../models/qwen2_5_omni/qwen2_5_omni.py       |  52 ++++-
 2 files changed, 233 insertions(+), 19 deletions(-)

diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
index e6e29b039c4..4897509a22e 100644
--- a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
+++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
@@ -117,32 +117,196 @@ def test_positions_past_thinker_reply_keep_codec_embedding():
     assert torch.equal(out_embeds[1:], base[1:])
 
 
-def test_one_token_prefill_tail_routes_to_prefill():
-    # Chunked prefill can end in a span_len == 1 tail. With the runner
-    # metadata saying is_prefill, the row must go down the prefill path and
-    # must not touch the decode-side reply buffer (the span heuristic would
-    # have misrouted it to decode and consumed one row).
-    model = _make_minimal_model()
-    sentinel = (torch.tensor([0]), torch.zeros(1, HIDDEN), {"via": "prefill"})
-    calls = []
+THINKER_PROMPT_LEN = 6
+TALKER_PROMPT_LEN = THINKER_PROMPT_LEN + 2  # thinker prompt + speaker/pad row + BOS row
+CODEC_PAD, CODEC_START, CODEC_MASK = 8292, 8293, 8296
 
-    def fake_prefill_process(input_ids, input_embeds, payload):
-        calls.append(input_ids)
-        return sentinel
 
-    model.thinker_to_talker_process = fake_prefill_process
+class _FakeTalker(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = SimpleNamespace(hidden_size=HIDDEN)
+        self.anchor = torch.nn.Parameter(torch.zeros(1), requires_grad=False)
 
-    out = model.talker_preprocess(
+    def embed_input_ids(self, ids: torch.Tensor) -> torch.Tensor:
+        return _fake_embed_input_ids(ids)
+
+
+def _make_prefill_model():
+    """Minimal model on which the real ``_thinker_to_talker_prefill`` runs."""
+    cls = _qwen2_5_omni_mod().Qwen2_5OmniForConditionalGeneration
+    model = cls.__new__(cls)
+    torch.nn.Module.__init__(model)
+    model.talker = _FakeTalker()
+    model.model = model.talker
+    model.talker_config = SimpleNamespace(
+        tts_codec_pad_token_id=CODEC_PAD,
+        tts_codec_start_token_id=CODEC_START,
+        tts_codec_mask_token_id=CODEC_MASK,
+    )
+    model.voice_type = "Chelsie"
+    model.embed_text_bos_token = torch.full((1, HIDDEN), 7.0)
+    model.embed_text_spk_tokens = {}
+    model.embed_codec_pad_token = torch.full((1, HIDDEN), 0.5)
+    model.embed_codec_bos_token = torch.full((1, HIDDEN), 0.25)
+    return model
+
+
+def _thinker_payload(num_output_rows: int = 5) -> dict:
+    # Thinker prompt row i is 1000 * (i + 1); thinker output row k is
+    # 10 * (k + 1), so ``thinker_reply`` (output rows 1:) starts at 20.
+    prompt = (torch.arange(1, THINKER_PROMPT_LEN + 1, dtype=torch.float32).reshape(-1, 1) * 1000.0).expand(-1, HIDDEN)
+    output = (torch.arange(1, num_output_rows + 1, dtype=torch.float32).reshape(-1, 1) * 10.0).expand(-1, HIDDEN)
+    return dict(
+        embed={"prefill": prompt.contiguous()},
+        hidden_states={"output": output.contiguous()},
+        ids={"prompt": list(range(100, 100 + THINKER_PROMPT_LEN)), "output": list(range(200, 200 + num_output_rows))},
+    )
+
+
+def _full_talker_prompt(model, payload):
+    """The complete prompt as reconstructed for an unchunked prefill."""
+    ids, embeds, _ = model.talker_preprocess(
+        torch.zeros(TALKER_PROMPT_LEN, dtype=torch.long),
+        None,
+        **payload,
+        _omni_is_prefill=True,
+        _omni_prompt_len=TALKER_PROMPT_LEN,
+        _omni_num_computed_tokens=0,
+    )
+    assert ids.shape[0] == embeds.shape[0] == TALKER_PROMPT_LEN
+    return ids, embeds
+
+
+@pytest.mark.parametrize(("num_computed", "span_len"), [(7, 1), (2, 3), (0, 4)])
+def test_chunked_prefill_span_gets_rows_at_its_offset(num_computed, span_len):
+    # The runner copies the returned rows into the positions it scheduled,
+    # starting at num_computed_tokens. A chunk at a non-zero offset (e.g. a
+    # 1-token tail at offset 7) must receive those prompt rows, not row 0.
+    model = _make_prefill_model()
+    payload = _thinker_payload()
+    full_ids, full_embeds = _full_talker_prompt(model, payload)
+
+    out_ids, out_embeds, update = model.talker_preprocess(
+        torch.zeros(span_len, dtype=torch.long),
+        None,
+        **payload,
+        _omni_is_prefill=True,
+        _omni_prompt_len=TALKER_PROMPT_LEN,
+        _omni_num_computed_tokens=num_computed,
+    )
+
+    end = num_computed + span_len
+    assert torch.equal(out_ids, full_ids[num_computed:end])
+    assert torch.equal(out_embeds, full_embeds[num_computed:end])
+    assert torch.equal(update["embed"]["thinker_reply"], payload["hidden_states"]["output"][1:])
+
+
+def test_prefill_tail_row_is_first_thinker_output_plus_bos():
+    # The last prompt row conditions on thinker output row 0 plus codec BOS;
+    # a 1-token tail at offset 7 must receive exactly that row.
+    model = _make_prefill_model()
+    payload = _thinker_payload()
+
+    out_ids, out_embeds, _ = model.talker_preprocess(
         torch.tensor([5], dtype=torch.long),
         None,
-        embed={"thinker_reply": _thinker_reply(4)},
+        **payload,
         _omni_is_prefill=True,
-        _omni_prompt_len=8,
-        _omni_num_computed_tokens=7,
+        _omni_prompt_len=TALKER_PROMPT_LEN,
+        _omni_num_computed_tokens=TALKER_PROMPT_LEN - 1,
     )
 
-    assert len(calls) == 1
-    assert out == sentinel
+    assert out_ids.tolist() == [CODEC_START]
+    expected = payload["hidden_states"]["output"][:1] + model.embed_codec_bos_token
+    assert torch.equal(out_embeds, expected)
+
+
+def _original_decode_embeds(model, payload, generated_ids):
+    """Embeddings the decode steps produced before preemption, one step at a time."""
+    _, _, update = model.talker_preprocess(
+        torch.zeros(TALKER_PROMPT_LEN, dtype=torch.long),
+        None,
+        **payload,
+        _omni_is_prefill=True,
+        _omni_prompt_len=TALKER_PROMPT_LEN,
+        _omni_num_computed_tokens=0,
+    )
+    rows = []
+    for k, token in enumerate(generated_ids):
+        _, embeds, _ = model.talker_preprocess(
+            torch.tensor([token], dtype=torch.long),
+            None,
+            embed={**payload["embed"], **update["embed"]},
+            _omni_is_prefill=False,
+            _omni_prompt_len=TALKER_PROMPT_LEN,
+            _omni_num_computed_tokens=TALKER_PROMPT_LEN + k,
+        )
+        rows.append(embeds)
+    return torch.cat(rows)
+
+
+@pytest.mark.parametrize("num_computed", [0, 5])
+def test_replay_span_crossing_prompt_boundary_conditions_generated_suffix(num_computed):
+    # After preemption the scheduler resets num_computed_tokens to 0 but keeps
+    # the generated tokens, and may schedule the prompt plus those tokens in
+    # one span (or, with chunking, a span starting inside the prompt). The
+    # prompt rows come from the reconstructed prompt at the offset; the
+    # generated suffix must get the same thinker-conditioned embeddings as the
+    # original decode steps, which a prompt-only result would leave unset.
+    model = _make_prefill_model()
+    payload = _thinker_payload()
+    full_ids, full_embeds = _full_talker_prompt(model, payload)
+    generated = [301, 302, 303]
+    span_ids = torch.cat(
+        [
+            torch.zeros(TALKER_PROMPT_LEN - num_computed, dtype=torch.long),
+            torch.tensor(generated, dtype=torch.long),
+        ]
+    )
+
+    out_ids, out_embeds, _ = model.talker_preprocess(
+        span_ids,
+        None,
+        **payload,
+        _omni_is_prefill=True,
+        _omni_prompt_len=TALKER_PROMPT_LEN,
+        _omni_num_computed_tokens=num_computed,
+    )
+
+    num_prompt = TALKER_PROMPT_LEN - num_computed
+    assert out_ids.shape[0] == out_embeds.shape[0] == span_ids.shape[0]
+    assert torch.equal(out_ids[:num_prompt], full_ids[num_computed:])
+    assert torch.equal(out_embeds[:num_prompt], full_embeds[num_computed:])
+    assert out_ids[num_prompt:].tolist() == generated
+    assert torch.equal(out_embeds[num_prompt:], _original_decode_embeds(model, payload, generated))
+
+
+def test_repeated_replay_span_is_identical():
+    # A request can be preempted and replayed more than once; each replay of
+    # the same span must rebuild identical embeddings.
+    model = _make_prefill_model()
+    payload = _thinker_payload()
+    span_ids = torch.cat([torch.zeros(TALKER_PROMPT_LEN, dtype=torch.long), torch.tensor([301, 302])])
+    meta = dict(_omni_is_prefill=True, _omni_prompt_len=TALKER_PROMPT_LEN, _omni_num_computed_tokens=0)
+
+    _, first, _ = model.talker_preprocess(span_ids, None, **payload, **meta)
+    _, second, _ = model.talker_preprocess(span_ids, None, **payload, **meta)
+
+    assert torch.equal(first, second)
+
+
+def test_prefill_without_runner_metadata_returns_full_prompt():
+    # Legacy callers without the runner metadata keep receiving the complete
+    # reconstructed prompt, as before.
+    model = _make_prefill_model()
+    payload = _thinker_payload()
+    full_ids, full_embeds = _full_talker_prompt(model, payload)
+
+    out_ids, out_embeds, _ = model.talker_preprocess(torch.zeros(TALKER_PROMPT_LEN, dtype=torch.long), None, **payload)
+
+    assert torch.equal(out_ids, full_ids)
+    assert torch.equal(out_embeds, full_embeds)
 
 
 def test_missing_runner_metadata_falls_back_to_legacy_single_step():
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
index 44a1ceebc3a..a3f34fedce8 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
@@ -722,9 +722,59 @@ def talker_preprocess(
             # Ensure we have base embeddings when only ids are provided
             if input_embeds is None and input_ids is not None:
                 input_embeds = self.talker.embed_input_ids(input_ids)
-            return self.thinker_to_talker_process(input_ids, input_embeds, payload)
+            return self.thinker_to_talker_prefill_span(input_ids, input_embeds, payload)
         return self.thinker_to_talker_decode(input_ids, input_embeds, payload)
 
+    def thinker_to_talker_prefill_span(self, input_ids, input_embeds, payload: OmniPayload):
+        """Assemble embeddings for a prefill span that starts at the runner offset.
+
+        :meth:`thinker_to_talker_process` reconstructs the complete Talker
+        prompt, while the runner copies the returned rows into the positions
+        it scheduled, starting at ``_omni_num_computed_tokens``. The prompt
+        rows are therefore sliced at that offset, so a later chunked-prefill
+        chunk (e.g. a 1-token tail at offset 7) receives its own rows rather
+        than the beginning of the prompt.
+
+        After preemption the scheduler resets ``num_computed_tokens`` to 0
+        while keeping the generated tokens, and may schedule the prompt and
+        those tokens in one span. Positions past the prompt are assembled by
+        the stateless :meth:`thinker_to_talker_decode`, using the
+        ``thinker_reply`` reconstructed by this call, so the replay matches
+        the embeddings of the original decode steps.
+
+        Without the runner metadata (legacy out-of-tree callers), the complete
+        reconstructed prompt is returned unchanged.
+        """
+        req_input_ids, req_embeds, update_dict = self.thinker_to_talker_process(input_ids, input_embeds, payload)
+        try:
+            offset = max(0, int(payload["_omni_num_computed_tokens"]))
+            prompt_len = int(payload["_omni_prompt_len"])
+        except (KeyError, TypeError, ValueError):
+            return req_input_ids, req_embeds, update_dict
+
+        span_len = int(input_ids.shape[0])
+        prompt_end = min(offset + span_len, prompt_len, int(req_embeds.shape[0]))
+        num_prompt = max(0, prompt_end - offset)
+        out_ids = req_input_ids[offset : offset + num_prompt]
+        out_embeds = req_embeds[offset : offset + num_prompt]
+
+        num_suffix = span_len - num_prompt
+        if num_suffix > 0:
+            embed = dict(payload.get("embed", {}))
+            embed.update(update_dict.get("embed", {}))
+            suffix_payload = {
+                **payload,
+                "embed": embed,
+                "_omni_num_computed_tokens": offset + num_prompt,
+                "_omni_prompt_len": prompt_len,
+            }
+            suffix_ids, suffix_embeds, _ = self.thinker_to_talker_decode(
+                input_ids[num_prompt:], input_embeds[num_prompt:], suffix_payload
+            )
+            out_ids = torch.cat([out_ids, suffix_ids.to(device=out_ids.device, dtype=out_ids.dtype)])
+            out_embeds = torch.cat([out_embeds, suffix_embeds.to(device=out_embeds.device, dtype=out_embeds.dtype)])
+        return out_ids, out_embeds, update_dict
+
     def thinker_to_talker_process(
         self,
         input_ids: torch.Tensor,

From 196b3f4aff85333a40a5b53f30feb5ad446c29b4 Mon Sep 17 00:00:00 2001
From: LJX-xixi <lijinxi413@icloud.com>
Date: Sun, 18 Oct 2026 23:37:46 +0800
Subject: [PATCH 3/5] [Bugfix] Qwen2.5-Omni: assemble talker spans only from
 runner positions

docs/design/feature/preprocess_phase_contract.md makes _omni_prompt_len,
_omni_num_computed_tokens and _omni_is_prefill a stable runner-owned
interface, and asks models not to fall back silently to the span length
when they need chunking correctness. Every in-tree runner (V1, MRv2, NPU)
provides these fields, so the span-length heuristic, the full-prompt return
and the consuming single-step decode path kept for callers without them were
unreachable, and kept a second, stateful way to condition decode positions,
including a model-level thinker_reply_part pool that is filled during
profiling and consumed across requests.

Talker preprocessing now derives the phase and every row from the span's
absolute positions: prompt positions take rows of the reconstructed prompt,
generated positions take thinker_reply[offset + i], and a replay crossing the
prompt boundary uses both. Missing metadata, or a reconstructed prompt whose
length differs from the scheduled prompt, raises instead of silently
misaligning rows. The stateful decode path and the shared pool are removed.

Also use the vLLM-Omni SPDX copyright line in the test file, as required by
the check-spdx-header pre-commit hook.

Signed-off-by: LJX-xixi <lijinxi413@icloud.com>
---
 .../test_qwen2_5_omni_talker_preprocess.py    |  65 +++---
 .../models/qwen2_5_omni/qwen2_5_omni.py       | 205 ++++--------------
 2 files changed, 74 insertions(+), 196 deletions(-)

diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
index 4897509a22e..481f1cc252a 100644
--- a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
+++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
@@ -1,15 +1,13 @@
 # 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 the Qwen2.5-Omni talker preprocess phase routing and the
-stateless decode-span embedding assembly.
-
-Regression tests for the span-length prefill/decode heuristic: a
-speculative-decoding verify step schedules num_draft_tokens + 1 decode
-positions in one span, and a chunked prefill can end in a 1-token tail —
-both misroute under ``span_len > 1``. The decode path must also assemble
-one ``thinker_reply`` row per position, statelessly, so a preempt/resume
-replay span reconstructs exactly the embeddings the original steps used.
+Unit tests for Qwen2.5-Omni talker span assembly from the runner's positions.
+
+A chunked-prefill chunk, a speculative-decoding verify step and a preemption
+replay all schedule spans whose length does not identify their phase. Each
+position must receive the row for its absolute position: a prompt row of the
+reconstructed Talker prompt, or ``thinker_reply[k] + codec_embedding`` for
+generation offset ``k``. Replaying the same positions must rebuild the same rows.
 """
 
 import functools
@@ -296,37 +294,26 @@ def test_repeated_replay_span_is_identical():
     assert torch.equal(first, second)
 
 
-def test_prefill_without_runner_metadata_returns_full_prompt():
-    # Legacy callers without the runner metadata keep receiving the complete
-    # reconstructed prompt, as before.
+def test_missing_runner_metadata_is_rejected():
+    # Every in-tree runner supplies the span positions (#3662); without them
+    # the span cannot be assembled correctly, so preprocessing fails loudly.
     model = _make_prefill_model()
-    payload = _thinker_payload()
-    full_ids, full_embeds = _full_talker_prompt(model, payload)
-
-    out_ids, out_embeds, _ = model.talker_preprocess(torch.zeros(TALKER_PROMPT_LEN, dtype=torch.long), None, **payload)
 
-    assert torch.equal(out_ids, full_ids)
-    assert torch.equal(out_embeds, full_embeds)
-
-
-def test_missing_runner_metadata_falls_back_to_legacy_single_step():
-    # Out-of-tree callers that predate the #3662 runner metadata must keep the
-    # stock consuming single-step behavior, bit for bit.
-    model = _make_minimal_model()
-    sentinel = (torch.tensor([9]), torch.ones(1, HIDDEN), {"via": "legacy"})
-    calls = []
+    with pytest.raises(ValueError, match="runner metadata"):
+        model.talker_preprocess(torch.zeros(TALKER_PROMPT_LEN, dtype=torch.long), None, **_thinker_payload())
 
-    def fake_one_step(input_ids, input_embeds, payload):
-        calls.append(input_ids)
-        return sentinel
 
-    model.thinker_to_talker_decode_one_step = fake_one_step
-
-    out = model.talker_preprocess(
-        torch.tensor([9], dtype=torch.long),
-        None,
-        embed={"thinker_reply": _thinker_reply(4)},
-    )
+def test_prompt_length_mismatch_is_rejected():
+    # If the reconstructed Talker prompt and the scheduled prompt disagree in
+    # length, slicing by position would silently misalign every row.
+    model = _make_prefill_model()
 
-    assert len(calls) == 1
-    assert out == sentinel
+    with pytest.raises(ValueError, match="Reconstructed Talker prompt"):
+        model.talker_preprocess(
+            torch.zeros(4, dtype=torch.long),
+            None,
+            **_thinker_payload(),
+            _omni_is_prefill=True,
+            _omni_prompt_len=TALKER_PROMPT_LEN + 1,
+            _omni_num_computed_tokens=0,
+        )
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
index a3f34fedce8..8ac7deffbd0 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
@@ -341,7 +341,6 @@ def forward(
             # mock data for profile
             if input_ids is None:
                 input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
-                self.thinker_reply_part = torch.zeros_like(inputs_embeds)
 
             # TODO(Peiqi): temporal hack here to support voice_type.
             if not hasattr(self, "voice_type"):
@@ -689,87 +688,48 @@ def talker_preprocess(
         input_embeds: torch.Tensor,
         **info_dict: object,
     ):
-        # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
-        # Rules:
-        # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
-        # - Decode segments may span one position (sequential decoding) or several
-        #   (speculative-decoding verification, preempt/resume KV replay).
-        # - If additional_information is provided (can be a list split by request or a
-        #   concatenated tensor plus a list of shapes), then for each request, reconstruct
-        #   the thinker→talker input embeddings for the Prefill segments;
-        # - For Decode segments, add the request's ``thinker_reply`` row for each
-        #   position; past the end of the thinker reply, keep the codec embedding.
-
+        # The runner reports, for every scheduled span, the request's prompt length
+        # and the absolute position of the span's first token (#3662). The span is
+        # assembled from these positions, never from its length: a chunked-prefill
+        # chunk, a speculative-decoding verify step and a preemption replay all
+        # schedule spans whose length does not identify their phase.
+        # - Prompt positions take rows of the reconstructed Talker prompt.
+        # - Generated positions add the request's ``thinker_reply`` row for their
+        #   generation offset; past the end of the reply, the codec embedding is kept.
         payload: OmniPayload = info_dict
-
-        span_len = int(input_ids.shape[0])
-        # Phase must come from the runner's per-request metadata (#3662), not from
-        # the span length: a speculative-decoding verify step schedules span_len ==
-        # num_draft_tokens + 1 decode positions, and a chunked prefill can end in a
-        # 1-token tail — both misroute under the span_len > 1 heuristic.
-        is_prefill_raw = payload.get("_omni_is_prefill")
-        if isinstance(is_prefill_raw, bool):
-            is_prefill = is_prefill_raw
-        else:
-            try:
-                is_prefill = int(payload["_omni_num_computed_tokens"]) < int(payload["_omni_prompt_len"])
-            except (KeyError, TypeError, ValueError):
-                # Legacy callers without the runner metadata: fall back to the
-                # span-length heuristic (pre-#3662 behavior).
-                is_prefill = span_len > 1
-
-        if is_prefill:
-            # Ensure we have base embeddings when only ids are provided
-            if input_embeds is None and input_ids is not None:
-                input_embeds = self.talker.embed_input_ids(input_ids)
-            return self.thinker_to_talker_prefill_span(input_ids, input_embeds, payload)
-        return self.thinker_to_talker_decode(input_ids, input_embeds, payload)
-
-    def thinker_to_talker_prefill_span(self, input_ids, input_embeds, payload: OmniPayload):
-        """Assemble embeddings for a prefill span that starts at the runner offset.
-
-        :meth:`thinker_to_talker_process` reconstructs the complete Talker
-        prompt, while the runner copies the returned rows into the positions
-        it scheduled, starting at ``_omni_num_computed_tokens``. The prompt
-        rows are therefore sliced at that offset, so a later chunked-prefill
-        chunk (e.g. a 1-token tail at offset 7) receives its own rows rather
-        than the beginning of the prompt.
-
-        After preemption the scheduler resets ``num_computed_tokens`` to 0
-        while keeping the generated tokens, and may schedule the prompt and
-        those tokens in one span. Positions past the prompt are assembled by
-        the stateless :meth:`thinker_to_talker_decode`, using the
-        ``thinker_reply`` reconstructed by this call, so the replay matches
-        the embeddings of the original decode steps.
-
-        Without the runner metadata (legacy out-of-tree callers), the complete
-        reconstructed prompt is returned unchanged.
-        """
-        req_input_ids, req_embeds, update_dict = self.thinker_to_talker_process(input_ids, input_embeds, payload)
         try:
-            offset = max(0, int(payload["_omni_num_computed_tokens"]))
+            start = int(payload["_omni_num_computed_tokens"])
             prompt_len = int(payload["_omni_prompt_len"])
-        except (KeyError, TypeError, ValueError):
-            return req_input_ids, req_embeds, update_dict
-
-        span_len = int(input_ids.shape[0])
-        prompt_end = min(offset + span_len, prompt_len, int(req_embeds.shape[0]))
-        num_prompt = max(0, prompt_end - offset)
-        out_ids = req_input_ids[offset : offset + num_prompt]
-        out_embeds = req_embeds[offset : offset + num_prompt]
-
-        num_suffix = span_len - num_prompt
-        if num_suffix > 0:
-            embed = dict(payload.get("embed", {}))
-            embed.update(update_dict.get("embed", {}))
-            suffix_payload = {
-                **payload,
-                "embed": embed,
-                "_omni_num_computed_tokens": offset + num_prompt,
-                "_omni_prompt_len": prompt_len,
-            }
+        except (KeyError, TypeError, ValueError) as exc:
+            raise ValueError(
+                "Qwen2.5-Omni talker preprocessing requires the runner metadata "
+                "_omni_num_computed_tokens and _omni_prompt_len."
+            ) from exc
+
+        if start >= prompt_len:
+            thinker_reply = payload.get("embed", {}).get("thinker_reply")
+            return self.thinker_to_talker_decode(input_ids, input_embeds, thinker_reply, start - prompt_len)
+
+        # Ensure we have base embeddings when only ids are provided
+        if input_embeds is None and input_ids is not None:
+            input_embeds = self.talker.embed_input_ids(input_ids)
+        req_input_ids, req_embeds, update_dict = self.thinker_to_talker_process(input_ids, input_embeds, payload)
+        if req_embeds.shape[0] != prompt_len:
+            raise ValueError(
+                f"Reconstructed Talker prompt has {req_embeds.shape[0]} rows, but the "
+                f"scheduled prompt has {prompt_len} tokens."
+            )
+        num_prompt = min(int(input_ids.shape[0]), prompt_len - start)
+        out_ids = req_input_ids[start : start + num_prompt]
+        out_embeds = req_embeds[start : start + num_prompt]
+        if num_prompt < input_ids.shape[0]:
+            # After preemption the scheduler resets num_computed_tokens to 0 while
+            # keeping the generated tokens, and may schedule the prompt and those
+            # tokens in one span. The generated suffix uses the thinker_reply
+            # reconstructed above, so it matches the original decode steps.
+            thinker_reply = update_dict.get("embed", {}).get("thinker_reply")
             suffix_ids, suffix_embeds, _ = self.thinker_to_talker_decode(
-                input_ids[num_prompt:], input_embeds[num_prompt:], suffix_payload
+                input_ids[num_prompt:], input_embeds[num_prompt:], thinker_reply, 0
             )
             out_ids = torch.cat([out_ids, suffix_ids.to(device=out_ids.device, dtype=out_ids.dtype)])
             out_embeds = torch.cat([out_embeds, suffix_embeds.to(device=out_embeds.device, dtype=out_embeds.dtype)])
@@ -864,99 +824,30 @@ def _thinker_to_talker_prefill(
             )
         return prompt_token_ids_processed, prompt_embeds
 
-    def thinker_to_talker_decode(self, input_ids, input_embeds, payload: OmniPayload):
-        """Assemble decode-side embeddings for a span of one or more positions.
-
-        Each decode position needs ``thinker_reply[k] + codec_embedding`` where
-        ``k`` is the position's generation offset. The offset is derived
-        statelessly from the runner metadata (``_omni_num_computed_tokens -
-        _omni_prompt_len``) instead of consuming one queue row per call, so the
-        same call handles the sequential single-token step, the multi-position
-        span of a speculative-decoding verify step, and the multi-token replay
-        span submitted when a preempted request resumes — a replay reconstructs
-        exactly the embeddings the original steps used. Positions past the end
-        of the thinker reply keep the plain codec embedding.
-
-        When the runner metadata is absent (legacy out-of-tree callers), the
-        single-position stateful path :meth:`thinker_to_talker_decode_one_step`
-        is used unchanged.
-        """
-        try:
-            offset = max(
-                0,
-                int(payload["_omni_num_computed_tokens"]) - int(payload["_omni_prompt_len"]),
-            )
-        except (KeyError, TypeError, ValueError):
-            # Pre-#3662 callers: keep the consuming single-step behavior.
-            if input_embeds is None and input_ids is not None:
-                input_embeds = self.talker.embed_input_ids(input_ids)
-            return self.thinker_to_talker_decode_one_step(input_ids, input_embeds, payload)
+    def thinker_to_talker_decode(self, input_ids, input_embeds, thinker_reply, offset: int):
+        """Assemble embeddings for generated positions starting at generation ``offset``.
 
+        Each position needs ``thinker_reply[offset + i] + codec_embedding``. The
+        rows are selected by position and ``thinker_reply`` is never consumed, so
+        a single decode step, a speculative-decoding verify span and a
+        preemption replay of the same positions all receive the same rows.
+        """
         span_len = int(input_ids.shape[0])
         # Build the span embeddings from the codec ids rather than trusting the
         # incoming buffer: in a mixed batch the runner hands decode requests a
-        # slice of a ``torch.empty`` scratch tensor (uninitialized memory), and
-        # this is the only embedding build for the span (the entry-point one is
-        # prefill-only).
+        # slice of a ``torch.empty`` scratch tensor (uninitialized memory).
         out_dtype = input_embeds.dtype if input_embeds is not None else None
         input_embeds = self.talker.embed_input_ids(input_ids)
         if out_dtype is not None and input_embeds.dtype != out_dtype:
             input_embeds = input_embeds.to(out_dtype)
 
-        embed = payload.get("embed", {})
-        q = embed.get("thinker_reply", None)
-        if isinstance(q, torch.Tensor) and q.numel() > 0:
-            rows = q[offset : offset + span_len]
+        if isinstance(thinker_reply, torch.Tensor) and thinker_reply.numel() > 0:
+            rows = thinker_reply[offset : offset + span_len]
             n = int(rows.shape[0])
             if n > 0:
                 input_embeds[:n] = input_embeds[:n] + rows.to(dtype=input_embeds.dtype, device=input_embeds.device)
-        # Stateless: ``thinker_reply`` is not consumed, so no update to merge.
         return input_ids, input_embeds, {}
 
-    def thinker_to_talker_decode_one_step(self, input_ids, input_embeds, payload: OmniPayload):
-        embed = payload.get("embed", {})
-
-        update_dict = {}
-        # choose step vector in priority order
-        step_vec = None
-        q = embed.get("thinker_reply", None)
-        if isinstance(q, torch.Tensor) and q.numel() > 0:
-            step_vec = q[0:1]
-            new_q = q[1:].detach().to("cpu").contiguous()
-            update_dict.setdefault("embed", {})["thinker_reply"] = new_q
-        else:
-            # B) per-request provided decode vector (optional)
-            dv = embed.get("decode")
-            if isinstance(dv, torch.Tensor) and dv.numel() > 0:
-                step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
-            elif (
-                hasattr(self, "thinker_reply_part")
-                and isinstance(self.thinker_reply_part, torch.Tensor)
-                and self.thinker_reply_part.numel() > 0
-            ):
-                # C) fallback shared pool
-                step_vec = self.thinker_reply_part[0:1]
-                self.thinker_reply_part = self.thinker_reply_part[1:]
-
-        if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
-            one_id = input_ids[0:1]
-            _, one_embed = self._thinker_to_talker_decode_one_step(
-                output_prompt_embeds=step_vec.to(input_embeds.dtype).to(self._module_device(self.model)),
-                output_token_ids=one_id,
-            )
-            input_embeds[0] = one_embed[0]
-        return input_ids[0:1], input_embeds[0:1], update_dict
-
-    def _thinker_to_talker_decode_one_step(
-        self,
-        output_prompt_embeds,
-        output_token_ids,
-    ):
-        processed_output_token_embeds = output_prompt_embeds + self.talker.embed_input_ids(
-            output_token_ids
-        )  # for decode
-        return output_token_ids, processed_output_token_embeds
-
     def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, **kwargs: object) -> torch.Tensor | None:
         # Handle OmniOutput type
         if isinstance(hidden_states, OmniOutput):

From 6ed38c272a6c4a694b53b512c8be9468b2b973e7 Mon Sep 17 00:00:00 2001
From: LJX-xixi <lijinxi413@icloud.com>
Date: Tue, 27 Oct 2026 01:48:34 +0800
Subject: [PATCH 4/5] =?UTF-8?q?=E6=B5=8B=E8=AF=95:=20=E9=80=82=E9=85=8DQwe?=
 =?UTF-8?q?n2.5=E6=96=B0=E5=A2=9E=E6=96=87=E6=9C=AC=E6=9D=A1=E4=BB=B6?=
 =?UTF-8?q?=E8=AF=AD=E4=B9=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: LJX-xixi <lijinxi413@icloud.com>
---
 .../test_qwen2_5_omni_talker_preprocess.py    | 25 ++++++++++++++-----
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
index 481f1cc252a..34eaaf38bb8 100644
--- a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
+++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
@@ -137,6 +137,8 @@ def _make_prefill_model():
     torch.nn.Module.__init__(model)
     model.talker = _FakeTalker()
     model.model = model.talker
+    model.thinker_config = SimpleNamespace()
+    model.thinker_embedding = torch.nn.Embedding.from_pretrained(_fake_embed_input_ids(torch.arange(256)))
     model.talker_config = SimpleNamespace(
         tts_codec_pad_token_id=CODEC_PAD,
         tts_codec_start_token_id=CODEC_START,
@@ -145,14 +147,17 @@ def _make_prefill_model():
     model.voice_type = "Chelsie"
     model.embed_text_bos_token = torch.full((1, HIDDEN), 7.0)
     model.embed_text_spk_tokens = {}
+    model.embed_text_eos_token = torch.full((1, HIDDEN), 8.0)
+    model.embed_text_pad_token = torch.full((1, HIDDEN), 9.0)
     model.embed_codec_pad_token = torch.full((1, HIDDEN), 0.5)
     model.embed_codec_bos_token = torch.full((1, HIDDEN), 0.25)
     return model
 
 
 def _thinker_payload(num_output_rows: int = 5) -> dict:
-    # Thinker prompt row i is 1000 * (i + 1); thinker output row k is
-    # 10 * (k + 1), so ``thinker_reply`` (output rows 1:) starts at 20.
+    # Hidden rows and token embeddings are both position-distinguishing:
+    # prompt hidden row i is 1000 * (i + 1), output hidden row k is 10 * (k + 1),
+    # and each text-token embedding adds its token ID.
     prompt = (torch.arange(1, THINKER_PROMPT_LEN + 1, dtype=torch.float32).reshape(-1, 1) * 1000.0).expand(-1, HIDDEN)
     output = (torch.arange(1, num_output_rows + 1, dtype=torch.float32).reshape(-1, 1) * 10.0).expand(-1, HIDDEN)
     return dict(
@@ -197,12 +202,16 @@ def test_chunked_prefill_span_gets_rows_at_its_offset(num_computed, span_len):
     end = num_computed + span_len
     assert torch.equal(out_ids, full_ids[num_computed:end])
     assert torch.equal(out_embeds, full_embeds[num_computed:end])
-    assert torch.equal(update["embed"]["thinker_reply"], payload["hidden_states"]["output"][1:])
+    reply_conditioning = payload["hidden_states"]["output"] + _fake_embed_input_ids(
+        torch.tensor(payload["ids"]["output"])
+    )
+    expected_reply = torch.cat([reply_conditioning[1:], model.embed_text_eos_token, model.embed_text_pad_token])
+    assert torch.equal(update["embed"]["thinker_reply"], expected_reply)
 
 
 def test_prefill_tail_row_is_first_thinker_output_plus_bos():
-    # The last prompt row conditions on thinker output row 0 plus codec BOS;
-    # a 1-token tail at offset 7 must receive exactly that row.
+    # The last prompt row conditions on thinker output row 0 plus its text-token
+    # embedding and codec BOS; a 1-token tail at offset 7 must receive that row.
     model = _make_prefill_model()
     payload = _thinker_payload()
 
@@ -216,7 +225,11 @@ def test_prefill_tail_row_is_first_thinker_output_plus_bos():
     )
 
     assert out_ids.tolist() == [CODEC_START]
-    expected = payload["hidden_states"]["output"][:1] + model.embed_codec_bos_token
+    expected = (
+        payload["hidden_states"]["output"][:1]
+        + _fake_embed_input_ids(torch.tensor(payload["ids"]["output"][:1]))
+        + model.embed_codec_bos_token
+    )
     assert torch.equal(out_embeds, expected)
 
 

From de22e99417e0ba028c4a50c569a0745a0bd7e48c Mon Sep 17 00:00:00 2001
From: LJX-xixi <lijinxi413@icloud.com>
Date: Tue, 27 Oct 2026 02:10:33 +0800
Subject: [PATCH 5/5] =?UTF-8?q?=E4=BF=AE=E5=A4=8D:=20Qwen2.5=E6=96=87?=
 =?UTF-8?q?=E6=9C=AC=E7=BB=93=E6=9D=9F=E5=90=8E=E6=8C=81=E7=BB=AD=E4=BF=9D?=
 =?UTF-8?q?=E7=95=99PAD=E6=9D=A1=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: LJX-xixi <lijinxi413@icloud.com>
---
 .../test_qwen2_5_omni_talker_preprocess.py    | 30 +++++++++++++------
 .../models/qwen2_5_omni/qwen2_5_omni.py       |  7 ++++-
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
index 34eaaf38bb8..67fe52a6169 100644
--- a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
+++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_preprocess.py
@@ -93,10 +93,13 @@ def test_decode_span_is_replayable():
     assert torch.equal(first, second)
 
 
-def test_positions_past_thinker_reply_keep_codec_embedding():
-    # Near the end of the thinker text the span can extend past the last
-    # reply row; those positions keep the plain codec embedding (the verifier
-    # rejects drafts there naturally).
+@pytest.mark.parametrize(
+    ("offset", "expected_rows"),
+    [(4, [4, 5, 5]), (5, [5, 5, 5]), (6, [5, 5, 5]), (9, [5, 5, 5])],
+)
+def test_positions_past_thinker_reply_repeat_final_pad(offset, expected_rows):
+    # Prefill appends text EOS and PAD as the final two reply rows. The HF
+    # consumer retains the last PAD row for every remaining audio position.
     model = _make_minimal_model()
     input_ids = torch.tensor([1, 2, 3], dtype=torch.long)
     q = _thinker_reply(6)
@@ -107,12 +110,12 @@ def test_positions_past_thinker_reply_keep_codec_embedding():
         embed={"thinker_reply": q},
         _omni_is_prefill=False,
         _omni_prompt_len=2,
-        _omni_num_computed_tokens=7,  # offset 5: one reply row left
+        _omni_num_computed_tokens=2 + offset,
     )
 
     base = _fake_embed_input_ids(input_ids)
-    assert torch.equal(out_embeds[0], base[0] + q[5])
-    assert torch.equal(out_embeds[1:], base[1:])
+    assert torch.equal(out_embeds, base + q[expected_rows])
+    assert torch.equal(q, _thinker_reply(6))
 
 
 THINKER_PROMPT_LEN = 6
@@ -268,7 +271,7 @@ def test_replay_span_crossing_prompt_boundary_conditions_generated_suffix(num_co
     model = _make_prefill_model()
     payload = _thinker_payload()
     full_ids, full_embeds = _full_talker_prompt(model, payload)
-    generated = [301, 302, 303]
+    generated = list(range(301, 309))  # Includes EOS and several PAD-conditioned positions.
     span_ids = torch.cat(
         [
             torch.zeros(TALKER_PROMPT_LEN - num_computed, dtype=torch.long),
@@ -291,6 +294,15 @@ def test_replay_span_crossing_prompt_boundary_conditions_generated_suffix(num_co
     assert torch.equal(out_embeds[:num_prompt], full_embeds[num_computed:])
     assert out_ids[num_prompt:].tolist() == generated
     assert torch.equal(out_embeds[num_prompt:], _original_decode_embeds(model, payload, generated))
+    # Check the reference tail rule independently of our sequential decoder.
+    num_reply = payload["hidden_states"]["output"].shape[0] - 1
+    terminal_ids = torch.tensor(generated[num_reply:])
+    terminal_conditioning = torch.cat(
+        [model.embed_text_eos_token, model.embed_text_pad_token.expand(len(terminal_ids) - 1, -1)]
+    )
+    assert torch.equal(
+        out_embeds[num_prompt + num_reply :], _fake_embed_input_ids(terminal_ids) + terminal_conditioning
+    )
 
 
 def test_repeated_replay_span_is_identical():
@@ -298,7 +310,7 @@ def test_repeated_replay_span_is_identical():
     # the same span must rebuild identical embeddings.
     model = _make_prefill_model()
     payload = _thinker_payload()
-    span_ids = torch.cat([torch.zeros(TALKER_PROMPT_LEN, dtype=torch.long), torch.tensor([301, 302])])
+    span_ids = torch.cat([torch.zeros(TALKER_PROMPT_LEN, dtype=torch.long), torch.arange(301, 309)])
     meta = dict(_omni_is_prefill=True, _omni_prompt_len=TALKER_PROMPT_LEN, _omni_num_computed_tokens=0)
 
     _, first, _ = model.talker_preprocess(span_ids, None, **payload, **meta)
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
index cc0d255c0e4..90069c1482d 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
@@ -746,7 +746,7 @@ def talker_preprocess(
         # schedule spans whose length does not identify their phase.
         # - Prompt positions take rows of the reconstructed Talker prompt.
         # - Generated positions add the request's ``thinker_reply`` row for their
-        #   generation offset; past the end of the reply, the codec embedding is kept.
+        #   generation offset, repeating its final text-PAD row after text ends.
         payload: OmniPayload = info_dict
         try:
             start = int(payload["_omni_num_computed_tokens"])
@@ -901,6 +901,8 @@ def thinker_to_talker_decode(self, input_ids, input_embeds, thinker_reply, offse
         rows are selected by position and ``thinker_reply`` is never consumed, so
         a single decode step, a speculative-decoding verify span and a
         preemption replay of the same positions all receive the same rows.
+        Prefill appends text EOS and PAD; the final PAD conditions every
+        remaining audio position, matching the reference decoder.
         """
         span_len = int(input_ids.shape[0])
         # Build the span embeddings from the codec ids rather than trusting the
@@ -916,6 +918,9 @@ def thinker_to_talker_decode(self, input_ids, input_embeds, thinker_reply, offse
             n = int(rows.shape[0])
             if n > 0:
                 input_embeds[:n] = input_embeds[:n] + rows.to(dtype=input_embeds.dtype, device=input_embeds.device)
+            if n < span_len:
+                pad = thinker_reply[-1:].to(dtype=input_embeds.dtype, device=input_embeds.device)
+                input_embeds[n:] = input_embeds[n:] + pad
         return input_ids, input_embeds, {}
 
     def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, **kwargs: object) -> torch.Tensor | None: