[Feature] Define the runner preprocess phase contract and regression coverage (#7136)

Signed-off-by: xiaobaixia222 <289982752@qq.com>

.buildkite/cuda/test-merge.yml

@@ -79,6 +79,8 @@ steps:
     steps:
       - label: "TTS · Qwen3-TTS CustomVoice Test"
         source_file_dependencies:
+          - tests/e2e/features/prefill_phase/
+          - vllm_omni/worker/gpu_model_runner.py
           - tests/helpers/media.py
           - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
           - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
@@ -96,7 +98,7 @@ steps:
             timeout 40m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
               export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
-              pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+              pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py tests/e2e/features/prefill_phase/test_chunked_prefill_tail.py -m 'advanced_model and cuda' --run-level 'advanced_model'
             "
         mirror_hardwares: l4_1
 

docs/.nav.yml

@@ -127,6 +127,7 @@ nav:
         - Async Chunk: design/feature/async_chunk.md
         - Async Diffusion Output: design/feature/async_diffusion_output.md
         - Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
+        - Prefill/Decode Phase Contract: design/feature/preprocess_phase_contract.md
         - Automatic Prefix Caching: design/feature/prefix_caching.md
         - Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
       - Communication:

docs/contributing/model/adding_omni_model.md

@@ -222,6 +222,16 @@ __all__ = ["Qwen3OmniMoeForConditionalGeneration"]
 
 ## Key Components
 
+### Preprocess phase metadata
+
+Models with `has_preprocess=True`, including out-of-tree plugins, can rely on
+the runner's stable `_omni_is_prefill`, `_omni_prompt_len`, and
+`_omni_num_computed_tokens` fields. Use request progress to distinguish prompt
+prefill from decode: a one-token span can still be a prefill tail. See the
+[runner-to-model phase contract](../../design/feature/preprocess_phase_contract.md)
+for types, batched-hook semantics, MTP routing guarantees, and the correctness
+workaround for vLLM-Omni 0.20.x and earlier.
+
 ### 1. Model Interfaces
 
 Your model should implement the appropriate interfaces:

docs/design/feature/preprocess_phase_contract.md

@@ -0,0 +1,172 @@
+# Runner-to-model prefill/decode phase contract
+
+## Overview
+
+Autoregressive models with a custom `preprocess` hook need to distinguish
+prompt processing from decoding. A one-token span can be either a complete
+one-token prompt, the last chunk of a longer prompt, or a decode step. Using
+span length as the phase signal can replace a prompt embedding with a generated
+audio embedding and advance teacher-forced text too early without an exception.
+
+[Issue #4382](https://github.com/vllm-project/vllm-omni/issues/4382) requests a
+stable contract for the metadata and routing guards introduced by
+[PR #3662](https://github.com/vllm-project/vllm-omni/pull/3662).
+The shared runner already implements those guards. This design documents the
+existing interface and adds runner-level regression coverage rather than
+introducing a second phase detector or changing scheduling policy.
+
+## Scope and objectives
+
+- Make the three existing phase keys safe to depend on in out-of-tree models.
+- Preserve prompt embeddings and model state on a one-token prefill tail.
+- Keep normal and batched decode routing consistent in mixed batches.
+- Explain the correctness workaround required on older releases.
+
+This contract does not enable chunked prefill for models whose other state
+handling does not support it. It does not change the scheduler, introduce new
+kwargs, or standardize other `_omni_*` fields. The implementation adds
+documentation and regression coverage; no hot-path allocation is added.
+
+## Design
+
+### Ownership and data flow
+
+```text
+Scheduler / current input batch
+  prompt_token_ids + num_computed_tokens_cpu (before this step)
+                         |
+            OmniGPUModelRunner._preprocess
+                         |
+            per-request phase metadata
+                  /                 \
+       prefill / other spans     single-token decode
+          preprocess             preprocess or preprocess_decode_batch
+              |                              |
+       prompt embedding              runner-managed talker_mtp
+                  \                 /
+                    model forward
+```
+
+The producer is `OmniGPUModelRunner._preprocess` in
+`vllm_omni/worker/gpu_model_runner.py`. `GPUARModelRunner` and
+`GPUGenerationModelRunner` inherit it. The current `OmniNPUModelRunner` also
+inherits this method; device-specific execution still needs its own hardware
+validation. A plugin runner that overrides preprocessing is responsible for
+preserving this interface.
+
+### Stable model-facing fields
+
+The following names, types, and meanings are a supported interface for in-tree
+and out-of-tree model hooks despite the leading underscore. Models should treat
+them as read-only runner-owned values. Changes to their meanings or removal
+require an explicit compatibility/migration plan.
+
+| Key | Python type | Meaning |
+| --- | --- | --- |
+| `_omni_prompt_len` | `int` | Length of the current request state's `prompt_token_ids`. This is the stage's prompt length, not the current chunk size or raw text length. |
+| `_omni_num_computed_tokens` | `int` | The row's `input_batch.num_computed_tokens_cpu` value **before** processing the scheduled span. It includes already-computed/cached progress, not just tokens computed in the current call. |
+| `_omni_is_prefill` | `bool` | Exactly `_omni_num_computed_tokens < _omni_prompt_len`. Equality and greater progress mean decode. |
+
+The runner refreshes the fields after current request state has been gathered,
+immediately before selecting the per-row hook. Existing values in intermediate
+model state cannot override the current phase. Row indices follow the current
+input batch (including any reordering), not request arrival order. A resumed
+or streaming request is classified from its **current** prompt/progress values;
+the phase is not latched for the lifetime of the request.
+
+The normal hook receives these fields as keyword arguments:
+
+```python
+def preprocess(self, input_ids, input_embeds, **info):
+    if info["_omni_is_prefill"]:
+        offset = info["_omni_num_computed_tokens"]
+        # Select the scheduled slice of the model's prompt embeddings.
+        ...
+    else:
+        # Advance decode-only text/audio state.
+        ...
+```
+
+For `preprocess_decode_batch(input_ids=..., req_infos=...)`, each `req_infos`
+entry contains the same fields, aligned with its corresponding input row.
+Every such row has `_omni_is_prefill == False`. The optional earlier
+`preprocess_batch` hook is a separate whole-batch preparation hook; these freshly
+computed per-row fields are not promised at that earlier boundary.
+
+### Routing guarantee
+
+With `has_preprocess=True`, the batched decode hook is eligible only when it is
+callable, the runner has `talker_mtp`, the scheduled span is one token, and
+`_omni_is_prefill` is false. Otherwise the normal `preprocess` hook runs.
+
+Runner-managed MTP is likewise restricted to single-token decode rows. A prefill
+row, including a one-token tail, is never added to that MTP batch: the runner
+does not consume its `mtp_inputs`, generate codes for it, or overwrite its
+preprocessed embedding with an MTP embedding. Decode rows in the same step
+continue to receive their MTP outputs at their own token offsets. Models may
+still implement their own model-specific work inside `preprocess`.
+
+| Prompt length | Computed before step | Scheduled span | Phase | Runner MTP eligible |
+| --- | --- | --- | --- | --- |
+| 5 | 0 | 4 | Prefill | No |
+| 5 | 4 | 1 | Prefill tail | No |
+| 5 | 5 | 1 | Decode | Yes |
+| 5 | 6 | 1 | Decode | Yes |
+| 1 | 0 | 1 | Prefill | No |
+| 5 | 5 | 2 | Decode | No (not a single-token span) |
+
+### Compatibility
+
+!!! warning "vLLM-Omni 0.20.x and earlier"
+    AR talker models with custom preprocessing that infer phase from span length
+    must use `enable_chunked_prefill: false` for correctness on these releases.
+    They lack the phase metadata and the runner's prefill-tail MTP protection.
+    A model-only heuristic cannot prevent the older runner from overwriting a
+    one-token prompt embedding. Disabling chunking avoids split tails; it does
+    not make span length a reliable phase detector for every possible prompt.
+
+The mechanism landed in #3662 (v0.21.0rc2 / v0.22.0 and later); this document
+formalizes that existing behavior. Prefer upgrading the matching vLLM and
+vLLM-Omni release pair and consuming the explicit phase keys. Do not silently
+fall back to `input_ids.numel() > 1` when the model requires chunking correctness.
+For a custom or backported runtime, verify both metadata and MTP guards.
+
+## Correctness and testing plan
+
+Runner tests call the production `_preprocess` and `_talker_mtp_forward` methods
+with CPU buffers and a deterministic MTP model. They inspect kwargs at the model
+boundary, resulting embeddings, and per-request generated codes. Coverage must
+include a multi-step prompt split with a one-token tail followed by decode,
+mixed prefill/decode batches in different orders, normal/batched decode hooks,
+adjacent decode requests in the same hook call with exact per-request codes,
+single-token prompts, cached progress, and decode spans longer than one token.
+No test should pass merely because MTP was disabled for every row.
+
+Real-weight validation uses Qwen3-TTS CustomVoice through the full talker and
+code2wav pipeline. Record the actual scheduled prompt progress, assert that at
+least one one-token prefill tail was observed, check its hook metadata and MTP
+exclusion, and require finite non-empty output audio. Compare to an unchunked
+run with the same real input, greedy sampling, and `VLLM_BATCH_INVARIANT=1`.
+Greedy sampling alone does not guarantee identical codes across different
+prefill shapes: floating-point reduction differences can change an argmax and
+diverge in subsequent autoregressive steps. Invariant kernels keep this exact
+comparison focused on the phase contract. The two engines must be initialized
+separately because their chunking settings and token budgets differ. GPU tests belong under
+`tests/e2e/features/`; existing model tests remain useful adjacent regression
+coverage. Functional audio checks do not establish throughput or all-model
+semantic parity; this change makes no performance claim.
+
+## Open questions and discussions
+
+There are no new field names or execution modes to choose. Future runner
+implementations, including Model Runner V2, should preserve these semantics or
+provide an explicit migration path. Hardware validation beyond the tested
+platform remains separate from the shared Python routing tests.
+
+## References
+
+- [Issue #4382](https://github.com/vllm-project/vllm-omni/issues/4382)
+- [Original metadata and guards, #3662](https://github.com/vllm-project/vllm-omni/pull/3662)
+- [Qwen2.5-Omni consumer fix, #5019](https://github.com/vllm-project/vllm-omni/pull/5019)
+- [Adding an Omni model](../../contributing/model/adding_omni_model.md)
+- [Test writing guide](../../contributing/ci/test_writing_guide.md)

docs/design/index.md

@@ -24,6 +24,7 @@ implementation contract; it is not, by itself, a general support claim.
 - [Async Chunk](feature/async_chunk.md)
 - [Async Diffusion Output](feature/async_diffusion_output.md)
 - [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
+- [Runner-to-model Prefill/Decode Phase Contract](feature/preprocess_phase_contract.md)
 - [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
 - [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
 

pyproject.toml

@@ -53,7 +53,9 @@ dev = [
     "soundfile==0.14.0",
     "imageio[ffmpeg]==2.37.4",
     "opencv-python==5.0.0.93",
-    "mooncake-transfer-engine-cuda13==0.3.11.post1",
+    # 0.3.12 includes acknowledged TCP writes (Mooncake #2850): completion
+    # must mean the receive buffer is ready before zero-copy consumers read it.
+    "mooncake-transfer-engine-cuda13==0.3.12.post1",
     "av==18.0.0", # for ComfyUI tests
     "openpyxl==3.1.5", # for nightly CI
     "pyttsx3==2.99",

tests/distributed/omni_connectors/test_mooncake_tcp_completion.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""TCP completion must mean that the destination has applied the write."""
+
+import multiprocessing
+import os
+import signal
+import threading
+import time
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import TransferEngine
+
+pytestmark = [
+    pytest.mark.parallel,
+    pytest.mark.core_model,
+    pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not installed"),
+    pytest.mark.skipif(os.name != "posix", reason="Requires POSIX process suspension"),
+]
+
+
+def _tcp_receiver(pipe):
+    engine = TransferEngine()
+    assert engine.initialize("127.0.0.1", "P2PHANDSHAKE", "tcp", "") == 0
+    destination = torch.zeros(65536, dtype=torch.uint8).pin_memory()
+    assert engine.register_memory(destination.data_ptr(), destination.numel()) == 0
+    pipe.send((engine.get_rpc_port(), destination.data_ptr()))
+    while True:
+        value = pipe.recv()
+        if value is None:
+            return
+        deadline = time.monotonic() + 10
+        while not bool((destination == value).all()):
+            assert time.monotonic() < deadline, "Destination did not receive the expected bytes"
+        pipe.send(True)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_tcp_write_completion_waits_for_receiver(monkeypatch):
+    # Also exercise TCP on hosts where Mooncake would otherwise auto-select RDMA.
+    monkeypatch.setenv("MC_FORCE_TCP", "1")
+    context = multiprocessing.get_context("spawn")
+    parent, child = context.Pipe()
+    receiver = context.Process(target=_tcp_receiver, args=(child,))
+    receiver.start()
+    receiver_pid = receiver.pid
+    assert receiver_pid is not None
+    stopped = False
+    try:
+        assert parent.poll(120), "Receiver did not initialize"
+        port, address = parent.recv()
+        engine = TransferEngine()
+        assert engine.initialize("127.0.0.1", "P2PHANDSHAKE", "tcp", "") == 0
+        source = torch.ones(65536, dtype=torch.uint8).pin_memory()
+        assert engine.register_memory(source.data_ptr(), source.numel()) == 0
+        remote = f"127.0.0.1:{port}"
+
+        # Establish the data connection before suspending the receiver, so the
+        # second call tests data completion rather than RPC/connection setup.
+        assert engine.batch_transfer_sync_write(remote, [source.data_ptr()], [address], [source.numel()]) == 0
+        parent.send(1)
+        assert parent.poll(30) and parent.recv() is True
+        os.kill(receiver_pid, signal.SIGSTOP)
+        stopped = True
+        _, status = os.waitpid(receiver_pid, os.WUNTRACED)
+        assert os.WIFSTOPPED(status)
+        source.fill_(173)
+        completed = threading.Event()
+        result = []
+
+        def transfer():
+            result.append(engine.batch_transfer_sync_write(remote, [source.data_ptr()], [address], [source.numel()]))
+            completed.set()
+
+        thread = threading.Thread(target=transfer, daemon=True)
+        thread.start()
+        completed_while_stopped = completed.wait(0.5)
+        os.kill(receiver_pid, signal.SIGCONT)
+        stopped = False
+        thread.join(30)
+        assert result == [0], "Write did not complete successfully after resuming the receiver"
+        parent.send(173)
+        assert parent.poll(30) and parent.recv() is True
+        assert not completed_while_stopped, "TCP reported completion before destination memory could be updated"
+    finally:
+        if stopped:
+            os.kill(receiver_pid, signal.SIGCONT)
+        if receiver.is_alive():
+            parent.send(None)
+            receiver.join(10)
+        if receiver.is_alive():
+            receiver.terminate()
+            receiver.join(10)
+        parent.close()
+        child.close()

tests/e2e/features/prefill_phase/test_chunked_prefill_tail.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Real-weight regression: force a final one-token prefill chunk, then decode."""
+
+import json
+import os
+from pathlib import Path
+
+import pytest
+import soundfile as sf
+import torch
+from transformers import AutoTokenizer
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import Qwen3TTSPromptEmbedsBuilder
+
+pytestmark = [pytest.mark.advanced_model, pytest.mark.tts]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_one_token_prefill_tail_matches_unchunked_audio(tmp_path, monkeypatch):
+    # Greedy sampling alone does not make different prefill shapes numerically
+    # identical. Use invariant kernels for the exact codes/waveform comparison.
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
+    model = str(Path(os.environ.get("MODEL_PREFIX", "")) / "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice")
+    tokenizer = AutoTokenizer.from_pretrained(model)
+    config = Qwen3TTSConfig.from_pretrained(model)
+    info = {
+        "task_type": ["CustomVoice"],
+        "text": ["Hello, this is a test of speech generation."],
+        "language": ["English"],
+        "speaker": ["vivian"],
+        "max_new_tokens": [128],
+    }
+    prompt_len = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+        info,
+        task_type="CustomVoice",
+        tokenize_prompt=lambda text: tokenizer(text)["input_ids"],
+        codec_language_id=config.talker_config.codec_language_id,
+        spk_is_dialect=config.talker_config.spk_is_dialect,
+    )
+    assert prompt_len > 2
+    prompt = {"prompt_token_ids": [0] * prompt_len, "additional_information": info}
+    observations = []
+    audios = []
+    for chunked in (False, True):
+        deploy = modify_stage_config(
+            get_deploy_config_path("qwen3_tts.yaml"),
+            updates={
+                "async_chunk": False,
+                "stages": {
+                    0: {
+                        "enable_chunked_prefill": chunked,
+                        "enable_prefix_caching": False,
+                        "max_num_batched_tokens": prompt_len - 1 if chunked else 2048,
+                        "max_model_len": 2048,
+                        "max_num_seqs": 1,
+                        "enforce_eager": True,
+                        "async_scheduling": False,
+                        "default_sampling_params.temperature": 0.0,
+                        "default_sampling_params.max_tokens": 128,
+                        "subtalker_sampling_params": {"do_sample": False},
+                    },
+                    1: {"enforce_eager": True, "async_scheduling": False},
+                },
+            },
+        )
+        engine = AsyncOmni(
+            model=model,
+            deploy_config=deploy,
+            worker_extension_cls=("tests.e2e.features.prefill_phase.worker_extension.PhaseContractWorkerExtension"),
+            stage_init_timeout=600,
+            init_timeout=900,
+        )
+        try:
+            assert await engine.collective_rpc("start_phase_contract_probe", stage_ids=[0]) == [[True]]
+            audio_parts = []
+            async for output in engine.generate(prompt, request_id=f"phase-{chunked}"):
+                if output.final_output_type == "audio":
+                    multimodal = output.outputs[0].multimodal_output
+                    assert (torch.as_tensor(multimodal["sr"]) == 24000).all()
+                    audio = multimodal["audio"]
+                    audio_parts.extend(audio if isinstance(audio, list) else [audio])
+            assert audio_parts, "pipeline returned no audio"
+            waveform = torch.cat([part.detach().cpu().reshape(-1) for part in audio_parts])
+            assert waveform.numel() > 0 and torch.isfinite(waveform).all()
+            assert waveform.abs().max() > 0
+            sf.write(tmp_path / f"chunked-{chunked}.wav", waveform.float().numpy(), 24000)
+            audios.append(waveform)
+            probe = await engine.collective_rpc("get_phase_contract_probe", stage_ids=[0])
+            observations.append(probe[0][0])
+        finally:
+            engine.shutdown()
+
+    baseline, chunked_probe = observations
+    evidence = {
+        "prompt_len": prompt_len,
+        "observations": observations,
+        "audio_samples": audios[0].numel(),
+        "chunked_audio_samples": audios[1].numel(),
+        "max_audio_difference": (
+            float((audios[1] - audios[0]).abs().max()) if audios[1].shape == audios[0].shape else None
+        ),
+    }
+    evidence_path = tmp_path / "phase_contract.json"
+    evidence_path.write_text(json.dumps(evidence, indent=2))
+    print(f"Phase contract evidence: {evidence_path}")
+
+    # Preserve both observations even when a parity assertion fails in CI.
+    assert [prompt_len, 0, prompt_len, True] in baseline["events"]
+    assert [prompt_len, 0, prompt_len - 1, True] in chunked_probe["events"]
+    assert [prompt_len, prompt_len - 1, 1, True] in chunked_probe["events"]
+    assert baseline["codes"] and chunked_probe["codes"]
+    assert chunked_probe["codes"] == baseline["codes"]
+    torch.testing.assert_close(audios[1], audios[0], rtol=1e-4, atol=1e-5)

tests/e2e/features/prefill_phase/worker_extension.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Observe the real runner/model boundary without replacing inference."""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+
+
+class PhaseContractWorkerExtension:
+    model_runner: "OmniGPUModelRunner"
+
+    def start_phase_contract_probe(self):
+        runner = self.model_runner
+        self.phase_contract_events: list[tuple[int, int, int, bool]] = []
+        self.phase_contract_codes: list[list[list[int]]] = []
+        tail_embeddings = {}
+        original_preprocess = runner.model.preprocess
+        original_mtp = runner._talker_mtp_forward
+
+        def preprocess(input_ids, input_embeds, **info):
+            rid = info["request_id"]
+            index = runner.input_batch.req_id_to_index[rid]
+            prompt_len = len(runner.requests[rid].prompt_token_ids)
+            computed = int(runner.input_batch.num_computed_tokens_cpu[index])
+            assert info["_omni_prompt_len"] == prompt_len
+            assert info["_omni_num_computed_tokens"] == computed
+            assert info["_omni_is_prefill"] is (computed < prompt_len)
+            result = original_preprocess(input_ids=input_ids, input_embeds=input_embeds, **info)
+            span = input_ids.numel()
+            self.phase_contract_events.append((prompt_len, computed, span, info["_omni_is_prefill"]))
+            if computed == prompt_len - 1 and span == 1:
+                offset = int(runner.query_start_loc.cpu[index])
+                tail_embeddings[offset] = result[1].detach().clone()
+            return result
+
+        def mtp(decode_req_ids, inputs_embeds, start_offsets=None):
+            for rid in decode_req_ids:
+                index = runner.input_batch.req_id_to_index[rid]
+                assert int(runner.input_batch.num_computed_tokens_cpu[index]) >= len(
+                    runner.requests[rid].prompt_token_ids
+                ), "prefill row was routed through MTP"
+            original_mtp(decode_req_ids, inputs_embeds, start_offsets)
+            for offset, expected in tail_embeddings.items():
+                torch.testing.assert_close(inputs_embeds[offset : offset + 1], expected, rtol=0, atol=0)
+            tail_embeddings.clear()
+            for rid in decode_req_ids:
+                self.phase_contract_codes.append(runner.model_intermediate_buffer[rid]["codes"]["audio"].tolist())
+
+        runner.model.preprocess = preprocess
+        setattr(runner, "_talker_mtp_forward", mtp)
+        return True
+
+    def get_phase_contract_probe(self):
+        return {"events": self.phase_contract_events, "codes": self.phase_contract_codes}

tests/worker/test_omni_gpu_model_runner.py

@@ -6,6 +6,7 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
 from vllm.v1.cudagraph_dispatcher import CUDAGraphMode
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
@@ -108,7 +109,8 @@ class DummyInputBatch:
 class DummyReqState:
     """A minimal request state container."""
 
-    pass
+    mm_features: list[str]
+    additional_information_cpu: dict
 
 
 def test_model_forward_passes_request_ids_to_decode_metadata(monkeypatch):
@@ -878,3 +880,151 @@ def test_maybe_attach_mimo_audio_req_infos_no_req_state_returns_input():
 
     # When no req_state, helper should be a no-op.
     assert result is req_infos
+
+
+def _make_phase_runner(monkeypatch, rows, *, batched_decode):
+    """Exercise real preprocessing and MTP routing with CPU model buffers.
+
+    Each row is (request id, prompt length, computed tokens, scheduled tokens).
+    Only model computation and the device forward context are replaced.
+    """
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod, "get_pp_group", lambda: SimpleNamespace(is_first_rank=True))
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+    runner = _make_runner(req_ids=tuple(row[0] for row in rows))
+    total = sum(row[3] for row in rows)
+    runner.supports_mm_inputs = False
+    runner.enable_prompt_embeds = False
+    runner.uses_mrope = False
+    runner.uses_xdrope_dim = 0
+    runner.has_talker_mtp = True
+    runner.model_config = SimpleNamespace(is_encoder_decoder=False)
+    runner.vllm_config.model_config.async_chunk = True
+    runner._init_model_kwargs = dict
+    runner.input_ids = DummyBuffer(torch.arange(total, dtype=torch.int64))
+    runner.inputs_embeds = DummyBuffer(torch.zeros(total, 4))
+    runner.positions = torch.arange(total)
+    runner.input_batch.num_computed_tokens_cpu = [row[2] for row in rows]
+    offsets = [0]
+    for rid, prompt_len, _, scheduled in rows:
+        runner.requests[rid].prompt_token_ids = list(range(prompt_len))
+        offsets.append(offsets[-1] + scheduled)
+        # Runner-owned metadata must override stale model/intermediate state.
+        runner.model_intermediate_buffer[rid] = {
+            "_omni_is_prefill": "stale",
+            "_omni_prompt_len": -1,
+            "_omni_num_computed_tokens": -1,
+        }
+    runner.query_start_loc.cpu = torch.tensor(offsets, dtype=torch.int32)
+    calls = []
+    batch_calls = []
+
+    def preprocess(input_ids, input_embeds, **info):
+        calls.append(("normal", dict(info)))
+        embeds = input_ids.float().view(-1, 1).expand(-1, 4).clone()
+        updates = {}
+        if not info["_omni_is_prefill"]:
+            updates["mtp_inputs"] = (torch.zeros_like(embeds), torch.zeros_like(embeds))
+        return input_ids, embeds, updates
+
+    def preprocess_decode_batch(input_ids, req_infos):
+        batch_calls.append([info["request_id"] for info in req_infos])
+        calls.extend(("batch", dict(info)) for info in req_infos)
+        embeds = input_ids.float().view(-1, 1).expand(-1, 4).clone()
+        return input_ids, embeds, torch.zeros_like(embeds), torch.zeros_like(embeds), [{} for _ in req_infos]
+
+    runner.model = SimpleNamespace(
+        has_preprocess=True,
+        preprocess=preprocess,
+        talker_mtp_output_key=("codes", "audio"),
+    )
+    if batched_decode:
+        runner.model.preprocess_decode_batch = preprocess_decode_batch
+    output = SchedulerOutput(
+        total_num_scheduled_tokens=total,
+        num_scheduled_tokens={row[0]: row[3] for row in rows},
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        scheduled_spec_decode_tokens={},
+        scheduled_encoder_inputs={},
+        num_common_prefix_blocks=[],
+        finished_req_ids=set(),
+        free_encoder_mm_hashes=[],
+    )
+    return runner, output, calls, batch_calls
+
+
+@pytest.mark.parametrize("batched_decode", [False, True], ids=["normal", "batch"])
+@pytest.mark.parametrize("order", ["interleaved", "reversed", "adjacent"])
+def test_preprocess_phase_contract_mixed_batch(monkeypatch, batched_decode, order):
+    rows = [
+        ("decode-before", 5, 5, 1),
+        ("tail", 5, 4, 1),
+        ("decode-after", 5, 6, 1),
+        ("single-prompt", 1, 0, 1),
+        ("cached-prefill", 8, 4, 3),
+        ("multi-token-decode", 5, 5, 2),
+    ]
+    if order == "reversed":
+        rows.reverse()
+    elif order == "adjacent":
+        rows = [rows[1], rows[0], rows[2], *rows[3:]]
+    runner, scheduled, calls, batch_calls = _make_phase_runner(monkeypatch, rows, batched_decode=batched_decode)
+    _, embeds, *_ = runner._preprocess(scheduled, scheduled.total_num_scheduled_tokens)
+
+    if batched_decode:
+        decode_order = [rid for rid, _, _, _ in rows if rid.startswith("decode-")]
+        expected_batches = [decode_order] if order == "adjacent" else [[rid] for rid in decode_order]
+        assert batch_calls == expected_batches
+    else:
+        assert batch_calls == []
+    assert len(calls) == len(rows)
+    offset = 0
+    code_index = 0
+    for (route, info), (rid, prompt_len, computed, span) in zip(calls, rows, strict=True):
+        assert info["request_id"] == rid
+        assert type(info["_omni_prompt_len"]) is int
+        assert type(info["_omni_num_computed_tokens"]) is int
+        assert type(info["_omni_is_prefill"]) is bool
+        assert info["_omni_prompt_len"] == prompt_len
+        assert info["_omni_num_computed_tokens"] == computed
+        assert info["_omni_is_prefill"] == (computed < prompt_len)
+        mtp_eligible = rid in {"decode-before", "decode-after"}
+        assert route == ("batch" if batched_decode and mtp_eligible else "normal")
+        expected = torch.arange(offset, offset + span).float().view(-1, 1).expand(-1, 4)
+        torch.testing.assert_close(embeds[offset : offset + span], expected + int(mtp_eligible))
+        assert ("codes" in runner.model_intermediate_buffer[rid]) == mtp_eligible
+        if mtp_eligible:
+            torch.testing.assert_close(
+                runner.model_intermediate_buffer[rid]["codes"]["audio"],
+                torch.tensor([[code_index]], dtype=torch.int64),
+                rtol=0,
+                atol=0,
+            )
+            code_index += 1
+        offset += span
+
+
+@pytest.mark.parametrize("batched_decode", [False, True], ids=["normal", "batch"])
+def test_preprocess_one_token_chunked_prefill_tail_then_decode(monkeypatch, batched_decode):
+    # Replay a five-token prompt split at a four-token scheduling budget.
+    runner, scheduled, calls, _ = _make_phase_runner(monkeypatch, [("r", 5, 0, 4)], batched_decode=batched_decode)
+    runner._preprocess(scheduled, 4)
+    assert calls[-1][1]["_omni_is_prefill"] is True
+    assert "codes" not in runner.model_intermediate_buffer["r"]
+
+    for computed, is_prefill in [(4, True), (5, False)]:
+        runner.input_batch.num_computed_tokens_cpu = [computed]
+        scheduled.total_num_scheduled_tokens = 1
+        scheduled.num_scheduled_tokens = {"r": 1}
+        runner.query_start_loc.cpu = torch.tensor([0, 1], dtype=torch.int32)
+        runner.input_ids.gpu[0] = 42
+        _, embeds, *_ = runner._preprocess(scheduled, 1)
+        route, info = calls[-1]
+        assert info["_omni_is_prefill"] is is_prefill
+        assert info["_omni_num_computed_tokens"] == computed
+        assert info["_omni_prompt_len"] == 5
+        assert route == ("batch" if batched_decode and not is_prefill else "normal")
+        torch.testing.assert_close(embeds, torch.full((1, 4), 42.0 if is_prefill else 43.0))
+        assert ("codes" in runner.model_intermediate_buffer["r"]) == (not is_prefill)