diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index d2a3ffb02..6bad98f79 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.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 diff --git a/docs/.nav.yml b/docs/.nav.yml index 8a03c7a4f..aa84f2819 100644 --- a/docs/.nav.yml +++ b/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: diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 1f1de15be..2a2f62c4a 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/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: diff --git a/docs/design/feature/preprocess_phase_contract.md b/docs/design/feature/preprocess_phase_contract.md new file mode 100644 index 000000000..d795c9c23 --- /dev/null +++ b/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) diff --git a/docs/design/index.md b/docs/design/index.md index c5b28130c..cff0b1eb3 100644 --- a/docs/design/index.md +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 73d21cbe8..cd7544af8 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/distributed/omni_connectors/test_mooncake_tcp_completion.py b/tests/distributed/omni_connectors/test_mooncake_tcp_completion.py new file mode 100644 index 000000000..c8a6312e8 --- /dev/null +++ b/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() diff --git a/tests/e2e/features/prefill_phase/test_chunked_prefill_tail.py b/tests/e2e/features/prefill_phase/test_chunked_prefill_tail.py new file mode 100644 index 000000000..8c59e3d23 --- /dev/null +++ b/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) diff --git a/tests/e2e/features/prefill_phase/worker_extension.py b/tests/e2e/features/prefill_phase/worker_extension.py new file mode 100644 index 000000000..eb0e9b83c --- /dev/null +++ b/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} diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index 81d2a6b8a..4155e5ffb 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/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)