From f64a51902e7575f0c801ed5479fde124c6f1216a Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Tue, 25 Aug 2026 19:21:42 +0800 Subject: [PATCH 01/41] refactor(stage_input_processors): RFC #4872 P3/P8a shared helpers + P1/P7 contracts (Phase 1-2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - P3/P8a: 10 模块委托 _common 共享辅助函数，新增 _dispatch.py 契约层 - P1: stage_engine_core_client.py / orchestrator.py 两调用点收敛 - P7: __init__.py 包导出；新增 golden 77 测试（conftest + test_common_helpers_golden） Signed-off-by: RyanYun09 --- .../stage_input_processors/conftest.py | 126 +++++ .../test_common_helpers_golden.py | 531 ++++++++++++++++++ vllm_omni/diffusion/output_formatter.py | 6 +- vllm_omni/engine/orchestrator.py | 26 +- vllm_omni/engine/stage_engine_core_client.py | 22 +- .../stage_input_processors/__init__.py | 42 ++ .../stage_input_processors/_common.py | 427 ++++++++++++++ .../stage_input_processors/_dispatch.py | 304 ++++++++++ .../stage_input_processors/audex.py | 8 +- .../stage_input_processors/cosyvoice3.py | 40 +- .../stage_input_processors/glm_tts.py | 15 +- .../stage_input_processors/higgs_audio_v2.py | 37 +- .../stage_input_processors/higgs_audio_v3.py | 40 +- .../stage_input_processors/ming_flash_omni.py | 14 +- .../stage_input_processors/qwen3_omni.py | 12 +- .../stage_input_processors/step_audio2.py | 15 +- 16 files changed, 1533 insertions(+), 132 deletions(-) create mode 100644 tests/model_executor/stage_input_processors/conftest.py create mode 100644 tests/model_executor/stage_input_processors/test_common_helpers_golden.py create mode 100644 vllm_omni/model_executor/stage_input_processors/_common.py create mode 100644 vllm_omni/model_executor/stage_input_processors/_dispatch.py diff --git a/tests/model_executor/stage_input_processors/conftest.py b/tests/model_executor/stage_input_processors/conftest.py new file mode 100644 index 00000000000..2eb57c78f60 --- /dev/null +++ b/tests/model_executor/stage_input_processors/conftest.py @@ -0,0 +1,126 @@ +"""Local-dev import shim for ``vllm_omni`` tests (macOS). + +vllm does not install on macOS, so importing real ``vllm_omni`` submodules +fails at ``import vllm``. This conftest installs a *test-only* shim that: + +1. Bypasses ``vllm_omni/__init__.py`` (so its ``patch`` / ``transformers`` + machinery is never executed) by pre-inserting a synthetic ``vllm_omni`` + package into ``sys.modules`` whose ``__path__`` points at the real + package directory. +2. Stubs the entire ``vllm`` package tree with permissive ``_Dummy`` classes + so module-level ``from vllm... import ...`` statements resolve. + +The shim is active **only** when real ``vllm`` is not importable (local dev). +On CI (vllm installed) it is a no-op and the real package is imported, which +guards against the shim masking genuine import problems. + +Scope: only pure-logic / unit-level tests (no model execution). Any test that +needs a real vLLM runtime must be marked to run on GPU/NPU CI instead. +""" + +import importlib.abc +import importlib.machinery +import os +import sys +import types + +# --------------------------------------------------------------------------- +# Locate the repo root (this file lives at: +# /tests/model_executor/stage_input_processors/conftest.py) +# --------------------------------------------------------------------------- +_REPO_ROOT = os.path.dirname( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) +) +_VLLM_OMNI_DIR = os.path.join(_REPO_ROOT, "vllm_omni") + + +class _Meta(type): + """Metaclass so that *class-level* attribute access on a _Dummy also + yields _Dummy (e.g. ``SomeStubClass.foo`` used in module-level code).""" + + def __getattr__(cls, name): + return _Dummy + + +class _Dummy(metaclass=_Meta): + """Permissive stand-in for any vllm symbol. + + Supports subclassing (``__mro_entries__``), ``total=False`` TypedDict + style (``__init_subclass__``), ``X | None`` unions, and arbitrary + attribute / call / item access — enough to get module *import* past the + vllm dependency while keeping pure helper functions callable. + """ + + def __init__(self, *args, **kwargs): + pass + + def __call__(self, *args, **kwargs): + return _Dummy() + + def __getattr__(self, name): + return _Dummy + + def __class_getitem__(cls, item): + return cls + + def __mro_entries__(self, bases): + return (_Dummy,) + + def __init_subclass__(cls, **kwargs): + pass + + def __or__(self, other): + return _Dummy + + def __ror__(self, other): + return _Dummy + + +class _StubLoader(importlib.abc.Loader): + """Creates an empty module with a wildcard __getattr__ for any vllm.*.""" + + def create_module(self, spec): + module = types.ModuleType(spec.name) + module.__path__ = [] + module.__getattr__ = lambda attr: _Dummy + return module + + def exec_module(self, module): + pass + + +class _StubFinder(importlib.abc.MetaPathFinder): + """Auto-creates a stub package for ``vllm`` and any ``vllm.*`` submodule.""" + + def find_spec(self, fullname, path=None, target=None): + if fullname == "vllm" or fullname.startswith("vllm."): + return importlib.machinery.ModuleSpec( + fullname, _StubLoader(), is_package=True + ) + return None + + +def _install_shim(): + """Bypass the real vllm_omni/__init__ + stub vllm.""" + if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + # Synthetic vllm_omni parent package (skip the real __init__). + pkg = types.ModuleType("vllm_omni") + pkg.__path__ = [_VLLM_OMNI_DIR] + pkg.__package__ = "vllm_omni" + pkg._SHIM_ACTIVE = True + sys.modules.setdefault("vllm_omni", pkg) + # Wildcard vllm stub finder (meta path -> takes precedence for vllm.*). + sys.meta_path.insert(0, _StubFinder()) + + +def _activate_if_needed(): + try: + import vllm # noqa: F401 + except ModuleNotFoundError: + _install_shim() + + +_activate_if_needed() diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py new file mode 100644 index 00000000000..eec93d2ba0d --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -0,0 +1,531 @@ +"""Golden tests that lock the *current* per-module behavior of duplicated +helper functions in ``stage_input_processors``. + +Phase 1a (RFC #4872 Proposal 3): capture behaviour BEFORE introducing +``_common.py``. Each case here encodes a deliberately-observed semantic +delta between modules (None handling, tensor->list, ``ConstantList._x`` +unwrapping, tuple handling, codec-frame validity masks, delay-pattern +strictness, ...). ``_common.ensure_list`` (Phase 1b) must pass the same +cases; do NOT silently "pick one" semantics across variants. + +Modules that cannot be imported in this environment (e.g. ``cosyvoice3`` / +``step_audio2`` on macOS because of missing deps) are skipped locally and +run on CI where the real dependencies exist. +""" + +import importlib +from typing import Any, Callable + +import pytest +import torch + +_PREFIX = "vllm_omni.model_executor.stage_input_processors." + + +def _import(name: str) -> Any: + try: + if name == "diffusion.output_formatter": + # The real module lives under ``vllm_omni.diffusion``, not under + # ``stage_input_processors`` (which only hosts the golden alias). + return importlib.import_module("vllm_omni.diffusion.output_formatter") + return importlib.import_module(_PREFIX + name) + except Exception as exc: # pragma: no cover - env dependent + pytest.skip(f"module {name!r} not importable here: {type(exc).__name__}: {exc}") + raise AssertionError("unreachable") # keep type checkers happy + + +class _ConstantList: + """Mimics the ``ConstantList._x`` attribute seen across modules.""" + + def __init__(self, values): + self._x = list(values) + + +def _norm(value: Any) -> Any: + """Normalize a result for comparison: tensors -> nested lists.""" + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + return value + + +def _is_tensor(value: Any) -> bool: + return isinstance(value, torch.Tensor) + + +# =========================================================================== +# _ensure_list +# =========================================================================== + + +def _ensure_list_cases(): + return [ + # (module, input, expected-or-None-if-tensor, expect_typeerror) + ("qwen3_omni", [1, 2], [1, 2], False), + ("qwen3_omni", _ConstantList([1, 2]), [1, 2], False), + # qwen3_omni returns non-list values *unchanged* (incl. None / tuple). + ("qwen3_omni", None, None, False), + ("qwen3_omni", (1, 2), (1, 2), False), + # qwen3_omni has NO tolist() branch: tensor comes back unchanged. + ("qwen3_omni", "TENSOR", None, False), + ("ming_flash_omni", [1, 2], [1, 2], False), + ("ming_flash_omni", _ConstantList([1, 2]), [1, 2], False), + ("ming_flash_omni", "TENSOR", [1, 2], False), + ("ming_flash_omni", (1, 2), [1, 2], False), + # ming_flash_omni: list(None) raises TypeError. + ("ming_flash_omni", None, None, True), + ("audex", None, [], False), + ("audex", "TENSOR", [1, 2, 3, 4], False), + ("audex", [1, 2], [1, 2], False), + ("audex", (1, 2), [1, 2], False), + # audex: list(5) raises TypeError. + ("audex", 5, None, True), + ("cosyvoice3", None, [], False), + ("cosyvoice3", (1, 2), [1, 2], False), + ("cosyvoice3", 5, [5], False), # TypeError -> [x] + ("cosyvoice3", _ConstantList([1, 2]), [1, 2], False), + ("cosyvoice3", [1, 2], [1, 2], False), + ("step_audio2", _ConstantList([1, 2]), [1, 2], False), + ("step_audio2", "TENSOR", [1, 2], False), + ("step_audio2", (1, 2), [1, 2], False), + # step_audio2: non-iterable scalar -> [x]; None is not iterable -> [None]. + ("step_audio2", 5, [5], False), + ("step_audio2", None, [None], False), + ("diffusion.output_formatter", [1, 2], [1, 2], False), + ("diffusion.output_formatter", None, [], False), + ("diffusion.output_formatter", 5, [5], False), + ] + + +@pytest.mark.parametrize("module,inp,expected,expect_typeerror", + _ensure_list_cases()) +def test_ensure_list_golden(module, inp, expected, expect_typeerror): + mod = _import(module) + fn = getattr(mod, "_ensure_list") + if inp == "TENSOR": + inp = torch.tensor([[1, 2], [3, 4]]) if module == "audex" else torch.tensor([1, 2]) + if expect_typeerror: + with pytest.raises(TypeError): + fn(inp) + return + result = fn(inp) + if _is_tensor(result): + # qwen3_omni may return a tensor unchanged. + assert _is_tensor(expected) or expected is None + if expected is None: + assert torch.equal(result, torch.tensor([1, 2])) + return + assert _norm(result) == expected + + +# =========================================================================== +# _extract_last_frame +# =========================================================================== + + +def _codes_audio(tensor): + return {"codes": {"audio": tensor}} + + +@pytest.mark.parametrize("module,payload,expected,expect_error", [ + # higgs_audio_v2: codes.audio key path; range filter [-0, _NUM_REAL_CODES). + ("higgs_audio_v2", + _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), + [4, 5], False), + ("higgs_audio_v2", + _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 1024]])), # >= _NUM_REAL_CODES + None, False), + ("higgs_audio_v2", + _codes_audio(torch.tensor([[0, 1], [2, 3], [-1, 5]])), + None, False), + ("higgs_audio_v2", + _codes_audio(torch.tensor([1, 2, 3])), + [1, 2, 3], False), + ("higgs_audio_v2", + _codes_audio(torch.empty(0, 4)), + None, False), + ("higgs_audio_v2", + _codes_audio(torch.zeros(2, 2, 2)), + None, True), # 3-D -> ValueError + ("higgs_audio_v2", + {"codes": {}}, + None, False), + # fish_speech: TOP-LEVEL audio_codes key; audio_code_valid mask; cpu+long. + ("fish_speech", + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), + "audio_code_valid": torch.tensor([1, 1, 1])}, + [4, 5], False), + ("fish_speech", + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), + "audio_code_valid": torch.tensor([1, 1, 0])}, + None, False), + ("fish_speech", + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), + "audio_code_valid": True}, + [4, 5], False), + ("fish_speech", + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), + "audio_code_valid": False}, + None, False), + ("fish_speech", + {"audio_codes": torch.tensor([[0, 1], [2, 3], [0, 0]])}, + None, False), # no valid key -> frame.any() False + ("fish_speech", + {"audio_codes": torch.tensor([1, 2, 3])}, + [1, 2, 3], False), + ("fish_speech", + {"audio_codes": torch.zeros(2, 2, 2)}, + None, True), + ("fish_speech", + {}, + None, False), + # qwen3_tts: codes.audio; frame.any() gate; long (no cpu). + ("qwen3_tts", + _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), + [4, 5], False), + ("qwen3_tts", + _codes_audio(torch.tensor([[0, 1], [2, 3], [0, 0]])), + None, False), + ("qwen3_tts", + _codes_audio(torch.tensor([1, 2, 3])), + [1, 2, 3], False), + ("qwen3_tts", + _codes_audio(torch.zeros(2, 2, 2)), + None, True), + ("qwen3_tts", + {"codes": {}}, + None, False), + # voxtral_tts: for a single tensor it flattens the WHOLE tensor (not just + # the last row); for a list of tensors it flattens the last tensor. + ("voxtral_tts", + _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), + [0, 1, 2, 3, 4, 5], False), + ("voxtral_tts", + {"codes": {"audio": [torch.tensor([1, 2]), torch.tensor([3, 4])]}}, + [3, 4], False), + ("voxtral_tts", + {"codes": {}}, + None, False), +]) +def test_extract_last_frame_golden(module, payload, expected, expect_error): + mod = _import(module) + fn = getattr(mod, "_extract_last_frame") + if expect_error: + with pytest.raises(ValueError): + fn(payload) + return + result = fn(payload) + assert _norm(result) == expected + + +# =========================================================================== +# _revert_delay_pattern (higgs_audio_v2 lenient vs higgs_audio_v3 strict) +# =========================================================================== + + +def test_revert_delay_pattern_higgs_v2_golden(): + mod = _import("higgs_audio_v2") + fn = getattr(mod, "_revert_delay_pattern") + inp = torch.tensor([ + [10, 1, 2, 3, 4], + [11, 12, 5, 6, 7], + [13, 14, 15, 8, 9], + ]) # [Q=3, T=5], seq_len = 3 + out = fn(inp) + assert out.shape == (3, 3) + assert out.tolist() == [[10, 1, 2], [12, 5, 6], [15, 8, 9]] + # Lenient: strictly T < Q returns input unchanged. + short = torch.tensor([[1], [2]]) # Q=2, T=1 -> t < q True + assert torch.equal(fn(short), short) + # Non-lenient: T >= Q processes (even when the result is narrower). + proc = torch.tensor([[1, 2], [3, 4]]) # Q=2, T=2 -> seq_len = 1 + assert fn(proc).tolist() == [[1], [4]] + + +def test_revert_delay_pattern_higgs_v3_golden(): + mod = _import("higgs_audio_v3") + fn = getattr(mod, "_revert_delay_pattern") + q = int(mod._NUM_CODEBOOKS) + # Build a [Q, T] input with T = q + 2 (seq_len = 3). + t = q + 2 + inp = torch.arange(q * t, dtype=torch.long).reshape(q, t) + out = fn(inp) + assert out.shape == (q, 3) + expected = torch.cat( + [inp[i : i + 1, i : 3 + i] for i in range(q)], dim=0 + ) + assert torch.equal(out, expected) + # Strict: wrong codebook count raises. + with pytest.raises(ValueError): + fn(torch.zeros(q + 1, t, dtype=torch.long)) + # Strict: T < Q raises. + with pytest.raises(ValueError): + fn(torch.zeros(q, q - 1, dtype=torch.long)) + + +# =========================================================================== +# _filter_real_code_frames (higgs v2 [frames,Q] vs v3 [Q,frames]) +# =========================================================================== + + +def test_filter_real_code_frames_higgs_v2_golden(): + mod = _import("higgs_audio_v2") + fn = getattr(mod, "_filter_real_code_frames") + nrc = int(mod._NUM_REAL_CODES) + inp = torch.tensor([ + [0, 1, 2], + [nrc, 5, 6], # invalid (>= _NUM_REAL_CODES) + [7, 8, 9], + [-1, 0, 1], # invalid (< 0) + ]) + out = fn(inp) + assert out.tolist() == [[0, 1, 2], [7, 8, 9]] + # Empty input passes through. + empty = torch.empty(0, 3, dtype=torch.long) + assert torch.equal(fn(empty), empty) + + +def test_filter_real_code_frames_higgs_v3_golden(): + mod = _import("higgs_audio_v3") + fn = getattr(mod, "_filter_real_code_frames") + nrc = int(mod._NUM_REAL_CODES) + # Input is [Q, frames]: 4 frames x 3 codebooks -> transpose to [Q=3, frames=4]. + inp = torch.tensor([ + [0, 1, 2], # frame 0 (valid) + [nrc, 5, 6], # frame 1 (invalid: >= _NUM_REAL_CODES) + [7, 8, 9], # frame 2 (valid) + [-1, 0, 1], # frame 3 (invalid: < 0) + ]).t().contiguous() + out = fn(inp) # keeps frames 0 and 2 + assert out.shape == (3, 2) + assert out.tolist() == [[0, 7], [1, 8], [2, 9]] + empty = torch.empty(3, 0, dtype=torch.long) + assert torch.equal(fn(empty), empty) + + +# =========================================================================== +# _to_cpu_tensor (glm_tts) +# =========================================================================== + + +@pytest.mark.parametrize("inp,expected", [ + ("TENSOR", "TENSOR"), + ([], None), + ("LIST_TENSOR", "LIST_TENSOR"), + ([1, 2], None), + (None, None), + (5, None), +]) +def test_to_cpu_tensor_glm_tts_golden(inp, expected): + mod = _import("glm_tts") + fn = getattr(mod, "_to_cpu_tensor") + t = torch.tensor([1.0, 2.0]) + if inp == "TENSOR": + value = t + elif inp == "LIST_TENSOR": + value = [t] + else: + value = inp + result = fn(value) + if expected == "TENSOR": + assert torch.equal(result, t) + elif expected == "LIST_TENSOR": + assert torch.equal(result, t) + else: + assert result is expected + + +# =========================================================================== +# _to_token_id_list (dynin_omni; cosyvoice3 recursive variant) +# =========================================================================== + + +@pytest.mark.parametrize("inp,expected", [ + ("SCALAR_TENSOR", [5]), + ("VEC_TENSOR", [1, 2, 3]), + ("MAT_TENSOR", [1, 2]), # 2-D -> first row only + ("NESTED_LIST", [1, 2]), # nested list -> first nested row only + ([1, 2, 3], [1, 2, 3]), + (None, []), + (5, [5]), +]) +def test_to_token_id_list_dynin_omni_golden(inp, expected): + mod = _import("dynin_omni") + fn = getattr(mod, "_to_token_id_list") + if inp == "SCALAR_TENSOR": + value = torch.tensor(5) + elif inp == "VEC_TENSOR": + value = torch.tensor([1, 2, 3]) + elif inp == "MAT_TENSOR": + value = torch.tensor([[1, 2], [3, 4]]) + elif inp == "NESTED_LIST": + value = [[1, 2], [3, 4]] + else: + value = inp + assert fn(value) == expected + + +def test_to_token_id_list_cosyvoice3_recursive_golden(): + mod = _import("cosyvoice3") + fn = getattr(mod, "_to_token_id_list") + # cosyvoice3 recursively flattens ALL nesting (dynin takes first row only). + assert fn([[1, 2], [3, 4]]) == [1, 2, 3, 4] + assert fn(None) == [] + assert fn(torch.tensor([[1, 2], [3, 4]])) == [1, 2, 3, 4] + assert fn(5) == [5] + + +# =========================================================================== +# P8a: placeholder length helper (qwen3_omni._compute_talker_prompt_ids_length) +# =========================================================================== + + +def test_compute_talker_prompt_ids_length_golden(): + mod = _import("qwen3_omni") + fn = getattr(mod, "_compute_talker_prompt_ids_length") + # im_start markers at 0 and 6; user segment len 6, last assistant +9. + prompt = [151644, 872, 10, 11, 12, 13, 151644, 77091, 20, 21] + info = {"ids": {"all": prompt, "prompt": prompt}} + assert fn(info, device="cpu") == 15 + + +# =========================================================================== +# Phase 1b: _common.py must reproduce the golden-locked behaviour. +# =========================================================================== + + +def _common_import() -> Any: + try: + import vllm_omni.model_executor.stage_input_processors._common as c + return c + except Exception as exc: # pragma: no cover - env dependent + pytest.skip(f"_common not importable here: {type(exc).__name__}: {exc}") + + +def _shim_active() -> bool: + """True when the local-dev vllm stub (no real vllm) is in effect.""" + try: + import vllm_omni # noqa: F401 + return bool(getattr(vllm_omni, "_SHIM_ACTIVE", False)) + except Exception: + return False + + +def test_common_ensure_list_matches_golden(): + c = _common_import() + # Canonical default: list / tuple / None -> [] / scalar -> [x]. + assert c.ensure_list([1, 2]) == [1, 2] + assert c.ensure_list((1, 2)) == [1, 2] + assert c.ensure_list(None) == [] + assert c.ensure_list(_ConstantList([1, 2])) == [1, 2] + assert c.ensure_list(5) == [5] + # Tensor -> .tolist() preserves dims (ming_flash / step_audio2 semantics). + assert c.ensure_list(torch.tensor([[1, 2], [3, 4]])) == [[1, 2], [3, 4]] + # Legacy qwen3_omni: non-list returned unchanged (incl. None / tuple). + assert c.ensure_list_unchanged(None) is None + assert c.ensure_list_unchanged((1, 2)) == (1, 2) + assert c.ensure_list_unchanged([1, 2]) == [1, 2] + # Strict: non-iterable scalar raises (ming_flash / audex legacy). + with pytest.raises(TypeError): + c.ensure_list_strict(5) + + +def test_common_to_cpu_tensor_matches_golden(): + c = _common_import() + t = torch.tensor([1.0, 2.0]) + assert torch.equal(c.to_cpu_tensor(t), t) + assert c.to_cpu_tensor([]) is None + assert torch.equal(c.to_cpu_tensor([t]), t) + assert c.to_cpu_tensor([1, 2]) is None + assert c.to_cpu_tensor(None) is None + + +def test_common_to_token_id_list_matches_golden(): + c = _common_import() + # dynin semantics: first nested row only (non-recursive). + assert c.to_token_id_list(torch.tensor(5)) == [5] + assert c.to_token_id_list(torch.tensor([1, 2, 3])) == [1, 2, 3] + assert c.to_token_id_list(torch.tensor([[1, 2], [3, 4]])) == [1, 2] + assert c.to_token_id_list([[1, 2], [3, 4]]) == [1, 2] + assert c.to_token_id_list(None) == [] + assert c.to_token_id_list(5) == [5] + # cosyvoice3 semantics: recursive flatten. + assert c.to_token_id_list([[1, 2], [3, 4]], recursive=True) == [1, 2, 3, 4] + assert c.to_token_id_list( + torch.tensor([[1, 2], [3, 4]]), recursive=True + ) == [1, 2, 3, 4] + + +def test_common_revert_delay_pattern_matches_golden(): + c = _common_import() + inp = torch.tensor([ + [10, 1, 2, 3, 4], + [11, 12, 5, 6, 7], + [13, 14, 15, 8, 9], + ]) # [Q=3, T=5] + assert c.revert_delay_pattern(inp, allow_short=True).tolist() == [ + [10, 1, 2], [12, 5, 6], [15, 8, 9] + ] + # Lenient (v2): T < Q returns input unchanged. + short = torch.tensor([[1], [2]]) + assert torch.equal(c.revert_delay_pattern(short, allow_short=True), short) + # Strict (v3): T < Q raises; wrong codebook count raises. + with pytest.raises(ValueError): + c.revert_delay_pattern(short, allow_short=False) + with pytest.raises(ValueError): + c.revert_delay_pattern(torch.zeros(5, 6), expected_codebooks=3) + + +def test_common_filter_real_code_frames_matches_golden(): + c = _common_import() + nrc = 1024 + # frames-first (v2 layout). + frames_first = torch.tensor([ + [0, 1, 2], [nrc, 5, 6], [7, 8, 9], [-1, 0, 1], + ]) + assert c.filter_real_code_frames( + frames_first, num_real_codes=nrc, layout="frames_first" + ).tolist() == [[0, 1, 2], [7, 8, 9]] + # codebooks-first (v3 layout). + codebooks_first = torch.tensor([ + [0, 1, 2], [nrc, 5, 6], [7, 8, 9], [-1, 0, 1], + ]).t().contiguous() # [Q=3, frames=4] + assert c.filter_real_code_frames( + codebooks_first, num_real_codes=nrc, layout="codebooks_first" + ).tolist() == [[0, 7], [1, 8], [2, 9]] + + +def test_common_compute_placeholder_prompt_len_matches_golden(): + c = _common_import() + prompt = [151644, 872, 10, 11, 12, 13, 151644, 77091, 20, 21] + # full mode == qwen3_omni._compute_talker_prompt_ids_length golden (15). + assert c.compute_placeholder_prompt_len( + ids_or_prompt={"ids": {"all": prompt, "prompt": prompt}}, mode="full" + ) == 15 + # stage0_only mode: just the stage-0 input prompt length (prewarm). + assert c.compute_placeholder_prompt_len( + ids_or_prompt=prompt, mode="stage0_only" + ) == len(prompt) + + +def test_common_pack_placeholder_prompt(): + c = _common_import() + prompt = c.pack_placeholder_prompt(prompt_len=4, voice_metadata={"speaker": 1}) + empty_meta = c.pack_placeholder_prompt(prompt_len=1) + if _shim_active(): + # Under the local vllm stub, OmniTokensPrompt's real constructor is + # bypassed, so field access returns stub values. Field assertions run + # on CI where real vllm constructs the model properly. + pytest.skip("OmniTokensPrompt field construction needs real vllm (CI)") + # Real vllm (0.27) builds OmniTokensPrompt as a dict subclass (MRO + # ['OmniTokensPrompt', 'dict', 'object']): attribute access is not + # available, so read the packed fields via dict indexing. + assert prompt["prompt_token_ids"] == [0, 0, 0, 0] + assert prompt["additional_information"] == {"speaker": 1} + assert empty_meta["prompt_token_ids"] == [0] + # When additional_information is None the key may be absent from the + # dict, so use .get() for the empty-metadata case. + assert empty_meta.get("additional_information") is None + + +# Keep a reference to the case builder so linters don't flag dead code. +_EXTRA_ALIAS = _ensure_list_cases diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index 4074806e1d6..0d1e81f5484 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -12,6 +12,7 @@ from vllm_omni.diffusion.registry import DiffusionModelRegistry from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniPromptType +from vllm_omni.model_executor.stage_input_processors import _common from vllm_omni.outputs import OmniRequestOutput from vllm_omni.outputs.output_metadata import ( DiffusionMetadata, @@ -202,9 +203,8 @@ def _stream_metadata_from_diffusion_output(diffusion_output: DiffusionOutput) -> def _ensure_list(outputs: DiffusionPayloadValue) -> list[DiffusionPayloadValue]: - if isinstance(outputs, list): - return outputs - return [outputs] if outputs is not None else [] + """Delegate to the canonical shared helper (RFC #4872 Step 1c).""" + return _common.ensure_list(outputs) def _primary_payload(postprocess_output: DiffusionPostprocessOutput) -> DiffusionPayloadValue | None: diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 800bc0ad2ed..76bfeecb49a 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -2595,20 +2595,22 @@ async def _forward_to_next_stage_unguarded( if next_client is not None and next_client.custom_process_input_func is not None: _t_ar2d = _time.perf_counter() _fn = next_client.custom_process_input_func - _extra_kwargs: dict[str, Any] = {} - # TODO: replace signature probe with explicit kwarg contract. - try: - import inspect as _inspect + # RFC #4872 Phase 2 (P1): dispatch through the shared contract + # layer, replacing the ad-hoc `sampling_params` name probe. + from vllm_omni.model_executor.stage_input_processors._dispatch import ( + OrchestratorInputContext, + invoke_orchestrator_processor, + ) - if "sampling_params" in _inspect.signature(_fn).parameters: - _extra_kwargs["sampling_params"] = params - except (TypeError, ValueError): - pass - diffusion_prompt = _fn( + _ctx = OrchestratorInputContext( + prompt=req_state.prompt, + requires_multimodal_data=requires_multimodal_data, + sampling_params=params, + ) + diffusion_prompt = invoke_orchestrator_processor( + _fn, diffusion_source_outputs, - req_state.prompt, - requires_multimodal_data, - **_extra_kwargs, + _ctx, ) _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000 req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index a8ab3678e9b..d2dc0ded71b 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -9,7 +9,6 @@ from __future__ import annotations -import inspect import os import socket from typing import TYPE_CHECKING, Any @@ -464,10 +463,29 @@ def process_engine_inputs( and the original prompt. """ if self.custom_process_input_func is not None: - return self._call_custom_process_input(source_outputs, prompt, streaming_context) + if self.custom_process_input_func is not None: + # RFC #4872 Phase 2 (P1): dispatch through the shared contract layer, + # replacing the ad-hoc arity probe (`len(signature.parameters) >= 4`). + from vllm_omni.model_executor.stage_input_processors._dispatch import ( + OrchestratorInputContext, + invoke_orchestrator_processor, + ) + + ctx = OrchestratorInputContext( + prompt=prompt, + requires_multimodal_data=self.requires_multimodal_data, + streaming_context=streaming_context, + ) + return invoke_orchestrator_processor( + self.custom_process_input_func, + source_outputs, + ctx, + ) if not self.engine_input_source: raise ValueError(f"engine_input_source empty for stage {self.stage_id}") + # The default path below is itself a C0-shape processor; it is kept out of + # the contract dispatch to preserve its exact fallback semantics unchanged. return _default_process_engine_inputs(source_outputs, prompt, self.requires_multimodal_data) def _call_custom_process_input( diff --git a/vllm_omni/model_executor/stage_input_processors/__init__.py b/vllm_omni/model_executor/stage_input_processors/__init__.py index e69de29bb2d..562047e98d1 100644 --- a/vllm_omni/model_executor/stage_input_processors/__init__.py +++ b/vllm_omni/model_executor/stage_input_processors/__init__.py @@ -0,0 +1,42 @@ +"""Stage input processors: per-model builders for inter-stage transitions. + +Module-naming convention (RFC #4872): + - ``*_full_payload`` : producer-side (worker) connector payload builder + (``FullPayloadProducer``; kwarg ``pooling_output`` load-bearing). + - ``*_async_chunk`` : producer-side (worker) streaming chunk builder + (``AsyncChunkProducer``; kwarg ``multimodal_output`` load-bearing). + - ``*_token_only`` : consumer-side (orchestrator) length-only placeholder + builder for the non-async-chunk path. + - no suffix (legacy) : consumer-side (orchestrator) sync builder; may be a + placeholder builder or a diffusion input builder. + +Orchestrator-facing roles are normalized through +``wrap_orchestrator_processor`` / ``invoke_orchestrator_processor``: + - placeholder : upstream outputs -> next-stage token prompts + (``PlaceholderPromptBuilder``). + - diffusion : upstream outputs -> diffusion payload(s) + (``DiffusionInputBuilder``). + +This package only re-exports the dispatch contract; it imports no other +processor modules, so importing it never pulls in model-specific code. +""" + +from vllm_omni.model_executor.stage_input_processors._dispatch import ( + AsyncChunkProducer, + DiffusionInputBuilder, + FullPayloadProducer, + OrchestratorInputContext, + PlaceholderPromptBuilder, + invoke_orchestrator_processor, + wrap_orchestrator_processor, +) + +__all__ = [ + "OrchestratorInputContext", + "PlaceholderPromptBuilder", + "DiffusionInputBuilder", + "FullPayloadProducer", + "AsyncChunkProducer", + "invoke_orchestrator_processor", + "wrap_orchestrator_processor", +] diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py new file mode 100644 index 00000000000..30bb02e1d8e --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -0,0 +1,427 @@ +"""Canonical, shared helpers for ``stage_input_processors``. + +RFC #4872 Proposal 3 / P8a. These are consolidated implementations of +helpers that were previously duplicated across model modules with subtle, +behavioral differences (see ``tests/model_executor/stage_input_processors/ +test_common_helpers_golden.py``, which locks the observed per-module +behaviour). + +**Consolidation rule:** where two legacy variants disagreed, we do NOT +silently pick one. The default implementation follows the most-complete +semantics; divergent legacy behaviour is preserved through explicit named +variants (e.g. ``ensure_list_unchanged``) or parameters (e.g. +``filter_real_code_frames(..., layout=...)``). A module may only switch to +``_common`` after its legacy behaviour is golden-locked and the matching +variant is used. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from vllm_omni.inputs.data import OmniTokensPrompt + + +# =========================================================================== +# ensure_list +# =========================================================================== + +def _unwrap_constant_list(x: Any) -> Any: + """Unwrap a ``ConstantList``-like object exposing ``_x``.""" + if hasattr(x, "_x"): + return x._x + return x + + +def ensure_list(x: Any) -> list[Any]: + """Canonical ``ensure_list``. + + Semantics (the union of legacy variants): + - ``ConstantList``-like objects (``_x``) are unwrapped. + - ``list`` / ``tuple`` are converted to a fresh ``list``. + - ``torch.Tensor`` is converted via ``.tolist()`` (all dims flattened). + - ``None`` becomes ``[]``. + - Any other iterable is converted via ``list(x)``; a non-iterable scalar + becomes ``[x]``. + """ + x = _unwrap_constant_list(x) + if isinstance(x, torch.Tensor): + return x.detach().cpu().tolist() + if x is None: + return [] + if isinstance(x, (list, tuple)): + return list(x) + try: + return list(x) + except TypeError: + return [x] + + +def ensure_list_or_empty(x: Any) -> list[Any]: + """Explicit alias: always returns a list; ``None`` -> ``[]``. + + Matches ``audex`` / ``cosyvoice3`` ``_ensure_list`` semantics. + """ + return ensure_list(x) + + +def ensure_list_unchanged(x: Any) -> Any: + """Legacy ``qwen3_omni._ensure_list`` semantics. + + ``ConstantList`` is unwrapped; a non-``list`` value (including ``None``, + a ``tuple``, or a ``torch.Tensor``) is returned **unchanged**. + """ + x = _unwrap_constant_list(x) + if not isinstance(x, list): + return x + return list(x) + + +def ensure_list_strict(x: Any) -> list[Any]: + """Strict ``ensure_list``: no ``None`` special-case, non-iterables raise. + + Matches ``ming_flash_omni._ensure_list``: ``None`` and non-iterable + scalars both fall through to ``list(x)`` and raise ``TypeError``; a + ``torch.Tensor`` is converted via ``.tolist()``. + """ + x = _unwrap_constant_list(x) + if isinstance(x, torch.Tensor): + return x.detach().cpu().tolist() + if isinstance(x, (list, tuple)): + return list(x) + return list(x) + + +def ensure_list_flatten(x: Any) -> list[Any]: + """``ensure_list`` that **flattens** tensors with ``reshape(-1).tolist()``. + + Matches ``audex._ensure_list``: ``None`` -> ``[]``; a ``torch.Tensor`` is + flattened to a 1-D list; a non-iterable scalar raises ``TypeError``. + """ + x = _unwrap_constant_list(x) + if x is None: + return [] + if isinstance(x, torch.Tensor): + return x.detach().cpu().reshape(-1).tolist() + if isinstance(x, (list, tuple)): + return list(x) + return list(x) + + +def ensure_list_preserve_none(x: Any) -> list[Any]: + """``ensure_list`` that **preserves** ``None`` as ``[None]``. + + Matches ``step_audio2._ensure_list``: + - ``ConstantList``-like objects (``_x``) are unwrapped. + - ``torch.Tensor`` is converted via ``.tolist()`` (dimensions preserved). + - ``list`` / ``tuple`` are converted to a fresh ``list``. + - Any other **iterable** is converted via ``list(x)``. + - A non-iterable value (including ``None`` and a scalar like ``5``) is + wrapped as ``[x]`` (so ``None`` -> ``[None]``, ``5`` -> ``[5]``). + """ + if hasattr(x, "_x"): + return list(x._x) + if isinstance(x, torch.Tensor): + return x.detach().cpu().tolist() + if isinstance(x, (list, tuple)): + return list(x) + return list(x) if hasattr(x, "__iter__") else [x] + + +# =========================================================================== +# to_cpu_tensor +# =========================================================================== + +def to_cpu_tensor(value: Any) -> torch.Tensor | None: + """Convert a value to a CPU tensor when possible; else ``None``. + + Matches ``glm_tts._to_cpu_tensor``: + - a list is unwrapped to its first element first; + - empty list -> ``None``; + - only ``torch.Tensor`` inputs yield a result (``detach().cpu()``). + """ + if isinstance(value, list): + if not value: + return None + value = value[0] + if isinstance(value, torch.Tensor): + return value.detach().cpu() + return None + + +# =========================================================================== +# to_token_id_list +# =========================================================================== + +def to_token_id_list(value: Any, *, recursive: bool = False) -> list[int]: + """Convert a token-ish value into a flat ``list[int]``. + + ``recursive=False`` (default) matches ``dynin_omni._to_token_id_list``: + for a 2-D tensor or a nested list only the **first** row / element is + used. ``recursive=True`` matches ``cosyvoice3._to_token_id_list``: all + nesting is flattened. + """ + if isinstance(value, torch.Tensor): + value = value.detach().to("cpu") + if value.ndim == 0: + return [int(value.item())] + if value.ndim > 1: + if recursive: + value = value.reshape(-1) + else: + value = value[0] + return [int(x) for x in value.tolist()] + if isinstance(value, list): + if not value: + return [] + if isinstance(value[0], list): + if recursive: + out: list[int] = [] + for item in value: + out.extend(to_token_id_list(item, recursive=True)) + return out + return [int(x) for x in value[0]] + return [int(x) for x in value] + if value is None: + return [] + return [int(value)] + + +# =========================================================================== +# revert_delay_pattern (higgs_audio_v2 lenient vs higgs_audio_v3 strict) +# =========================================================================== + +def revert_delay_pattern( + audio_codes_qt: torch.Tensor, + *, + expected_codebooks: int | None = None, + allow_short: bool = False, +) -> torch.Tensor: + """Reverse the MusicGen-style delay pattern. + + Input ``[Q, T]`` (``T = seq_len + Q - 1``) -> output ``[Q, seq_len]``. + + ``allow_short=True`` matches ``higgs_audio_v2``: when ``T < Q`` the input + is returned **unchanged**. ``allow_short=False`` (default, matches + ``higgs_audio_v3``) raises on ``T < Q``. When ``expected_codebooks`` is + given, a row-count mismatch raises ``ValueError`` (v3 checks ``Q == + _NUM_CODEBOOKS``; v2 does not). + """ + if audio_codes_qt.ndim != 2: + raise ValueError( + f"_revert_delay_pattern expects [Q, T] input; " + f"got {tuple(audio_codes_qt.shape)}" + ) + q, t = audio_codes_qt.shape + if expected_codebooks is not None and q != expected_codebooks: + raise ValueError( + f"Expected exactly {expected_codebooks} codebook rows, got {q}. " + f"Input shape: [{q}, {t}]" + ) + if t < q: + if allow_short: + return audio_codes_qt + raise ValueError(f"Not enough frames to revert delay pattern: T={t} < Q={q}") + seq_len = t - q + 1 + out_l = [audio_codes_qt[i : i + 1, i : seq_len + i] for i in range(q)] + return torch.cat(out_l, dim=0) + + +# =========================================================================== +# filter_real_code_frames (higgs v2 [frames,Q] vs v3 [Q,frames]) +# =========================================================================== + +def filter_real_code_frames( + audio_codes: torch.Tensor, + *, + num_real_codes: int, + layout: str = "frames_first", +) -> torch.Tensor: + """Keep only frames whose codebook values are all in ``[0, num_real_codes)``. + + ``layout="frames_first"`` matches ``higgs_audio_v2``: input is + ``[num_frames, num_codebooks]`` and the same layout is returned. + ``layout="codebooks_first"`` matches ``higgs_audio_v3``: input is + ``[num_codebooks, num_frames]``; frames are filtered and the layout is + restored (``[num_codebooks, kept_frames]``). + """ + if audio_codes.numel() == 0: + return audio_codes + if layout == "frames_first": + if audio_codes.ndim != 2: + raise ValueError( + f"expected [num_frames, num_codebooks] audio_codes; " + f"got shape {tuple(audio_codes.shape)}" + ) + valid = (audio_codes >= 0).all(dim=1) & ( + audio_codes < num_real_codes + ).all(dim=1) + return audio_codes[valid] + if layout == "codebooks_first": + frames = audio_codes.t() + valid = (frames >= 0).all(dim=1) & (frames < num_real_codes).all(dim=1) + return frames[valid].t().contiguous() + raise ValueError(f"unknown layout: {layout!r}") + + +# =========================================================================== +# extract_last_codec_frame +# =========================================================================== + +def extract_last_codec_frame( + payload: Any, + *, + key_path: tuple[str, ...] = ("codes", "audio"), + validate: str | None = None, + to_cpu: bool = False, + to_long: bool = True, +) -> torch.Tensor | None: + """Extract the last inter-stage codec frame from a payload. + + Parameterized canonical implementation covering the four legacy + ``_extract_last_frame`` variants: + + - ``key_path``: ``("codes", "audio")`` (higgs_v2 / qwen3_tts / voxtral) + vs ``("audio_codes",)`` (fish_speech top-level). + - ``validate``: + - ``None``: only empty checks (higgs_v2 / voxtral); + - ``"any"``: drop a frame whose values are all zero (qwen3_tts); + - ``"valid_mask"``: consult ``audio_code_valid`` then fall back to + ``any`` (fish_speech). + - ``to_cpu`` / ``to_long``: fish_speech moves to CPU and casts to long; + higgs_v2 / qwen3_tts cast to long without moving to CPU; voxtral does + neither (plain ``flatten``). + + The returned tensor is always 1-D. + """ + if not isinstance(payload, dict): + return None + audio_codes: Any = payload + for key in key_path: + audio_codes = audio_codes.get(key) if isinstance(audio_codes, dict) else None + if audio_codes is None: + return None + if isinstance(audio_codes, list): + if not audio_codes: + return None + audio_codes = audio_codes[-1] + if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: + return None + if audio_codes.ndim == 2: + frame: torch.Tensor = audio_codes[-1] + elif audio_codes.ndim == 1: + frame = audio_codes + else: + raise ValueError( + f"unexpected audio_codes shape: {tuple(audio_codes.shape)}" + ) + if frame.numel() == 0: + return None + if validate in ("any", "valid_mask"): + if validate == "valid_mask": + valid = payload.get("audio_code_valid") + if isinstance(valid, torch.Tensor) and valid.numel() > 0: + is_valid = bool(valid.reshape(-1)[-1].item()) + elif valid is not None: + is_valid = bool(valid) + else: + is_valid = bool(frame.any().item()) + else: + is_valid = bool(frame.any().item()) + if not is_valid: + return None + frame = frame.flatten() + if to_cpu: + frame = frame.detach().cpu() + if to_long: + frame = frame.to(torch.long) + return frame + + +# =========================================================================== +# P8a: placeholder prompt length + packing +# =========================================================================== + +# Qwen chat-template sentinel ids (single source within this module; the +# codec token ids themselves are centralized separately in P6). +_IM_START_TOKEN_ID = 151644 +_SYSTEM_TOKEN_ID = 8948 +_USER_TOKEN_ID = 872 +_ASSISTANT_TOKEN_ID = 77091 +_ASSISTANT_TAIL_LEN = 9 # 3 + 4 + 1 + 1 + + +def compute_placeholder_prompt_len( + *, + ids_or_prompt: Any, + mode: str = "full", + device: torch.device | str = "cpu", +) -> int: + """Compute the downstream placeholder prompt length. + + ``mode="full"`` replicates the Qwen chat-template scan of + ``qwen3_omni._compute_talker_prompt_ids_length``: the input is an + ``OmniPayload``-like dict with ``ids: {"all": [...], "prompt": [...]}``, + and the length is ``sum(user segments) + 9`` (the assistant tail). + + ``mode="stage0_only"`` is the async-chunk **prewarm** estimate: there is + no upstream output yet, so the length is simply the stage-0 input prompt + length (``ids_or_prompt`` is a sequence of token ids). + """ + if mode == "stage0_only": + return len(list(ids_or_prompt)) + if mode != "full": + raise ValueError(f"unknown mode: {mode!r}") + + ids = (ids_or_prompt or {}).get("ids", {}) + thinker_sequences = torch.tensor( + ids["all"], dtype=torch.long, device=device + ).unsqueeze(0) # [1, T] + input_ids = torch.tensor( + ids["prompt"], dtype=torch.long, device=device + ).unsqueeze(0) # [1, T] + + im_start_indexes = torch.cat( + [ + torch.nonzero(input_ids[0] == _IM_START_TOKEN_ID).squeeze(1), + torch.tensor( + [thinker_sequences.shape[-1]], + device=input_ids.device, + dtype=input_ids.dtype, + ), + ], + dim=0, + ) + + sum_user_len = 0 + assistant_len = 0 + for i in range(len(im_start_indexes) - 1): + s = int(im_start_indexes[i].item()) + e = int(im_start_indexes[i + 1].item()) + role = int(input_ids[0, s + 1].item()) + if role == _SYSTEM_TOKEN_ID: + continue + if role == _USER_TOKEN_ID: + sum_user_len += e - s + elif role == _ASSISTANT_TOKEN_ID and i == len(im_start_indexes) - 2: + assistant_len += _ASSISTANT_TAIL_LEN + return sum_user_len + assistant_len + + +def pack_placeholder_prompt( + *, + prompt_len: int, + voice_metadata: dict[str, Any] | None = None, +) -> OmniTokensPrompt: + """Pack a KV-slot placeholder prompt: ``[0] * prompt_len`` + voice meta. + + Shared by both the sync forward path (``*_token_only``) and the + async-chunk prewarm path (Proposal 8b). Bulk conditioning arrives + separately via the connector, so only the length matters here. + """ + return OmniTokensPrompt( + prompt_token_ids=[0] * max(1, int(prompt_len)), + additional_information=(voice_metadata if voice_metadata else None), + ) diff --git a/vllm_omni/model_executor/stage_input_processors/_dispatch.py b/vllm_omni/model_executor/stage_input_processors/_dispatch.py new file mode 100644 index 00000000000..ace4e02ad47 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/_dispatch.py @@ -0,0 +1,304 @@ +"""RFC #4872 Phase 2 (P1): orchestrator input-dispatch contract layer. + +This module defines the **consumer-side (orchestrator / stage-client) dispatch +contract** for ``stage_input_processors``. It replaces the two inconsistent +``inspect.signature`` probes in the orchestrator runtime with one shared +contract: + +* ``stage_engine_core_client.process_engine_inputs`` used an **arity** probe + (``len(signature.parameters) >= 4``) to pick a 3-arg vs 4-arg positional + call. +* ``orchestrator._forward_to_next_stage`` (diffusion branch) used a **name** + probe (``"sampling_params" in signature.parameters``) to decide whether to + inject the diffusion-stage sampling params as a keyword. + +Both are collapsed into a single normalization layer: + +* ``OrchestratorInputContext`` — the fixed transition context the runtime + passes to every consumer-side builder. +* ``PlaceholderPromptBuilder`` / ``DiffusionInputBuilder`` — the two + orchestrator-facing callable roles (C1-compatible ``(source_outputs, ctx)``). +* ``wrap_orchestrator_processor`` / ``invoke_orchestrator_processor`` — + processors that already speak the ``ctx`` contract are passed through + unchanged; legacy positional shapes (C0/C2/C3/C4 below) are adapted with a + ``DeprecationWarning``. + +**Producer roles are intentionally out of scope here.** ``FullPayloadProducer`` +(``*, transfer_manager, pooling_output, request, is_finished=...``) and +``AsyncChunkProducer`` (``*, transfer_manager, multimodal_output, request, +is_finished=False``) run inside workers and never receive an +``OrchestratorInputContext``; their ``pooling_output`` / ``multimodal_output`` +keyword names are load-bearing parts of the producer contract and are kept +as-is (RFC P1/P2 keep the producer kwargs contract unchanged). + +Legacy positional shapes normalized here (RFC P1 migration matrix): + +* C0 3-arg: ``(source_outputs, prompt, requires_multimodal_data)``. +* C2 placeholder: ``(source_outputs, prompt, requires_multimodal_data, + streaming_context=None)`` (the ``_streaming_context`` variant used by + ``forced_aligner.code2wav2aligner`` / ``minicpmo_4_5_omni.llm2tts`` is also + recognized). +* C3 diffusion: ``(source_outputs, prompt, requires_multimodal_data, + sampling_params=None)``. +* C4 legacy multi-source: ``(stage_list, engine_input_source, ...)`` — only + ``moss_tts.talker2codec`` remains today; normalized by + ``_adapt_moss_processor``, which binds ``stage_list=source_outputs`` and + ``engine_input_source=ctx.prompt`` and drops the trailing + ``prompt``/``requires_multimodal_data`` arguments (unused by that + implementation). +""" + +from __future__ import annotations + +import inspect +import warnings +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + +from vllm_omni.inputs.data import OmniTokensPrompt + +__all__ = [ + "OrchestratorInputContext", + "PlaceholderPromptBuilder", + "DiffusionInputBuilder", + "FullPayloadProducer", + "AsyncChunkProducer", + "wrap_orchestrator_processor", + "invoke_orchestrator_processor", +] + + +@dataclass(frozen=True) +class OrchestratorInputContext: + """Fixed transition context passed to orchestrator-facing processors. + + There is deliberately **no** ``model_config`` field: a processor that needs + a model config reads it through the upstream stage closure, never through + this context. + """ + + prompt: Any | None = None + requires_multimodal_data: bool = False + streaming_context: Any | None = None + sampling_params: Any | None = None + + +@runtime_checkable +class PlaceholderPromptBuilder(Protocol): + """C2-style sync builder: upstream outputs -> next-stage token prompts.""" + + def __call__( + self, + source_outputs: list[Any], + ctx: OrchestratorInputContext, + ) -> list[OmniTokensPrompt]: ... + + +@runtime_checkable +class DiffusionInputBuilder(Protocol): + """Diffusion-stage input builder: upstream outputs -> diffusion payload(s).""" + + def __call__( + self, + source_outputs: list[Any], + ctx: OrchestratorInputContext, + ) -> dict | list[dict] | None: ... + + +@runtime_checkable +class FullPayloadProducer(Protocol): + """Producer-side (worker) full-payload builder. + + ``pooling_output`` is a **load-bearing** keyword name: it is the connector + data-plane contract and must not be renamed or made positional. These + builders never receive an ``OrchestratorInputContext``. + """ + + def __call__( + self, + *, + transfer_manager: Any, + pooling_output: Any, + request: Any, + is_finished: bool = ..., + ) -> Any: ... + + +@runtime_checkable +class AsyncChunkProducer(Protocol): + """Producer-side (worker) async-chunk builder. + + ``multimodal_output`` is a **load-bearing** keyword name: it is the + connector data-plane contract and must not be renamed or made positional. + These builders never receive an ``OrchestratorInputContext``. + """ + + def __call__( + self, + *, + transfer_manager: Any, + multimodal_output: Any, + request: Any, + is_finished: bool = False, + ) -> Any: ... + + +# --------------------------------------------------------------------------- +# Signature inspection helpers +# --------------------------------------------------------------------------- + + +def _is_orchestrator_context_annotation(annotation: Any) -> bool: + """Whether ``annotation`` names ``OrchestratorInputContext``. + + Handles the class object, plain string annotations (``from __future__ + import annotations``) and fully-qualified dotted strings. + """ + if annotation is inspect.Parameter.empty: + return False + if annotation is OrchestratorInputContext: + return True + if isinstance(annotation, str): + text = annotation.strip().strip("'\"") + return text == "OrchestratorInputContext" or text.endswith("OrchestratorInputContext") + return False + + +def _accepts_orchestrator_input_context(fn: Any) -> bool: + """Whether ``fn`` already speaks the new ``ctx`` contract. + + True when the signature contains a parameter named ``ctx``, or any + parameter whose type annotation is ``OrchestratorInputContext``. + """ + try: + signature = inspect.signature(fn) + except (TypeError, ValueError): + return False + for name, param in signature.parameters.items(): + if name == "ctx": + return True + if _is_orchestrator_context_annotation(param.annotation): + return True + return False + + +def _positional_parameter_names(fn: Any) -> list[str] | None: + """Names of positional (or positional-or-keyword) parameters of ``fn``. + + Returns ``None`` when the signature cannot be inspected (e.g. builtins). + """ + try: + signature = inspect.signature(fn) + except (TypeError, ValueError): + return None + names: list[str] = [] + for param in signature.parameters.values(): + if param.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD): + names.append(param.name) + elif param.kind == inspect.Parameter.VAR_POSITIONAL: + names.append("*args") + return names + + +def _warn_legacy_contract(fn: Any) -> None: + """Emit a ``DeprecationWarning`` for a legacy positional processor shape.""" + name = getattr(fn, "__qualname__", None) or getattr(fn, "__name__", None) or repr(fn) + warnings.warn( + f"stage-input processor {name!r} uses the legacy positional contract; " + "migrate it to the OrchestratorInputContext contract " + "(RFC #4872 Phase 2).", + DeprecationWarning, + stacklevel=2, + ) + + +# --------------------------------------------------------------------------- +# Legacy adapters +# --------------------------------------------------------------------------- + + +def _adapt_moss_processor(fn: Any) -> Any: + """C4 legacy multi-source adapter (e.g. ``moss_tts.talker2codec``). + + The legacy multi-source sync processors take + ``(stage_list, engine_input_source, prompt=None, requires_multimodal_data=False)``. + The orchestrator contract dispatches ``(source_outputs, ctx)``; we bind + ``stage_list=source_outputs`` and ``engine_input_source=ctx.prompt`` and + drop the trailing ``prompt``/``requires_multimodal_data`` arguments (the + ``moss_tts.talker2codec`` implementation does not use them, so this is + behaviour-preserving w.r.t. the previous 4-arg positional call). + """ + + def _adapted(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: + return fn(source_outputs, ctx.prompt) + + _warn_legacy_contract(fn) + return _adapted + + +def wrap_orchestrator_processor(fn: Any) -> PlaceholderPromptBuilder | DiffusionInputBuilder: + """Return a C1-compatible callable ``(source_outputs, ctx)`` for ``fn``. + + - If ``fn`` already accepts ``ctx`` (``_accepts_orchestrator_input_context``), + it is returned unchanged. + - Otherwise the legacy positional shape is adapted: + * C4 ``(stage_list, engine_input_source, ...)`` -> ``_adapt_moss_processor``; + * C3 ``(source_outputs, prompt, requires_multimodal_data, sampling_params=...)`` + forwards ``ctx.sampling_params`` when the parameter exists; + * C2 ``(source_outputs, prompt, requires_multimodal_data, + streaming_context=...)`` (or ``_streaming_context``) forwards + ``ctx.streaming_context``; + * C0 3-arg ``(source_outputs, prompt, requires_multimodal_data)``. + Every legacy shape emits a ``DeprecationWarning``. + """ + if not callable(fn): + raise TypeError(f"stage-input processor must be callable, got {fn!r}") + + if _accepts_orchestrator_input_context(fn): + return fn # type: ignore[return-value] + + names = _positional_parameter_names(fn) + + # Uninspectable callable (e.g. builtin): fall back to the most conservative + # 3-arg C0 shape; the runtime previously treated it via the <4 arity path. + if names is None: + names = ["source_outputs", "prompt", "requires_multimodal_data"] + + # C4 legacy multi-source (e.g. moss_tts.talker2codec). + if len(names) >= 2 and names[0] == "stage_list" and names[1] == "engine_input_source": + return _adapt_moss_processor(fn) # type: ignore[return-value] + + # C3 diffusion: forwarding sampling_params (diffusion-stage params). + if "sampling_params" in names: + + def _c3(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: + return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data, + sampling_params=ctx.sampling_params) + + _warn_legacy_contract(fn) + return _c3 # type: ignore[return-value] + + # C2 placeholder: forwarding the streaming context. + if "streaming_context" in names or "_streaming_context" in names: + + def _c2(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: + return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data, ctx.streaming_context) + + _warn_legacy_contract(fn) + return _c2 # type: ignore[return-value] + + # C0 3-arg legacy. + def _c0(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: + return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data) + + _warn_legacy_contract(fn) + return _c0 # type: ignore[return-value] + + +def invoke_orchestrator_processor( + fn: Any, + source_outputs: list[Any], + ctx: OrchestratorInputContext, +) -> Any: + """Invoke an orchestrator-facing processor under the fixed ``(source_outputs, ctx)`` contract.""" + return wrap_orchestrator_processor(fn)(source_outputs, ctx) diff --git a/vllm_omni/model_executor/stage_input_processors/audex.py b/vllm_omni/model_executor/stage_input_processors/audex.py index 011ec904306..2f3589dd514 100644 --- a/vllm_omni/model_executor/stage_input_processors/audex.py +++ b/vllm_omni/model_executor/stage_input_processors/audex.py @@ -22,6 +22,7 @@ from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt +from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) @@ -44,11 +45,8 @@ def _ensure_list(value: Any) -> list[int]: - if value is None: - return [] - if isinstance(value, torch.Tensor): - return value.reshape(-1).tolist() - return list(value) + """RFC #4872 P3: delegate to the canonical ``_common.ensure_list_flatten``.""" + return _common.ensure_list_flatten(value) def _connector_extra(transfer_manager: Any) -> dict[str, Any]: diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index e8123fe8751..f34f6da86f1 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py @@ -18,6 +18,7 @@ from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.models.cosyvoice3.utils import unpad_prompt_conditioning +from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) @@ -37,35 +38,13 @@ def _build_prompt_embed_struct(prompt_payload: dict[str, Any]) -> EmbeddingsStru def _ensure_list(x: Any) -> list[Any]: - if hasattr(x, "_x"): - return list(x._x) - if isinstance(x, list): - return list(x) - if isinstance(x, tuple): - return list(x) - if x is None: - return [] - try: - return list(x) - except TypeError: - return [x] + """Delegate to the canonical shared helper (RFC #4872 Step 1c).""" + return _common.ensure_list(x) def _to_token_id_list(value: Any) -> list[int]: - if value is None: - return [] - if isinstance(value, torch.Tensor): - value = value.detach().to("cpu").reshape(-1).tolist() - token_ids: list[int] = [] - for item in _ensure_list(value): - if isinstance(item, torch.Tensor): - token_ids.extend(_to_token_id_list(item)) - continue - if isinstance(item, (list, tuple)): - token_ids.extend(_to_token_id_list(item)) - continue - token_ids.append(int(item)) - return token_ids + """Recursive token-id flattening (cosyvoice3 golden semantics, Step 1c).""" + return _common.to_token_id_list(value, recursive=True) def _strip_prompt_prefix(output_ids: list[Any], prefix_ids: list[Any]) -> list[Any]: @@ -84,13 +63,8 @@ def _prompt_speech_token_ids(multi_modal_data: dict[str, Any]) -> list[int]: def _to_cpu_tensor(x: Any) -> torch.Tensor | None: - if isinstance(x, list): - if not x: - return None - x = x[0] - if isinstance(x, torch.Tensor): - return x.detach().cpu() - return None + """Delegate to the canonical shared helper (RFC #4872 Step 1c).""" + return _common.to_cpu_tensor(x) def talker2code2wav_async_chunk( diff --git a/vllm_omni/model_executor/stage_input_processors/glm_tts.py b/vllm_omni/model_executor/stage_input_processors/glm_tts.py index af037aee77c..77d73ff9854 100644 --- a/vllm_omni/model_executor/stage_input_processors/glm_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/glm_tts.py @@ -18,6 +18,7 @@ OmniPayloadStruct, ) from vllm_omni.engine.serialization import deserialize_additional_information +from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) @@ -170,15 +171,11 @@ def _extract_last_speech_token(pooling_output: dict[str, Any]) -> int | None: def _to_cpu_tensor(value: Any) -> torch.Tensor | None: - """Convert value to CPU tensor if possible.""" - if isinstance(value, torch.Tensor): - return value.detach().cpu() - if isinstance(value, list): - if not value: - return None - if isinstance(value[0], torch.Tensor): - return value[0].detach().cpu() - return None + """Convert value to CPU tensor if possible. + + RFC #4872 P3: delegate to the canonical ``_common.to_cpu_tensor``. + """ + return _common.to_cpu_tensor(value) # --------------------------------------------------------------------------- diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py index 53f60ba4ba4..668c2e80930 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Stage-input processor for higgs-audio v2: Talker -> Code2Wav. Mirrors the qwen3_tts contract: @@ -29,6 +29,7 @@ OmniPayload, OmniPayloadStruct, ) +from vllm_omni.model_executor.stage_input_processors import _common __all__ = ["talker2code2wav", "talker2code2wav_async_chunk"] @@ -48,28 +49,24 @@ _NUM_REAL_CODES = _AUDIO_STREAM_BOS_ID # codes in [0, 1023] are real -def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: - """Reverse the delay-pattern shift, mirroring upstream's - ``boson_multimodal.model.higgs_audio.utils.revert_delay_pattern``. +def _filter_real_code_frames(audio_codes: torch.Tensor) -> torch.Tensor: + """Keep only frames whose codes are entirely in [0, 1023]. + + RFC #4872 P3: delegate to the canonical ``_common.filter_real_code_frames`` + with the higgs-v2 ``[num_frames, num_codebooks]`` layout. + """ + return _common.filter_real_code_frames( + audio_codes, num_real_codes=_NUM_REAL_CODES, layout="frames_first" + ) + - Input shape: ``[num_codebooks, seq_len + num_codebooks - 1]``. - Output shape: ``[num_codebooks, seq_len]``. +def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: + """Reverse the delay-pattern shift (higgs-v2 lenient variant). - For each codebook ``i``, slice rows ``[i:i+1]`` from columns - ``[i : seq_len + i]`` to remove ``i`` leading BOS pads and ``Q-1-i`` - trailing PAD entries. + RFC #4872 P3: delegate to the canonical ``_common.revert_delay_pattern`` + with ``allow_short=True`` (T < Q returns the input unchanged). """ - if audio_codes_qt.ndim != 2: - raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}") - q, t = audio_codes_qt.shape - if t < q: - # Not enough frames to revert delay pattern; return as-is. - return audio_codes_qt - seq_len = t - q + 1 - out_l = [] - for i in range(q): - out_l.append(audio_codes_qt[i : i + 1, i : seq_len + i]) - return torch.cat(out_l, dim=0) + return _common.revert_delay_pattern(audio_codes_qt, allow_short=True) def talker2code2wav( diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py index 8a7c66121aa..86eac44b652 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py @@ -34,6 +34,7 @@ OmniPayloadStruct, ) from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors import _common __all__ = ["talker2code2wav", "talker2code2wav_async_chunk"] @@ -60,40 +61,27 @@ def _empty_code2wav_prompt() -> Any: def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: - """Reverse the MusicGen-style delay pattern. + """Reverse the MusicGen-style delay pattern (higgs-v3 strict variant). - Input shape: [num_codebooks, seq_len + num_codebooks - 1]. - Output shape: [num_codebooks, seq_len]. - - For each codebook i, extract delayed[i, i : i + seq_len] to remove - the i leading BOC pads and Q-1-i trailing EOC entries. + RFC #4872 P3: delegate to the canonical ``_common.revert_delay_pattern`` + with the strict codebook-count and short-input checks. """ - if audio_codes_qt.ndim != 2: - raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}") - q, t = audio_codes_qt.shape - if q != _NUM_CODEBOOKS: - raise ValueError(f"Expected exactly {_NUM_CODEBOOKS} codebook rows, got {q}. Input shape: [{q}, {t}]") - if t < q: - raise ValueError(f"Not enough frames to revert delay pattern: T={t} < Q={q}") - seq_len = t - q + 1 - out_l = [] - for i in range(q): - out_l.append(audio_codes_qt[i : i + 1, i : seq_len + i]) - return torch.cat(out_l, dim=0) + return _common.revert_delay_pattern( + audio_codes_qt, + expected_codebooks=_NUM_CODEBOOKS, + allow_short=False, + ) def _filter_real_code_frames(audio_codes_qt: torch.Tensor) -> torch.Tensor: """Keep only frames where ALL codebook values are in [0, 1023]. - Input shape: [num_codebooks, num_frames]. - Called AFTER delay pattern reversal. + RFC #4872 P3: delegate to the canonical ``_common.filter_real_code_frames`` + with the higgs-v3 ``[num_codebooks, num_frames]`` layout. """ - if audio_codes_qt.numel() == 0: - return audio_codes_qt - # Transpose to [num_frames, num_codebooks] for per-frame filtering - frames = audio_codes_qt.t() - valid = (frames >= 0).all(dim=1) & (frames < _NUM_REAL_CODES).all(dim=1) - return frames[valid].t().contiguous() + return _common.filter_real_code_frames( + audio_codes_qt, num_real_codes=_NUM_REAL_CODES, layout="codebooks_first" + ) def _whole_utterance_codec_codes(audio_rows: torch.Tensor) -> torch.Tensor: diff --git a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py index 653e3477789..937e4d39a0c 100644 --- a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py @@ -13,6 +13,7 @@ from vllm.inputs import TextPrompt from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors import _common logger = logging.getLogger(__name__) @@ -204,14 +205,11 @@ def _resolve_image_patch_token_id(stage: Any) -> int: def _ensure_list(x) -> list[int]: - """Convert ConstantList / tensor-like to plain list.""" - if hasattr(x, "_x"): - return list(x._x) - if isinstance(x, list): - return x - if hasattr(x, "tolist"): - return x.tolist() - return list(x) + """Convert ConstantList / tensor-like to plain list. + + RFC #4872 P3: delegate to the canonical ``_common.ensure_list_strict``. + """ + return _common.ensure_list_strict(x) def _slice_patch_hidden( diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 3fee73ac3df..1a1cc089b7e 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -23,6 +23,7 @@ ) from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors import _common from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import compute_ramp_emit, parse_chunk_ramp from vllm_omni.model_executor.stage_input_processors.tts_utils import ( extract_language_from_prompt, @@ -157,12 +158,11 @@ def _compute_talker_prompt_ids_length(info: OmniPayload, device: torch.device | def _ensure_list(x): - """Convert ConstantList / tensor-like to Python list.""" - if hasattr(x, "_x"): - return list(x._x) - elif not isinstance(x, list): - return x - return list(x) + """Convert ConstantList / tensor-like to Python list. + + RFC #4872 P3: delegate to the canonical ``_common.ensure_list_unchanged``. + """ + return _common.ensure_list_unchanged(x) def _as_tensor_or_none(value: Any) -> torch.Tensor | None: diff --git a/vllm_omni/model_executor/stage_input_processors/step_audio2.py b/vllm_omni/model_executor/stage_input_processors/step_audio2.py index c0c0932b27a..f196e3a9af4 100644 --- a/vllm_omni/model_executor/stage_input_processors/step_audio2.py +++ b/vllm_omni/model_executor/stage_input_processors/step_audio2.py @@ -14,19 +14,18 @@ DEFAULT_STREAM_CONFIG, DEFAULT_TOKEN_CONFIG, ) +from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) def _ensure_list(x): - """Convert ConstantList / tensor / iterable to Python list.""" - if hasattr(x, "_x"): - return list(x._x) - elif isinstance(x, torch.Tensor): - return x.tolist() - elif not isinstance(x, list): - return list(x) if hasattr(x, "__iter__") else [x] - return list(x) + """Convert ConstantList / tensor / iterable to Python list. + + Preserves ``None`` as ``[None]`` (step_audio2 semantics, RFC #4872 + Step 1c): delegates to ``_common.ensure_list_preserve_none``. + """ + return _common.ensure_list_preserve_none(x) # ========================= From f7578020ba9d8543ebae9b69c9ecab9ea0a24b55 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Tue, 25 Aug 2026 19:48:13 +0800 Subject: [PATCH 02/41] feat(stage_input_processors): RFC #4872 P2 processor registry + startup validation (Phase 3a) Signed-off-by: RyanYun09 --- .../stage_input_processors/test_registry.py | 370 +++ .../chunk_transfer_adapter.py | 36 +- vllm_omni/engine/stage_init_utils.py | 23 +- .../stage_input_processors/__init__.py | 23 +- .../stage_input_processors/_registry.py | 679 +++++ .../omni_connector_model_runner_mixin.py | 2275 +++++++++++++++++ 6 files changed, 3396 insertions(+), 10 deletions(-) create mode 100644 tests/model_executor/stage_input_processors/test_registry.py create mode 100644 vllm_omni/model_executor/stage_input_processors/_registry.py diff --git a/tests/model_executor/stage_input_processors/test_registry.py b/tests/model_executor/stage_input_processors/test_registry.py new file mode 100644 index 00000000000..968d70d9e88 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_registry.py @@ -0,0 +1,370 @@ +"""Unit tests for the RFC #4872 Phase 3 (Part A) processor registry. + +Pure-logic, CPU-only: exercises kind inference, structural validation, +dead-processor hints and ``resolve_processor`` without executing any processor +logic or loading model weights. + +Subject under test: ``vllm_omni.model_executor.stage_input_processors._registry`` +""" + +import importlib +import warnings + +import pytest + +from vllm_omni.model_executor.stage_input_processors import ( + ProcessorKind, + ProcessorSpec, + ProcessorValidationError, + dead_processor_hint, + infer_kind, + register_processor, + resolve_processor, + validate_processor, +) +from vllm_omni.model_executor.stage_input_processors import _registry + + +# --------------------------------------------------------------------------- +# Dummy processors (simple callables used as validation subjects). +# --------------------------------------------------------------------------- + + +def placeholder_a_token_only(source_outputs, _prompt=None, _requires_multimodal_data=False): + return [] + + +def placeholder_b_token_only(prompt, ctx=None): + return [] + + +def placeholder_c_badname(source_outputs, _prompt=None): + return [] + + +def diffusion_a_ar2diffusion(source_outputs, _prompt=None, _requires_multimodal_data=False): + return {} + + +def producer_a_full_payload(transfer_manager, pooling_output, request): + return None + + +def producer_b_full_payload(transfer_manager, pooling_output, request, is_finished=False): + return None + + +def producer_c_full_payload(transfer_manager, multimodal_output, request, is_finished=False): + return None + + +def producer_a_async_chunk(transfer_manager, multimodal_output, request, is_finished=False): + return None + + +def producer_b_async_chunk(transfer_manager, multimodal_output, request): + return None + + +def producer_c_async_chunk(transfer_manager, multimodal_output, request, is_finished, /): + return None + + +def talker2codec_shape(stage_list, engine_input_source, prompt=None, requires_multimodal_data=False): + return [] + + +def _mod_path(fn_name: str) -> str: + """Dotted path to a callable defined in this test module.""" + return f"{__name__}.{fn_name}" + + +# --------------------------------------------------------------------------- +# Kind inference (suffix rules). +# --------------------------------------------------------------------------- + + +class TestInferKind: + @pytest.mark.parametrize( + "path,expected", + [ + ("pkg.mod.foo_token_only", "placeholder_prompt_builder"), + ("pkg.mod.foo_full_payload", "producer_full_payload"), + ("pkg.mod.foo_batch", "producer_full_payload"), + ("pkg.mod.foo_async_chunk", "producer_async_chunk"), + ("pkg.mod.ar2diffusion", "diffusion_input_builder"), + ("pkg.mod.ar2dit", "diffusion_input_builder"), + ("pkg.mod.thinker2imagegen", "diffusion_input_builder"), + ("pkg.mod.legacy_builder", "legacy_orchestrator_builder"), + ("pkg.mod.talker2code2wav", "legacy_orchestrator_builder"), + ("vllm_omni.model_executor.stage_input_processors.moss_tts.talker2codec", "legacy_multi_source"), + ("pkg.moss_tts.talker2codec", "legacy_multi_source"), + ( + "vllm_omni.model_executor.stage_input_processors.moss_tts.talker2codec_delay_async_chunk", + "producer_async_chunk", + ), + ], + ) + def test_suffix_rules(self, path: str, expected: ProcessorKind): + assert infer_kind(None, path=path) == expected + + def test_invalid_path_raises(self): + with pytest.raises(ValueError): + infer_kind(None, path="") + with pytest.raises(ValueError): + infer_kind(None, path="nodots") + + def test_register_processor_override(self): + register_processor("pkg.mod.custom_path", "diffusion_input_builder") + try: + assert infer_kind(None, path="pkg.mod.custom_path") == "diffusion_input_builder" + finally: + _registry._KIND_OVERRIDES.clear() + + def test_register_processor_invalid_kind(self): + with pytest.raises(ValueError): + register_processor("pkg.mod.custom_path", "bogus_kind") # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- +# Structural validation. +# --------------------------------------------------------------------------- + + +class TestValidateProcessor: + def test_full_payload_ok_no_warning(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor( + producer_b_full_payload, + kind="producer_full_payload", + path="pkg.mod.producer_b_full_payload", + ) + + def test_full_payload_missing_is_finished_warns(self): + with pytest.warns(RuntimeWarning, match="is_finished"): + validate_processor( + producer_a_full_payload, + kind="producer_full_payload", + path="pkg.mod.producer_a_full_payload", + ) + + def test_full_payload_cross_kind_second_param_warns(self): + with pytest.warns(RuntimeWarning, match="multimodal_output"): + validate_processor( + producer_c_full_payload, + kind="producer_full_payload", + path="pkg.mod.producer_c_full_payload", + ) + + def test_full_payload_missing_transfer_manager_fails(self): + def bad(manager, pooling_output, request): + return None + + with pytest.raises(ProcessorValidationError) as excinfo: + validate_processor(bad, kind="producer_full_payload", path="pkg.mod.producer_full_payload") + assert excinfo.value.rule == "transfer_manager" + + def test_async_chunk_ok_no_warning(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor( + producer_a_async_chunk, + kind="producer_async_chunk", + path="pkg.mod.producer_a_async_chunk", + ) + + def test_async_chunk_missing_is_finished_hard_fail(self): + with pytest.raises(ProcessorValidationError) as excinfo: + validate_processor( + producer_b_async_chunk, + kind="producer_async_chunk", + path="pkg.mod.producer_b_async_chunk", + ) + assert excinfo.value.rule == "is_finished_required" + assert "producer_b_async_chunk" in str(excinfo.value) + + def test_async_chunk_positional_only_is_finished_hard_fail(self): + with pytest.raises(ProcessorValidationError) as excinfo: + validate_processor( + producer_c_async_chunk, + kind="producer_async_chunk", + path="pkg.mod.producer_c_async_chunk", + ) + assert excinfo.value.rule == "is_finished_positional_only" + + def test_placeholder_ok_no_warning(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor( + placeholder_a_token_only, + kind="placeholder_prompt_builder", + path="pkg.mod.placeholder_a_token_only", + ) + + def test_placeholder_ctx_contract_ok(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor( + placeholder_b_token_only, + kind="placeholder_prompt_builder", + path="pkg.mod.placeholder_b_token_only", + ) + + def test_placeholder_suffix_kind_conflict_warns(self): + with pytest.warns(RuntimeWarning, match="suffix"): + validate_processor( + placeholder_c_badname, + kind="placeholder_prompt_builder", + path="pkg.mod.placeholder_c_badname", + ) + + def test_diffusion_ok_no_warning(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor( + diffusion_a_ar2diffusion, + kind="diffusion_input_builder", + path="pkg.mod.diffusion_a_ar2diffusion", + ) + + def test_legacy_multi_source_allowlist_ok(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor( + talker2codec_shape, + kind="legacy_multi_source", + path="vllm_omni.model_executor.stage_input_processors.moss_tts.talker2codec", + ) + + def test_legacy_multi_source_allowlist_rejects_non_moss(self): + with pytest.raises(ProcessorValidationError) as excinfo: + validate_processor( + talker2codec_shape, + kind="legacy_multi_source", + path="pkg.mod.talker2codec", + ) + assert excinfo.value.rule == "legacy_multi_source_allowlist" + + def test_uninspectable_callable_fails(self): + with pytest.raises(ProcessorValidationError) as excinfo: + validate_processor( + object(), # not callable / not inspectable + kind="producer_full_payload", + path="pkg.mod.producer_full_payload", + ) + assert excinfo.value.rule == "inspectable_signature" + + +# --------------------------------------------------------------------------- +# Dead-processor hint (gate three-state). +# --------------------------------------------------------------------------- + + +class TestDeadProcessor: + def test_async_chunk_receiving_downstream_is_dead(self): + assert ( + dead_processor_hint( + "placeholder_prompt_builder", + async_chunk=True, + downstream_receives_async_chunks=True, + has_sync=False, + ) + is True + ) + + def test_async_chunk_non_receiving_downstream_alive(self): + assert ( + dead_processor_hint( + "diffusion_input_builder", + async_chunk=True, + downstream_receives_async_chunks=False, + has_sync=False, + ) + is False + ) + + def test_sync_mode_alive_unless_overridden(self): + assert ( + dead_processor_hint( + "diffusion_input_builder", + async_chunk=False, + downstream_receives_async_chunks=False, + has_sync=False, + ) + is False + ) + assert ( + dead_processor_hint( + "diffusion_input_builder", + async_chunk=False, + downstream_receives_async_chunks=False, + has_sync=True, + ) + is True + ) + + def test_producer_never_dead(self): + for kind in ("producer_full_payload", "producer_async_chunk"): + assert ( + dead_processor_hint( + kind, + async_chunk=True, + downstream_receives_async_chunks=True, + has_sync=False, + ) + is False + ) + + +# --------------------------------------------------------------------------- +# Resolution entry point. +# --------------------------------------------------------------------------- + + +class TestResolveProcessor: + def test_resolve_normal(self): + spec = resolve_processor( + _mod_path("placeholder_a_token_only"), + expected_kind="placeholder_prompt_builder", + ) + assert isinstance(spec, ProcessorSpec) + assert spec.kind == "placeholder_prompt_builder" + assert spec.path == _mod_path("placeholder_a_token_only") + assert spec.fn is placeholder_a_token_only + + def test_resolve_full_payload_without_expected_kind(self): + spec = resolve_processor(_mod_path("producer_b_full_payload")) + assert spec.kind == "producer_full_payload" + + def test_resolve_drop_in_equivalence(self): + """resolve_processor().fn must be the same object as the legacy lookup.""" + spec = resolve_processor( + _mod_path("producer_b_full_payload"), + expected_kind="producer_full_payload", + ) + legacy = getattr(importlib.import_module(__name__), "producer_b_full_payload") + assert spec.fn is legacy + + def test_expected_kind_mismatch_raises(self): + with pytest.raises(ProcessorValidationError) as excinfo: + resolve_processor( + _mod_path("placeholder_a_token_only"), + expected_kind="producer_full_payload", + ) + assert excinfo.value.rule == "expected_kind" + assert "placeholder_prompt_builder" in str(excinfo.value) + + def test_missing_symbol_raises(self): + with pytest.raises(AttributeError): + resolve_processor(_mod_path("does_not_exist")) + + def test_package_exports(self): + import vllm_omni.model_executor.stage_input_processors as sip + + assert sip.resolve_processor is resolve_processor + assert sip.ProcessorSpec is ProcessorSpec + assert sip.ProcessorValidationError is ProcessorValidationError + assert sip.infer_kind is infer_kind + assert sip.dead_processor_hint is dead_processor_hint + assert sip.register_processor is register_processor diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index cdaf2720df0..5b683c4bce8 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import copy -import importlib import inspect import threading import time @@ -16,6 +15,28 @@ from vllm.v1.utils import ConstantList from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload +from vllm_omni.model_executor.stage_input_processors import resolve_processor + +from ..adapter import construct_next_stage_streaming_input_prompt +from ..factory import OmniConnectorFactory +from ..utils.config import ConnectorSpec, stage_receives_chunks +from ..utils.initialization import resolve_connector_spec +from ..utils.kv_utils import get_local_tp_rank, get_omni_replica_id +from ..utils.logging import get_connector_logger +from .base import OmniTransferAdapterBase + +import time +from collections import defaultdict, deque +from collections.abc import Callable, Iterable, Mapping +from typing import Any + +import torch +from vllm.v1.metrics.stats import PrefillStats +from vllm.v1.request import Request, RequestStatus +from vllm.v1.utils import ConstantList + +from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload +from vllm_omni.model_executor.stage_input_processors import resolve_processor from ..adapter import construct_next_stage_streaming_input_prompt from ..connectors.shm_connector import SharedMemoryConnector @@ -181,10 +202,15 @@ def __init__(self, vllm_config: Any): self.custom_process_next_stage_input_func: Callable[..., OmniPayloadStruct | None] | None = None custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None) if custom_process_next_stage_input_func: - module_path, func_name = custom_process_next_stage_input_func.rsplit(".", 1) - module = importlib.import_module(module_path) - self.custom_process_next_stage_input_func = getattr(module, func_name) - self._processor_accepts_step_tokens: dict[Callable[..., Any], bool] = {} + # Resolve through the registry with the async-chunk producer + # contract. ``_select_processor_funcs`` already merged the + # async_chunk field into this path, so it must be a real + # ``*_async_chunk`` producer; ``is_finished`` is required (the + # scheduler always passes it below) — a missing/positional-only + # declaration is a hard configuration error (ProcessorValidationError). + spec = resolve_processor(custom_process_next_stage_input_func, expected_kind="producer_async_chunk") + self.custom_process_next_stage_input_func = spec.fn + # mapping for request id and chunk id self.put_req_chunk: dict[str, int] = defaultdict(int) self.get_req_chunk: dict[str, int] = defaultdict(int) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 4413cd01aad..36c51f9e8f2 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -57,7 +57,8 @@ from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams -from vllm_omni.inputs.preprocess import build_omni_renderer, omni_renderer_cls +from vllm_omni.inputs.preprocess import OmniInputPreprocessor, build_omni_renderer, omni_renderer_cls +from vllm_omni.model_executor.stage_input_processors import resolve_processor from vllm_omni.outputs.output_processor import MultimodalOutputProcessor from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.inc_config import OmniINCConfig @@ -639,8 +640,24 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: custom_process_input_func: Callable | None = None _cpif_path = _get_attr_or_item(stage_config, "custom_process_input_func") if _cpif_path: - mod_path, fn_name = _cpif_path.rsplit(".", 1) - custom_process_input_func = getattr(importlib.import_module(mod_path), fn_name) + # RFC #4872 Phase 3 (Part A): resolve through the stage-input processor + # registry. expected_kind is deliberately None here — orchestrator-side + # processors carry no suffix contract, so we rely on suffix inference + # only (compatible with existing configs). A bad path must not kill + # startup (M0): warn and continue without the hook. + try: + custom_process_input_func = resolve_processor( + _cpif_path, + expected_kind=None, + stage_config=stage_config, + ).fn + except Exception as exc: + logger.warning( + "[stage_init] Failed to resolve custom_process_input_func %r; " + "continuing without it (RFC #4872 M0: warn, not fail): %s", + _cpif_path, + exc, + ) prompt_transform_func: Callable | None = None _ptf_path = _get_attr_or_item(stage_config, "prompt_transform_func") diff --git a/vllm_omni/model_executor/stage_input_processors/__init__.py b/vllm_omni/model_executor/stage_input_processors/__init__.py index 562047e98d1..e7e2bcce0a6 100644 --- a/vllm_omni/model_executor/stage_input_processors/__init__.py +++ b/vllm_omni/model_executor/stage_input_processors/__init__.py @@ -17,8 +17,9 @@ - diffusion : upstream outputs -> diffusion payload(s) (``DiffusionInputBuilder``). -This package only re-exports the dispatch contract; it imports no other -processor modules, so importing it never pulls in model-specific code. +This package only re-exports the dispatch contract and the processor registry; +it imports no other processor modules, so importing it never pulls in +model-specific code. """ from vllm_omni.model_executor.stage_input_processors._dispatch import ( @@ -30,6 +31,16 @@ invoke_orchestrator_processor, wrap_orchestrator_processor, ) +from vllm_omni.model_executor.stage_input_processors._registry import ( + ProcessorKind, + ProcessorSpec, + ProcessorValidationError, + dead_processor_hint, + infer_kind, + register_processor, + resolve_processor, + validate_processor, +) __all__ = [ "OrchestratorInputContext", @@ -39,4 +50,12 @@ "AsyncChunkProducer", "invoke_orchestrator_processor", "wrap_orchestrator_processor", + "ProcessorKind", + "ProcessorSpec", + "ProcessorValidationError", + "resolve_processor", + "register_processor", + "infer_kind", + "validate_processor", + "dead_processor_hint", ] diff --git a/vllm_omni/model_executor/stage_input_processors/_registry.py b/vllm_omni/model_executor/stage_input_processors/_registry.py new file mode 100644 index 00000000000..50057b74b7f --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/_registry.py @@ -0,0 +1,679 @@ +"""RFC #4872 Phase 3 (Part A): stage-input processor registry. + +This module is the **structural validation** layer behind the three ``importlib`` +resolution points that RFC P2 calls out as independent symbol lookups: + +* ``engine/stage_init_utils.extract_legacy_stage_metadata`` — consumer-side + (orchestrator) input processors (``custom_process_input_func`` / + ``sync_process_input_func``). +* ``worker/omni_connector_model_runner_mixin._load_custom_func`` — the + producer-side full-payload builder (``FullPayloadProducer``). +* ``distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.__init__`` + — the producer-side async-chunk builder (``AsyncChunkProducer``). + +Scope / responsibility split +---------------------------- +The registry **only performs signature-level structural checks** +(``inspect.signature``); it never executes processor logic and never loads +model weights. Behavioural guarantees (payload shape, flush semantics, golden +outputs) belong to the golden / parity tests, not here — do not conflate the +two responsibilities. + +Producer kwargs are contract. ``FullPayloadProducer`` keeps ``pooling_output`` +(and ``is_finished`` is **optional** — the worker retries without it when the +processor rejects the kwarg); ``AsyncChunkProducer`` keeps ``multimodal_output`` +and **must** accept ``is_finished`` — the scheduler always passes it on the +async-chunk path (see ``chunk_transfer_adapter._send_single_request``). + +Kind inference is name-driven (suffix-based), following the RFC P2 naming +convention documented in :mod:`vllm_omni.model_executor.stage_input_processors`: + +* ``_token_only`` -> ``placeholder_prompt_builder`` +* ``_full_payload`` / ``_batch`` -> ``producer_full_payload`` +* ``_async_chunk`` -> ``producer_async_chunk`` +* ``ar2diffusion`` / ``ar2dit`` / ``thinker2imagegen`` (and friends) -> + ``diffusion_input_builder`` +* no suffix -> ``legacy_orchestrator_builder`` +* ``moss_tts.talker2codec`` (legacy multi-source shape) -> + ``legacy_multi_source`` + +``resolve_processor`` is a drop-in replacement for the legacy +``getattr(importlib.import_module(mod_path), fn_name)`` lookups: it returns a +:class:`ProcessorSpec` whose ``fn`` is the **same callable object** the legacy +code produced, so callers keep identical runtime behaviour while gaining +validated kind metadata. +""" + +from __future__ import annotations + +import importlib +import inspect +import warnings +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Literal + +ProcessorKind = Literal[ + "placeholder_prompt_builder", + "diffusion_input_builder", + "legacy_orchestrator_builder", + "producer_full_payload", + "producer_async_chunk", + "legacy_multi_source", +] + +#: Suffixes that identify a producer-side full-payload builder. +_FULL_PAYLOAD_SUFFIXES: tuple[str, ...] = ("_full_payload", "_batch") +#: Suffix that identifies a producer-side async-chunk builder. +_ASYNC_CHUNK_SUFFIX: str = "_async_chunk" +#: Suffix that identifies a consumer-side placeholder builder. +_TOKEN_ONLY_SUFFIX: str = "_token_only" + +#: Diffusion-stage input-builder name stems (suffix match). +_DIFFUSION_SUFFIXES: tuple[str, ...] = ( + "ar2diffusion", + "ar2dit", + "thinker2imagegen", + "ar2acoustic", + "ar_to_dit", + "llm2audio_vae", + "thinker2token2wav", +) + +#: Module hint for the only remaining legacy multi-source processor family. +_MOSS_MULTI_SOURCE_MODULE_HINT: str = "moss_tts" +#: Exact function names in ``moss_tts`` with the multi-source shape. +_MOSS_MULTI_SOURCE_NAMES: frozenset[str] = frozenset({"talker2codec"}) + +#: Producer kinds are never orchestrator-side, so the dead-processor hint never +#: applies to them. +_PRODUCER_KINDS: frozenset[str] = frozenset({"producer_full_payload", "producer_async_chunk"}) + +#: All valid :class:`ProcessorKind` values (used to validate manual overrides). +_PROCESSOR_KINDS: frozenset[str] = frozenset( + { + "placeholder_prompt_builder", + "diffusion_input_builder", + "legacy_orchestrator_builder", + "producer_full_payload", + "producer_async_chunk", + "legacy_multi_source", + } +) + +#: Soft-check scalar return annotations that contradict a stage-input role. +_SCALAR_RETURN_ANNOTATIONS: frozenset[Any] = frozenset({int, str, bool, float}) + +__all__ = [ + "ProcessorKind", + "ProcessorSpec", + "ProcessorValidationError", + "import_symbol", + "infer_kind", + "validate_processor", + "resolve_processor", + "dead_processor_hint", + "register_processor", +] + + +@dataclass(frozen=True) +class ProcessorSpec: + """Resolved and validated stage-input processor. + + Attributes: + path: The dotted ``module.attr`` path the processor was resolved from. + kind: The inferred (or overridden) processor kind. + fn: The callable itself — identical to what the legacy + ``getattr(importlib.import_module(...), ...)`` lookup returned. + """ + + path: str + kind: ProcessorKind + fn: Any + + +class ProcessorValidationError(Exception): + """A stage-input processor failed structural validation. + + Carries enough context to locate the offending stage / config line: + ``stage_id`` (when a ``stage_config`` was supplied), ``path``, ``kind`` and + the specific ``rule`` that was violated. + """ + + def __init__( + self, + *, + stage_id: int | None = None, + path: str | None = None, + kind: ProcessorKind | None = None, + rule: str | None = None, + message: str | None = None, + ) -> None: + self.stage_id = stage_id + self.path = path + self.kind = kind + self.rule = rule + self.message = message + super().__init__(self._format()) + + def _format(self) -> str: + parts: list[str] = [] + if self.stage_id is not None: + parts.append(f"stage_id={self.stage_id}") + if self.path: + parts.append(f"path={self.path!r}") + if self.kind: + parts.append(f"kind={self.kind!r}") + if self.rule: + parts.append(f"rule={self.rule}") + detail = ", ".join(parts) or "processor validation failed" + if self.message: + return f"{detail}: {self.message}" + return detail + + +# --------------------------------------------------------------------------- +# Small path / config helpers +# --------------------------------------------------------------------------- + + +def _fn_name(path: str) -> str: + return path.rsplit(".", 1)[-1] + + +def _module_part(path: str) -> str: + return path.rsplit(".", 1)[0] + + +def _stage_id_of(stage_config: Any) -> int | None: + """Best-effort stage id extraction from an arbitrary config object.""" + if stage_config is None: + return None + sid = getattr(stage_config, "stage_id", None) + if sid is None: + return None + try: + return int(sid) + except (TypeError, ValueError): + return None + + +# --------------------------------------------------------------------------- +# Symbol resolution +# --------------------------------------------------------------------------- + + +def import_symbol(path: str) -> Any: + """Import a dotted ``module.attr`` symbol (importlib + ``getattr``). + + This mirrors the two-step shape used by every legacy resolution point: + ``importlib.import_module(module_path)`` then ``getattr(module, attr)``. + A missing attribute or module raises the underlying + ``ImportError``/``AttributeError``; callers that need a best-effort lookup + (e.g. the worker candidate chain) catch these themselves. + """ + if not isinstance(path, str) or not path: + raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") + module_path, attr_name = path.rsplit(".", 1) + module = importlib.import_module(module_path) + return getattr(module, attr_name) + + +# --------------------------------------------------------------------------- +# Kind inference (name / suffix driven) +# --------------------------------------------------------------------------- + +#: Manual kind overrides registered via :func:`register_processor`. +_KIND_OVERRIDES: dict[str, ProcessorKind] = {} + + +def register_processor(path: str, kind: ProcessorKind) -> None: + """Register a manual kind override for a processor *path*. + + This is an escape hatch for models whose function names do not follow the + RFC P2 suffix convention. ``infer_kind`` consults these overrides before + applying the suffix rules. Overrides are validated eagerly so typos fail + fast. + """ + if not isinstance(path, str) or not path: + raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") + if kind not in _PROCESSOR_KINDS: + raise ValueError(f"unknown processor kind {kind!r}") + _KIND_OVERRIDES[path] = kind + + +def infer_kind(fn: Any, *, path: str) -> ProcessorKind: + """Infer the processor kind from the *suffix* of ``path``. + + ``fn`` is accepted for API symmetry with :func:`validate_processor` but the + inference is purely name-driven (RFC P2 naming convention); the callable + object itself is not inspected here. + """ + del fn # inference is name-based; kept for a symmetric API + if not isinstance(path, str) or not path or "." not in path: + raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") + + override = _KIND_OVERRIDES.get(path) + if override is not None: + return override + + fn_name = _fn_name(path) + + if _MOSS_MULTI_SOURCE_MODULE_HINT in _module_part(path) and fn_name in _MOSS_MULTI_SOURCE_NAMES: + return "legacy_multi_source" + + if fn_name.endswith(_TOKEN_ONLY_SUFFIX): + return "placeholder_prompt_builder" + if fn_name.endswith(_FULL_PAYLOAD_SUFFIXES): + return "producer_full_payload" + if fn_name.endswith(_ASYNC_CHUNK_SUFFIX): + return "producer_async_chunk" + for suffix in _DIFFUSION_SUFFIXES: + if fn_name.endswith(suffix): + return "diffusion_input_builder" + return "legacy_orchestrator_builder" + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +def _warn(path: str, kind: ProcessorKind | None, rule: str, message: str) -> None: + """Emit a soft ``RuntimeWarning`` for a non-fatal structural mismatch.""" + warnings.warn( + f"stage-input processor {path!r} (kind={kind!r}, rule={rule}): {message}", + RuntimeWarning, + stacklevel=3, + ) + + +def _fail( + *, + path: str, + kind: ProcessorKind | None, + rule: str, + stage_config: Any, + message: str, +) -> None: + raise ProcessorValidationError( + stage_id=_stage_id_of(stage_config), + path=path, + kind=kind, + rule=rule, + message=message, + ) + + +def _second_parameter_name(params: Mapping[str, inspect.Parameter]) -> str | None: + """Name of the second (non-var) parameter, mirroring the producer shape.""" + names = [ + name + for name, param in params.items() + if param.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) + ] + return names[1] if len(names) >= 2 else None + + +def _check_orchestrator_first_param( + params: Mapping[str, inspect.Parameter], + *, + path: str, + kind: ProcessorKind, + stage_config: Any, +) -> None: + """First param must be ``source_outputs`` or the signature must carry ``ctx``.""" + names = [ + name + for name, param in params.items() + if param.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) + ] + first = names[0] if names else None + if first != "source_outputs" and "ctx" not in params: + _warn( + path, + kind, + "orchestrator_first_param", + f"expected first parameter 'source_outputs' or a 'ctx' parameter, got first={first!r}", + ) + + +def _check_suffix( + path: str, + kind: ProcessorKind, + *, + suffix_pool: tuple[str, ...], + stage_config: Any, + what: str, +) -> None: + fn_name = _fn_name(path) + if not any(fn_name.endswith(suffix) for suffix in suffix_pool): + _warn(path, kind, "suffix_kind_mismatch", f"name {fn_name!r} does not match expected {what}") + + +def _check_full_payload_is_finished( + params: Mapping[str, inspect.Parameter], + *, + path: str, + kind: ProcessorKind, + stage_config: Any, +) -> None: + """``is_finished`` is optional for the full-payload producer. + + The worker passes it best-effort and retries without it on + ``TypeError`` (see ``omni_connector_model_runner_mixin._build_custom_process_payload``). + """ + if "is_finished" not in params: + _warn( + path, + kind, + "is_finished_optional", + "full-payload producer does not declare is_finished; worker retries without it (best-effort)", + ) + return + if params["is_finished"].kind == inspect.Parameter.POSITIONAL_ONLY: + _warn( + path, + kind, + "is_finished_positional_only", + "full-payload producer declares is_finished positional-only; worker retries without it", + ) + + +def _check_async_chunk_is_finished( + params: Mapping[str, inspect.Parameter], + *, + path: str, + kind: ProcessorKind, + stage_config: Any, +) -> None: + """``is_finished`` is **required** for the async-chunk producer. + + The scheduler always passes ``is_finished`` (see + ``chunk_transfer_adapter._send_single_request``), so a processor that cannot + accept it is a hard configuration error. + """ + if "is_finished" not in params: + _fail( + path=path, + kind=kind, + rule="is_finished_required", + stage_config=stage_config, + message="async-chunk producer must declare 'is_finished' (the scheduler always passes it)", + ) + return # unreachable; kept for type-checkers + is_finished_param = params["is_finished"] + if is_finished_param.kind == inspect.Parameter.POSITIONAL_ONLY: + _fail( + path=path, + kind=kind, + rule="is_finished_positional_only", + stage_config=stage_config, + message="async-chunk producer declares 'is_finished' positional-only; " + "it must be keyword-only or have a default", + ) + return # unreachable + if is_finished_param.kind != inspect.Parameter.KEYWORD_ONLY and is_finished_param.default is inspect.Parameter.empty: + _warn( + path, + kind, + "is_finished_default", + "async-chunk producer declares a required positional-or-keyword 'is_finished'; " + "prefer keyword-only or a default", + ) + + +def _check_return_annotation( + signature: inspect.Signature, + *, + path: str, + kind: ProcessorKind, +) -> None: + """Soft return-annotation check: mismatches warn, never fail.""" + ann = signature.return_annotation + if ann is inspect.Signature.empty or isinstance(ann, str): + # Unannotated or postponed-string annotations are not verifiable here. + return + if ann in _SCALAR_RETURN_ANNOTATIONS: + _warn( + path, + kind, + "return_annotation", + f"return annotation {ann!r} is a scalar, unlikely a stage-input processor result", + ) + + +def validate_processor(fn: Any, *, kind: ProcessorKind, path: str, stage_config: Any = None) -> None: + """Structurally validate *fn* against the expected *kind*. + + Only ``inspect.signature`` metadata is examined; the processor body is never + executed. Hard contract violations raise :class:`ProcessorValidationError`; + soft mismatches emit a ``RuntimeWarning``. + + Args: + fn: The processor callable. + kind: The expected :class:`ProcessorKind`. + path: The dotted ``module.attr`` path (used for reporting). + stage_config: Optional config object; its ``stage_id`` is attached to + any raised :class:`ProcessorValidationError`. + """ + if not isinstance(path, str) or not path: + raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") + + try: + signature = inspect.signature(fn) + except (TypeError, ValueError) as exc: + _fail( + path=path, + kind=kind, + rule="inspectable_signature", + stage_config=stage_config, + message=f"cannot inspect signature: {exc}", + ) + return # unreachable + + params = signature.parameters + + if kind == "placeholder_prompt_builder": + _check_orchestrator_first_param(params, path=path, kind=kind, stage_config=stage_config) + _check_suffix( + path, + kind, + suffix_pool=(_TOKEN_ONLY_SUFFIX,), + stage_config=stage_config, + what="a *_token_only suffix", + ) + + elif kind == "diffusion_input_builder": + _check_orchestrator_first_param(params, path=path, kind=kind, stage_config=stage_config) + _check_suffix( + path, + kind, + suffix_pool=_DIFFUSION_SUFFIXES, + stage_config=stage_config, + what="a diffusion-builder suffix (ar2diffusion/ar2dit/thinker2imagegen/...)", + ) + + elif kind == "producer_full_payload": + if "transfer_manager" not in params: + _fail( + path=path, + kind=kind, + rule="transfer_manager", + stage_config=stage_config, + message="full-payload producer must declare 'transfer_manager'", + ) + if "request" not in params: + _fail( + path=path, + kind=kind, + rule="request", + stage_config=stage_config, + message="full-payload producer must declare 'request'", + ) + second = _second_parameter_name(params) + if second == "multimodal_output": + _warn( + path, + kind, + "cross_kind_second_param", + "second parameter is 'multimodal_output' (async-chunk producer shape) " + "but kind is producer_full_payload; expected 'pooling_output'", + ) + elif second is not None and second != "pooling_output": + _warn( + path, + kind, + "second_param", + f"expected second parameter 'pooling_output', got {second!r}", + ) + _check_full_payload_is_finished(params, path=path, kind=kind, stage_config=stage_config) + + elif kind == "producer_async_chunk": + if "transfer_manager" not in params: + _fail( + path=path, + kind=kind, + rule="transfer_manager", + stage_config=stage_config, + message="async-chunk producer must declare 'transfer_manager'", + ) + if "request" not in params: + _fail( + path=path, + kind=kind, + rule="request", + stage_config=stage_config, + message="async-chunk producer must declare 'request'", + ) + second = _second_parameter_name(params) + if second == "pooling_output": + _warn( + path, + kind, + "cross_kind_second_param", + "second parameter is 'pooling_output' (full-payload producer shape) " + "but kind is producer_async_chunk; expected 'multimodal_output'", + ) + elif second is not None and second != "multimodal_output": + _warn( + path, + kind, + "second_param", + f"expected second parameter 'multimodal_output', got {second!r}", + ) + _check_async_chunk_is_finished(params, path=path, kind=kind, stage_config=stage_config) + + elif kind == "legacy_multi_source": + if _MOSS_MULTI_SOURCE_MODULE_HINT not in _module_part(path): + _fail( + path=path, + kind=kind, + rule="legacy_multi_source_allowlist", + stage_config=stage_config, + message=f"legacy multi-source processors are allowlisted to moss_tts.*, got {path!r}", + ) + + elif kind == "legacy_orchestrator_builder": + # No structural contract: the orchestrator normalizes legacy positional + # shapes (C0/C2/C3/C4) through wrap_orchestrator_processor, and the + # first-parameter name is not reliable for legacy builders (configs may + # legitimately point at C-callables / operator builtins as test doubles). + # Deliberately no first-param check here. + pass + + else: # pragma: no cover - defensive against future literal additions + raise ValueError(f"unknown processor kind {kind!r}") + + _check_return_annotation(signature, path=path, kind=kind) + + +# --------------------------------------------------------------------------- +# Resolution entry point +# --------------------------------------------------------------------------- + + +def resolve_processor( + path: str, + *, + expected_kind: ProcessorKind | None = None, + stage_config: Any = None, +) -> ProcessorSpec: + """Resolve a dotted *path* to a validated :class:`ProcessorSpec`. + + Flow: ``import_symbol`` -> ``infer_kind`` -> ``validate_processor`` -> + (optional) ``expected_kind`` check. The returned ``spec.fn`` is the same + callable the legacy ``getattr(importlib.import_module(...), ...)`` lookup + produced, so callers keep identical runtime behaviour. + + Args: + path: Dotted ``module.attr`` path of the processor. + expected_kind: When given, ``infer_kind`` must match; a mismatch raises + :class:`ProcessorValidationError` (rule ``expected_kind``). + stage_config: Optional config whose ``stage_id`` is attached to errors. + """ + fn = import_symbol(path) + if not callable(fn): + _fail( + path=path, + kind=expected_kind, + rule="callable", + stage_config=stage_config, + message=f"resolved symbol is not callable: {fn!r}", + ) + raise AssertionError("unreachable") # pragma: no cover + + kind = infer_kind(fn, path=path) + validate_processor(fn, kind=kind, path=path, stage_config=stage_config) + + if expected_kind is not None and kind != expected_kind: + _fail( + path=path, + kind=kind, + rule="expected_kind", + stage_config=stage_config, + message=f"inferred kind {kind!r} does not match expected_kind {expected_kind!r}", + ) + raise AssertionError("unreachable") # pragma: no cover + + return ProcessorSpec(path=path, kind=kind, fn=fn) + + +# --------------------------------------------------------------------------- +# Dead-processor hint (gate three-state) +# --------------------------------------------------------------------------- + + +def dead_processor_hint( + kind: ProcessorKind, + *, + async_chunk: bool, + downstream_receives_async_chunks: bool, + has_sync: bool, +) -> bool: + """Whether an orchestrator-side ``custom_process_input_func`` is never invoked. + + This mirrors the orchestrator gate in ``engine/orchestrator._route_output`` + (``(not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1))``): + + * ``async_chunk=True`` and the downstream stage **receives** async chunks -> + the orchestrator skips ``process_engine_inputs`` for that transition, so + the input processor may be **dead** (returns ``True``). + * ``async_chunk=True`` and the downstream stage does **not** receive async + chunks -> the transition still runs forward, so the processor is alive + (returns ``False``). + * ``async_chunk=False`` -> forward path is active; the processor is only + dead when a ``sync_process_input_func`` overrides it (``has_sync=True``). + + Producer kinds are never orchestrator-side, so they are never "dead" under + this definition. + + This is a **pure decision helper** for warnings/tests only; the orchestrator + does not enforce it at runtime (M0 phase is warn-first). + """ + if kind in _PRODUCER_KINDS: + return False + if async_chunk: + return bool(downstream_receives_async_chunks) + return bool(has_sync) diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 3f94c16a725..4c89972379c 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -10,6 +10,15 @@ needs_omni_connector, should_accumulate_full_payload_output, ) +from vllm_omni.model_executor.stage_input_processors import ( + ProcessorValidationError, + resolve_processor, +) +from vllm_omni.outputs import OmniConnectorOutput +from vllm_omni.worker.payload_span import ( + get_tensor_span, + merge_tensor_spans, +) __all__ = [ "OmniConnectorModelRunnerMixin", @@ -25,3 +34,2269 @@ class OmniConnectorModelRunnerMixin(_OmniConnectorPayloadTransportMixin): Runtime ownership and payload transport live under ``distributed.omni_connectors``. """ +def should_accumulate_full_payload_output(model_config, custom_process_func) -> bool: + """Producer-side structural gate. + + Fires iff the stage explicitly declares a downstream full-payload + producer hook via ``custom_process_next_stage_input_func``. Consumer + stages may have ``custom_process_input_func`` values that can be + mechanically derived to ``*_full_payload`` helper names in the same + module; those are intentionally not enough to make the stage a producer. + """ + if custom_process_func is None: + return False + if getattr(model_config, "async_chunk", False): + return False + if getattr(model_config, "final_output", False): + return False + next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) + if not isinstance(next_stage_func, str) or not next_stage_func: + return False + return getattr(model_config, "model_stage", None) is not None + + +class OmniConnectorModelRunnerMixin: + """Unified data-plane communication mixin for Model Runners. + + Provides three transfer modes through a single pair of bg I/O threads: + - **full_payload_mode**: ``recv_full_payload_inputs`` / ``send_full_payload_outputs`` + - **Streaming (async_chunk)**: ``recv_chunk`` / ``send_chunk`` + - **KV cache**: ``send_kv_cache`` / ``recv_kv_cache`` (delegates to + the existing ``OmniKVTransferManager``) + + The mixin owns connector instances and background threads. It never + touches scheduling queues -- readiness is communicated to the Scheduler + via ``OmniConnectorOutput``. + """ + + # ------------------------------------------------------------------ # + # Init / Shutdown + # ------------------------------------------------------------------ # + + def init_omni_connectors( + self, + model_config: OmniModelConfig, + kv_transfer_manager: OmniKVTransferManager | None = None, + ) -> None: + """Initialize connectors and background threads. + + Args: + model_config: Stage-level model config with connector settings. + kv_transfer_manager: Existing KV transfer manager to delegate to. + """ + self._omni_connector: OmniConnectorBase | None = ( + self._create_connector(model_config) if _should_create_payload_connector(model_config) else None + ) + self._kv_transfer_manager = kv_transfer_manager + + self._async_chunk: bool = getattr(model_config, "async_chunk", False) + self._model_mode: str = getattr(model_config, "worker_type", "ar") + stage_id = getattr(model_config, "stage_id", 0) + if isinstance(stage_id, str): + stage_id = int(stage_id) + self._stage_id: int = stage_id if isinstance(stage_id, int) else 0 + + self._custom_process_func_path, self._custom_process_func = self._load_custom_func(model_config) + self._custom_process_supports_is_finished = self._custom_process_supports_is_finished_kwarg() + logger.debug( + "[Stage-%s] init_omni_connectors: async_chunk=%s, custom_process_func=%s, connector=%s, func_path=%s", + self._stage_id, + self._async_chunk, + self._custom_process_func, + type(self._omni_connector).__name__ if self._omni_connector else None, + self._custom_process_func_path, + ) + + # -- next stage ID (from connector config or default stage_id + 1) -- + self._next_stage_id: int = self._resolve_next_stage_id(model_config) + + # -- heterogeneous TP rank support -- + rank_cfg = self._parse_rank_mapping(model_config) + if self._kv_transfer_manager is not None: + topology = getattr(self._kv_transfer_manager, "tp_topology", None) + effective_mapping = ( + getattr(topology, "source_tp_size", None), + getattr(topology, "target_tp_size", None), + getattr(topology, "local_rank", None), + ) + if all(isinstance(value, int) for value in effective_mapping): + rank_cfg = { + "from_tp": effective_mapping[0], + "to_tp": effective_mapping[1], + "local_rank": effective_mapping[2], + } + self._from_tp: int = rank_cfg["from_tp"] + self._to_tp: int = rank_cfg["to_tp"] + self._local_rank: int = rank_cfg["local_rank"] + if self._kv_transfer_manager is not None: + self._kv_transfer_manager.kv_send_key_builder = self.get_rank_aware_kv_send_keys + self._kv_transfer_manager.kv_recv_key_builder = self.get_rank_aware_kv_keys + self._kv_transfer_manager.kv_payload_merger = self._merge_rank_sharded_kv_payloads + self._kv_transfer_manager.kv_payload_slicer = self._slice_rank_sharded_kv_payload + + # -- chunk index tracking (ported from OmniChunkTransferAdapter) -- + self._put_req_chunk: dict[str, int] = defaultdict(int) + self._get_req_chunk: dict[str, int] = defaultdict(int) + # Segment-local chunk counter: incremented alongside _put_req_chunk + # and popped at request cleanup. Note: the mixin path (uniproc mode) + # does not have segment boundary infrastructure; multi-segment support + # is only available via chunk_transfer_adapter (distributed path). + self._ramp_chunk_count: dict[str, int] = defaultdict(int) + # Send-side async accumulation / staging buffer. Receive-side payload + # ownership lives in ``_local_stage_payload_cache``. + self._send_side_request_payload: dict[str, dict[str, Any]] = {} + self._code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list) + self._cached_ic: dict[str, int] = {} + self._request_ids_mapping: dict[str, str] = {} + + # -- async I/O state (shared by chunk + full_payload_mode) -- + self._pending_load_reqs: dict[str, Any] = {} + self._finished_load_reqs: set[str] = set() + self._pending_save_reqs: dict[str, deque] = {} + self._pending_save_counts: dict[str, int] = defaultdict(int) + self._deferred_send_cleanup: set[str] = set() + # -- per-cycle output accumulator -- + self._chunk_ready_req_ids: set[str] = set() + self._chunk_finished_req_ids: set[str] = set() + self._stage_recv_req_ids: set[str] = set() + self._full_payload_pending_broadcast_req_ids: set[str] = set() + self._async_chunk_updated_req_ids: set[str] = set() + + # -- Model Runner local payload cache (RFC §2.4) -- + # Full stage payloads land here first on the recv side. We + # intentionally do not write connector recv results straight into + # `model_intermediate_buffer`: runner-owned runtime state is + # materialized later by `_sync_local_stage_payloads()` on the + # model thread. This keeps recv timing separate from execute-step + # visibility and avoids mixing connector I/O with model runtime + # ownership. + self._local_stage_payload_cache: dict[str, dict[str, Any]] = {} + # Lightweight scheduling metadata pending delivery to the Scheduler. + self._local_request_metadata: dict[str, dict[str, Any]] = {} + + # -- persistent set of request IDs whose chunk stream is complete -- + # Prevents re-registration after the finish sentinel has been received. + self._chunk_stream_completed: set[str] = set() + + # -- full_payload_mode: accumulate latest pooler_output per request, + # send only when the request finishes (next-cycle flush) -- + self._pending_full_payload_send: dict[str, tuple[Any, ...]] = {} + + # -- KV sent accumulator -- + self._kv_sent_req_ids: list[str] = [] + + # -- KV transfer lifecycle (absorbed from scheduler) -- + # Requests marked for KV transfer: {req_id: {seq_len, block_ids}} + self._kv_pending_transfers: dict[str, dict[str, Any]] = {} + # Requests whose KV transfer has been submitted but not yet acked + self._kv_active_transfers: set[str] = set() + # Requests whose KV transfer is complete (acked by kv_extracted_req_ids) + self._kv_completed_transfers: set[str] = set() + # Dedup guard: requests that have already triggered KV transfer + self._kv_triggered_requests: set[str] = set() + + self._lock = threading.Lock() + self._stop_event = threading.Event() + self._work_available = threading.Event() + + # Start background threads only when there's a connector + self._recv_thread: threading.Thread | None = None + self._save_thread: threading.Thread | None = None + if self._omni_connector is not None: + self._recv_thread = threading.Thread( + target=self._recv_loop, + daemon=True, + name="omni-mixin-recv", + ) + self._recv_thread.start() + self._save_thread = threading.Thread( + target=self._save_loop, + daemon=True, + name="omni-mixin-save", + ) + self._save_thread.start() + + # Explicit "fully initialised" marker so other parts of the runner + # (e.g. _update_states cleanup) can branch on a stable contract + # instead of probing for private mixin attribute names. Must be set + # only after every field above has been bound, so a partially + # constructed mixin is never observable as initialised. + self._omni_connector_initialized = True + + def shutdown_omni_connectors(self) -> None: + """Stop background threads and release connector resources.""" + self._stop_event.set() + if self._recv_thread is not None: + self._recv_thread.join(timeout=5) + if self._save_thread is not None: + self._save_thread.join(timeout=5) + if self._omni_connector is not None: + try: + self._omni_connector.close() + except Exception: + pass + + def cleanup_finished_request(self, req_id: str) -> None: + """Clean up per-request state after a request is fully finished. + + Call this when a request is freed from the model runner to prevent + memory leaks in the mixin's tracking dicts/sets. + + Two senders use different keys: ``send_chunk`` keys per-request + state under the EXTERNAL id (after mapping resolution), while + ``send_full_payload_outputs`` keys under the INTERNAL id. To cover + both modes (and forward compat with id-rename scenarios) we attempt + cleanup against both keys; the entry that doesn't exist for the + active mode is a no-op pop. Only the key that actually has pending + saves is added to ``_deferred_send_cleanup`` so the bg save's + decrement path drains it without leaving orphans. + """ + # Force-flush any pending full-payload accumulator entry before + # cleanup proceeds. Without this, finished requests with no + # downstream consumer (e.g. text-only on multi-modal arch) leave + # the entry orphaned in _pending_full_payload_send across requests, + # which empirically destabilises subsequent thinker forwards by + # making prefix-cache reuse observe stale accumulator state. The + # flush is idempotent when the entry has already been flushed by the + # scheduler-driven path, but this cleanup path runs for every request, + # so skip it entirely when the request never accumulated a payload. + if req_id in self._pending_full_payload_send: + try: + self.flush_full_payload_outputs({req_id}) + except Exception: + # Cleanup must still proceed regardless of flush errors here -- + # we already gated on ``_omni_connector_initialized`` upstream, + # so any exception here reflects a real connector-side issue + # (shared memory corruption, background thread crash) worth + # surfacing rather than silently swallowing. + logger.warning( + "flush_full_payload_outputs(%s) raised during cleanup; continuing tear-down.", + req_id, + exc_info=True, + ) + + ext_id = self._request_ids_mapping.pop(req_id, None) + keys_to_clean: list[str] = [req_id] + if ext_id is not None and ext_id != req_id: + keys_to_clean.append(ext_id) + + with self._lock: + keys_pending = [k for k in keys_to_clean if self._pending_save_counts.get(k, 0)] + for k in keys_pending: + self._deferred_send_cleanup.add(k) + for k in keys_to_clean: + if k in keys_pending: + continue + self._put_req_chunk.pop(k, None) + self._send_side_request_payload.pop(k, None) + self._code_prompt_token_ids.pop(k, None) + self._cached_ic.pop(k, None) + self._ramp_chunk_count.pop(k, None) + self._kv_pending_transfers.pop(req_id, None) + self._kv_active_transfers.discard(req_id) + self._kv_completed_transfers.discard(req_id) + self._kv_triggered_requests.discard(req_id) + self._cleanup_recv_delivery_state(req_id) + + def drop_inactive_request_delivery_state(self, req_id: str) -> None: + """Clear recv-side state for inactive requests.""" + ext_id = self._request_ids_mapping.pop(req_id, None) + if hasattr(self, "_lock"): + with self._lock: + self._drop_send_side_payload_state(req_id, ext_id) + else: + self._drop_send_side_payload_state(req_id, ext_id) + self._cleanup_recv_delivery_state(req_id) + + def _drop_send_side_payload_state(self, req_id: str, ext_id: str | None) -> None: + if ext_id is not None: + self._send_side_request_payload.pop(ext_id, None) + self._cached_ic.pop(ext_id, None) + self._send_side_request_payload.pop(req_id, None) + self._cached_ic.pop(req_id, None) + + def _cleanup_recv_delivery_state(self, req_id: str) -> None: + """Clear recv-side delivery-cycle state.""" + if hasattr(self, "_lock"): + with self._lock: + self._clear_recv_delivery_state(req_id) + else: + self._clear_recv_delivery_state(req_id) + + def _clear_recv_delivery_state(self, req_id: str) -> None: + self._get_req_chunk.pop(req_id, None) + self._pending_load_reqs.pop(req_id, None) + self._finished_load_reqs.discard(req_id) + self._chunk_ready_req_ids.discard(req_id) + self._chunk_finished_req_ids.discard(req_id) + self._chunk_stream_completed.discard(req_id) + self._stage_recv_req_ids.discard(req_id) + self._full_payload_pending_broadcast_req_ids.discard(req_id) + self._async_chunk_updated_req_ids.discard(req_id) + self._local_stage_payload_cache.pop(req_id, None) + self._local_request_metadata.pop(req_id, None) + + def prune_inactive_requests(self, active_req_ids: Any) -> set[str]: + """Drop connector state for requests that no longer exist locally. + + Preempted / unscheduled requests are expected to stay in + ``self.requests`` and therefore remain untouched. This only prunes + stale request IDs that have already fallen out of the active request + map, preventing background recv/send bookkeeping from outliving the + request lifecycle. + """ + if active_req_ids is None: + return set() + + active_req_ids = set(active_req_ids) + pending_req_ids = set(getattr(self, "_pending_load_reqs", {}).keys()) + received_req_ids = set(getattr(self, "_stage_recv_req_ids", set())) + received_req_ids.update(getattr(self, "_full_payload_pending_broadcast_req_ids", set())) + received_req_ids.update(getattr(self, "_local_request_metadata", {}).keys()) + # Pending recv requests may not yet be in the caller's active set + # (e.g. WAITING_FOR_CHUNK requests live in the coordinator's internal + # queues, not in model runner self.requests). Protect them so that + # legitimate waiting requests are not pruned. + # + # Likewise, a full payload can arrive on the background recv thread + # after the scheduler_output snapshot for the current execute_model() + # cycle was already materialized. Those requests may briefly live only + # in recv-side buffers/local cache until the next scheduler cycle wakes + # them up; pruning them here drops the payload before stage_recv can be + # published. + active_req_ids.update(pending_req_ids) + active_req_ids.update(received_req_ids) + stale_req_ids: set[str] = set() + + # NOTE: _pending_load_reqs is excluded from the scan list because + # all its entries are unconditionally protected above. The mixin + # cannot distinguish a legitimately-waiting pending recv from an + # orphaned one (only the coordinator/scheduler knows). + # + # Requests with freshly received full payloads / local stage payloads + # are also protected above. Their scheduler wake-up may lag the recv + # thread by one execute_model() cycle, especially when the request was + # added after the current scheduler_output snapshot. + # + # Orphaned pending recv entries (e.g. from upstream stage crash) are + # handled by collect_timed_out_request_ids() -- on + # OmniSchedulingCoordinator for full-payload requests, and on + # OmniChunkTransferAdapter for async-chunk ones -- which detect + # wait-time violations. The scheduler then removes the request from + # its queues, sets FINISHED_ERROR, and calls _free_request() which + # ultimately triggers cleanup_finished_request() here. + for attr_name in ( + "_request_ids_mapping", + "_get_req_chunk", + "_finished_load_reqs", + "_chunk_ready_req_ids", + "_chunk_finished_req_ids", + "_chunk_stream_completed", + "_stage_recv_req_ids", + "_full_payload_pending_broadcast_req_ids", + "_async_chunk_updated_req_ids", + "_local_stage_payload_cache", + "_local_request_metadata", + "_kv_pending_transfers", + "_kv_active_transfers", + "_kv_completed_transfers", + "_kv_triggered_requests", + ): + state = getattr(self, attr_name, None) + if isinstance(state, dict): + stale_req_ids.update(req_id for req_id in state if req_id not in active_req_ids) + elif isinstance(state, set): + stale_req_ids.update(req_id for req_id in state if req_id not in active_req_ids) + + for req_id in stale_req_ids: + self.cleanup_finished_request(req_id) + + return stale_req_ids + + # ------------------------------------------------------------------ # + # Local payload cache (RFC §2.4 – Model Runner ownership) + # ------------------------------------------------------------------ # + + def put_local_stage_payload(self, req_id: str, payload: OmniPayload) -> None: + """Store a full stage payload in the local cache.""" + self._local_stage_payload_cache[req_id] = payload + + def get_local_stage_payload(self, req_id: str) -> OmniPayload | None: + """Read a stage payload without removing it.""" + return self._local_stage_payload_cache.get(req_id) + + def pop_local_stage_payload(self, req_id: str) -> OmniPayload | None: + """Remove and return a stage payload (consume after use).""" + return self._local_stage_payload_cache.pop(req_id, None) + + def put_local_request_metadata(self, req_id: str, metadata: dict[str, Any]) -> None: + """Store lightweight scheduling metadata for a request.""" + self._local_request_metadata[req_id] = metadata + + def get_local_request_metadata(self, req_id: str) -> dict[str, Any] | None: + """Retrieve scheduling metadata for a request.""" + return self._local_request_metadata.get(req_id) + + # ------------------------------------------------------------------ # + # Scheduling metadata extraction + # ------------------------------------------------------------------ # + + @classmethod + def _extract_scheduling_metadata(cls, payload: OmniPayload) -> dict[str, Any]: + """Extract only the fields the scheduler needs from a full payload.""" + extracted: dict[str, Any] = {} + meta = payload.get("meta") if isinstance(payload, dict) else None + meta = meta if isinstance(meta, dict) else {} + + if "next_stage_prompt_len" in meta: + extracted["next_stage_prompt_len"] = meta["next_stage_prompt_len"] + elif "next_stage_prompt_len" in payload: + logger.warning_once( + "legacy flat 'next_stage_prompt_len' key in payload; expected 'meta.next_stage_prompt_len'" + ) + extracted["next_stage_prompt_len"] = payload["next_stage_prompt_len"] + + audio_codes = cls._payload_audio_codes(payload) + if audio_codes is not None: + extracted["code_predictor_codes"] = audio_codes + + if "left_context_size" in meta: + extracted["left_context_size"] = meta["left_context_size"] + elif "left_context_size" in payload: + logger.warning_once("legacy flat 'left_context_size' key in payload; expected 'meta.left_context_size'") + + return extracted + + _NON_CONSUMABLE_PAYLOAD_KEYS: set[tuple[str, str]] = { + ("meta", "finished"), + ("meta", "override_keys"), + ("meta", "next_stage_prompt_len"), + ("meta", "left_context_size"), + ("ids", "output"), + ("embed", "decode_token_start"), + ("embed", "decode_token_end"), + } + + @staticmethod + def _payload_value_has_content(value: Any) -> bool: + if value is None: + return False + if isinstance(value, torch.Tensor): + return value.numel() > 0 + if isinstance(value, (list, tuple, dict, set)): + return len(value) > 0 + return True + + @staticmethod + def _payload_finished(payload: Any) -> bool: + if not isinstance(payload, dict): + return False + if "finished" in payload: + logger.warning_once("legacy flat 'finished' key in payload; expected 'meta.finished'") + meta = payload.get("meta") + if not isinstance(meta, dict) or "finished" not in meta: + return False + flag = meta["finished"] + if isinstance(flag, torch.Tensor): + return flag.numel() == 1 and bool(flag.item()) + return bool(flag) + + @staticmethod + def _payload_audio_codes(payload: Any) -> Any: + if not isinstance(payload, dict): + return None + if "code_predictor_codes" in payload: + logger.warning_once("legacy flat 'code_predictor_codes' key in payload; expected 'codes.audio'") + codes = payload.get("codes") + if isinstance(codes, dict): + return codes.get("audio") + return None + + @classmethod + def _payload_is_consumable(cls, payload: OmniPayload | None) -> bool: + """Return True when an async payload can drive a real forward step. + + Metadata-only wake-ups should not transition WAITING_FOR_CHUNK requests + back to schedulable state. In particular, a widened token horizon without + any newly visible thinker decode embeds should not force a placeholder-only + talker decode step. + """ + if not isinstance(payload, dict) or not payload: + return False + + embed = payload.get("embed") + if isinstance(embed, dict): + decode_embeddings = embed.get("decode") + if isinstance(decode_embeddings, torch.Tensor): + if decode_embeddings.ndim == 0: + return True + return decode_embeddings.numel() > 0 and decode_embeddings.shape[0] > 0 + + audio_codes = cls._payload_audio_codes(payload) + if audio_codes is not None: + if isinstance(audio_codes, torch.Tensor): + return audio_codes.numel() > 0 + if hasattr(audio_codes, "__len__"): + return len(audio_codes) > 0 + return True + + for key, value in payload.items(): + if isinstance(value, dict): + for sk, sv in value.items(): + if (key, sk) in cls._NON_CONSUMABLE_PAYLOAD_KEYS: + continue + if cls._payload_value_has_content(sv): + return True + continue + if cls._payload_value_has_content(value): + return True + return False + + @staticmethod + def _get_local_tp_group() -> Any | None: + """Return the local TP group when tensor parallelism is initialized.""" + try: + return get_tp_group() + except Exception: + return None + + def _recv_ordinary_stage_result( + self, + connector: OmniConnectorBase, + from_stage: str, + to_stage: str, + connector_get_key: str, + ) -> Any: + """Receive one ordinary non-KV stage payload on the local leader rank only.""" + tp_group = self._get_local_tp_group() + if tp_group is None or getattr(tp_group, "world_size", 1) <= 1: + return connector.get(from_stage, to_stage, connector_get_key) + if not self.is_data_transfer_rank(): + return None + return connector.get(from_stage, to_stage, connector_get_key) + + def _recv_full_payload_result( + self, + connector: OmniConnectorBase, + from_stage: str, + to_stage: str, + connector_get_key: str, + ) -> Any: + """Receive one full-payload transfer on the local leader rank only.""" + return self._recv_ordinary_stage_result( + connector, + from_stage, + to_stage, + connector_get_key, + ) + + def _recv_async_chunk_result( + self, + connector: OmniConnectorBase, + from_stage: str, + to_stage: str, + connector_get_key: str, + ) -> Any: + """Receive one ordinary async chunk on the local leader rank only.""" + return self._recv_ordinary_stage_result( + connector, + from_stage, + to_stage, + connector_get_key, + ) + + @staticmethod + def _snapshot_payload(payload: Any) -> Any: + if isinstance(payload, dict): + return dict(payload) + return payload + + def _broadcast_tp_payload_packet(self, packet: Any) -> Any: + """Broadcast one ordinary payload packet from TP rank 0 when TP is active.""" + tp_group = self._get_local_tp_group() + if tp_group is None or getattr(tp_group, "world_size", 1) <= 1: + return packet + leader_packet = packet if self.is_data_transfer_rank() else None + return tp_group.broadcast_object(leader_packet, src=0) + + def _apply_staged_payloads_locked(self, staged_payloads: dict[str, Any]) -> None: + for req_id, payload in staged_payloads.items(): + self._local_stage_payload_cache[req_id] = self._snapshot_payload(payload) + + def _collect_full_payload_results_locked(self) -> dict[str, Any] | None: + if not self._full_payload_pending_broadcast_req_ids: + return None + results: dict[str, Any] = {} + missing_req_ids: list[str] = [] + for req_id in tuple(self._full_payload_pending_broadcast_req_ids): + payload = self._local_stage_payload_cache.get(req_id) + if payload is None: + missing_req_ids.append(req_id) + continue + results[req_id] = self._snapshot_payload(payload) + self._full_payload_pending_broadcast_req_ids.discard(req_id) + if missing_req_ids: + logger.warning( + "[Stage-%s] _collect_full_payload_results_locked: " + "pending full-payload reqs missing from local cache: %s", + self._stage_id, + missing_req_ids, + ) + return results or None + + def _collect_async_chunk_fanout_packet_locked(self) -> dict[str, Any] | None: + payload_req_ids = set(self._async_chunk_updated_req_ids) + payload_req_ids.update(self._finished_load_reqs) + payload_req_ids.update(self._chunk_finished_req_ids) + payload_req_ids.update(self._local_request_metadata) + if not ( + payload_req_ids or self._finished_load_reqs or self._chunk_finished_req_ids or self._local_request_metadata + ): + return None + + staged_payloads = { + req_id: self._snapshot_payload(self._local_stage_payload_cache[req_id]) + for req_id in payload_req_ids + if req_id in self._local_stage_payload_cache + } + packet = { + "staged_payloads": staged_payloads, + "request_metadata": dict(self._local_request_metadata), + "newly_finished": set(self._finished_load_reqs), + "chunk_finished": set(self._chunk_finished_req_ids), + } + + self._async_chunk_updated_req_ids.clear() + self._finished_load_reqs.clear() + self._chunk_finished_req_ids.clear() + self._local_request_metadata.clear() + + for req_id in packet["chunk_finished"]: + if req_id not in self._local_stage_payload_cache: + continue + ext_req_id = self._request_ids_mapping.get(req_id, req_id) + self._send_side_request_payload.pop(ext_req_id, None) + if ext_req_id != req_id: + self._send_side_request_payload.pop(req_id, None) + + return packet + + def _apply_async_chunk_fanout_packet(self, packet: dict[str, Any]) -> None: + staged_payloads = packet.get("staged_payloads", {}) + chunk_finished = set(packet.get("chunk_finished", ())) + with self._lock: + self._apply_staged_payloads_locked(staged_payloads) + for req_id in chunk_finished: + self._pending_load_reqs.pop(req_id, None) + self._chunk_stream_completed.add(req_id) + + # ------------------------------------------------------------------ # + # full_payload_mode (recv_full_payload_inputs / send_full_payload_outputs) + # ------------------------------------------------------------------ # + + def recv_full_payload_inputs(self, scheduler_output: Any) -> dict[str, Any] | None: + """Check for incoming full_payload_mode stage inputs (non-blocking). + + Returns a dict mapping ``request_id -> engine_inputs`` for data + that has arrived, or ``None`` if nothing is ready. Stores full + payloads in the local cache and extracts scheduling metadata. + """ + # Fast path: when TP is trivial (no peer ranks waiting on a broadcast) + # and the bg recv thread has not staged anything, skip the lock + TP + # broadcast cycle entirely. _broadcast_tp_payload_packet already + # returns its input unchanged under the same world_size<=1 condition, + # so the original code path was a no-op here on every empty step. + tp_group = self._get_local_tp_group() + if ( + tp_group is None or getattr(tp_group, "world_size", 1) <= 1 + ) and not self._full_payload_pending_broadcast_req_ids: + return None + with self._lock: + results = self._collect_full_payload_results_locked() if self.is_data_transfer_rank() else None + results = self._broadcast_tp_payload_packet(results) + if not results: + return None + with self._lock: + self._stage_recv_req_ids.update(results.keys()) + for req_id in results: + self._pending_load_reqs.pop(req_id, None) + self._apply_staged_payloads_locked(results) + for req_id, payload in results.items(): + self._local_request_metadata[req_id] = self._extract_scheduling_metadata(payload) + logger.debug( + "[Stage-%s] recv_full_payload_inputs: consumed %s reqs: %s, stage_recv_req_ids now=%s", + self._stage_id, + len(results), + list(results.keys()), + self._stage_recv_req_ids, + ) + return results + + def _get_model_config(self) -> Any: + model_config = getattr(self, "model_config", None) + if model_config is not None: + return model_config + return getattr(getattr(self, "vllm_config", None), "model_config", None) + + def _should_accumulate_full_payload_output(self) -> bool: + """Gate send-side full-payload output accumulation only. + + Cached per instance: the result depends only on model_config / + _custom_process_func, both of which are set at init time. Avoid + the per-step dynamic import inside the model decode loop. + """ + if getattr(self, "_omni_connector", None) is None: + # No connector at all: send_full_payload_outputs would no-op. + # Skip the per-step accumulator+build that would otherwise be + # silently discarded. Defends against a terminal stage whose + # custom_process_input_func has a *_full_payload derivative in + # the same module (e.g. dynin stage 2 token2image_to_token2audio + # in pipelines that don't configure any connector at all). + # + # Known limitation: a *terminal-consumer* stage that has a + # connector configured for receiving upstream input is NOT + # caught here -- ``_omni_connector`` is non-None for it, and + # ``_load_custom_func`` may still resolve a ``*_full_payload`` + # derivative from this stage's ``custom_process_input_func``. + # In that case the accumulator builds payloads that + # ``send_full_payload_outputs`` later drops via its own + # connector-side checks (wasted CPU, not a functional bug). + # A topology-aware gate (explicit producer field or pipeline + # is_terminal info) would close the gap; that change is out + # of scope for this PR. + self._should_accumulate_full_payload_output_cached = False + return False + cached = getattr(self, "_should_accumulate_full_payload_output_cached", None) + if cached is not None: + return cached + model_config = self._get_model_config() + if model_config is None: + self._should_accumulate_full_payload_output_cached = False + return False + result = should_accumulate_full_payload_output( + model_config, + getattr(self, "_custom_process_func", None), + ) + self._should_accumulate_full_payload_output_cached = result + return result + + @staticmethod + def _new_full_payload_accumulator(output: dict[str, Any]): + chunks: dict[str, list[torch.Tensor]] = {} + latest: dict[str, Any] = {} + rows: dict[str, int] = {} + for k, v in output.items(): + if isinstance(v, torch.Tensor) and v.dim() >= 2: + chunks[k] = [v] + rows[k] = int(v.shape[0]) + else: + latest[k] = v + return chunks, latest, rows + + @staticmethod + def _materialize_full_payload_entry(entry): + if len(entry) == 2: + return entry + chunks, latest, _rows, request = entry + output = dict(latest) + for k, tensors in chunks.items(): + if tensors: + output[k] = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0) + return output, request + + def _resolve_full_payload_replace_keys(self) -> frozenset: + """Per-model REPLACE-key set for the full-payload accumulator. + + Looked up from the stage-input-processor module that ships the model's sync builder + (`model_config.custom_process_input_func.__module__`). The module + declares ``_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str]``; if absent, + returns the empty set. + + Cached per instance. Keys in this set use REPLACE semantics in the + accumulator (subsequent emissions discard prior chunks) instead of + the default CONCAT semantics. Use for tensors that carry the full + result so far rather than per-step deltas (e.g. ``model_outputs``). + """ + cached = getattr(self, "_full_payload_replace_keys_cached", None) + if cached is not None: + return cached + proc = getattr(self, "_custom_process_func", None) + if proc is None: + self._full_payload_replace_keys_cached = frozenset() + return self._full_payload_replace_keys_cached + module_name = getattr(proc, "__module__", None) + if module_name is None: + self._full_payload_replace_keys_cached = frozenset() + return self._full_payload_replace_keys_cached + try: + import sys as _sys + + mod = _sys.modules.get(module_name) or importlib.import_module(module_name) + keys = getattr(mod, "_FULL_PAYLOAD_REPLACE_KEYS", frozenset()) + except ImportError: + logger.debug( + "Could not import stage input processor module %s while resolving " + "_FULL_PAYLOAD_REPLACE_KEYS; using CONCAT semantics for all keys.", + module_name, + exc_info=True, + ) + keys = frozenset() + if not isinstance(keys, (frozenset, set)): + logger.debug( + "Ignoring non-set _FULL_PAYLOAD_REPLACE_KEYS from %s: %s", + module_name, + type(keys).__name__, + ) + keys = frozenset() + self._full_payload_replace_keys_cached = frozenset(keys) + logger.debug( + "Resolved _FULL_PAYLOAD_REPLACE_KEYS for %s: %s", + module_name, + sorted(self._full_payload_replace_keys_cached), + ) + return self._full_payload_replace_keys_cached + + def accumulate_full_payload_output( + self, + req_id: str, + pooler_output: Any, + request: Any, + ) -> None: + """Accumulate pooler_output for a request across steps (full_payload_mode). + + Per-token tensors (2-D+, matching trailing dims) are concatenated + along dim-0. Scalar / global tensors (1-D or 0-D) are replaced + with the latest value. + + Note: codec rows are NOT filtered for zero placeholders here. The + downstream consumer ``_extract_qwen3_full_payload_codec_rows`` crops + codec rows using ``output_token_ids`` as the authoritative source, + which makes any sender-side zero filtering redundant. Skipping the + sender-side ``t.any()`` scan also avoids a per-tensor GPU->CPU device + sync that stalled the decode pipeline. + + The data is actually sent when ``flush_full_payload_outputs`` is called + with the finished request IDs from the next scheduler cycle. + """ + replace_keys = self._resolve_full_payload_replace_keys() + existing = self._pending_full_payload_send.get(req_id) + + if existing is None: + chunks, latest, rows = self._new_full_payload_accumulator(pooler_output) + self._pending_full_payload_send[req_id] = (chunks, latest, rows, request) + return + + if len(existing) == 2: + chunks, latest, rows = self._new_full_payload_accumulator(existing[0]) + else: + chunks, latest, rows, _ = existing + + for k, v in pooler_output.items(): + if v is None: + continue + if k in replace_keys: + # Explicit REPLACE semantics: the new value supersedes any + # prior chunks (e.g. `model_outputs` carries the full result + # so far, not an appendable per-step delta). + latest.pop(k, None) + if isinstance(v, torch.Tensor) and v.dim() >= 2: + chunks[k] = [v] + rows[k] = int(v.shape[0]) + else: + chunks.pop(k, None) + rows.pop(k, None) + latest[k] = v + continue + if isinstance(v, torch.Tensor) and v.dim() >= 2: + if k in chunks and chunks[k] and v.shape[1:] == chunks[k][0].shape[1:]: + chunks[k].append(v) + rows[k] += int(v.shape[0]) + else: + latest.pop(k, None) + chunks[k] = [v] + rows[k] = int(v.shape[0]) + else: + chunks.pop(k, None) + rows.pop(k, None) + latest[k] = v + + self._pending_full_payload_send[req_id] = (chunks, latest, rows, request) + + def flush_full_payload_outputs(self, finished_req_ids: set[str]) -> None: + """Send accumulated full_payload outputs for requests that just finished.""" + pending_req_ids = set(self._pending_full_payload_send.keys()) + if not (finished_req_ids & pending_req_ids): + return + + logger.debug( + "[Stage-%s] flush_full_payload_outputs: finished_req_ids=%s, pending=%s", + self._stage_id, + finished_req_ids, + list(self._pending_full_payload_send.keys()), + ) + to_send: dict[str, tuple[Any, Any]] = {} + for req_id in finished_req_ids: + entry = self._pending_full_payload_send.pop(req_id, None) + if entry is not None: + to_send[req_id] = self._materialize_full_payload_entry(entry) + logger.debug("[Stage-%s] flush_full_payload_outputs: to_send=%s", self._stage_id, list(to_send.keys())) + if to_send: + self.send_full_payload_outputs(scheduler_output=None, outputs=to_send) + + def send_full_payload_outputs( + self, + scheduler_output: Any, + outputs: dict[str, tuple[Any, Any] | Any], + ) -> list[str]: + """Send full_payload stage outputs to the next stage via connector. + + Args: + outputs: Mapping of ``req_id`` to either a + ``(pooling_output, request)`` tuple (preferred) or a raw + payload dict. When a tuple is supplied the request object + is forwarded to ``custom_process_stage_input_func``. + + Returns list of request IDs successfully enqueued. + """ + if self._omni_connector is None: + logger.debug("[Stage-%s] send_full_payload_outputs: connector is None, skip", self._stage_id) + return [] + if not self.is_data_transfer_rank(): + logger.debug( + "[Stage-%s] send_full_payload_outputs: not data_transfer_rank (rank=%s), skip", + self._stage_id, + self._local_rank, + ) + return list(outputs.keys()) + sent_ids: list[str] = [] + next_stage_id = self._next_stage_id + for req_id, value in outputs.items(): + if isinstance(value, tuple) and len(value) == 2: + raw_output, request = value + else: + raw_output, request = value, None + + payload = raw_output + if self._custom_process_func is not None: + payload = self._build_custom_process_payload( + request_id=req_id, + request=request, + pooling_output=raw_output, + ) + if payload is None: + continue + if payload is None: + logger.debug("[Stage-%s] send_full_payload_outputs: payload is None for %s", self._stage_id, req_id) + continue + if isinstance(payload, dict): + audio_codes = self._payload_audio_codes(payload) + if isinstance(audio_codes, torch.Tensor): + code_len = int(audio_codes.numel()) + elif hasattr(audio_codes, "__len__"): + code_len = len(audio_codes) + else: + code_len = None + meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {} + logger.debug( + "[Stage-%s] send_full_payload_outputs: req=%s payload_keys=%s code_len=%s left_context_size=%s", + self._stage_id, + req_id, + sorted(payload.keys()), + code_len, + meta.get("left_context_size"), + ) + + external_req_id = self._resolve_external_req_id(request, req_id) + chunk_id = self._put_req_chunk[req_id] + self._put_req_chunk[req_id] += 1 + self._ramp_chunk_count[req_id] += 1 + connector_put_key = f"{external_req_id}_{self._stage_id}_{chunk_id}" + + logger.debug( + "[Stage-%s] send_full_payload_outputs: enqueue req=%s put_key=%s next_stage=%s", + self._stage_id, + req_id, + connector_put_key, + next_stage_id, + ) + task = { + "stage_id": self._stage_id, + "next_stage_id": next_stage_id, + "put_key": connector_put_key, + "data": payload, + "request_id": req_id, + } + with self._lock: + self._pending_save_reqs.setdefault(req_id, deque()).append(task) + self._pending_save_counts[req_id] += 1 + sent_ids.append(req_id) + if sent_ids: + self._work_available.set() + return sent_ids + + # ------------------------------------------------------------------ # + # Streaming chunk mode (recv_chunk / send_chunk) + # ------------------------------------------------------------------ # + + def register_chunk_recv(self, request: Any) -> None: + """Register a request for async chunk retrieval by the bg thread. + + Stage-0 has no upstream producer so this is a no-op there. + Skips requests whose batch data has already been received to + prevent the bg thread from polling for non-existent chunks. + """ + if self._stage_id == 0: + return + request_id = request.request_id + # Explicit external_req_id=None must fall back to request_id; + # otherwise recv keys become `None__` and collide + # across requests. + ext = getattr(request, "external_req_id", None) + self._request_ids_mapping[request_id] = ext if ext is not None else request_id + with self._lock: + if request_id in self._stage_recv_req_ids: + return + # Don't re-register if the finish sentinel was already received + if request_id in self._chunk_stream_completed: + return + self._pending_load_reqs[request_id] = request + self._work_available.set() + + def recv_chunk(self) -> dict[str, Any]: + """Collect chunks received by the bg thread since last call. + + Returns a dict ``{request_id: chunk_payload}`` for newly arrived + chunks. Empty dict when nothing is ready. + + This method reads from ``_finished_load_reqs`` without clearing + it -- ``get_omni_connector_output()`` is the sole consumer that + drains and resets ``_finished_load_reqs`` at the end of each + ``execute_model`` cycle. + + Returns **shallow copies** of the cached payloads so that the + caller can read them without racing against the background recv + thread, which may concurrently mutate the live cache entries via + ``dict.update()``. + """ + with self._lock: + finished = set(self._finished_load_reqs) + if not finished: + return {} + # Snapshot the payloads under the lock to avoid racing with + # _poll_single_request which does existing.update(payload_data) + # on the same dict objects. + result = {} + for rid in finished: + payload = self._local_stage_payload_cache.get(rid) + result[rid] = dict(payload) if isinstance(payload, dict) else payload + + self._chunk_ready_req_ids.update(finished) + return result + + def send_chunk( + self, + request: Any, + pooling_output: Any | None = None, + ) -> bool: + """Derive and enqueue one chunk for async sending. + + Payload extraction runs in the caller thread (via + ``custom_process_stage_input_func``); the actual + ``connector.put()`` is done by the background save thread. + Non-KV data is identical across TP ranks; only rank 0 sends. + """ + if self._omni_connector is None: + logger.warning("[Stage-%s] send_chunk: connector is None", self._stage_id) + return False + if not self.is_data_transfer_rank(): + return True + raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None) + request_id = self._resolve_external_req_id(request, raw_req_id) + # Cache the internal→external mapping so that finish sentinels can + # resolve the external ID even after the request is freed. + if raw_req_id and raw_req_id != request_id: + self._request_ids_mapping.setdefault(raw_req_id, request_id) + chunk_id = self._put_req_chunk[request_id] + + payload_data = self._build_custom_process_payload( + request_id=request_id, + request=request, + pooling_output=pooling_output, + ) + if payload_data is None: + if chunk_id == 0: + logger.warning( + "[Stage-%s] send_chunk: payload is None for req=%s chunk=%s (process_func=%s)", + self._stage_id, + request_id, + chunk_id, + self._custom_process_func, + ) + return False + + self._put_req_chunk[request_id] += 1 + self._ramp_chunk_count[request_id] += 1 + next_stage_id = self._next_stage_id + connector_put_key = f"{request_id}_{self._stage_id}_{chunk_id}" + + if chunk_id == 0: + logger.debug( + "[Stage-%s] send_chunk: first chunk enqueued, req=%s key=%s", + self._stage_id, + request_id, + connector_put_key, + ) + + task = { + "stage_id": self._stage_id, + "next_stage_id": next_stage_id, + "put_key": connector_put_key, + "data": payload_data, + "request_id": request_id, + } + with self._lock: + self._pending_save_reqs.setdefault(request_id, deque()).append(task) + self._pending_save_counts[request_id] += 1 + self._work_available.set() + return True + + # ------------------------------------------------------------------ # + # KV cache (delegates to OmniKVTransferManager) + # ------------------------------------------------------------------ # + + def send_kv_cache( + self, + finished_reqs: dict[str, dict[str, Any]], + kv_caches: list[torch.Tensor], + block_size: int, + cache_dtype: str, + request_id_resolver: Any | None = None, + ) -> list[str]: + """Send KV cache for finished requests. + + Delegates to the existing ``OmniKVTransferManager``. + """ + if self._kv_transfer_manager is None: + return list(finished_reqs.keys()) if finished_reqs else [] + result = self._kv_transfer_manager.handle_finished_requests_kv_transfer( + finished_reqs=finished_reqs, + kv_caches=kv_caches, + block_size=block_size, + cache_dtype=cache_dtype, + request_id_resolver=request_id_resolver, + ) + if result: + self._kv_sent_req_ids.extend(result) + return result + + def recv_kv_cache( + self, + request_id: str, + target_device: torch.device | None = None, + ) -> tuple[dict[str, Any] | None, int]: + """Receive KV cache for a request. + + Delegates to the existing ``OmniKVTransferManager``. + """ + if self._kv_transfer_manager is None: + return None, 0 + return self._kv_transfer_manager.receive_kv_cache_for_request( + request_id=request_id, + target_device=target_device, + ) + + def receive_cfg_companion_kv_payloads( + self, + cfg_request_ids: dict[str, str], + target_device: torch.device | None = None, + ) -> dict[str, tuple[dict[str, Any] | None, int]]: + """Receive raw CFG companion KV payloads keyed by role.""" + return { + role: self.recv_kv_cache(companion_rid, target_device=target_device) + for role, companion_rid in cfg_request_ids.items() + } + + def receive_multi_kv_cache( + self, + req: Any, + cfg_kv_collect_func: Any | None = None, + target_device: torch.device | None = None, + ) -> bool: + """Receive primary and optional companion KV caches for a request. + + The mixin owns the runner-facing orchestration: primary KV receive, + companion payload fetch, and applying any model-specific CFG fields back + onto ``req.sampling_params``. + """ + if self._kv_transfer_manager is None: + return False + + request_id = getattr(req, "request_id", None) + if not request_id: + logger.warning("Request has no ID, cannot receive KV cache") + return False + + active_requests = getattr(self, "requests", None) + if active_requests is not None and request_id not in active_requests: + logger.debug("Skip receiving KV cache for inactive request %s", request_id) + return False + + primary_ok = False + data, _size = self.recv_kv_cache(request_id, target_device=target_device) + if data: + self._kv_transfer_manager.apply_kv_cache_to_request(req, data) + primary_ok = True + + cfg_ids = getattr(getattr(req, "sampling_params", None), "cfg_kv_request_ids", None) + if cfg_ids and cfg_kv_collect_func: + try: + cfg_role_payloads = self.receive_cfg_companion_kv_payloads( + cfg_ids, + target_device=target_device, + ) + cfg_kvs = cfg_kv_collect_func(request_id, cfg_role_payloads) + if cfg_kvs and hasattr(req, "sampling_params") and req.sampling_params is not None: + for key, value in cfg_kvs.items(): + setattr(req.sampling_params, key, value) + logger.debug("Applied CFG KV caches: %s", list(cfg_kvs.keys())) + except Exception: + logger.exception("Failed to collect CFG KV caches for %s", request_id) + + return primary_ok + + # ------------------------------------------------------------------ # + # Rank-aware KV transfer routing + # ------------------------------------------------------------------ # + + def get_rank_aware_kv_keys( + self, + req_id: str, + from_stage: int, + to_stage: int | None = None, + chunk_id: int = 0, + ) -> list[str]: + """Build recv-side connector keys for all remote ranks this rank needs. + + For heterogeneous TP receive, the local rank is the target rank and must + fetch one or more source-rank shards keyed as ``from_rank -> to_rank``. + """ + if self._from_tp <= 1 and self._to_tp <= 1: + resolved_to_stage = self._next_stage_id if to_stage is None else to_stage + return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] + + remote_ranks = self.get_kv_remote_ranks() + return [ + self.get_kv_connector_key( + req_id=req_id, + from_stage=from_stage, + chunk_id=chunk_id, + from_rank=remote_rank, + to_rank=self._local_rank, + ) + for remote_rank in remote_ranks + ] + + def get_kv_target_ranks_for_send(self) -> list[int]: + """Determine which target ranks this local rank should send KV shards to.""" + self._validate_kv_tp_topology() + if self._from_tp == self._to_tp: + return [self._local_rank] + if self._from_tp > self._to_tp: + tp_ratio = self._from_tp // self._to_tp + return [self._local_rank // tp_ratio] + tp_ratio = self._to_tp // self._from_tp + base_rank = self._local_rank * tp_ratio + return [base_rank + i for i in range(tp_ratio)] + + def get_rank_aware_kv_send_keys( + self, + req_id: str, + from_stage: int, + to_stage: int | None = None, + chunk_id: int = 0, + ) -> list[str]: + """Build send-side connector keys for this rank's KV shard(s).""" + if self._from_tp <= 1 and self._to_tp <= 1: + resolved_to_stage = self._next_stage_id if to_stage is None else to_stage + return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] + + target_ranks = self.get_kv_target_ranks_for_send() + return [ + self.get_kv_connector_key( + req_id=req_id, + from_stage=from_stage, + chunk_id=chunk_id, + from_rank=self._local_rank, + to_rank=target_rank, + ) + for target_rank in target_ranks + ] + + @staticmethod + def _merge_rank_sharded_kv_payloads(payloads: list[dict[str, Any]]) -> dict[str, Any] | None: + """Merge multiple source-rank KV shards for one target rank.""" + payloads = [payload for payload in payloads if isinstance(payload, dict)] + if not payloads: + return None + if len(payloads) == 1: + return payloads[0] + + merged = dict(payloads[0]) + layer_blocks = merged.get("layer_blocks") + if not isinstance(layer_blocks, dict): + return merged + + def _merge_tensor_lists(name: str) -> list[torch.Tensor | None]: + merged_list: list[torch.Tensor | None] = [] + cache_lists = [payload.get("layer_blocks", {}).get(name, []) for payload in payloads] + max_len = max((len(cache_list) for cache_list in cache_lists), default=0) + for idx in range(max_len): + tensors = [cache_list[idx] for cache_list in cache_lists if idx < len(cache_list)] + tensors = [tensor for tensor in tensors if isinstance(tensor, torch.Tensor)] + if not tensors: + merged_list.append(None) + elif len(tensors) == 1: + merged_list.append(tensors[0]) + else: + merged_list.append(torch.cat(tensors, dim=-2).contiguous()) + return merged_list + + merged["layer_blocks"] = { + "key_cache": _merge_tensor_lists("key_cache"), + "value_cache": _merge_tensor_lists("value_cache"), + } + metadata = dict(merged.get("metadata", {})) + metadata["merged_remote_rank_count"] = len(payloads) + merged["metadata"] = metadata + return merged + + def _slice_rank_sharded_kv_payload(self, payload: dict[str, Any] | None) -> dict[str, Any] | None: + """Slice a duplicated source-rank KV shard for ``from_tp < to_tp`` cases.""" + if payload is None or self._from_tp >= self._to_tp: + return payload + + tp_ratio = self._to_tp // self._from_tp + shard_index = self._local_rank % tp_ratio + layer_blocks = payload.get("layer_blocks") if isinstance(payload, dict) else None + if not isinstance(layer_blocks, dict): + return payload + + def _slice_tensor_list(name: str) -> list[torch.Tensor | None]: + sliced: list[torch.Tensor | None] = [] + for tensor in layer_blocks.get(name, []): + if not isinstance(tensor, torch.Tensor) or tensor.ndim < 2: + sliced.append(tensor) + continue + head_dim = tensor.shape[-2] + if head_dim % tp_ratio != 0: + sliced.append(tensor) + continue + per_rank = head_dim // tp_ratio + start = shard_index * per_rank + sliced.append(tensor.narrow(-2, start, per_rank).contiguous()) + return sliced + + payload = dict(payload) + payload["layer_blocks"] = { + "key_cache": _slice_tensor_list("key_cache"), + "value_cache": _slice_tensor_list("value_cache"), + } + metadata = dict(payload.get("metadata", {})) + metadata["sliced_for_local_rank"] = self._local_rank + payload["metadata"] = metadata + return payload + + def should_replicate_payload(self) -> bool: + """Whether non-KV payloads should be replicated across ranks. + + Data payloads (stage inputs, chunks) are identical after all-gather, + so only rank 0 transfers them. KV payloads are rank-specific and + all ranks participate. + """ + return self._local_rank != 0 + + def get_kv_rank_mapping(self) -> dict[str, Any]: + """Return the current rank mapping configuration. + + Useful for debugging and for downstream code that needs to know + the TP topology without re-parsing model config. + """ + return { + "from_tp": self._from_tp, + "to_tp": self._to_tp, + "local_rank": self._local_rank, + "remote_ranks": self.get_kv_remote_ranks(), + "is_data_transfer_rank": self.is_data_transfer_rank(), + } + + # ------------------------------------------------------------------ # + # KV transfer lifecycle (RFC – mixin-owned) + # ------------------------------------------------------------------ # + + def mark_kv_transfer( + self, + req_id: str, + seq_len: int, + block_ids: list[int], + custom_metadata: dict[str, Any] | None = None, + ) -> None: + """Mark a request as needing KV cache transfer. + + Called by the scheduler when a transfer trigger fires. The mixin + owns the lifecycle from this point: pending → active → completed. + """ + if req_id in self._kv_pending_transfers: + return + self._kv_triggered_requests.add(req_id) + transfer = { + "seq_len": seq_len, + "block_ids": block_ids, + } + if custom_metadata is not None: + transfer["custom_metadata"] = custom_metadata + self._kv_pending_transfers[req_id] = transfer + + def drain_pending_kv_transfers(self) -> dict[str, dict[str, Any]]: + """Drain pending KV transfers and move them to active. + + Returns ``{req_id: {seq_len, block_ids}}`` for the model runner + to submit to ``send_kv_cache``. + """ + if not self._kv_pending_transfers: + return {} + pending = dict(self._kv_pending_transfers) + self._kv_active_transfers.update(pending.keys()) + self._kv_pending_transfers.clear() + return pending + + def ack_kv_transfers(self, req_ids: list[str] | set[str]) -> None: + """Acknowledge completed KV transfers (from kv_extracted_req_ids). + + Moves requests from active to completed so the scheduler can + safely free their blocks. + """ + for req_id in req_ids: + self._kv_active_transfers.discard(req_id) + self._kv_completed_transfers.add(req_id) + + def drain_completed_kv_transfers(self) -> set[str]: + """Drain and return completed KV transfer request IDs. + + The scheduler calls this to know which requests' blocks can be freed. + """ + completed = set(self._kv_completed_transfers) + self._kv_completed_transfers.clear() + return completed + + def is_kv_transfer_triggered(self, req_id: str) -> bool: + """Check if a request has already triggered KV transfer.""" + return req_id in self._kv_triggered_requests + + def has_pending_kv_work(self) -> bool: + """True if any KV transfers are pending, active, or awaiting ack.""" + return bool(self._kv_pending_transfers or self._kv_active_transfers or self._kv_completed_transfers) + + # Output aggregation + # ------------------------------------------------------------------ # + + def get_omni_connector_output(self) -> OmniConnectorOutput: + """Collect and reset transfer results for this execute_model cycle. + + ``request_metadata`` carries only lightweight scheduling metadata. + Full payloads remain owned by the Model Runner local cache for all + paths. + """ + if not hasattr(self, "_lock"): + return OmniConnectorOutput() + + tp_group = self._get_local_tp_group() + if self._async_chunk and tp_group is not None and getattr(tp_group, "world_size", 1) > 1: + if self.is_data_transfer_rank(): + with self._lock: + fanout_packet = self._collect_async_chunk_fanout_packet_locked() + else: + fanout_packet = None + fanout_packet = self._broadcast_tp_payload_packet(fanout_packet) + if fanout_packet is None: + newly_finished = set() + chunk_finished = set() + request_metadata = {} + else: + if not self.is_data_transfer_rank(): + self._apply_async_chunk_fanout_packet(fanout_packet) + newly_finished = set(fanout_packet["newly_finished"]) + chunk_finished = set(fanout_packet["chunk_finished"]) + request_metadata = dict(fanout_packet["request_metadata"]) + else: + with self._lock: + newly_finished = set(self._finished_load_reqs) + self._finished_load_reqs.clear() + chunk_finished = set(self._chunk_finished_req_ids) + self._chunk_finished_req_ids.clear() + request_metadata = dict(self._local_request_metadata) + self._local_request_metadata.clear() + # _send_side_request_payload is the async accumulation buffer for + # future recv chunks. Clearing it on every consumable wake-up drops + # intermediate + # thinker decode spans before the model side can consume them. + # Only terminal chunk_finished requests may release that buffer. + for req_id in chunk_finished: + if req_id not in self._local_stage_payload_cache: + continue + ext_req_id = self._request_ids_mapping.get(req_id, req_id) + self._send_side_request_payload.pop(ext_req_id, None) + if ext_req_id != req_id: + self._send_side_request_payload.pop(req_id, None) + self._chunk_ready_req_ids.update(newly_finished) + + output = OmniConnectorOutput( + chunk_ready_req_ids=set(self._chunk_ready_req_ids), + chunk_finished_req_ids=chunk_finished, + request_metadata=request_metadata, + kv_sent_req_ids=list(self._kv_sent_req_ids), + stage_recv_req_ids=set(self._stage_recv_req_ids), + has_pending_kv_work=self.has_pending_kv_work(), + ) + if output.stage_recv_req_ids or chunk_finished or newly_finished: + logger.debug( + "[Stage-%s] get_omni_connector_output: stage_recv=%s, chunk_finished=%s, chunk_ready=%s", + self._stage_id, + output.stage_recv_req_ids, + chunk_finished, + output.chunk_ready_req_ids, + ) + self._chunk_ready_req_ids.clear() + self._kv_sent_req_ids.clear() + self._stage_recv_req_ids.clear() + return output + + @staticmethod + def _connector_output_has_signals(output: OmniConnectorOutput) -> bool: + return bool( + output.chunk_ready_req_ids + or output.chunk_finished_req_ids + or output.request_metadata + or output.kv_sent_req_ids + or output.stage_recv_req_ids + or output.has_pending_kv_work + ) + + def attach_omni_connector_output(self, result: Any | None) -> Any: + omni_output = self.get_omni_connector_output() + if not self._connector_output_has_signals(omni_output): + return result + + from copy import copy + + from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT + + wrapped = copy(result if result is not None else EMPTY_MODEL_RUNNER_OUTPUT) + wrapped.omni_connector_output = omni_output + return wrapped + + # ------------------------------------------------------------------ # + # Properties for compatibility with custom_process funcs that access + # transfer_manager.put_req_chunk / request_payload / code_prompt_token_ids + # ------------------------------------------------------------------ # + + @property + def put_req_chunk(self) -> dict[str, int]: + return self._put_req_chunk + + @property + def ramp_chunk_count(self) -> dict[str, int]: + return self._ramp_chunk_count + + @property + def request_payload(self) -> dict[str, dict[str, Any]]: + return self._send_side_request_payload + + @request_payload.setter + def request_payload(self, value: dict[str, dict[str, Any]]) -> None: + self._send_side_request_payload = value + + @property + def code_prompt_token_ids(self) -> dict[str, list[list[int]]]: + return self._code_prompt_token_ids + + @property + def connector(self) -> Any | None: + return self._omni_connector + + # ------------------------------------------------------------------ # + # Background I/O threads + # ------------------------------------------------------------------ # + + def _recv_loop(self) -> None: + """Background thread: poll connector for incoming data.""" + _recv_poll_count = 0 + while not self._stop_event.is_set(): + with self._lock: + pending_ids = list(self._pending_load_reqs.keys()) + + if not pending_ids: + self._work_available.wait(timeout=0.01) + self._work_available.clear() + continue + + _recv_poll_count += 1 + if _recv_poll_count % 5000 == 1: + logger.debug( + "[Stage-%s] _recv_loop: polling %s pending reqs: %s (poll#%s)", + self._stage_id, + len(pending_ids), + pending_ids[:5], + _recv_poll_count, + ) + + made_progress = False + for req_id in pending_ids: + if self._stop_event.is_set(): + break + try: + made_progress = self._poll_single_request(req_id) or made_progress + except Exception: + logger.warning("Error receiving data for %s", req_id, exc_info=True) + + if not made_progress and not self._stop_event.is_set(): + self._work_available.wait(timeout=0.005) + self._work_available.clear() + + _MAX_SEND_RETRIES = 3 + + def _save_loop(self) -> None: + """Background thread: send outgoing data via connector.""" + while not self._stop_event.is_set(): + task = None + with self._lock: + for req_id in list(self._pending_save_reqs.keys()): + dq = self._pending_save_reqs[req_id] + if dq: + task = dq.popleft() + if not dq: + del self._pending_save_reqs[req_id] + break + del self._pending_save_reqs[req_id] + + if task is not None: + success = False + try: + success = self._send_single_request(task) + except Exception: + logger.error( + "Error saving data for %s", + task.get("request_id"), + exc_info=True, + ) + if not success: + self._requeue_or_drop_failed_send(task) + continue + + self._work_available.wait(timeout=0.01) + self._work_available.clear() + + def _requeue_or_drop_failed_send(self, task: dict) -> None: + """Re-enqueue a failed send task or drop it after max retries.""" + retry_count = task.get("_retry_count", 0) + 1 + req_id = task.get("request_id") + if retry_count <= self._MAX_SEND_RETRIES: + task["_retry_count"] = retry_count + logger.warning( + "[Stage-%s] Re-enqueuing failed send for %s (retry %d/%d)", + getattr(self, "_stage_id", "?"), + req_id, + retry_count, + self._MAX_SEND_RETRIES, + ) + with self._lock: + dq = self._pending_save_reqs.setdefault(req_id, deque()) + dq.appendleft(task) + else: + logger.error( + "[Stage-%s] Giving up on send for %s after %d retries", + getattr(self, "_stage_id", "?"), + req_id, + self._MAX_SEND_RETRIES, + ) + self._decrement_pending_save_count(req_id) + + # ------------------------------------------------------------------ # + # Chunk-level poll / send (ported from OmniChunkTransferAdapter) + # ------------------------------------------------------------------ # + + def _poll_single_request(self, req_id: str) -> bool: + """Poll connector for one chunk of a request (non-blocking).""" + connector = self._omni_connector + if connector is None: + return False + + if self._async_chunk and self._model_mode != "ar": + with self._lock: + staged_payload = self._local_stage_payload_cache.get(req_id) + metadata_in_flight = req_id in self._local_request_metadata + scheduler_wakeup_pending = req_id in self._finished_load_reqs + if self._payload_is_consumable(staged_payload) or metadata_in_flight or scheduler_wakeup_pending: + logger.debug( + "[Stage-%s] delaying recv for req=%s until staged async payload is handed to scheduler", + self._stage_id, + req_id, + ) + return False + + target_stage_id = self._stage_id - 1 + chunk_id = self._get_req_chunk[req_id] + external_req_id = self._request_ids_mapping.get(req_id, req_id) + connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" + + if self._async_chunk: + result = self._recv_async_chunk_result( + connector, + str(target_stage_id), + str(self._stage_id), + connector_get_key, + ) + else: + result = self._recv_full_payload_result( + connector, + str(target_stage_id), + str(self._stage_id), + connector_get_key, + ) + + if result is None: + return False + + payload_data, _size = result + if not payload_data: + return False + if isinstance(payload_data, dict): + logger.debug( + "[Stage-%s] recv_chunk_result: req=%s ext=%s key=%s keys=%s finished=%s", + self._stage_id, + req_id, + external_req_id, + connector_get_key, + sorted(payload_data.keys()), + self._payload_finished(payload_data), + ) + + self._get_req_chunk[req_id] += 1 + + if self._async_chunk: + is_finished = self._payload_finished(payload_data) + incoming_payload_consumable = self._payload_is_consumable(payload_data) + + if self._model_mode == "ar": + payload_data = self._accumulate_payload(external_req_id, payload_data) + payload_consumable = incoming_payload_consumable + else: + new_ids = self._payload_audio_codes(payload_data) or [] + if not new_ids and not is_finished: + return False + payload_consumable = self._payload_is_consumable(payload_data) + + with self._lock: + if is_finished: + self._chunk_finished_req_ids.add(req_id) + self._chunk_stream_completed.add(req_id) + # Local cache (RFC §2.4) — merge, don't replace, so that + # earlier chunk keys (e.g. thinker_prefill_embeddings from + # chunk 0) are not overwritten by later chunks. + existing = self._local_stage_payload_cache.get(req_id) + if existing is not None and isinstance(existing, dict) and isinstance(payload_data, dict): + existing.update(payload_data) + else: + self._local_stage_payload_cache[req_id] = payload_data + staged_payload = self._local_stage_payload_cache[req_id] + self._async_chunk_updated_req_ids.add(req_id) + self.put_local_request_metadata(req_id, self._extract_scheduling_metadata(staged_payload)) + # A finish-only sentinel still needs one terminal wake-up so + # the downstream stage can sync the merged local payload and + # flush/finish even when the last recv carries no new + # consumable chunk bytes. + if payload_consumable or is_finished: + self._finished_load_reqs.add(req_id) + if is_finished and not payload_consumable: + logger.debug( + "[Stage-%s] finish sentinel arrived for req=%s without new consumable payload", + self._stage_id, + req_id, + ) + elif not payload_consumable: + logger.debug( + "[Stage-%s] req=%s received metadata-only / non-consumable async payload; delaying wake-up", + self._stage_id, + req_id, + ) + if is_finished: + self._pending_load_reqs.pop(req_id, None) + else: + # full_payload_mode: the complete payload arrives in a single get(), + # so always unregister immediately. + if isinstance(payload_data, dict): + engine_inputs = payload_data.get("engine_inputs", payload_data) + else: + engine_inputs = payload_data + with self._lock: + self._local_stage_payload_cache[req_id] = self._snapshot_payload(engine_inputs) + # Publish full-payload readiness only after the aligned TP broadcast + # path in recv_full_payload_inputs() has materialized the payload on all + # local ranks. Publishing metadata / stage_recv from the background recv + # thread can let the scheduler observe a request before the payload is + # actually visible to the model thread. + self._full_payload_pending_broadcast_req_ids.add(req_id) + self._pending_load_reqs.pop(req_id, None) + logger.debug( + "[Stage-%s] full_payload recv complete: req=%s key=%s payload_type=%s", + self._stage_id, + req_id, + connector_get_key, + type(engine_inputs).__name__, + ) + + logger.debug("[Stage-%s] Received data for key %s", self._stage_id, connector_get_key) + return True + + def _build_custom_process_payload( + self, + request_id: str | None, + request: Any | None, + pooling_output: Any | None, + ) -> Any | None: + """Run the custom process hook with a best-effort finished kwarg.""" + if self._custom_process_func is None: + return None + + kwargs = { + "transfer_manager": self, + "pooling_output": pooling_output, + "request": request, + } + supports_is_finished = getattr( + self, + "_custom_process_supports_is_finished", + self._custom_process_supports_is_finished_kwarg(), + ) + is_finished_fn = getattr(request, "is_finished", None) + if callable(is_finished_fn): + try: + if supports_is_finished is not False: + kwargs["is_finished"] = bool(is_finished_fn()) + except Exception: + logger.debug("request.is_finished() failed for %s", request_id, exc_info=True) + + try: + return self._custom_process_func(**kwargs) + except TypeError as exc: + if "is_finished" not in kwargs or not self._is_unexpected_is_finished_kwarg_error(exc): + logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) + return None + kwargs.pop("is_finished", None) + try: + return self._custom_process_func(**kwargs) + except Exception: + logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) + return None + except Exception: + logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) + return None + + def _custom_process_supports_is_finished_kwarg(self) -> bool | None: + """Return whether the custom process hook accepts `is_finished`.""" + if self._custom_process_func is None: + return None + try: + signature = inspect.signature(self._custom_process_func) + except (TypeError, ValueError): + return None + + for param in signature.parameters.values(): + if param.kind == inspect.Parameter.VAR_KEYWORD: + return True + + is_finished_param = signature.parameters.get("is_finished") + if is_finished_param is None: + return False + return is_finished_param.kind in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + + @staticmethod + def _is_unexpected_is_finished_kwarg_error(exc: TypeError) -> bool: + message = str(exc) + return ( + "unexpected keyword argument 'is_finished'" in message + or 'unexpected keyword argument "is_finished"' in message + or "positional-only arguments passed as keyword arguments: 'is_finished'" in message + ) + + def _send_single_request(self, task: dict) -> bool: + """Send one queued task via connector.put(). + + Returns True on success. On failure (put() raises or returns + ``success=False``), returns False **without** decrementing + ``_pending_save_counts`` so the caller can retry or clean up. + """ + connector = self._omni_connector + if connector is None: + return True + + request_id = task.get("request_id") + payload_data = task.get("data") + if payload_data is None and task.get("request") is not None: + payload_data = self._build_custom_process_payload( + request_id=request_id, + request=task.get("request"), + pooling_output=task.get("pooling_output"), + ) + put_key = task.get("put_key") + + success, _size, _metadata = connector.put( + from_stage=str(task["stage_id"]), + to_stage=str(task["next_stage_id"]), + put_key=put_key, + data=payload_data, + ) + logger.debug( + "[Stage-%s] _send_single_request: put_key=%s success=%s size=%s", + task["stage_id"], + put_key, + success, + _size, + ) + + if not success: + return False + + self._decrement_pending_save_count(request_id) + return True + + def _decrement_pending_save_count(self, request_id: str) -> None: + """Decrement pending save count and run deferred cleanup if zero.""" + cleanup_req_id = None + with self._lock: + remaining = self._pending_save_counts.get(request_id, 0) + if remaining > 1: + self._pending_save_counts[request_id] = remaining - 1 + elif remaining == 1: + self._pending_save_counts.pop(request_id, None) + if request_id in self._deferred_send_cleanup: + self._deferred_send_cleanup.remove(request_id) + cleanup_req_id = request_id + if cleanup_req_id is not None: + self._put_req_chunk.pop(cleanup_req_id, None) + self._send_side_request_payload.pop(cleanup_req_id, None) + self._code_prompt_token_ids.pop(cleanup_req_id, None) + self._cached_ic.pop(cleanup_req_id, None) + self._ramp_chunk_count.pop(cleanup_req_id, None) + + # ------------------------------------------------------------------ # + # Payload accumulation (ported from OmniChunkTransferAdapter) + # ------------------------------------------------------------------ # + + def _accumulate_payload(self, req_id: str, payload_data: OmniPayload) -> OmniPayload: + """Accumulate chunk payloads (concat tensors, extend lists).""" + if req_id not in self._send_side_request_payload: + self._send_side_request_payload[req_id] = dict(payload_data) + return dict(self._send_side_request_payload[req_id]) + + origin = self._send_side_request_payload[req_id] + merged = dict(origin) + raw_ok = payload_data.get("meta", {}).get("override_keys", []) if isinstance(payload_data, dict) else [] + override_keys = {tuple(k) if isinstance(k, list) else k for k in raw_ok} + + for key, value in payload_data.items(): + if isinstance(value, dict): + origin_sub = origin.get(key) + merged_sub = dict(origin_sub) if isinstance(origin_sub, dict) else {} + span_handled: set[str] = set() + if key == "embed" and isinstance(origin_sub, dict): + for tk, sk, ek in _EMBED_SPAN_GROUPS: + if tk not in value or (key, tk) in override_keys: + continue + span = merge_tensor_spans( + get_tensor_span(origin_sub, tensor_key=tk, start_key=sk, end_key=ek), + get_tensor_span(value, tensor_key=tk, start_key=sk, end_key=ek), + ) + if span is None: + continue + t, s, e = span + merged_sub[tk] = t + merged_sub[sk] = s + merged_sub[ek] = e + span_handled |= {tk, sk, ek} + for qual, qval in value.items(): + if qual in span_handled: + continue + if key == "meta" and qual == "finished": + merged_sub[qual] = qval + continue + if (key, qual) in override_keys: + merged_sub[qual] = qval + continue + osv = merged_sub.get(qual) + if isinstance(qval, torch.Tensor) and isinstance(osv, torch.Tensor): + merged_sub[qual] = torch.cat([osv, qval], dim=0) + elif isinstance(qval, list) and isinstance(osv, list): + merged_sub[qual] = osv + qval + else: + merged_sub[qual] = qval + merged[key] = merged_sub + else: + if key in override_keys: + merged[key] = value + continue + ov = origin.get(key) + if isinstance(value, torch.Tensor) and isinstance(ov, torch.Tensor): + merged[key] = torch.cat([ov, value], dim=0) + elif isinstance(value, list) and isinstance(ov, list): + merged[key] = ov + value + else: + merged[key] = value + + self._send_side_request_payload[req_id] = merged + return dict(merged) + + def drop_inactive_request_runtime_state(self, req_id: str) -> None: + """Clear inactive request state used by both the runner and mixin. + + This centralizes the model-runner-side cleanup pattern so + ``OmniGPUModelRunner`` can reuse it instead of open-coding the same + inactive-request state mutations. + """ + if hasattr(self, "model_intermediate_buffer"): + self.model_intermediate_buffer.pop(req_id, None) + self.drop_inactive_request_delivery_state(req_id) + + # ------------------------------------------------------------------ # + # Helpers + # ------------------------------------------------------------------ # + + @staticmethod + def _create_connector(model_config: Any) -> OmniConnectorBase | None: + """Create a connector from model_config, or None if unconfigured.""" + connector_config = getattr(model_config, "stage_connector_config", None) + if connector_config is None: + return None + + if not isinstance(connector_config, dict): + connector_config = { + "name": getattr(connector_config, "name", None), + "extra": getattr(connector_config, "extra", None), + } + + name = connector_config.get("name") + if not isinstance(name, str) or not name.strip(): + raise RuntimeError("Invalid stage connector config: missing connector name") + name = name.strip() + + extra = connector_config.get("extra") + if extra is None: + extra = {} + elif not isinstance(extra, dict): + raise RuntimeError(f"Invalid extra config for connector {name}: expected dict, got {type(extra).__name__}") + + spec = ConnectorSpec(name=name, extra=extra) + try: + return OmniConnectorFactory.create_connector(spec) + except Exception as exc: + raise RuntimeError(f"Failed to create connector {name}") from exc + + @staticmethod + def _load_custom_func(model_config: Any) -> tuple[str | None, Any | None]: + """Load the connector payload builder for the downstream stage. + + Preferred source is ``custom_process_next_stage_input_func``. Some + full_payload_mode configs (async_chunk=false) only expose the next-stage prompt builder via + ``custom_process_input_func`` (for example ``thinker2talker``), while the + connector payload builder lives beside it as ``thinker2talker_full_payload``. + In that case, derive the full_payload_mode builder path automatically. + """ + candidates: list[str] = [] + + next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) + if isinstance(next_stage_func, str) and next_stage_func: + candidates.append(next_stage_func) + + if not getattr(model_config, "async_chunk", False): + input_func = getattr(model_config, "custom_process_input_func", None) + if isinstance(input_func, str) and input_func: + try: + module_path, func_name = input_func.rsplit(".", 1) + if func_name.endswith("_full_payload") or func_name.endswith("_batch"): + candidates.append(f"{module_path}.{func_name}") + else: + candidates.append(f"{module_path}.{func_name}_full_payload") + candidates.append(f"{module_path}.{func_name}_batch") + candidates.append(input_func) + except ValueError: + candidates.append(input_func) + + tried: set[str] = set() + for func_path in candidates: + if func_path in tried: + continue + tried.add(func_path) + try: + # RFC #4872 Phase 3 (Part A): resolve through the registry with + # the full-payload producer contract. is_finished is optional + # here (missing only warns); a structurally incompatible + # candidate is skipped so the fallback chain keeps working. + spec = resolve_processor(func_path, expected_kind="producer_full_payload") + return spec.path, spec.fn + except ProcessorValidationError as exc: + logger.debug( + "Skipping incompatible connector payload hook %s: %s", + func_path, + exc, + ) + continue + except Exception: + logger.warning("Failed to load custom func: %s", func_path, exc_info=True) + + return None, None + + @staticmethod + def _is_connector_payload_builder(func: Any) -> bool: + """Whether *func* matches the mixin payload-builder contract.""" + try: + signature = inspect.signature(func) + except (TypeError, ValueError): + return False + + params = signature.parameters + if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()): + return True + + required = {"transfer_manager", "pooling_output", "request"} + supported = { + name + for name, param in params.items() + if param.kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + } + return required.issubset(supported) + + def _resolve_external_req_id(self, request: Any, fallback_req_id: str) -> str: + """Resolve the external request ID consistently. + + Checks ``_request_ids_mapping`` first (populated by + ``register_chunk_recv``), then falls back to the request's + ``external_req_id`` attribute, and finally to the given + ``fallback_req_id``. + """ + mapped = self._request_ids_mapping.get(fallback_req_id) + if mapped is not None: + return mapped + if request is not None: + # external_req_id may be explicitly None; fall back. + ext = getattr(request, "external_req_id", None) + if ext is not None: + return ext + return fallback_req_id + + def _resolve_next_stage_id(self, model_config: Any) -> int: + """Determine the downstream stage ID from connector config. + + Falls back to ``stage_id + 1`` when the config does not specify + a ``to_stage`` explicitly. + """ + connector_config = getattr(model_config, "stage_connector_config", None) + if connector_config is not None: + if isinstance(connector_config, dict): + to_stage = connector_config.get("to_stage") + else: + to_stage = getattr(connector_config, "to_stage", None) + if isinstance(to_stage, int): + return to_stage + if isinstance(to_stage, str) and to_stage.strip(): + return int(to_stage) + return self._stage_id + 1 + + @staticmethod + def _parse_rank_mapping(model_config: Any) -> dict[str, int]: + """Parse rank_mapping from connector config (optional). + + Returns ``{"from_tp": int, "to_tp": int, "local_rank": int}``. + When ``rank_mapping`` is absent, assumes 1:1 homogeneous mapping. + """ + connector_config = getattr(model_config, "stage_connector_config", None) + if connector_config is not None and not isinstance(connector_config, dict): + connector_config = getattr(connector_config, "__dict__", {}) + + rank_mapping: dict = {} + if isinstance(connector_config, dict): + rank_mapping = connector_config.get("rank_mapping", {}) + + from_tp = int(rank_mapping.get("from_tp", 1)) + to_tp = int(rank_mapping.get("to_tp", 1)) + + local_rank = 0 + try: + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + except (ValueError, TypeError): + pass + + return {"from_tp": from_tp, "to_tp": to_tp, "local_rank": local_rank} + + # ------------------------------------------------------------------ # + # Heterogeneous TP rank support + # ------------------------------------------------------------------ # + + def _validate_kv_tp_topology(self) -> None: + """Reject heterogeneous TP mappings that cannot be routed losslessly.""" + if self._from_tp <= 0 or self._to_tp <= 0: + raise ValueError(f"Invalid KV TP mapping: from_tp={self._from_tp}, to_tp={self._to_tp}") + larger = max(self._from_tp, self._to_tp) + smaller = min(self._from_tp, self._to_tp) + if larger % smaller != 0: + raise ValueError( + f"KV TP mapping must be divisible for rank-aware routing: from_tp={self._from_tp}, to_tp={self._to_tp}" + ) + + def get_kv_remote_ranks(self) -> list[int]: + """Determine which remote ranks this local rank exchanges KV with. + + Follows vLLM's ``TpKVTopology.get_target_remote_ranks()`` pattern: + - ``from_tp > to_tp``: each to-rank reads from multiple from-ranks + - ``from_tp < to_tp``: multiple to-ranks read from the same from-rank + - ``from_tp == to_tp``: 1:1 mapping + """ + self._validate_kv_tp_topology() + if self._from_tp == self._to_tp: + return [self._local_rank] + + if self._from_tp > self._to_tp: + tp_ratio = self._from_tp // self._to_tp + return [self._local_rank * tp_ratio + i for i in range(tp_ratio)] + else: + tp_ratio = self._to_tp // self._from_tp + return [self._local_rank // tp_ratio] + + def is_data_transfer_rank(self) -> bool: + """Whether this rank should participate in data (non-KV) transfer. + + Ordinary stage payloads are TP-identical, so exactly one TP rank + should talk to the connector. When TP is initialized, use TP rank 0 + so the connector leader matches TP-local broadcast source rank. + Otherwise fall back to LOCAL_RANK==0 for the single-rank case. + """ + tp_group = self._get_local_tp_group() + if tp_group is not None and getattr(tp_group, "world_size", 1) > 1: + return getattr(tp_group, "rank_in_group", 0) == 0 + return self._local_rank == 0 + + def get_kv_connector_key( + self, + req_id: str, + from_stage: int, + chunk_id: int, + from_rank: int, + to_rank: int, + ) -> str: + """Build connector key that includes rank info for KV transfers.""" + return f"{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}" From 1ba009b9f6977f1089585f2d3960af187cddf71f Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Tue, 25 Aug 2026 20:18:19 +0800 Subject: [PATCH 03/41] feat(stage_input_processors): RFC #4872 P8b unify async-chunk placeholder prompt builder (Phase 3b) qwen3_omni dual-entry (build_forward/prewarm), _prewarm_async_chunk_stages routed through registered builder + fallback, adapter removes hardcoded Qwen token id, sync_process_input_func config/metadata chain, parity test. Signed-off-by: RyanYun09 --- .../test_orchestrator_stage_input_bridge.py | 81 ++++++ .../test_placeholder_parity.py | 249 ++++++++++++++++++ vllm_omni/config/stage_config.py | 7 + .../distributed/omni_connectors/adapter.py | 38 ++- vllm_omni/engine/orchestrator.py | 192 ++++++++++---- vllm_omni/engine/stage_engine_core_client.py | 4 + vllm_omni/engine/stage_init_utils.py | 25 +- .../stage_input_processors/qwen3_omni.py | 119 +++++++-- 8 files changed, 618 insertions(+), 97 deletions(-) create mode 100644 tests/model_executor/stage_input_processors/test_placeholder_parity.py diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index a156077e77a..e580c296efd 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -302,6 +302,87 @@ async def test_async_prewarm_skips_stage_with_custom_process_input_func() -> Non assert 2 in req_state.stage_submit_ts +async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: + """RFC #4872 P8b: async-chunk prewarm resolves ``sync_process_input_func`` + (``*_token_only``) to ``build_prewarm_placeholder`` and uses its estimate. + + The placeholder length is ``stage0_only`` = ``len(stage0 prompt_token_ids)`` + (best-effort; the connector fixup path replaces it later). + """ + orchestrator = object.__new__(Orchestrator) + stage0 = FakePrewarmPool("sender") + stage1 = FakePrewarmPool("receiver") + # The real qwen3_omni sync placeholder builder exposes build_prewarm_placeholder + # off the resolved fn (see test_placeholder_parity.py). + stage1.stage_client = SimpleNamespace( + sync_process_input_func=( + "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only" + ), + ) + orchestrator.stage_pools = [stage0, stage1] + orchestrator._emit_tx_edge = lambda **_kwargs: None + orchestrator._record_duplex_stage_submission = MagicMock() + req_state = OrchestratorRequestState( + request_id="req-prewarm-resolved", + prompt={"prompt_token_ids": [1, 2]}, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=1, + ) + + prewarmed = await orchestrator._prewarm_async_chunk_stages( + "req-prewarm-resolved", + SimpleNamespace(prompt_token_ids=[1, 2], resumable=True), + req_state, + ) + + assert prewarmed is True + assert len(stage1.submitted) == 1 + # build_prewarm_placeholder uses stage0_only = len(stage0 prompt) == 2. + assert stage1.submitted[0].prompt_token_ids == [0, 0] + orchestrator._record_duplex_stage_submission.assert_called_once_with( + 1, + "req-prewarm-resolved", + 0, + req_state, + ) + + +@pytest.mark.asyncio +async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None: + port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission() + prewarm.return_value = True + + result = await port.submit(submission) + + assert result.stage_id == 0 + stage_pools[0].submit_initial.assert_awaited_once() + prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"]) + + +@pytest.mark.asyncio +async def test_duplex_submit_bails_out_when_prewarm_failed_the_request() -> None: + """A failed prewarm has already aborted the request and popped its state. + + Returning a success result here would hand the control plane a replica for a + request that no longer exists, and the trailing bookkeeping would re-register + a running counter the cleanup just released -- a leak that never decrements. + """ + port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission() + counter = MagicMock() + port._running_counter = counter + prewarm.return_value = False + request_state = request_states["req-duplex"] + + with pytest.raises(RuntimeError, match="prewarm failed"): + await port.submit(submission) + + prewarm.assert_awaited_once() + assert request_state.duplex_stage_fences == {} + assert request_state.stage_submit_ts == {} + assert request_state.running_counter_registered is False + counter.increment.assert_not_called() + + @pytest.mark.asyncio async def test_async_route_forwards_to_outgoing_only_stage() -> None: orchestrator = object.__new__(Orchestrator) diff --git a/tests/model_executor/stage_input_processors/test_placeholder_parity.py b/tests/model_executor/stage_input_processors/test_placeholder_parity.py new file mode 100644 index 00000000000..354a11682ab --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_placeholder_parity.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""P8b parity tests: async-chunk prewarm placeholder vs sync forward placeholder. + +RFC #4872 Proposal 8b: ``*_token_only`` placeholder builders expose a dual entry +— ``build_forward_placeholder(source_outputs, ctx)`` for the sync forward path +and ``build_prewarm_placeholder(*, stage0_prompt, ctx, downstream_stage_id)`` for +the async-chunk prewarm path (no ``source_outputs`` yet). Both must be fed by +the same ``_common`` length / packing helpers so async and sync placeholders stay +consistent, and the streaming fixup helper +(``adapter.compute_talker_prompt_ids_length``) must use the same length helper. + +Length semantics +---------------- +- forward (``mode="full"``) is the Qwen chat-template scan (golden-locked at 15 + for the canonical prompt ``test_common_helpers_golden.py``). +- prewarm (``mode="stage0_only"``) is a best-effort estimate equal to + ``len(stage0_prompt)``; the connector fixup path replaces it with the real + length once the upstream chunk arrives. For a chat-template prompt these + differ (15 vs 10); for a single user segment (no trailing assistant marker) + they coincide (6 == 6), which is the exact parity the prewarm targets. + +These tests are CPU-only (no model loading). They run under the local vllm +stub (shim) and on CI where real vllm builds ``OmniTokensPrompt`` as a dict +subclass; length capture uses a ``pack_placeholder_prompt`` spy so it works in +both environments. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from vllm_omni.model_executor.stage_input_processors import _common +from vllm_omni.model_executor.stage_input_processors import qwen3_omni as q3 +from vllm_omni.model_executor.stage_input_processors._dispatch import OrchestratorInputContext +from vllm_omni.model_executor.stage_input_processors._registry import resolve_processor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +#: Canonical Qwen chat-template prompt used by the golden tests: im_start at 0/6, +#: one user segment of 6 tokens + one trailing assistant segment (+9) -> full 15. +GOLDEN_PROMPT = [151644, 872, 10, 11, 12, 13, 151644, 77091, 20, 21] + +#: Single user segment (no trailing assistant marker): full == stage0_only == 6. +SINGLE_USER_PROMPT = [151644, 872, 10, 11, 12, 13] + + +def _shim_active() -> bool: + """True when the local-dev vllm stub (no real vllm) is in effect.""" + try: + import vllm_omni # noqa: F401 + return bool(getattr(vllm_omni, "_SHIM_ACTIVE", False)) + except Exception: + return False + + +def _source_outputs(prompt_token_ids, output_ids=()): + return [ + SimpleNamespace( + request_id="req-1", + prompt_token_ids=list(prompt_token_ids), + outputs=[SimpleNamespace(cumulative_token_ids=list(output_ids))], + ) + ] + + +def _ctx(prompt=None): + return OrchestratorInputContext( + prompt=prompt, + requires_multimodal_data=False, + streaming_context=None, + ) + + +def _capture_pack(monkeypatch): + """Spy on ``_common.pack_placeholder_prompt``; returns a recorder dict. + + ``pack_placeholder_prompt`` is the single packing chokepoint shared by both + builders, so capturing ``prompt_len`` / ``voice_metadata`` verifies the + builders' length logic without depending on how ``OmniTokensPrompt`` is + constructed (shim vs real vllm). + """ + + recorded: dict[str, Any] = {} + + def _spy(*, prompt_len: int, voice_metadata: dict[str, Any] | None = None) -> Any: + recorded["prompt_len"] = prompt_len + recorded["voice_metadata"] = voice_metadata + return "PLACEHOLDER" + + monkeypatch.setattr(_common, "pack_placeholder_prompt", _spy) + return recorded + + +# --------------------------------------------------------------------------- +# Forward builder (sync path) locks the golden full-mode length. +# --------------------------------------------------------------------------- + + +def test_build_forward_placeholder_length_matches_golden(monkeypatch): + recorded = _capture_pack(monkeypatch) + out = q3.build_forward_placeholder(_source_outputs(GOLDEN_PROMPT), _ctx()) + assert len(out) == 1 + assert out[0] == "PLACEHOLDER" + # Golden chat-template scan: sum(user segment) + 9 == 15. + assert recorded["prompt_len"] == 15 + + +def test_build_forward_placeholder_preserves_voice_metadata(monkeypatch): + recorded = _capture_pack(monkeypatch) + prompt = {"additional_information": {"speaker": ["ethan"], "language": ["English"]}} + out = q3.build_forward_placeholder(_source_outputs([1, 2], [3]), _ctx(prompt=prompt)) + assert len(out) == 1 + assert recorded["voice_metadata"] == {"speaker": ["ethan"], "language": ["English"]} + + +# --------------------------------------------------------------------------- +# Prewarm builder (async path) is a stage0-only best-effort estimate. +# --------------------------------------------------------------------------- + + +def test_build_prewarm_placeholder_length_from_stage0(monkeypatch): + recorded = _capture_pack(monkeypatch) + out = q3.build_prewarm_placeholder( + stage0_prompt=GOLDEN_PROMPT, + ctx=_ctx(), + downstream_stage_id=1, + ) + assert len(out) == 1 + assert out[0] == "PLACEHOLDER" + assert recorded["prompt_len"] == len(GOLDEN_PROMPT) # 10 (best-effort estimate) + # Voice metadata is intentionally not forwarded (matches pre-existing prewarm). + assert recorded["voice_metadata"] is None + + +def test_build_prewarm_placeholder_accepts_request_like_stage0(monkeypatch): + recorded = _capture_pack(monkeypatch) + stage0_request = SimpleNamespace(prompt_token_ids=GOLDEN_PROMPT) + q3.build_prewarm_placeholder(stage0_prompt=stage0_request, ctx=_ctx(), downstream_stage_id=1) + assert recorded["prompt_len"] == len(GOLDEN_PROMPT) + + +# --------------------------------------------------------------------------- +# Parity: forward and prewarm must agree where the prompt shape allows. +# --------------------------------------------------------------------------- + + +def test_forward_prewarm_parity_single_user_segment(monkeypatch): + """Single user segment: forward (full) == prewarm (stage0_only) == 6.""" + recorded = _capture_pack(monkeypatch) + q3.build_forward_placeholder(_source_outputs(SINGLE_USER_PROMPT), _ctx()) + forward_len = recorded["prompt_len"] + recorded.clear() + q3.build_prewarm_placeholder(stage0_prompt=SINGLE_USER_PROMPT, ctx=_ctx(), downstream_stage_id=1) + assert forward_len == recorded["prompt_len"] == 6 + + +def test_forward_prewarm_golden_relationship(monkeypatch): + """Documented relationship for the chat-template golden prompt. + + forward (full) == 15 (golden-locked); prewarm (stage0_only) == 10 + (``= len(stage0_prompt)``). The difference is intentional: prewarm is a + best-effort estimate and the connector fixup path + (``adapter.construct_next_stage_streaming_input_prompt``) replaces it with + the real length once the upstream chunk arrives. + """ + recorded = _capture_pack(monkeypatch) + q3.build_forward_placeholder(_source_outputs(GOLDEN_PROMPT), _ctx()) + forward_len = recorded["prompt_len"] + recorded.clear() + q3.build_prewarm_placeholder(stage0_prompt=GOLDEN_PROMPT, ctx=_ctx(), downstream_stage_id=1) + assert forward_len == 15 + assert recorded["prompt_len"] == len(GOLDEN_PROMPT) + assert forward_len != recorded["prompt_len"] # full vs stage0_only estimate + + +# --------------------------------------------------------------------------- +# Packed placeholder structure (real vllm asserts; shim skips field reads). +# --------------------------------------------------------------------------- + + +def test_pack_placeholder_prompt_structure(): + prompt = _common.pack_placeholder_prompt(prompt_len=4, voice_metadata={"speaker": ["ethan"]}) + empty_meta = _common.pack_placeholder_prompt(prompt_len=1) + if _shim_active(): + # Under the local vllm stub, OmniTokensPrompt's real constructor is + # bypassed, so field access returns stub values. Field assertions run + # on CI where real vllm constructs the model properly. + pytest.skip("OmniTokensPrompt field construction needs real vllm (CI)") + assert prompt["prompt_token_ids"] == [0, 0, 0, 0] + assert prompt["additional_information"] == {"speaker": ["ethan"]} + assert empty_meta["prompt_token_ids"] == [0] + assert empty_meta.get("additional_information") is None + + +# --------------------------------------------------------------------------- +# thinker2talker_token_only delegation keeps the public sync entry unchanged. +# --------------------------------------------------------------------------- + + +def test_thinker2talker_token_only_delegates_to_forward_placeholder(monkeypatch): + recorded = _capture_pack(monkeypatch) + prompt = {"additional_information": {"speaker": ["ethan"], "language": ["English"]}} + out = q3.thinker2talker_token_only( + _source_outputs([1, 2], [3]), + prompt=prompt, + requires_multimodal_data=False, + streaming_context=None, + ) + assert len(out) == 1 + assert out[0] == "PLACEHOLDER" + assert recorded["voice_metadata"] == {"speaker": ["ethan"], "language": ["English"]} + + +# --------------------------------------------------------------------------- +# Streaming fixup helper shares the same length helper as the forward builder. +# --------------------------------------------------------------------------- + + +def test_adapter_compute_length_consistent_with_common(): + # Lazy import keeps this module import-light under the shim. + from vllm_omni.distributed.omni_connectors import adapter + + assert adapter.compute_talker_prompt_ids_length(GOLDEN_PROMPT) == 15 + assert adapter.compute_talker_prompt_ids_length(GOLDEN_PROMPT) == _common.compute_placeholder_prompt_len( + ids_or_prompt={"ids": {"all": GOLDEN_PROMPT, "prompt": GOLDEN_PROMPT}}, + mode="full", + ) + assert adapter.compute_talker_prompt_ids_length(SINGLE_USER_PROMPT) == 6 + + +# --------------------------------------------------------------------------- +# Registry surface: the resolved sync fn exposes build_prewarm_placeholder, so +# the orchestrator's prewarm resolution path resolves without a module scan. +# --------------------------------------------------------------------------- + + +def test_resolved_sync_fn_exposes_build_prewarm_placeholder(): + spec = resolve_processor( + "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only", + expected_kind=None, + ) + assert callable(getattr(spec.fn, "build_prewarm_placeholder", None)) + assert callable(getattr(spec.fn, "build_forward_placeholder", None)) + # The attached builder is the same module-level object used by tests above. + assert spec.fn.build_prewarm_placeholder is q3.build_prewarm_placeholder diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 8739dec11d1..3a148befe40 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -1387,6 +1387,7 @@ def merge_pipeline_deploy( stage_type=stage_type, input_sources=list(ps.input_sources), custom_process_input_func=input_proc, + sync_process_input_func=ps.sync_process_input_func, final_output=ps.final_output, final_output_type=ps.final_output_type, worker_type=worker_type, @@ -1416,6 +1417,7 @@ class StageConfig: stage_type: StageType = StageType.LLM input_sources: list[int] = field(default_factory=list) custom_process_input_func: str | None = None + sync_process_input_func: str | None = None final_output: bool = False final_output_type: str | None = None worker_type: str | None = None @@ -1511,6 +1513,11 @@ def to_omegaconf(self) -> Any: if self.custom_process_input_func: config_dict["custom_process_input_func"] = self.custom_process_input_func + # RFC #4872 P8b: carry the raw ``sync_process_input_func`` path to the + # orchestrator so async-chunk prewarm can resolve the registered + # placeholder builder (``*_token_only``) without re-reading the YAML. + if self.sync_process_input_func: + config_dict["sync_process_input_func"] = self.sync_process_input_func # Pass through extra YAML fields (default_sampling_params, # output_connectors, input_connectors, tts_args, etc.) diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index dc3be166fc9..10c73d91681 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -193,34 +193,28 @@ def try_recv_via_connector( def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int: """Compute the length of the talker prompt ids. + RFC #4872 P8b: delegates to the canonical + ``_common.compute_placeholder_prompt_len(mode="full")`` so the streaming + fixup path (``construct_next_stage_streaming_input_prompt``) and the + forward placeholder builder (``build_forward_placeholder``) share one length + helper. Behaviour is golden-locked (``test_common_helpers_golden.py`` + reproduces the 15-token chat-template result for the canonical prompt). + Args: prompt_ids: The prompt ids tensor. Returns: The length of the talker prompt ids. """ - im_start_token_id = 151644 - system_token_id = 8948 - user_token_id = 872 - assistant_token_id = 77091 - im_start_indexes = [i for i in range(len(prompt_ids)) if prompt_ids[i] == im_start_token_id] - im_start_indexes.append(len(prompt_ids)) - sum_user_len = 0 - assistant_len = 0 - for i in range(len(im_start_indexes) - 1): - s = im_start_indexes[i] - e = im_start_indexes[i + 1] - role = prompt_ids[s + 1] - if role == system_token_id: - continue - elif role == user_token_id: - sum_user_len += e - s - elif role == assistant_token_id and i == len(im_start_indexes) - 2: - assistant_len += 9 # 3 + 4 + 1 + 1 - else: - pass - - return sum_user_len + assistant_len + from vllm_omni.model_executor.stage_input_processors import _common + + # A flat prompt list has no ``all``/``prompt`` split, so feed the same list + # to both roles; ``_common``'s terminal index then equals ``len(prompt_ids)`` + # exactly as the legacy inline scan did. + return _common.compute_placeholder_prompt_len( + ids_or_prompt={"ids": {"all": prompt_ids, "prompt": prompt_ids}}, + mode="full", + ) def construct_next_stage_streaming_input_prompt( diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 76bfeecb49a..7e0b007d090 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -3006,51 +3006,16 @@ async def _prewarm_async_chunk_stages( operation="async-chunk prewarm", ) else: - from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length - - try: - next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) - except Exception: - next_prompt_len = max(1, len(prompt_token_ids)) - # A stage may size its own prewarmed placeholder prompt via - # hf_overrides.async_chunk_prewarm_prompt_len (e.g. a talker - # whose engine positions must cover a speaker-prompt prefill; - # the model validates the value and reports the right one on - # mismatch). - prewarm_len = getattr( - getattr(next_pool.stage_vllm_config, "model_config", None), - "hf_config", - None, + # RFC #4872 P8b: build the prewarm placeholder through the + # registered placeholder builder's ``build_prewarm_placeholder`` + # (resolved from ``sync_process_input_func``), falling back to + # the legacy inline estimate when no builder is available. + base_input = self._build_prewarm_placeholder_input( + next_stage_id, + request_id, + prompt_token_ids, + req_state, ) - prewarm_len = getattr(prewarm_len, "async_chunk_prewarm_prompt_len", None) - if prewarm_len is not None and int(prewarm_len) > 1: - next_prompt_len = int(prewarm_len) - - original_prompt = req_state.prompt - if isinstance(original_prompt, dict): - from vllm_omni.engine.request_snapshot import copy_request_snapshot - - # The placeholder replaces the prompt and multimodal fields, - # so they are not copied (multimodal inputs can be megabytes). - # additional_information is serialized below, which copies - # tensor bytes; its containers are copied, its tensors shared. - memo: dict[int, Any] = {} - base_input = copy_request_snapshot( - {k: v for k, v in original_prompt.items() if k not in _PREWARM_UNCOPIED_PROMPT_KEYS}, - memo, - ) - if "additional_information" in original_prompt: - base_input["additional_information"] = copy_request_snapshot( - original_prompt["additional_information"], - memo, - shared_types=(torch.Tensor,), - ) - else: - base_input = {} - - base_input["prompt_token_ids"] = [0] * next_prompt_len - base_input["multi_modal_data"] = None - base_input["mm_processor_kwargs"] = None downstream_resumable = bool(getattr(entry_request, "resumable", req_state.streaming.enabled)) # Session-owned and resumable requests keep their downstream # request across turns, so they stay on the chunk-0 path. @@ -3181,6 +3146,145 @@ def _diffusion_submit_kwargs( remote_bootstrap_addr=bootstrap_addr_from_kv_transfer_config(config), ) } + def _get_prewarm_placeholder_builder(self, stage_id: int) -> Any: + """Resolve the registered ``build_prewarm_placeholder`` for a stage. + + RFC #4872 P8b: the async-chunk prewarm reuses ``sync_process_input_func`` + (the ``*_token_only`` path). Resolution order: + + 1. ``client.sync_process_input_func`` -> ``resolve_processor(...).fn`` + and read ``build_prewarm_placeholder`` off the resolved function + (with a module-level sibling fallback). + 2. The already-resolved ``custom_process_input_func`` fn carrying the + attribute (some models install the forward processor directly). + + Returns ``None`` when no builder is available so the caller keeps the + legacy inline estimate. The lookup is deliberately defensive — it never + raises (fake pools / minimal stage clients in tests and future + connectors may not expose ``stage_client``), so the prewarm path cannot + regress on builder resolution. + """ + from vllm_omni.model_executor.stage_input_processors import resolve_processor + + try: + pool = self.stage_pools[stage_id] + client = getattr(pool, "stage_client", None) + if client is None: + return None + + path = getattr(client, "sync_process_input_func", None) + if path: + try: + spec = resolve_processor(path, expected_kind=None) + builder = getattr(spec.fn, "build_prewarm_placeholder", None) + if not callable(builder): + # Module-level sibling (covers paths whose fn does not + # expose the attribute directly). + import importlib + + module = importlib.import_module(path.rsplit(".", 1)[0]) + builder = getattr(module, "build_prewarm_placeholder", None) + if callable(builder): + return builder + logger.warning( + "[Orchestrator] stage-%s sync_process_input_func=%r has no " + "build_prewarm_placeholder; using the inline estimate", + stage_id, + path, + ) + except Exception as exc: + logger.warning( + "[Orchestrator] failed to resolve build_prewarm_placeholder " + "for %r (stage-%s); using the inline estimate: %s", + path, + stage_id, + exc, + ) + + fn = getattr(client, "custom_process_input_func", None) + if callable(fn): + builder = getattr(fn, "build_prewarm_placeholder", None) + if callable(builder): + return builder + except Exception as exc: # pragma: no cover - defensive only + logger.warning( + "[Orchestrator] prewarm placeholder builder lookup for stage-%s " + "failed; using the inline estimate: %s", + stage_id, + exc, + ) + return None + + def _build_prewarm_placeholder_input( + self, + next_stage_id: int, + request_id: str, + prompt_token_ids: Any, + req_state: OrchestratorRequestState, + ) -> dict[str, Any]: + """Build the async-chunk prewarm ``base_input`` dict for a downstream stage. + + Uses the registered placeholder builder's ``build_prewarm_placeholder`` + when available; otherwise (and on any builder failure) falls back to the + legacy inline estimate (``max(1, compute_talker_prompt_ids_length(...))``) + with a warning — the prewarm path must never raise. + """ + from vllm_omni.model_executor.stage_input_processors import OrchestratorInputContext + + original_prompt = req_state.prompt + if isinstance(original_prompt, dict): + import copy + + base_input = copy.deepcopy(original_prompt) + else: + base_input = {} + + prewarm_builder = self._get_prewarm_placeholder_builder(next_stage_id) + if prewarm_builder is not None: + try: + ctx = OrchestratorInputContext( + prompt=req_state.prompt, + requires_multimodal_data=False, + streaming_context=None, + ) + placeholders = prewarm_builder( + stage0_prompt=prompt_token_ids, + ctx=ctx, + downstream_stage_id=next_stage_id, + ) + placeholder = placeholders[0] if placeholders else None + if placeholder is not None: + ph_ids = ( + placeholder["prompt_token_ids"] + if isinstance(placeholder, dict) + else getattr(placeholder, "prompt_token_ids", None) + ) + if ph_ids: + base_input["prompt_token_ids"] = list(ph_ids) + base_input["multi_modal_data"] = None + base_input["mm_processor_kwargs"] = None + return base_input + except Exception as exc: + logger.warning( + "[Orchestrator] req=%s: prewarm placeholder builder for " + "stage-%s raised; falling back to inline estimate: %s", + request_id, + next_stage_id, + exc, + ) + + # Legacy inline fallback (RFC #4872 P8b: warn-and-keep-old behaviour so + # the prewarm never regresses on models without a dual-entry builder). + from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length + + try: + next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) + except Exception: + next_prompt_len = max(1, len(prompt_token_ids)) + base_input["prompt_token_ids"] = [0] * next_prompt_len + base_input["multi_modal_data"] = None + base_input["mm_processor_kwargs"] = None + return base_input def _build_kv_sender_info( self, diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index d2dc0ded71b..ca8ac786b8f 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -153,6 +153,10 @@ def __init__( self.prompt_expand_func = metadata.prompt_expand_func self.custom_process_input_func = metadata.custom_process_input_func self.async_chunk_prewarm_payload_func = getattr(metadata, "async_chunk_prewarm_payload_func", None) + # RFC #4872 P8b: raw ``sync_process_input_func`` path so the + # orchestrator can resolve the async-chunk prewarm placeholder + # builder (``build_prewarm_placeholder``) at prewarm time. + self.sync_process_input_func = metadata.sync_process_input_func self.engine_outputs: Any = None self.client_addresses = dict(client_addresses or {}) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 36c51f9e8f2..d77bed2b3a7 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -572,7 +572,10 @@ class StageMetadata: custom_process_input_func: Callable | None model_stage: str | None runtime_cfg: Any - prompt_transform_func: Callable | None = None + # RFC #4872 P8b: raw ``sync_process_input_func`` path (the ``*_token_only`` + # placeholder builder) so the orchestrator can resolve + # ``build_prewarm_placeholder`` at async-chunk prewarm time. + sync_process_input_func: str | None = None prompt_expand_func: Callable | None = None cfg_kv_collect_func: Callable | None = None async_chunk_prewarm_payload_func: Callable | None = None @@ -659,11 +662,12 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: exc, ) - prompt_transform_func: Callable | None = None - _ptf_path = _get_attr_or_item(stage_config, "prompt_transform_func") - if _ptf_path: - _mod, _fn = _ptf_path.rsplit(".", 1) - prompt_transform_func = getattr(importlib.import_module(_mod), _fn) + sync_process_input_func: str | None = None + _spif_path = _get_attr_or_item(stage_config, "sync_process_input_func") + if _spif_path: + # Keep the raw dotted path (not a resolved callable): the orchestrator + # resolves ``build_prewarm_placeholder`` off the ``*_token_only`` fn. + sync_process_input_func = _spif_path prompt_expand_func: Callable | None = None _pef_path = _get_attr_or_item(stage_config, "prompt_expand_func") @@ -697,6 +701,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: final_output_type=final_output_type, default_sampling_params=default_params, custom_process_input_func=custom_process_input_func, + sync_process_input_func=sync_process_input_func, model_stage=model_stage, runtime_cfg=runtime_cfg, prompt_transform_func=prompt_transform_func, @@ -718,6 +723,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: final_output_type=final_output_type, default_sampling_params=default_params, custom_process_input_func=custom_process_input_func, + sync_process_input_func=sync_process_input_func, model_stage=model_stage, runtime_cfg=runtime_cfg, prompt_transform_func=prompt_transform_func, @@ -752,6 +758,11 @@ def extract_stage_metadata_from_omni_stage_config( sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams sampling_params = sampling_params_cls(**(stage_config.model_config.default_sampling_params or {})) custom_process_input_func = _resolve_omni_metadata_hook(stage_config.custom_process_input_func) + # RFC #4872 P8b: carry the raw sync path for async-chunk prewarm resolution. + sync_process_input_func: str | None = None + _spif_path = getattr(stage_config.stage_pipeline_config, "sync_process_input_func", None) + if _spif_path: + sync_process_input_func = _spif_path if stage_type == "diffusion": return StageMetadata( @@ -765,6 +776,7 @@ def extract_stage_metadata_from_omni_stage_config( final_output_type=stage_config.final_output_type, default_sampling_params=sampling_params, custom_process_input_func=custom_process_input_func, + sync_process_input_func=sync_process_input_func, model_stage=stage_config.model_stage, runtime_cfg=stage_config.runtime_config, prompt_transform_func=_resolve_omni_metadata_hook(stage_config.prompt_transform_func), @@ -782,6 +794,7 @@ def extract_stage_metadata_from_omni_stage_config( final_output_type=stage_config.final_output_type, default_sampling_params=sampling_params, custom_process_input_func=custom_process_input_func, + sync_process_input_func=sync_process_input_func, model_stage=stage_config.model_stage, runtime_cfg=stage_config.runtime_config, prompt_transform_func=_resolve_omni_metadata_hook(stage_config.prompt_transform_func), diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 1a1cc089b7e..6b2b162d378 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -24,6 +24,7 @@ from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.stage_input_processors import _common +from vllm_omni.model_executor.stage_input_processors._dispatch import OrchestratorInputContext from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import compute_ramp_emit, parse_chunk_ramp from vllm_omni.model_executor.stage_input_processors.tts_utils import ( extract_language_from_prompt, @@ -634,27 +635,25 @@ def thinker2talker_full_payload( return payload -def thinker2talker_token_only( +def build_forward_placeholder( source_outputs: list[Any], - prompt: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, - streaming_context: Any | None = None, + ctx: OrchestratorInputContext, ) -> list[OmniTokensPrompt]: - """Orchestrator-side placeholder builder for Stage-1 (Talker) when - ``async_chunk=False``. - - After the communication-layer refactor, this function only allocates a - placeholder ``prompt_token_ids`` of the correct length so the scheduler can - reserve KV-cache slots. It does **not** forward bulk tensors. - - Bulk talker conditioning is sent through the connector. Speaker and - language are also copied from the original prompt so they survive when - Stage-0 request metadata is unavailable to the connector payload. - - ``prompt`` / ``requires_multimodal_data`` are kept for call-site signature - compatibility with other orchestrator input processors; they are unused. + """Orchestrator-side sync forward placeholder builder (RFC #4872 P8b). + + Semantics match the legacy ``thinker2talker_token_only`` forward path: one + placeholder ``OmniTokensPrompt`` per upstream output, sized so the + scheduler can reserve KV-cache slots, with speaker / language copied from + the original prompt so they survive when Stage-0 request metadata is + unavailable to the connector payload. + + RFC #4872 P8b consolidation: the length computation is delegated to + ``_common.compute_placeholder_prompt_len(mode="full")`` (the Qwen + chat-template scan, golden-locked at 15 for the golden prompt) and the + placeholder packing to ``_common.pack_placeholder_prompt``. """ talker_inputs: list[OmniTokensPrompt] = [] + streaming_context = ctx.streaming_context for i, thinker_output in enumerate(source_outputs): output = thinker_output.outputs[0] req_id = str(getattr(thinker_output, "request_id", f"idx-{i}")) @@ -673,25 +672,95 @@ def thinker2talker_token_only( thinker_sequences = prompt_token_ids + output_ids thinker_input_ids = prompt_token_ids info_for_len: OmniPayload = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}} - prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu") + prompt_len = _common.compute_placeholder_prompt_len( + ids_or_prompt=info_for_len, + mode="full", + device="cpu", + ) # Keep this fallback until the connector reliably preserves voice metadata. additional_information = to_dict( OmniPayloadStruct( - speaker=extract_speaker_from_prompt(prompt, index=i), - language=extract_language_from_prompt(prompt, index=i), + speaker=extract_speaker_from_prompt(ctx.prompt, index=i), + language=extract_language_from_prompt(ctx.prompt, index=i), ) ) talker_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[0] * prompt_len, - additional_information=additional_information or None, - multi_modal_data=None, - mm_processor_kwargs=None, + _common.pack_placeholder_prompt( + prompt_len=prompt_len, + voice_metadata=additional_information or None, ) ) return talker_inputs +def build_prewarm_placeholder( + *, + stage0_prompt: Any, + ctx: OrchestratorInputContext, + downstream_stage_id: int, +) -> list[OmniTokensPrompt]: + """Best-effort async-chunk prewarm placeholder builder (RFC #4872 P8b). + + async-chunk mode has **no upstream ``source_outputs``** yet at prewarm + time, so this is a best-effort estimate: the placeholder length comes from + the stage-0 input prompt (``_common.compute_placeholder_prompt_len( + mode="stage0_only")``, i.e. ``len(stage0_prompt)``). The connector fixup + path (``adapter.construct_next_stage_streaming_input_prompt``) replaces + this estimate with the real length once the upstream chunk arrives. + + ``ctx`` / ``downstream_stage_id`` are accepted for contract uniformity; + voice metadata is intentionally **not** forwarded here (matching the + pre-existing prewarm behaviour — no new behaviour is introduced). Per-edge + refinement (``downstream_stage_id``) is RFC #4872 P8c and out of scope. + """ + del ctx, downstream_stage_id # P8b uses only the stage-0 input estimate. + if hasattr(stage0_prompt, "prompt_token_ids"): + stage0_prompt = stage0_prompt.prompt_token_ids + prompt_len = max( + 1, + _common.compute_placeholder_prompt_len( + ids_or_prompt=stage0_prompt, + mode="stage0_only", + ), + ) + return [_common.pack_placeholder_prompt(prompt_len=prompt_len, voice_metadata=None)] + + +def thinker2talker_token_only( + source_outputs: list[Any], + prompt: OmniTokensPrompt | TextPrompt | None = None, + requires_multimodal_data: bool = False, + streaming_context: Any | None = None, +) -> list[OmniTokensPrompt]: + """Orchestrator-side placeholder builder for Stage-1 (Talker) when + ``async_chunk=False``. + + After the communication-layer refactor, this function only allocates a + placeholder ``prompt_token_ids`` of the correct length so the scheduler can + reserve KV-cache slots. It does **not** forward bulk tensors. + + RFC #4872 P8b: thin delegation to :func:`build_forward_placeholder` — the + dual-entry design keeps the sync forward path and the async-chunk prewarm + path on the same shared ``_common`` length / packing helpers. ``prompt`` / + ``requires_multimodal_data`` / ``streaming_context`` are kept for call-site + signature compatibility; they are folded into an ``OrchestratorInputContext``. + """ + ctx = OrchestratorInputContext( + prompt=prompt, + requires_multimodal_data=requires_multimodal_data, + streaming_context=streaming_context, + ) + return build_forward_placeholder(source_outputs, ctx) + + +# RFC #4872 P8b: the orchestrator resolves ``sync_process_input_func`` (the +# ``*_token_only`` path) via ``resolve_processor(...).fn`` and reads +# ``build_prewarm_placeholder`` off the resolved function object. Expose both +# dual-entry builders here so that lookup works without a separate module scan. +thinker2talker_token_only.build_forward_placeholder = build_forward_placeholder +thinker2talker_token_only.build_prewarm_placeholder = build_prewarm_placeholder + + # ========================= # Talker -> Code2Wav # ========================= From cad699d7a2410a2fc20d85fb86f075160e739360 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Tue, 25 Aug 2026 20:41:22 +0800 Subject: [PATCH 04/41] refactor(stage_input_processors): RFC #4872 P6 centralize codec/stop_token constants (Phase 4a) _constants.py centralizes 2150 stop token, qwen3 codec 4196/4197/4198 startup assertions, qwen2_5 8294 same-source; three pipelines reference shared constants; add consistency test. Signed-off-by: RyanYun09 --- .../test_constants_consistency.py | 100 ++++++++++++++++++ .../models/aura_omni/pipeline.py | 5 +- .../models/qwen2_5_omni/pipeline.py | 4 +- .../models/qwen2_5_omni/qwen2_5_omni.py | 7 +- .../models/qwen3_omni/pipeline.py | 5 +- .../models/qwen3_omni/qwen3_omni.py | 20 +++- .../models/qwen3_tts/pipeline.py | 5 +- .../stage_input_processors/_constants.py | 19 ++++ .../stage_input_processors/qwen3_omni.py | 38 +++++++ 9 files changed, 194 insertions(+), 9 deletions(-) create mode 100644 tests/model_executor/stage_input_processors/test_constants_consistency.py create mode 100644 vllm_omni/model_executor/stage_input_processors/_constants.py diff --git a/tests/model_executor/stage_input_processors/test_constants_consistency.py b/tests/model_executor/stage_input_processors/test_constants_consistency.py new file mode 100644 index 00000000000..70f619cef91 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_constants_consistency.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RFC #4872 P6: constant-consistency guards. + +Verifies that magic special-token ids that were previously duplicated across +processors / pipelines / model implementations now resolve to a single value: + +* qwen3 codec ``4196/4197/4198``: processor ``_QWEN3_CODEC_*`` == model + ``TALKER_CODEC_*`` == HF config defaults (``codec_pad_id`` / + ``codec_bos_id`` / ``codec_eos_token_id``). +* qwen2_5 talker codec end ``8294``: processor ``TALKER_CODEC_END_TOKEN_ID`` + == model ``TALKER_CODEC_EOS_TOKEN_ID`` == pipeline ``stop_token_ids``. +* ``2150``: single-sourced in ``stage_input_processors._constants`` and shared + by the qwen3_omni / qwen3_tts / aura_omni pipelines. + +CPU-only: imports modules and reads constants; never loads model weights. +""" + +from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE +from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import QWEN2_5_OMNI_PIPELINE +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni import TALKER_CODEC_EOS_TOKEN_ID +from vllm_omni.model_executor.models.qwen3_omni.pipeline import QWEN3_OMNI_PIPELINE +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import ( + TALKER_CODEC_BOS_TOKEN_ID, + TALKER_CODEC_EOS_TOKEN_ID as QWEN3_MODEL_CODEC_EOS, + TALKER_CODEC_PAD_TOKEN_ID, +) +from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSTalkerConfig +from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE +from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID +from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import TALKER_CODEC_END_TOKEN_ID +from vllm_omni.model_executor.stage_input_processors.qwen3_omni import ( + _QWEN3_CODEC_BOS_TOKEN_ID, + _QWEN3_CODEC_EOS_TOKEN_ID, + _QWEN3_CODEC_PAD_TOKEN_ID, +) + +# Expected values locked by the baseline commit `731a771`. +QWEN3_CODEC_EXPECTED = (4196, 4197, 4198) # (pad, bos, eos) +QWEN2_5_CODEC_END_EXPECTED = 8294 +QWEN3_CODEC_STOP_EXPECTED = 2150 + + +def _pipeline_stop_token_ids(pipeline) -> list[int]: + """Extract the (unique) ``stop_token_ids`` across a pipeline's stages.""" + ids: list[int] = [] + for stage in pipeline.stages: + constraints = getattr(stage, "sampling_constraints", None) or {} + stage_ids = constraints.get("stop_token_ids") + if stage_ids: + ids.extend(stage_ids) + return ids + + +def test_qwen3_codec_token_ids_processor_matches_model(): + """4196/4197/4198: processor and model implementation must agree.""" + processor = ( + _QWEN3_CODEC_PAD_TOKEN_ID, + _QWEN3_CODEC_BOS_TOKEN_ID, + _QWEN3_CODEC_EOS_TOKEN_ID, + ) + model = (TALKER_CODEC_PAD_TOKEN_ID, TALKER_CODEC_BOS_TOKEN_ID, QWEN3_MODEL_CODEC_EOS) + assert processor == model == QWEN3_CODEC_EXPECTED + + +def test_qwen3_codec_token_ids_match_hf_config(): + """4196/4197/4198: processor constants must match the in-repo Qwen3-TTS + HF config defaults (single source of truth).""" + cfg = Qwen3TTSTalkerConfig() + assert (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) == QWEN3_CODEC_EXPECTED + assert ( + _QWEN3_CODEC_PAD_TOKEN_ID, + _QWEN3_CODEC_BOS_TOKEN_ID, + _QWEN3_CODEC_EOS_TOKEN_ID, + ) == (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) + + +def test_qwen3_omni_runtime_talker_config_matches_processor(): + """The transformers ``Qwen3OmniMoeTalkerConfig`` (the config qwen3-omni + actually runs with) must also carry 4196/4197/4198.""" + from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeTalkerConfig, + ) + + cfg = Qwen3OmniMoeTalkerConfig() + assert (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) == QWEN3_CODEC_EXPECTED + + +def test_qwen2_5_stop_token_consistent(): + """8294: processor / model implementation / pipeline must agree.""" + assert TALKER_CODEC_END_TOKEN_ID == QWEN2_5_CODEC_END_EXPECTED + assert TALKER_CODEC_EOS_TOKEN_ID == TALKER_CODEC_END_TOKEN_ID + assert _pipeline_stop_token_ids(QWEN2_5_OMNI_PIPELINE) == [TALKER_CODEC_END_TOKEN_ID] + + +def test_qwen3_codec_stop_token_shared_across_pipelines(): + """2150: the three pipelines must share the single-sourced constant.""" + assert QWEN3_CODEC_EOS_TOKEN_ID == QWEN3_CODEC_STOP_EXPECTED + for pipeline in (QWEN3_OMNI_PIPELINE, QWEN3_TTS_PIPELINE, AURA_OMNI_PIPELINE): + assert _pipeline_stop_token_ids(pipeline) == [QWEN3_CODEC_EOS_TOKEN_ID] diff --git a/vllm_omni/model_executor/models/aura_omni/pipeline.py b/vllm_omni/model_executor/models/aura_omni/pipeline.py index 71f451c001a..221c45a1eb3 100644 --- a/vllm_omni/model_executor/models/aura_omni/pipeline.py +++ b/vllm_omni/model_executor/models/aura_omni/pipeline.py @@ -17,6 +17,9 @@ StageExecutionType, StagePipelineConfig, ) +# RFC #4872 P6: talker stop token single-sourced in the shared constants +# module (also used by qwen3_omni / qwen3_tts pipelines). +from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID _AURA_PROC = "vllm_omni.model_executor.stage_input_processors.aura_omni" _QWEN3_TTS_PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts" @@ -66,7 +69,7 @@ async_chunk_process_next_stage_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_async_chunk", sampling_constraints={ "detokenize": False, - "stop_token_ids": [2150], + "stop_token_ids": [QWEN3_CODEC_EOS_TOKEN_ID], }, ), StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py index 2036c42aa07..4d12484fc3a 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py @@ -12,6 +12,8 @@ StageExecutionType, StagePipelineConfig, ) +# RFC #4872 P6: 8294 single-sourced in the stage input processor. +from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import TALKER_CODEC_END_TOKEN_ID _PROC = "vllm_omni.model_executor.stage_input_processors.qwen2_5_omni" @@ -45,7 +47,7 @@ custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload", sampling_constraints={ "detokenize": True, - "stop_token_ids": [8294], + "stop_token_ids": [TALKER_CODEC_END_TOKEN_ID], }, requires_full_payload_input=True, ), 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..d9baf3c515f 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 @@ -44,6 +44,9 @@ Qwen2_5OmniThinkerProcessingInfo, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges +# RFC #4872 P6: the talker codec end token id is single-sourced in the stage +# input processor (``TALKER_CODEC_END_TOKEN_ID``). +from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import TALKER_CODEC_END_TOKEN_ID from vllm_omni.platforms import current_omni_platform # `get_llm_pos_ids_for_vision` was removed from upstream vLLM (commit @@ -51,7 +54,9 @@ # OmniMRotaryEmbedding. Bind it at module level to keep call sites unchanged. get_llm_pos_ids_for_vision = OmniMRotaryEmbedding._get_llm_pos_ids_for_vision -TALKER_CODEC_EOS_TOKEN_ID = 8294 +# RFC #4872 P6: alias kept for backwards compatibility; the single source is +# ``TALKER_CODEC_END_TOKEN_ID`` in ``stage_input_processors.qwen2_5_omni``. +TALKER_CODEC_EOS_TOKEN_ID = TALKER_CODEC_END_TOKEN_ID TALKER_CODEC_BOS_TOKEN_ID = 8293 diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index 7a6c07326a5..7d7f37af6f6 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -16,6 +16,9 @@ StagePipelineConfig, pipeline_cfg_resolver, ) +# RFC #4872 P6: talker stop token single-sourced in the shared constants +# module (also used by qwen3_tts / aura_omni pipelines). +from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_omni" @@ -63,7 +66,7 @@ supports_native_mrv2_data_plane=True, sampling_constraints={ "detokenize": False, - "stop_token_ids": [2150], + "stop_token_ids": [QWEN3_CODEC_EOS_TOKEN_ID], }, requires_full_payload_input=True, ), diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 92b1d8c6df5..ce08d41c754 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -64,6 +64,15 @@ ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_decode_windows +# RFC #4872 P6: codec token ids are single-sourced in the stage input +# processor (``_QWEN3_CODEC_*``); the model keeps the ``TALKER_CODEC_*`` +# aliases for backwards compatibility. At runtime the talker actually reads +# them from ``config.talker_config.codec_{pad,bos,eos}``. +from vllm_omni.model_executor.stage_input_processors.qwen3_omni import ( + _QWEN3_CODEC_BOS_TOKEN_ID, + _QWEN3_CODEC_EOS_TOKEN_ID, + _QWEN3_CODEC_PAD_TOKEN_ID, +) from vllm_omni.platforms import current_omni_platform from vllm_omni.utils.device_copy import to_device_nonblocking @@ -95,10 +104,13 @@ def resolve_default_speaker(speaker_token_ids: dict[str, int]) -> str: TTS_EOS_TOKEN_ID = 151673 # (end of dialogue) TTS_BOS_SINGLE_TOKEN_ID = 151674 # -# Talker codec tokens (talker vocabulary, used for RVQ code generation) -TALKER_CODEC_PAD_TOKEN_ID = 4196 # Padding token -TALKER_CODEC_BOS_TOKEN_ID = 4197 # Beginning of speech -TALKER_CODEC_EOS_TOKEN_ID = 4198 # End of speech +# Talker codec tokens (talker vocabulary, used for RVQ code generation). +# RFC #4872 P6: single-sourced from ``stage_input_processors.qwen3_omni`` +# (``_QWEN3_CODEC_*``); at runtime the talker reads these ids from +# ``config.talker_config.codec_{pad,bos,eos}`` instead. +TALKER_CODEC_PAD_TOKEN_ID = _QWEN3_CODEC_PAD_TOKEN_ID # Padding token +TALKER_CODEC_BOS_TOKEN_ID = _QWEN3_CODEC_BOS_TOKEN_ID # Beginning of speech +TALKER_CODEC_EOS_TOKEN_ID = _QWEN3_CODEC_EOS_TOKEN_ID # End of speech TALKER_CODEC_NOTHINK_ID = 4203 # No-think mode TALKER_CODEC_THINK_BOS_ID = 4204 # Think mode start TALKER_CODEC_THINK_EOS_ID = 4205 # Think mode end diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index e1644f24215..ebe858080b3 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -14,6 +14,9 @@ TensorAccumulationStrategy, register_key_accumulation_strategy, ) +# RFC #4872 P6: talker stop token single-sourced in the shared constants +# module (also used by qwen3_omni / aura_omni pipelines). +from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts" @@ -47,7 +50,7 @@ custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload", sampling_constraints={ "detokenize": False, - "stop_token_ids": [2150], + "stop_token_ids": [QWEN3_CODEC_EOS_TOKEN_ID], }, ), StagePipelineConfig( diff --git a/vllm_omni/model_executor/stage_input_processors/_constants.py b/vllm_omni/model_executor/stage_input_processors/_constants.py new file mode 100644 index 00000000000..87d4e5261e6 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/_constants.py @@ -0,0 +1,19 @@ +"""RFC #4872 P6: canonical cross-model constant definitions. + +Special-token ids that are shared across model pipelines / processors / +model implementations live here so no module hard-codes a magic number that +could drift out of sync. This module is deliberately **dependency-free** +(no torch / vllm imports) so pure-config modules (e.g. pipeline topologies) +can import it without pulling in the runtime stack. + +**Scope note (conservative):** only constants genuinely shared across more +than one module belong here. Model-specific codec ids keep their module-top +home with the processor as source and the HF config as the runtime single +source of truth (see ``qwen3_omni._assert_codec_token_ids_consistent`` and +``qwen2_5_omni.TALKER_CODEC_END_TOKEN_ID``). +""" + +# Qwen3-family talker codec stop (EOS) token id. Used as the AR talker +# stage's ``stop_token_ids`` across the qwen3_omni / qwen3_tts / aura_omni +# pipelines (three pipeline files shared the literal ``2150`` before P6). +QWEN3_CODEC_EOS_TOKEN_ID = 2150 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 6b2b162d378..765f3280854 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -49,6 +49,44 @@ _QWEN3_CODEC_EOS_TOKEN_ID = 4198 +def _assert_codec_token_ids_consistent() -> None: + """Fail-fast guard (RFC #4872 P6): codec token ids must not drift. + + These ids are model-specific special tokens. The HF config classes are + the single source of truth at runtime -- the talker reads + ``codec_pad_id`` / ``codec_bos_id`` / ``codec_eos_token_id`` off + ``talker_config``. The processor keeps its own module-top constants + (used by ``talker2code2wav_*`` bookkeeping) and verifies them against + both the in-repo Qwen3-TTS config and the transformers + ``Qwen3OmniMoeTalkerConfig`` (the config qwen3-omni actually runs with). + """ + from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( + Qwen3TTSTalkerConfig, + ) + from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeTalkerConfig, + ) + + expected = ( + _QWEN3_CODEC_PAD_TOKEN_ID, + _QWEN3_CODEC_BOS_TOKEN_ID, + _QWEN3_CODEC_EOS_TOKEN_ID, + ) + for name, cfg in ( + ("qwen3_tts", Qwen3TTSTalkerConfig()), + ("qwen3_omni_moe", Qwen3OmniMoeTalkerConfig()), + ): + actual = (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) + assert actual == expected, ( + f"{name} codec token ids {actual} drift from the " + f"stage_input_processors.qwen3_omni constants {expected}" + ) + + +# Single source of truth: keep processor / model impl / HF config in sync. +_assert_codec_token_ids_consistent() + + def _layer_tensor(layers: dict[Any, Any], key: str) -> torch.Tensor | None: """Fetch layer tensor with tolerant key lookup (str/int/semantic).""" if not isinstance(layers, dict): From 51cfc7d56affb1cca035d3778c27dac1571533d1 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sun, 30 Aug 2026 12:48:44 +0800 Subject: [PATCH 05/41] fix(stage_input_processors): resolve precheck-pr blockers (ruff/spdx/test marks/mypy) Signed-off-by: RyanYun09 --- .../test_orchestrator_stage_input_bridge.py | 2 +- .../stage_input_processors/conftest.py | 17 +- .../test_common_helpers_golden.py | 312 +++++++++--------- .../test_constants_consistency.py | 10 +- .../test_placeholder_parity.py | 3 +- .../stage_input_processors/test_registry.py | 6 +- vllm_omni/engine/orchestrator.py | 3 +- .../models/aura_omni/pipeline.py | 1 + .../models/qwen2_5_omni/pipeline.py | 1 + .../models/qwen2_5_omni/qwen2_5_omni.py | 1 + .../models/qwen3_omni/pipeline.py | 1 + .../models/qwen3_omni/qwen3_omni.py | 1 + .../models/qwen3_tts/pipeline.py | 1 + .../stage_input_processors/__init__.py | 3 + .../stage_input_processors/_common.py | 43 +-- .../stage_input_processors/_constants.py | 3 + .../stage_input_processors/_dispatch.py | 6 +- .../stage_input_processors/_registry.py | 8 +- .../stage_input_processors/audex.py | 2 +- .../stage_input_processors/higgs_audio_v2.py | 4 +- .../stage_input_processors/higgs_audio_v3.py | 4 +- .../stage_input_processors/qwen3_omni.py | 10 +- 22 files changed, 235 insertions(+), 207 deletions(-) diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index e580c296efd..09e941c9f10 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -314,7 +314,7 @@ async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: stage1 = FakePrewarmPool("receiver") # The real qwen3_omni sync placeholder builder exposes build_prewarm_placeholder # off the resolved fn (see test_placeholder_parity.py). - stage1.stage_client = SimpleNamespace( + stage1.stage_client = SimpleNamespace( # type: ignore[attr-defined] sync_process_input_func=( "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only" ), diff --git a/tests/model_executor/stage_input_processors/conftest.py b/tests/model_executor/stage_input_processors/conftest.py index 2eb57c78f60..64ae9ade526 100644 --- a/tests/model_executor/stage_input_processors/conftest.py +++ b/tests/model_executor/stage_input_processors/conftest.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Local-dev import shim for ``vllm_omni`` tests (macOS). vllm does not install on macOS, so importing real ``vllm_omni`` submodules @@ -28,11 +31,7 @@ # Locate the repo root (this file lives at: # /tests/model_executor/stage_input_processors/conftest.py) # --------------------------------------------------------------------------- -_REPO_ROOT = os.path.dirname( - os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ) -) +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) _VLLM_OMNI_DIR = os.path.join(_REPO_ROOT, "vllm_omni") @@ -84,7 +83,7 @@ class _StubLoader(importlib.abc.Loader): def create_module(self, spec): module = types.ModuleType(spec.name) module.__path__ = [] - module.__getattr__ = lambda attr: _Dummy + module.__getattr__ = lambda attr: _Dummy # type: ignore[method-assign] return module def exec_module(self, module): @@ -96,9 +95,7 @@ class _StubFinder(importlib.abc.MetaPathFinder): def find_spec(self, fullname, path=None, target=None): if fullname == "vllm" or fullname.startswith("vllm."): - return importlib.machinery.ModuleSpec( - fullname, _StubLoader(), is_package=True - ) + return importlib.machinery.ModuleSpec(fullname, _StubLoader(), is_package=True) return None @@ -110,7 +107,7 @@ def _install_shim(): pkg = types.ModuleType("vllm_omni") pkg.__path__ = [_VLLM_OMNI_DIR] pkg.__package__ = "vllm_omni" - pkg._SHIM_ACTIVE = True + pkg._SHIM_ACTIVE = True # type: ignore[attr-defined] sys.modules.setdefault("vllm_omni", pkg) # Wildcard vllm stub finder (meta path -> takes precedence for vllm.*). sys.meta_path.insert(0, _StubFinder()) diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py index eec93d2ba0d..b646349ba4c 100644 --- a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Golden tests that lock the *current* per-module behavior of duplicated helper functions in ``stage_input_processors``. @@ -14,11 +17,13 @@ """ import importlib -from typing import Any, Callable +from typing import Any import pytest import torch +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + _PREFIX = "vllm_omni.model_executor.stage_input_processors." @@ -96,8 +101,7 @@ def _ensure_list_cases(): ] -@pytest.mark.parametrize("module,inp,expected,expect_typeerror", - _ensure_list_cases()) +@pytest.mark.parametrize("module,inp,expected,expect_typeerror", _ensure_list_cases()) def test_ensure_list_golden(module, inp, expected, expect_typeerror): mod = _import(module) fn = getattr(mod, "_ensure_list") @@ -126,86 +130,69 @@ def _codes_audio(tensor): return {"codes": {"audio": tensor}} -@pytest.mark.parametrize("module,payload,expected,expect_error", [ - # higgs_audio_v2: codes.audio key path; range filter [-0, _NUM_REAL_CODES). - ("higgs_audio_v2", - _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), - [4, 5], False), - ("higgs_audio_v2", - _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 1024]])), # >= _NUM_REAL_CODES - None, False), - ("higgs_audio_v2", - _codes_audio(torch.tensor([[0, 1], [2, 3], [-1, 5]])), - None, False), - ("higgs_audio_v2", - _codes_audio(torch.tensor([1, 2, 3])), - [1, 2, 3], False), - ("higgs_audio_v2", - _codes_audio(torch.empty(0, 4)), - None, False), - ("higgs_audio_v2", - _codes_audio(torch.zeros(2, 2, 2)), - None, True), # 3-D -> ValueError - ("higgs_audio_v2", - {"codes": {}}, - None, False), - # fish_speech: TOP-LEVEL audio_codes key; audio_code_valid mask; cpu+long. - ("fish_speech", - {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), - "audio_code_valid": torch.tensor([1, 1, 1])}, - [4, 5], False), - ("fish_speech", - {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), - "audio_code_valid": torch.tensor([1, 1, 0])}, - None, False), - ("fish_speech", - {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), - "audio_code_valid": True}, - [4, 5], False), - ("fish_speech", - {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), - "audio_code_valid": False}, - None, False), - ("fish_speech", - {"audio_codes": torch.tensor([[0, 1], [2, 3], [0, 0]])}, - None, False), # no valid key -> frame.any() False - ("fish_speech", - {"audio_codes": torch.tensor([1, 2, 3])}, - [1, 2, 3], False), - ("fish_speech", - {"audio_codes": torch.zeros(2, 2, 2)}, - None, True), - ("fish_speech", - {}, - None, False), - # qwen3_tts: codes.audio; frame.any() gate; long (no cpu). - ("qwen3_tts", - _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), - [4, 5], False), - ("qwen3_tts", - _codes_audio(torch.tensor([[0, 1], [2, 3], [0, 0]])), - None, False), - ("qwen3_tts", - _codes_audio(torch.tensor([1, 2, 3])), - [1, 2, 3], False), - ("qwen3_tts", - _codes_audio(torch.zeros(2, 2, 2)), - None, True), - ("qwen3_tts", - {"codes": {}}, - None, False), - # voxtral_tts: for a single tensor it flattens the WHOLE tensor (not just - # the last row); for a list of tensors it flattens the last tensor. - ("voxtral_tts", - _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), - [0, 1, 2, 3, 4, 5], False), - ("voxtral_tts", - {"codes": {"audio": [torch.tensor([1, 2]), torch.tensor([3, 4])]}}, - [3, 4], False), - ("voxtral_tts", - {"codes": {}}, - None, False), -]) +@pytest.mark.parametrize( + "module,payload,expected,expect_error", + [ + # higgs_audio_v2: codes.audio key path; range filter [-0, _NUM_REAL_CODES). + ("higgs_audio_v2", _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False), + ( + "higgs_audio_v2", + _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 1024]])), # >= _NUM_REAL_CODES + None, + False, + ), + ("higgs_audio_v2", _codes_audio(torch.tensor([[0, 1], [2, 3], [-1, 5]])), None, False), + ("higgs_audio_v2", _codes_audio(torch.tensor([1, 2, 3])), [1, 2, 3], False), + ("higgs_audio_v2", _codes_audio(torch.empty(0, 4)), None, False), + ("higgs_audio_v2", _codes_audio(torch.zeros(2, 2, 2)), None, True), # 3-D -> ValueError + ("higgs_audio_v2", {"codes": {}}, None, False), + # fish_speech: TOP-LEVEL audio_codes key; audio_code_valid mask; cpu+long. + ( + "fish_speech", + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": torch.tensor([1, 1, 1])}, + [4, 5], + False, + ), + ( + "fish_speech", + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": torch.tensor([1, 1, 0])}, + None, + False, + ), + ( + "fish_speech", + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": True}, + [4, 5], + False, + ), + ( + "fish_speech", + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": False}, + None, + False, + ), + ( + "fish_speech", + {"audio_codes": torch.tensor([[0, 1], [2, 3], [0, 0]])}, + None, + False, + ), # no valid key -> frame.any() False + ("fish_speech", {"audio_codes": torch.tensor([1, 2, 3])}, [1, 2, 3], False), + ("fish_speech", {"audio_codes": torch.zeros(2, 2, 2)}, None, True), + ("fish_speech", {}, None, False), + # qwen3_tts: codes.audio; frame.any() gate; long (no cpu). + ("qwen3_tts", _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False), + ("qwen3_tts", _codes_audio(torch.tensor([[0, 1], [2, 3], [0, 0]])), None, False), + ("qwen3_tts", _codes_audio(torch.tensor([1, 2, 3])), [1, 2, 3], False), + ("qwen3_tts", _codes_audio(torch.zeros(2, 2, 2)), None, True), + ("qwen3_tts", {"codes": {}}, None, False), + # voxtral_tts: for a single tensor it flattens the WHOLE tensor (not just + # the last row); for a list of tensors it flattens the last tensor. + ("voxtral_tts", _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [0, 1, 2, 3, 4, 5], False), + ("voxtral_tts", {"codes": {"audio": [torch.tensor([1, 2]), torch.tensor([3, 4])]}}, [3, 4], False), + ("voxtral_tts", {"codes": {}}, None, False), + ], +) def test_extract_last_frame_golden(module, payload, expected, expect_error): mod = _import(module) fn = getattr(mod, "_extract_last_frame") @@ -225,11 +212,13 @@ def test_extract_last_frame_golden(module, payload, expected, expect_error): def test_revert_delay_pattern_higgs_v2_golden(): mod = _import("higgs_audio_v2") fn = getattr(mod, "_revert_delay_pattern") - inp = torch.tensor([ - [10, 1, 2, 3, 4], - [11, 12, 5, 6, 7], - [13, 14, 15, 8, 9], - ]) # [Q=3, T=5], seq_len = 3 + inp = torch.tensor( + [ + [10, 1, 2, 3, 4], + [11, 12, 5, 6, 7], + [13, 14, 15, 8, 9], + ] + ) # [Q=3, T=5], seq_len = 3 out = fn(inp) assert out.shape == (3, 3) assert out.tolist() == [[10, 1, 2], [12, 5, 6], [15, 8, 9]] @@ -250,9 +239,7 @@ def test_revert_delay_pattern_higgs_v3_golden(): inp = torch.arange(q * t, dtype=torch.long).reshape(q, t) out = fn(inp) assert out.shape == (q, 3) - expected = torch.cat( - [inp[i : i + 1, i : 3 + i] for i in range(q)], dim=0 - ) + expected = torch.cat([inp[i : i + 1, i : 3 + i] for i in range(q)], dim=0) assert torch.equal(out, expected) # Strict: wrong codebook count raises. with pytest.raises(ValueError): @@ -271,12 +258,14 @@ def test_filter_real_code_frames_higgs_v2_golden(): mod = _import("higgs_audio_v2") fn = getattr(mod, "_filter_real_code_frames") nrc = int(mod._NUM_REAL_CODES) - inp = torch.tensor([ - [0, 1, 2], - [nrc, 5, 6], # invalid (>= _NUM_REAL_CODES) - [7, 8, 9], - [-1, 0, 1], # invalid (< 0) - ]) + inp = torch.tensor( + [ + [0, 1, 2], + [nrc, 5, 6], # invalid (>= _NUM_REAL_CODES) + [7, 8, 9], + [-1, 0, 1], # invalid (< 0) + ] + ) out = fn(inp) assert out.tolist() == [[0, 1, 2], [7, 8, 9]] # Empty input passes through. @@ -289,12 +278,18 @@ def test_filter_real_code_frames_higgs_v3_golden(): fn = getattr(mod, "_filter_real_code_frames") nrc = int(mod._NUM_REAL_CODES) # Input is [Q, frames]: 4 frames x 3 codebooks -> transpose to [Q=3, frames=4]. - inp = torch.tensor([ - [0, 1, 2], # frame 0 (valid) - [nrc, 5, 6], # frame 1 (invalid: >= _NUM_REAL_CODES) - [7, 8, 9], # frame 2 (valid) - [-1, 0, 1], # frame 3 (invalid: < 0) - ]).t().contiguous() + inp = ( + torch.tensor( + [ + [0, 1, 2], # frame 0 (valid) + [nrc, 5, 6], # frame 1 (invalid: >= _NUM_REAL_CODES) + [7, 8, 9], # frame 2 (valid) + [-1, 0, 1], # frame 3 (invalid: < 0) + ] + ) + .t() + .contiguous() + ) out = fn(inp) # keeps frames 0 and 2 assert out.shape == (3, 2) assert out.tolist() == [[0, 7], [1, 8], [2, 9]] @@ -307,14 +302,17 @@ def test_filter_real_code_frames_higgs_v3_golden(): # =========================================================================== -@pytest.mark.parametrize("inp,expected", [ - ("TENSOR", "TENSOR"), - ([], None), - ("LIST_TENSOR", "LIST_TENSOR"), - ([1, 2], None), - (None, None), - (5, None), -]) +@pytest.mark.parametrize( + "inp,expected", + [ + ("TENSOR", "TENSOR"), + ([], None), + ("LIST_TENSOR", "LIST_TENSOR"), + ([1, 2], None), + (None, None), + (5, None), + ], +) def test_to_cpu_tensor_glm_tts_golden(inp, expected): mod = _import("glm_tts") fn = getattr(mod, "_to_cpu_tensor") @@ -339,15 +337,18 @@ def test_to_cpu_tensor_glm_tts_golden(inp, expected): # =========================================================================== -@pytest.mark.parametrize("inp,expected", [ - ("SCALAR_TENSOR", [5]), - ("VEC_TENSOR", [1, 2, 3]), - ("MAT_TENSOR", [1, 2]), # 2-D -> first row only - ("NESTED_LIST", [1, 2]), # nested list -> first nested row only - ([1, 2, 3], [1, 2, 3]), - (None, []), - (5, [5]), -]) +@pytest.mark.parametrize( + "inp,expected", + [ + ("SCALAR_TENSOR", [5]), + ("VEC_TENSOR", [1, 2, 3]), + ("MAT_TENSOR", [1, 2]), # 2-D -> first row only + ("NESTED_LIST", [1, 2]), # nested list -> first nested row only + ([1, 2, 3], [1, 2, 3]), + (None, []), + (5, [5]), + ], +) def test_to_token_id_list_dynin_omni_golden(inp, expected): mod = _import("dynin_omni") fn = getattr(mod, "_to_token_id_list") @@ -396,6 +397,7 @@ def test_compute_talker_prompt_ids_length_golden(): def _common_import() -> Any: try: import vllm_omni.model_executor.stage_input_processors._common as c + return c except Exception as exc: # pragma: no cover - env dependent pytest.skip(f"_common not importable here: {type(exc).__name__}: {exc}") @@ -405,6 +407,7 @@ def _shim_active() -> bool: """True when the local-dev vllm stub (no real vllm) is in effect.""" try: import vllm_omni # noqa: F401 + return bool(getattr(vllm_omni, "_SHIM_ACTIVE", False)) except Exception: return False @@ -450,21 +453,19 @@ def test_common_to_token_id_list_matches_golden(): assert c.to_token_id_list(5) == [5] # cosyvoice3 semantics: recursive flatten. assert c.to_token_id_list([[1, 2], [3, 4]], recursive=True) == [1, 2, 3, 4] - assert c.to_token_id_list( - torch.tensor([[1, 2], [3, 4]]), recursive=True - ) == [1, 2, 3, 4] + assert c.to_token_id_list(torch.tensor([[1, 2], [3, 4]]), recursive=True) == [1, 2, 3, 4] def test_common_revert_delay_pattern_matches_golden(): c = _common_import() - inp = torch.tensor([ - [10, 1, 2, 3, 4], - [11, 12, 5, 6, 7], - [13, 14, 15, 8, 9], - ]) # [Q=3, T=5] - assert c.revert_delay_pattern(inp, allow_short=True).tolist() == [ - [10, 1, 2], [12, 5, 6], [15, 8, 9] - ] + inp = torch.tensor( + [ + [10, 1, 2, 3, 4], + [11, 12, 5, 6, 7], + [13, 14, 15, 8, 9], + ] + ) # [Q=3, T=5] + assert c.revert_delay_pattern(inp, allow_short=True).tolist() == [[10, 1, 2], [12, 5, 6], [15, 8, 9]] # Lenient (v2): T < Q returns input unchanged. short = torch.tensor([[1], [2]]) assert torch.equal(c.revert_delay_pattern(short, allow_short=True), short) @@ -479,32 +480,45 @@ def test_common_filter_real_code_frames_matches_golden(): c = _common_import() nrc = 1024 # frames-first (v2 layout). - frames_first = torch.tensor([ - [0, 1, 2], [nrc, 5, 6], [7, 8, 9], [-1, 0, 1], - ]) - assert c.filter_real_code_frames( - frames_first, num_real_codes=nrc, layout="frames_first" - ).tolist() == [[0, 1, 2], [7, 8, 9]] + frames_first = torch.tensor( + [ + [0, 1, 2], + [nrc, 5, 6], + [7, 8, 9], + [-1, 0, 1], + ] + ) + assert c.filter_real_code_frames(frames_first, num_real_codes=nrc, layout="frames_first").tolist() == [ + [0, 1, 2], + [7, 8, 9], + ] # codebooks-first (v3 layout). - codebooks_first = torch.tensor([ - [0, 1, 2], [nrc, 5, 6], [7, 8, 9], [-1, 0, 1], - ]).t().contiguous() # [Q=3, frames=4] - assert c.filter_real_code_frames( - codebooks_first, num_real_codes=nrc, layout="codebooks_first" - ).tolist() == [[0, 7], [1, 8], [2, 9]] + codebooks_first = ( + torch.tensor( + [ + [0, 1, 2], + [nrc, 5, 6], + [7, 8, 9], + [-1, 0, 1], + ] + ) + .t() + .contiguous() + ) # [Q=3, frames=4] + assert c.filter_real_code_frames(codebooks_first, num_real_codes=nrc, layout="codebooks_first").tolist() == [ + [0, 7], + [1, 8], + [2, 9], + ] def test_common_compute_placeholder_prompt_len_matches_golden(): c = _common_import() prompt = [151644, 872, 10, 11, 12, 13, 151644, 77091, 20, 21] # full mode == qwen3_omni._compute_talker_prompt_ids_length golden (15). - assert c.compute_placeholder_prompt_len( - ids_or_prompt={"ids": {"all": prompt, "prompt": prompt}}, mode="full" - ) == 15 + assert c.compute_placeholder_prompt_len(ids_or_prompt={"ids": {"all": prompt, "prompt": prompt}}, mode="full") == 15 # stage0_only mode: just the stage-0 input prompt length (prewarm). - assert c.compute_placeholder_prompt_len( - ids_or_prompt=prompt, mode="stage0_only" - ) == len(prompt) + assert c.compute_placeholder_prompt_len(ids_or_prompt=prompt, mode="stage0_only") == len(prompt) def test_common_pack_placeholder_prompt(): diff --git a/tests/model_executor/stage_input_processors/test_constants_consistency.py b/tests/model_executor/stage_input_processors/test_constants_consistency.py index 70f619cef91..abc8c15f1af 100644 --- a/tests/model_executor/stage_input_processors/test_constants_consistency.py +++ b/tests/model_executor/stage_input_processors/test_constants_consistency.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """RFC #4872 P6: constant-consistency guards. Verifies that magic special-token ids that were previously duplicated across @@ -16,15 +16,19 @@ CPU-only: imports modules and reads constants; never loads model weights. """ +import pytest + from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import QWEN2_5_OMNI_PIPELINE from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni import TALKER_CODEC_EOS_TOKEN_ID from vllm_omni.model_executor.models.qwen3_omni.pipeline import QWEN3_OMNI_PIPELINE from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import ( TALKER_CODEC_BOS_TOKEN_ID, - TALKER_CODEC_EOS_TOKEN_ID as QWEN3_MODEL_CODEC_EOS, TALKER_CODEC_PAD_TOKEN_ID, ) +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import ( + TALKER_CODEC_EOS_TOKEN_ID as QWEN3_MODEL_CODEC_EOS, +) from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSTalkerConfig from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID @@ -35,6 +39,8 @@ _QWEN3_CODEC_PAD_TOKEN_ID, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + # Expected values locked by the baseline commit `731a771`. QWEN3_CODEC_EXPECTED = (4196, 4197, 4198) # (pad, bos, eos) QWEN2_5_CODEC_END_EXPECTED = 8294 diff --git a/tests/model_executor/stage_input_processors/test_placeholder_parity.py b/tests/model_executor/stage_input_processors/test_placeholder_parity.py index 354a11682ab..7c652798472 100644 --- a/tests/model_executor/stage_input_processors/test_placeholder_parity.py +++ b/tests/model_executor/stage_input_processors/test_placeholder_parity.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """P8b parity tests: async-chunk prewarm placeholder vs sync forward placeholder. RFC #4872 Proposal 8b: ``*_token_only`` placeholder builders expose a dual entry @@ -52,6 +52,7 @@ def _shim_active() -> bool: """True when the local-dev vllm stub (no real vllm) is in effect.""" try: import vllm_omni # noqa: F401 + return bool(getattr(vllm_omni, "_SHIM_ACTIVE", False)) except Exception: return False diff --git a/tests/model_executor/stage_input_processors/test_registry.py b/tests/model_executor/stage_input_processors/test_registry.py index 968d70d9e88..c7085cee55b 100644 --- a/tests/model_executor/stage_input_processors/test_registry.py +++ b/tests/model_executor/stage_input_processors/test_registry.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for the RFC #4872 Phase 3 (Part A) processor registry. Pure-logic, CPU-only: exercises kind inference, structural validation, @@ -16,14 +19,15 @@ ProcessorKind, ProcessorSpec, ProcessorValidationError, + _registry, dead_processor_hint, infer_kind, register_processor, resolve_processor, validate_processor, ) -from vllm_omni.model_executor.stage_input_processors import _registry +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] # --------------------------------------------------------------------------- # Dummy processors (simple callables used as validation subjects). diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 7e0b007d090..1265a00995c 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -3208,8 +3208,7 @@ def _get_prewarm_placeholder_builder(self, stage_id: int) -> Any: return builder except Exception as exc: # pragma: no cover - defensive only logger.warning( - "[Orchestrator] prewarm placeholder builder lookup for stage-%s " - "failed; using the inline estimate: %s", + "[Orchestrator] prewarm placeholder builder lookup for stage-%s failed; using the inline estimate: %s", stage_id, exc, ) diff --git a/vllm_omni/model_executor/models/aura_omni/pipeline.py b/vllm_omni/model_executor/models/aura_omni/pipeline.py index 221c45a1eb3..ccdd7813867 100644 --- a/vllm_omni/model_executor/models/aura_omni/pipeline.py +++ b/vllm_omni/model_executor/models/aura_omni/pipeline.py @@ -17,6 +17,7 @@ StageExecutionType, StagePipelineConfig, ) + # RFC #4872 P6: talker stop token single-sourced in the shared constants # module (also used by qwen3_omni / qwen3_tts pipelines). from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py index 4d12484fc3a..0c4ffb131d2 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py @@ -12,6 +12,7 @@ StageExecutionType, StagePipelineConfig, ) + # RFC #4872 P6: 8294 single-sourced in the stage input processor. from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import TALKER_CODEC_END_TOKEN_ID 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 d9baf3c515f..6ce6b8bf03e 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 @@ -44,6 +44,7 @@ Qwen2_5OmniThinkerProcessingInfo, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges + # RFC #4872 P6: the talker codec end token id is single-sourced in the stage # input processor (``TALKER_CODEC_END_TOKEN_ID``). from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import TALKER_CODEC_END_TOKEN_ID diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index 7d7f37af6f6..cebc9cf1fb1 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -16,6 +16,7 @@ StagePipelineConfig, pipeline_cfg_resolver, ) + # RFC #4872 P6: talker stop token single-sourced in the shared constants # module (also used by qwen3_tts / aura_omni pipelines). from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index ce08d41c754..8f246ebdb1b 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -64,6 +64,7 @@ ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_decode_windows + # RFC #4872 P6: codec token ids are single-sourced in the stage input # processor (``_QWEN3_CODEC_*``); the model keeps the ``TALKER_CODEC_*`` # aliases for backwards compatibility. At runtime the talker actually reads diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index ebe858080b3..8233f5150aa 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -14,6 +14,7 @@ TensorAccumulationStrategy, register_key_accumulation_strategy, ) + # RFC #4872 P6: talker stop token single-sourced in the shared constants # module (also used by qwen3_omni / aura_omni pipelines). from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID diff --git a/vllm_omni/model_executor/stage_input_processors/__init__.py b/vllm_omni/model_executor/stage_input_processors/__init__.py index e7e2bcce0a6..3fc95647204 100644 --- a/vllm_omni/model_executor/stage_input_processors/__init__.py +++ b/vllm_omni/model_executor/stage_input_processors/__init__.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Stage input processors: per-model builders for inter-stage transitions. Module-naming convention (RFC #4872): diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py index 30bb02e1d8e..c83dafadef9 100644 --- a/vllm_omni/model_executor/stage_input_processors/_common.py +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Canonical, shared helpers for ``stage_input_processors``. RFC #4872 Proposal 3 / P8a. These are consolidated implementations of @@ -23,11 +26,11 @@ from vllm_omni.inputs.data import OmniTokensPrompt - # =========================================================================== # ensure_list # =========================================================================== + def _unwrap_constant_list(x: Any) -> Any: """Unwrap a ``ConstantList``-like object exposing ``_x``.""" if hasattr(x, "_x"): @@ -134,6 +137,7 @@ def ensure_list_preserve_none(x: Any) -> list[Any]: # to_cpu_tensor # =========================================================================== + def to_cpu_tensor(value: Any) -> torch.Tensor | None: """Convert a value to a CPU tensor when possible; else ``None``. @@ -155,6 +159,7 @@ def to_cpu_tensor(value: Any) -> torch.Tensor | None: # to_token_id_list # =========================================================================== + def to_token_id_list(value: Any, *, recursive: bool = False) -> list[int]: """Convert a token-ish value into a flat ``list[int]``. @@ -193,6 +198,7 @@ def to_token_id_list(value: Any, *, recursive: bool = False) -> list[int]: # revert_delay_pattern (higgs_audio_v2 lenient vs higgs_audio_v3 strict) # =========================================================================== + def revert_delay_pattern( audio_codes_qt: torch.Tensor, *, @@ -210,16 +216,10 @@ def revert_delay_pattern( _NUM_CODEBOOKS``; v2 does not). """ if audio_codes_qt.ndim != 2: - raise ValueError( - f"_revert_delay_pattern expects [Q, T] input; " - f"got {tuple(audio_codes_qt.shape)}" - ) + raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}") q, t = audio_codes_qt.shape if expected_codebooks is not None and q != expected_codebooks: - raise ValueError( - f"Expected exactly {expected_codebooks} codebook rows, got {q}. " - f"Input shape: [{q}, {t}]" - ) + raise ValueError(f"Expected exactly {expected_codebooks} codebook rows, got {q}. Input shape: [{q}, {t}]") if t < q: if allow_short: return audio_codes_qt @@ -233,6 +233,7 @@ def revert_delay_pattern( # filter_real_code_frames (higgs v2 [frames,Q] vs v3 [Q,frames]) # =========================================================================== + def filter_real_code_frames( audio_codes: torch.Tensor, *, @@ -251,13 +252,8 @@ def filter_real_code_frames( return audio_codes if layout == "frames_first": if audio_codes.ndim != 2: - raise ValueError( - f"expected [num_frames, num_codebooks] audio_codes; " - f"got shape {tuple(audio_codes.shape)}" - ) - valid = (audio_codes >= 0).all(dim=1) & ( - audio_codes < num_real_codes - ).all(dim=1) + raise ValueError(f"expected [num_frames, num_codebooks] audio_codes; got shape {tuple(audio_codes.shape)}") + valid = (audio_codes >= 0).all(dim=1) & (audio_codes < num_real_codes).all(dim=1) return audio_codes[valid] if layout == "codebooks_first": frames = audio_codes.t() @@ -270,6 +266,7 @@ def filter_real_code_frames( # extract_last_codec_frame # =========================================================================== + def extract_last_codec_frame( payload: Any, *, @@ -314,9 +311,7 @@ def extract_last_codec_frame( elif audio_codes.ndim == 1: frame = audio_codes else: - raise ValueError( - f"unexpected audio_codes shape: {tuple(audio_codes.shape)}" - ) + raise ValueError(f"unexpected audio_codes shape: {tuple(audio_codes.shape)}") if frame.numel() == 0: return None if validate in ("any", "valid_mask"): @@ -376,12 +371,8 @@ def compute_placeholder_prompt_len( raise ValueError(f"unknown mode: {mode!r}") ids = (ids_or_prompt or {}).get("ids", {}) - thinker_sequences = torch.tensor( - ids["all"], dtype=torch.long, device=device - ).unsqueeze(0) # [1, T] - input_ids = torch.tensor( - ids["prompt"], dtype=torch.long, device=device - ).unsqueeze(0) # [1, T] + thinker_sequences = torch.tensor(ids["all"], dtype=torch.long, device=device).unsqueeze(0) # [1, T] + input_ids = torch.tensor(ids["prompt"], dtype=torch.long, device=device).unsqueeze(0) # [1, T] im_start_indexes = torch.cat( [ @@ -423,5 +414,5 @@ def pack_placeholder_prompt( """ return OmniTokensPrompt( prompt_token_ids=[0] * max(1, int(prompt_len)), - additional_information=(voice_metadata if voice_metadata else None), + additional_information=(voice_metadata if voice_metadata else None), # type: ignore[typeddict-item] ) diff --git a/vllm_omni/model_executor/stage_input_processors/_constants.py b/vllm_omni/model_executor/stage_input_processors/_constants.py index 87d4e5261e6..88a9e9444ae 100644 --- a/vllm_omni/model_executor/stage_input_processors/_constants.py +++ b/vllm_omni/model_executor/stage_input_processors/_constants.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """RFC #4872 P6: canonical cross-model constant definitions. Special-token ids that are shared across model pipelines / processors / diff --git a/vllm_omni/model_executor/stage_input_processors/_dispatch.py b/vllm_omni/model_executor/stage_input_processors/_dispatch.py index ace4e02ad47..dea5f990df0 100644 --- a/vllm_omni/model_executor/stage_input_processors/_dispatch.py +++ b/vllm_omni/model_executor/stage_input_processors/_dispatch.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """RFC #4872 Phase 2 (P1): orchestrator input-dispatch contract layer. This module defines the **consumer-side (orchestrator / stage-client) dispatch @@ -272,8 +275,7 @@ def wrap_orchestrator_processor(fn: Any) -> PlaceholderPromptBuilder | Diffusion if "sampling_params" in names: def _c3(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: - return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data, - sampling_params=ctx.sampling_params) + return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data, sampling_params=ctx.sampling_params) _warn_legacy_contract(fn) return _c3 # type: ignore[return-value] diff --git a/vllm_omni/model_executor/stage_input_processors/_registry.py b/vllm_omni/model_executor/stage_input_processors/_registry.py index 50057b74b7f..3eff813e077 100644 --- a/vllm_omni/model_executor/stage_input_processors/_registry.py +++ b/vllm_omni/model_executor/stage_input_processors/_registry.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """RFC #4872 Phase 3 (Part A): stage-input processor registry. This module is the **structural validation** layer behind the three ``importlib`` @@ -414,7 +417,10 @@ def _check_async_chunk_is_finished( "it must be keyword-only or have a default", ) return # unreachable - if is_finished_param.kind != inspect.Parameter.KEYWORD_ONLY and is_finished_param.default is inspect.Parameter.empty: + if ( + is_finished_param.kind != inspect.Parameter.KEYWORD_ONLY + and is_finished_param.default is inspect.Parameter.empty + ): _warn( path, kind, diff --git a/vllm_omni/model_executor/stage_input_processors/audex.py b/vllm_omni/model_executor/stage_input_processors/audex.py index 2f3589dd514..3da58cf0209 100644 --- a/vllm_omni/model_executor/stage_input_processors/audex.py +++ b/vllm_omni/model_executor/stage_input_processors/audex.py @@ -21,8 +21,8 @@ from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct from vllm_omni.inputs.data import OmniTokensPrompt -from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt from vllm_omni.model_executor.stage_input_processors import _common +from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt logger = init_logger(__name__) diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py index 668c2e80930..c1b00947f36 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py @@ -55,9 +55,7 @@ def _filter_real_code_frames(audio_codes: torch.Tensor) -> torch.Tensor: RFC #4872 P3: delegate to the canonical ``_common.filter_real_code_frames`` with the higgs-v2 ``[num_frames, num_codebooks]`` layout. """ - return _common.filter_real_code_frames( - audio_codes, num_real_codes=_NUM_REAL_CODES, layout="frames_first" - ) + return _common.filter_real_code_frames(audio_codes, num_real_codes=_NUM_REAL_CODES, layout="frames_first") def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py index 86eac44b652..8b781499aa8 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py @@ -79,9 +79,7 @@ def _filter_real_code_frames(audio_codes_qt: torch.Tensor) -> torch.Tensor: RFC #4872 P3: delegate to the canonical ``_common.filter_real_code_frames`` with the higgs-v3 ``[num_codebooks, num_frames]`` layout. """ - return _common.filter_real_code_frames( - audio_codes_qt, num_real_codes=_NUM_REAL_CODES, layout="codebooks_first" - ) + return _common.filter_real_code_frames(audio_codes_qt, num_real_codes=_NUM_REAL_CODES, layout="codebooks_first") def _whole_utterance_codec_codes(audio_rows: torch.Tensor) -> torch.Tensor: diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 765f3280854..094f4aa5b57 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -60,13 +60,14 @@ def _assert_codec_token_ids_consistent() -> None: both the in-repo Qwen3-TTS config and the transformers ``Qwen3OmniMoeTalkerConfig`` (the config qwen3-omni actually runs with). """ - from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( - Qwen3TTSTalkerConfig, - ) from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeTalkerConfig, ) + from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( + Qwen3TTSTalkerConfig, + ) + expected = ( _QWEN3_CODEC_PAD_TOKEN_ID, _QWEN3_CODEC_BOS_TOKEN_ID, @@ -78,8 +79,7 @@ def _assert_codec_token_ids_consistent() -> None: ): actual = (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) assert actual == expected, ( - f"{name} codec token ids {actual} drift from the " - f"stage_input_processors.qwen3_omni constants {expected}" + f"{name} codec token ids {actual} drift from the stage_input_processors.qwen3_omni constants {expected}" ) From 7037100cd2ff722e29a0e6e742bfcc79c92ffce2 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sun, 30 Aug 2026 15:01:04 +0800 Subject: [PATCH 06/41] docs(design): document stage input processor contract (RFC #4872) Add a Stage Input Processor Contract section to the async chunk design document covering the four-role suffix convention, the minimal processor set required in both async and non-async modes, the OrchestratorInputContext C1 contract, the processor registry and its startup validation semantics, the P8b forward/prewarm placeholder dual entry, the three-state async gate, and the load-bearing producer keyword contract. Link the section from docs/design/index.md (RFC P7). Signed-off-by: RyanYun09 --- docs/design/feature/async_chunk.md | 170 ++++++++++++++++++++++++++++- docs/design/index.md | 1 + 2 files changed, 170 insertions(+), 1 deletion(-) diff --git a/docs/design/feature/async_chunk.md b/docs/design/feature/async_chunk.md index 1448cf55f52..b34e262f458 100644 --- a/docs/design/feature/async_chunk.md +++ b/docs/design/feature/async_chunk.md @@ -6,7 +6,8 @@ 2. [Performance](#performance) 3. [Architecture](#architecture) 4. [Configuration](#configuration) -5. [Related Files](#related-files) +5. [Stage Input Processor Contract](#stage-input-processor-contract) +6. [Related Files](#related-files) ## Overview @@ -284,6 +285,173 @@ stages: max_num_seqs: 64 # Enables batched audio generation ``` +## Stage Input Processor Contract + +RFC #4872 standardized the per-model builders that convert one stage's outputs +into the next stage's inputs. Each builder belongs to exactly one *role* +(consumer-side vs producer-side) and is dispatched by the runtime according to +the `async_chunk` mode. The naming convention and the validation rules below +are load-bearing: the processor registry infers a builder's kind from its name +suffix and enforces the matching signature at startup. + +### Naming convention and roles + +| Suffix | Registry kind | Runs where | Role | +|--------|---------------|------------|------| +| `*_full_payload` | `producer_full_payload` | Worker (producer-side) | Packs the accumulated stage output into an `OmniPayload` and ships it through the connector (`FullPayloadProducer`) | +| `*_async_chunk` | `producer_async_chunk` | Scheduler (producer-side) | Streams one chunk per call while `async_chunk` is enabled (`AsyncChunkProducer`) | +| `*_token_only` | `placeholder_prompt_builder` | Orchestrator (consumer-side) | Allocates downstream KV slots only; bulk tensors arrive via the connector (`PlaceholderPromptBuilder`) | +| (no suffix, legacy) | `legacy_orchestrator_builder` | Orchestrator (consumer-side) | Legacy sync builder — a placeholder or diffusion input builder, adapted through `wrap_orchestrator_processor` | + +Diffusion-stage edges use dedicated suffixes +(`ar2diffusion` / `ar2dit` / `thinker2imagegen` / ...) mapped to +`diffusion_input_builder`, and `moss_tts.talker2codec` keeps the legacy +`legacy_multi_source` shape. The canonical naming documentation lives in the +`vllm_omni/model_executor/stage_input_processors/__init__.py` module docstring. + +### Minimal processor set per edge + +A new model that must work in **both** modes implements, for each inter-stage +edge, the three processors below: + +| Mode | Producer-side | Consumer-side (orchestrator) | +|------|---------------|------------------------------| +| `async_chunk=false` | `*_full_payload` (`FullPayloadProducer`) | `*_token_only` (`PlaceholderPromptBuilder`) | +| `async_chunk=true` | `*_async_chunk` (`AsyncChunkProducer`) | `*_token_only` (prewarm via `build_prewarm_placeholder`) | + +The minimal union set is `*_full_payload`, `*_async_chunk` and `*_token_only`. +The `*_token_only` placeholder builder is required in both modes: in non-async +mode it builds the forward placeholder, and in async mode the orchestrator +reuses it to prewarm the downstream stage. + +### OrchestratorInputContext and the C1 contract + +Every orchestrator-facing builder is invoked under the fixed C1 contract +`(source_outputs, ctx)`: + +```python +from vllm_omni.model_executor.stage_input_processors import OrchestratorInputContext + + +def my_token_only( + source_outputs: list[Any], + ctx: OrchestratorInputContext, +) -> list[OmniTokensPrompt]: + """Upstream outputs -> next-stage token prompts (C1 contract).""" + ... +``` + +`OrchestratorInputContext` carries the transition metadata and deliberately has +no `model_config` field (a processor that needs the model config reads it +through the upstream stage closure, never through this context): + +```python +@dataclass(frozen=True) +class OrchestratorInputContext: + prompt: Any | None = None + requires_multimodal_data: bool = False + streaming_context: Any | None = None + sampling_params: Any | None = None +``` + +Processors that already accept `ctx` are used unchanged. Legacy positional +shapes (C0 3-arg, C2 placeholder with `streaming_context`, C3 diffusion with +`sampling_params`, C4 `moss_tts.talker2codec` multi-source) are adapted by +`wrap_orchestrator_processor` / `invoke_orchestrator_processor` and emit a +`DeprecationWarning`. + +### Registry and startup validation + +The registry performs **signature-level structural checks only** — it never +executes processor logic and never loads model weights. Kind inference is +name-driven (suffix-based): + +- `register_processor(path, kind)` — manual kind override for names that do not + follow the suffix convention (escape hatch; overrides are validated eagerly). +- `infer_kind(fn, *, path)` — suffix rules: `_token_only` -> + `placeholder_prompt_builder`, `_full_payload` / `_batch` -> + `producer_full_payload`, `_async_chunk` -> `producer_async_chunk`, diffusion + suffixes -> `diffusion_input_builder`, no suffix -> + `legacy_orchestrator_builder`, `moss_tts.talker2codec` -> + `legacy_multi_source`. +- `validate_processor(fn, *, kind, path, stage_config=None)` — hard contract + violations raise `ProcessorValidationError`; soft mismatches emit a + `RuntimeWarning`. +- `resolve_processor(path, *, expected_kind=None, stage_config=None)` — the + drop-in replacement for the legacy `getattr(importlib.import_module(...), ...)` + lookups: imports, infers, validates, optionally checks `expected_kind`, and + returns a `ProcessorSpec` whose `fn` is the same callable the legacy lookup + produced. + +### P8b dual entry: forward and prewarm placeholders + +The `*_token_only` placeholder builder is exposed through two entry points so +the sync forward path and the async-chunk prewarm path share the same `_common` +length / packing helpers: + +- `build_forward_placeholder(source_outputs, ctx)` — the non-async forward + path. One placeholder `OmniTokensPrompt` per upstream output, sized by + `_common.compute_placeholder_prompt_len(mode="full")`. +- `build_prewarm_placeholder(*, stage0_prompt, ctx, downstream_stage_id)` — + async-chunk mode has no upstream `source_outputs` yet at prewarm time, so the + length is a best-effort estimate from the stage-0 input prompt + (`mode="stage0_only"`, i.e. `len(stage0_prompt)`). The connector fixup path + (`adapter.construct_next_stage_streaming_input_prompt`) replaces the estimate + with the real length once the upstream chunk arrives. + +The orchestrator routes prewarm through `_prewarm_async_chunk_stages`, which +reads `build_prewarm_placeholder` off the resolved `*_token_only` function +object (both builders are attached as attributes on the exported function). + +### Three-state async gate + +The orchestrator only forwards via `process_engine_inputs` when the transition +is not served by the async-chunk data plane: + +```python +if ( + (finished or segment_finished) + and stage_id < req_state.final_stage_id + and (not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1)) + and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled) +): +``` + +`_stage_receives_async_chunks(stage_id)` reports whether a stage's connector +supplies its runtime inputs. When `async_chunk=true` and the downstream stage +receives async chunks, the orchestrator skips `process_engine_inputs` for that +transition, so the corresponding input processor may be dead for the duration +of the request. `dead_processor_hint` is a pure decision helper (for warnings +and tests only); the runtime is warn-first in the M0 phase. + +### Producer keyword contract + +Producer-side builders never receive an `OrchestratorInputContext`; their +keyword-only parameters are load-bearing parts of the connector data plane and +must not be renamed or made positional: + +```python +def my_full_payload( + *, + transfer_manager: Any, + pooling_output: Any, + request: Any, + is_finished: bool = ..., # optional: the worker retries without it +) -> Any: ... + +def my_async_chunk( + *, + transfer_manager: Any, + multimodal_output: Any, + request: Any, + is_finished: bool = False, # required: the scheduler always passes it +) -> Any: ... +``` + +`pooling_output` (full payload) and `multimodal_output` (async chunk) are +cross-checked by `validate_processor`; using the wrong keyword name for a kind +is flagged as a structural warning. + ## Related Files - `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`: Chunk processing functions (receive `transfer_manager` as first param) diff --git a/docs/design/index.md b/docs/design/index.md index 64144a5dfe6..e7d2acb5acd 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -22,6 +22,7 @@ implementation contract; it is not, by itself, a general support claim. - [Disaggregated Inference](feature/disaggregated_inference.md) - [Host Weight Runtime](feature/host_weight_runtime.md) - [Async Chunk](feature/async_chunk.md) +- [Stage Input Processor Contract](feature/async_chunk.md#stage-input-processor-contract) — RFC #4872 `stage_input_processors` refactor - [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) From ac30134bfc58b34c533bbc3f20d71c09636e1280 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sun, 30 Aug 2026 15:01:05 +0800 Subject: [PATCH 07/41] docs(contributing): align model guides with stage input processor contract (RFC #4872) Upgrade the stage input processor examples in adding_omni_model.md to the C1 signature (OrchestratorInputContext), document register_processor / kind / resolve_processor / validate_processor usage and the minimal processor set for both modes, and fix the broken rfc_stage_input_processors_refactor.md link to point at the new contract section in docs/design/feature/async_chunk.md. In adding_tts_model.md, add the *_token_only placeholder builder (sync_process_input_func) and minimal processor set guidance, document the C1 contract and registration API, and rename the ar2decoder_async_chunk second parameter from pooling_output to multimodal_output to match registry validation. Signed-off-by: RyanYun09 --- docs/contributing/model/adding_omni_model.md | 67 +++++++++++++++++--- docs/contributing/model/adding_tts_model.md | 59 ++++++++++++++++- 2 files changed, 113 insertions(+), 13 deletions(-) diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 2a2f62c4a82..421d5795515 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -506,35 +506,48 @@ thinker_hidden_states = output.multimodal_output["24"] Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`. Each inter-stage edge should provide a **coherent processor set** rather than a single monolithic function: -| Suffix | Role | Runs when | -| -------- | ------ | ----------- | -| `*_full_payload` | Worker-side payload producer | `async_chunk=false`; accumulates tensors and ships via connector | -| `*_async_chunk` | Scheduler-side streaming producer | `async_chunk=true`; emits per-chunk payloads | -| `*_token_only` | Orchestrator placeholder builder | `async_chunk=false`; allocates downstream prompt slots only | +| Suffix | Registry kind | Runs where | Role | +|--------|---------------|------------|------| +| `*_full_payload` | `producer_full_payload` | Worker (producer-side) | `async_chunk=false`; accumulates tensors and ships via connector | +| `*_async_chunk` | `producer_async_chunk` | Scheduler (producer-side) | `async_chunk=true`; emits per-chunk payloads | +| `*_token_only` | `placeholder_prompt_builder` | Orchestrator (consumer-side) | Allocates downstream prompt slots only (both modes) | + +A model that supports **both** `async_chunk=false` and `async_chunk=true` must +implement all three processors for every inter-stage edge — `*_full_payload`, +`*_async_chunk` and `*_token_only`. The `*_token_only` placeholder builder is +required in both modes: in non-async mode it builds the forward placeholder, in +async mode the orchestrator reuses it to prewarm the downstream stage. + +Orchestrator-facing builders follow the C1 contract +`(source_outputs, ctx: OrchestratorInputContext)`: ```python -# qwen3_omni.py (Thinker → Talker, non-async path) +# qwen3_omni.py (Thinker → Talker) + +from vllm_omni.model_executor.stage_input_processors import OrchestratorInputContext + def thinker2talker_token_only( source_outputs: list[Any], - prompt: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, - streaming_context: Any | None = None, + ctx: OrchestratorInputContext, ) -> list[OmniTokensPrompt]: """Allocate talker prefill slots; bulk tensors arrive via the connector.""" ... def thinker2talker_full_payload( + *, transfer_manager: Any, pooling_output: dict[str, Any], request: OmniEngineCoreRequest, + is_finished: bool = ..., # optional: the worker retries without it ) -> dict[str, Any] | None: """Pack accumulated thinker hidden states into OmniPayload for stage-1.""" ... def thinker2talker_async_chunk( + *, transfer_manager: Any, multimodal_output: OmniPayload | dict[str, Any], request: OmniEngineCoreRequest, @@ -544,6 +557,13 @@ def thinker2talker_async_chunk( ... ``` +`OrchestratorInputContext` carries `prompt`, `requires_multimodal_data`, +`streaming_context` and `sampling_params` (and deliberately has no +`model_config` field). The producer-side builders (`*_full_payload`, +`*_async_chunk`) are keyword-only and never receive an +`OrchestratorInputContext`; `pooling_output` / `multimodal_output` are +load-bearing keyword names of the connector data plane. + Wire these in `pipeline.py`: ```python @@ -560,7 +580,34 @@ StagePipelineConfig( ), ``` -Do **not** add a no-suffix `thinker2talker` when a `sync_process_input_func` is already declared — `_select_processor_funcs()` always prefers the `*_token_only` hook in non-async mode, so the bare function would never run. See `docs/design/rfc_stage_input_processors_refactor.md` for the full contract. +### Registry registration and validation + +The processor registry infers a processor's kind from its name suffix and +structurally validates the signature at startup (it never executes processor +logic). Names that do not follow the suffix convention can register an explicit +kind override once, e.g. at module import time: + +```python +from vllm_omni.model_executor.stage_input_processors import register_processor + +register_processor(f"{_PROC}.thinker2talker_token_only", "placeholder_prompt_builder") +``` + +At startup the runtime resolves each processor through +`resolve_processor(path, expected_kind=...)` — import -> `infer_kind` -> +`validate_processor` -> expected-kind check — and returns a `ProcessorSpec` +whose `fn` is the same callable a legacy `getattr(importlib.import_module(...), +...)` lookup would produce. Hard contract violations raise +`ProcessorValidationError`; soft mismatches (e.g. a `*_full_payload` builder +declaring `multimodal_output` instead of `pooling_output`) emit a +`RuntimeWarning`. New model processors should therefore match the suffix +convention and the corresponding signature exactly. + +Do **not** add a no-suffix `thinker2talker` when a `sync_process_input_func` is +already declared — `_select_processor_funcs()` always prefers the `*_token_only` +hook in non-async mode, so the bare function would never run. See the +[Stage Input Processor Contract](../../design/feature/async_chunk.md#stage-input-processor-contract) +section of the async chunk design document for the full contract. ## Testing diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 806bec849d1..205165ab84b 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -481,6 +481,7 @@ YOUR_TTS_PIPELINE = PipelineConfig( execution_type=StageExecutionType.LLM_GENERATION, input_sources=(0,), model_arch="YourTTSDecoder", + sync_process_input_func=f"{_PROC}.ar2decoder_token_only", engine_output_type="audio", final_output=True, final_output_type="audio", @@ -521,6 +522,57 @@ Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours See `stage_input_processors/qwen3_tts.py` for the full reference implementation. +### Processor set and naming convention + +Each inter-stage edge provides a **coherent processor set** rather than a single +monolithic function. The suffix convention and the matching registry kind are +load-bearing: + +| Suffix | Registry kind | Runs where | Role | +|--------|---------------|------------|------| +| `*_full_payload` | `producer_full_payload` | Worker (producer-side) | `async_chunk=false`; packs the accumulated Stage 0 output and ships it via the connector | +| `*_async_chunk` | `producer_async_chunk` | Scheduler (producer-side) | `async_chunk=true`; streams one chunk per AR decode step | +| `*_token_only` | `placeholder_prompt_builder` | Orchestrator (consumer-side) | Allocates downstream KV slots only; bulk tensors arrive via the connector (both modes) | + +A model that supports **both** `async_chunk=false` and `async_chunk=true` must +implement all three processors per edge — `ar2decoder` (full payload), +`ar2decoder_async_chunk` and `ar2decoder_token_only`. The `*_token_only` +placeholder builder is wired as the downstream stage's `sync_process_input_func` +and follows the C1 contract `(source_outputs, ctx: OrchestratorInputContext)`: + +```python +from vllm_omni.model_executor.stage_input_processors import OrchestratorInputContext + + +def ar2decoder_token_only( + source_outputs: list[Any], + ctx: OrchestratorInputContext, +) -> list[OmniTokensPrompt]: + """Allocate decoder prefill slots; bulk audio codes arrive via the connector.""" + ... +``` + +`OrchestratorInputContext` carries `prompt`, `requires_multimodal_data`, +`streaming_context` and `sampling_params` (and deliberately has no +`model_config` field). Legacy positional shapes are adapted through +`wrap_orchestrator_processor` with a `DeprecationWarning`. + +Names that do not follow the suffix convention can register an explicit kind +override at module import time: + +```python +from vllm_omni.model_executor.stage_input_processors import register_processor + +register_processor(f"{_PROC}.ar2decoder_token_only", "placeholder_prompt_builder") +``` + +At startup the runtime resolves each processor through +`resolve_processor(path, expected_kind=...)` (import -> `infer_kind` -> +`validate_processor` -> expected-kind check). Hard contract violations raise +`ProcessorValidationError`; soft mismatches emit a `RuntimeWarning`. See the +[Stage Input Processor Contract](../../design/feature/async_chunk.md#stage-input-processor-contract) +section of the async chunk design document for the full contract. + ### Data structures Understanding what's available in stage outputs: @@ -563,8 +615,9 @@ have accumulated. The function signature follows the `OmniChunkTransferAdapter` ```python def ar2decoder_async_chunk( + *, transfer_manager: Any, - pooling_output: dict[str, Any] | None, + multimodal_output: dict[str, Any] | None, request: Any, is_finished: bool = False, ) -> dict[str, Any] | None: @@ -573,8 +626,8 @@ def ar2decoder_async_chunk( finished = bool(is_finished or request.is_finished()) # Extract and buffer the latest frame - if isinstance(pooling_output, dict): - frame = extract_frame(pooling_output) + if isinstance(multimodal_output, dict): + frame = extract_frame(multimodal_output) if frame is not None: transfer_manager.code_prompt_token_ids[request_id].append( frame.cpu().tolist() From 356007a35489f81b8fc9c7fd05201d7dc0de9434 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sun, 30 Aug 2026 15:31:54 +0800 Subject: [PATCH 08/41] style(docs): apply ruff format to markdown code blocks (RFC #4872) Signed-off-by: RyanYun09 --- docs/contributing/model/adding_omni_model.md | 34 ++++---------- docs/contributing/model/adding_tts_model.md | 48 +++++++++----------- docs/design/feature/async_chunk.md | 1 + 3 files changed, 33 insertions(+), 50 deletions(-) diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 421d5795515..64c280cdeb8 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -74,11 +74,8 @@ The thinker stage handles multimodal understanding. Key features: from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM -class Qwen3OmniMoeThinkerForConditionalGeneration( - Qwen3MoeForCausalLM, - SupportsMultiModal, - SupportsPP -): + +class Qwen3OmniMoeThinkerForConditionalGeneration(Qwen3MoeForCausalLM, SupportsMultiModal, SupportsPP): """Thinker stage: multimodal understanding → text generation.""" def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -93,10 +90,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( The talker stage converts text embeddings to codec codes: ```python -class Qwen3OmniMoeTalkerForConditionalGeneration( - Qwen3MoeForCausalLM, - SupportsPP -): +class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3MoeForCausalLM, SupportsPP): """Talker stage: text embeddings → RVQ codec codes.""" def __init__(self, vllm_config, talker_config, prefix): @@ -156,8 +150,7 @@ class Qwen3OmniMoeForConditionalGeneration( if self.model_stage == "thinker": # Initialize thinker model thinker_vllm_config = vllm_config.with_hf_config( - config.thinker_config, - architectures=["Qwen3OmniMoeThinkerForConditionalGeneration"] + config.thinker_config, architectures=["Qwen3OmniMoeThinkerForConditionalGeneration"] ) self.thinker = init_vllm_registered_model( vllm_config=thinker_vllm_config, @@ -170,8 +163,7 @@ class Qwen3OmniMoeForConditionalGeneration( elif self.model_stage == "talker": # Initialize talker model talker_vllm_config = vllm_config.with_hf_config( - config.talker_config, - architectures=["Qwen3OmniMoeTalkerForConditionalGeneration"] + config.talker_config, architectures=["Qwen3OmniMoeTalkerForConditionalGeneration"] ) self.talker = init_vllm_registered_model( vllm_config=talker_vllm_config, @@ -184,8 +176,7 @@ class Qwen3OmniMoeForConditionalGeneration( elif self.model_stage == "code2wav": # Initialize code2wav model code2wav_vllm_config = vllm_config.with_hf_config( - config.code2wav_config, - architectures=["Qwen3OmniMoeCode2Wav"] + config.code2wav_config, architectures=["Qwen3OmniMoeCode2Wav"] ) self.code2wav = init_vllm_registered_model( vllm_config=code2wav_vllm_config, @@ -196,8 +187,7 @@ class Qwen3OmniMoeForConditionalGeneration( self.model = self.code2wav else: raise ValueError( - f"Invalid model_stage: {self.model_stage}. " - f"Must be one of: 'thinker', 'talker', 'code2wav'" + f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'talker', 'code2wav'" ) ``` @@ -308,11 +298,10 @@ Register your model in `vllm_omni/model_executor/models/registry.py`: ```python _OMNI_MODELS = { # ... existing models ... - # Your new model "YourModelForConditionalGeneration": ( - "your_model_name", # Module folder name - "your_model", # Module file name (without .py) + "your_model_name", # Module folder name + "your_model", # Module file name (without .py) "YourModelForConditionalGeneration", # Class name ), "YourModelThinkerForConditionalGeneration": ( @@ -410,10 +399,7 @@ The stage transition process follows these steps: else: # Custom transition function (YOUR CODE HERE) return self.custom_process_input_func( - stage_list, - self.engine_input_source, - prompt, - self.requires_multimodal_data + stage_list, self.engine_input_source, prompt, self.requires_multimodal_data ) ``` - If `custom_process_input_func` is configured, it calls that function diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 205165ab84b..68c78572ce3 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -238,15 +238,15 @@ Do not wrap the HuggingFace model directly - that bypasses PagedAttention and fu from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer -class YourTTSARStage(nn.Module): +class YourTTSARStage(nn.Module): def __init__(self, config, vllm_config, prefix): - self.layers = nn.ModuleList([ - Qwen3DecoderLayer( - config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}" - ) - for i in range(config.num_hidden_layers) - ]) + self.layers = nn.ModuleList( + [ + Qwen3DecoderLayer(config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}") + for i in range(config.num_hidden_layers) + ] + ) self.lm_head = ParallelLMHead(config.codec_size, config.hidden_size) ``` @@ -257,8 +257,7 @@ See `qwen3_tts_code_predictor_vllm.py` for the full implementation. Implement `forward()` to return an `OmniOutput` with intermediate data for Stage 1: ```python -def forward(self, input_ids, positions, intermediate_tensors=None, - inputs_embeds=None, **kwargs) -> OmniOutput: +def forward(self, input_ids, positions, intermediate_tensors=None, inputs_embeds=None, **kwargs) -> OmniOutput: hidden_states = self.run_layers(input_ids, positions, intermediate_tensors, inputs_embeds) logits = self.lm_head(hidden_states) @@ -301,8 +300,8 @@ needed). Implement `chunked_decode_streaming()` to support async_chunk streaming ```python # your_model_decoder.py -class YourTTSDecoder(nn.Module): +class YourTTSDecoder(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() # Initialize your audio decoder (SpeechTokenizer, HiFiGAN, etc.) @@ -310,14 +309,13 @@ class YourTTSDecoder(nn.Module): def forward(self, codes: torch.Tensor, **kwargs) -> torch.Tensor: return self.decoder(codes) - def chunked_decode_streaming(self, codes, chunk_size=25, - left_context_size=25) -> torch.Tensor: + def chunked_decode_streaming(self, codes, chunk_size=25, left_context_size=25) -> torch.Tensor: """Decode with a sliding context window for smooth chunk boundaries.""" end_index = codes.shape[-1] context_size = 0 if end_index <= chunk_size else left_context_size wav_chunk = self(codes) # Trim left context to avoid duplicate audio - return wav_chunk[..., context_size * self.total_upsample:] + return wav_chunk[..., context_size * self.total_upsample :] ``` ### Step 3: Implement the Unified Model Class @@ -327,8 +325,8 @@ The unified class dispatches to the correct stage based on `model_stage` in the ```python # your_model.py -class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP): +class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.model_stage = vllm_config.model_config.model_stage @@ -411,15 +409,18 @@ Register all stage classes in `vllm_omni/model_executor/models/registry.py`: _OMNI_MODELS = { # (package_name, module_name, class_name) "YourTTSModelForConditionalGeneration": ( - "your_model_name", "your_model", + "your_model_name", + "your_model", "YourTTSModelForConditionalGeneration", ), "YourTTSARStageForConditionalGeneration": ( - "your_model_name", "your_model_ar_stage", + "your_model_name", + "your_model_ar_stage", "YourTTSARStageForConditionalGeneration", ), "YourTTSDecoder": ( - "your_model_name", "your_model_decoder", + "your_model_name", + "your_model_decoder", "YourTTSDecoder", ), } @@ -601,9 +602,7 @@ def ar2decoder( for output in stage_list[source_id].engine_outputs: result = output.outputs[0] codes = result.multimodal_output["audio_codes"].cpu() - decoder_inputs.append( - OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist()) - ) + decoder_inputs.append(OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist())) return decoder_inputs ``` @@ -629,9 +628,7 @@ def ar2decoder_async_chunk( if isinstance(multimodal_output, dict): frame = extract_frame(multimodal_output) if frame is not None: - transfer_manager.code_prompt_token_ids[request_id].append( - frame.cpu().tolist() - ) + transfer_manager.code_prompt_token_ids[request_id].append(frame.cpu().tolist()) elif not finished: return None @@ -914,11 +911,10 @@ any CUDA allocations. Stream via a per-request generator stored in an instance d ```python class YourModelForCausalLM(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() - self._lm: nn.Module | None = None # populated in load_weights() - self._stream_gens: dict[str, Any] = {} # request_key → generator + self._lm: nn.Module | None = None # populated in load_weights() + self._stream_gens: dict[str, Any] = {} # request_key → generator def load_weights(self, weights): # Load self._lm here, after vLLM distributed init diff --git a/docs/design/feature/async_chunk.md b/docs/design/feature/async_chunk.md index b34e262f458..bd602345773 100644 --- a/docs/design/feature/async_chunk.md +++ b/docs/design/feature/async_chunk.md @@ -439,6 +439,7 @@ def my_full_payload( is_finished: bool = ..., # optional: the worker retries without it ) -> Any: ... + def my_async_chunk( *, transfer_manager: Any, From c42c6c50f0364fd5ffdc844d1a359061c238fab6 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Mon, 31 Aug 2026 16:04:56 +0800 Subject: [PATCH 09/41] fix(stage_input_processors): restore diffusion output_formatter wrap-only ensure_list (RFC #4872 P3) output_formatter._ensure_list is not a stage-input helper. It used to wrap a non-list primary payload as [x]; reusing _common.ensure_list flattens tensors, iterates dict keys, and walks a PIL Image by rows, which would corrupt the images=outputs list in format_diffusion_outputs. - Add _common.ensure_list_wrap_only (list passthrough; non-list -> [x]; None -> []), performing no flattening / key iteration / row walk. - Delegate output_formatter._ensure_list to it, preserving the exact legacy wrap-only semantics. - Add test_output_formatter_ensure_list_wrap_only locking the no-flatten / no-key iteration / no-row-walk behaviour (golden). Reviewer: amy-why-3459 (Blocking 1) / hsliuustc0106 (blocking). Signed-off-by: RyanYun09 --- .../test_common_helpers_golden.py | 37 +++++++++++++++++++ vllm_omni/diffusion/output_formatter.py | 13 ++++++- .../stage_input_processors/_common.py | 14 +++++++ 3 files changed, 62 insertions(+), 2 deletions(-) diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py index b646349ba4c..835ebba47ec 100644 --- a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -121,6 +121,43 @@ def test_ensure_list_golden(module, inp, expected, expect_typeerror): assert _norm(result) == expected +def test_output_formatter_ensure_list_wrap_only(): + """output_formatter._ensure_list keeps wrap-only semantics (RFC #4872 P3). + + The diffusion formatter is NOT a stage-input helper: a non-list primary + payload is wrapped as ``[x]`` verbatim — no tensor flattening, no dict-key + iteration, no row-wise iteration of an iterable (e.g. a PIL ``Image``). + ``_common.ensure_list`` (the processor flatten) must NOT be reused here. + """ + mod = _import("diffusion.output_formatter") + fn = getattr(mod, "_ensure_list") + assert fn([1, 2]) == [1, 2] + assert fn(None) == [] + assert fn(5) == [5] + # A tensor is wrapped whole, NOT flattened via .tolist(). + t = torch.tensor([1, 2]) + (wrapped,) = fn(t) + assert wrapped is t + # A dict is wrapped whole, NOT iterated over its keys. + d = {"a": 1} + (wrapped_dict,) = fn(d) + assert wrapped_dict is d + + # A row-iterable (PIL.Image-like) is wrapped whole, NOT walked row by row. + class _RowIterable: + def __iter__(self): + yield from [[1, 2], [3, 4]] + + img = _RowIterable() + (wrapped_img,) = fn(img) + assert wrapped_img is img + # Contrast: the canonical processor flatten would have flattened each. + c = _common_import() + assert c.ensure_list(t) == [1, 2] # flattened, not [t] + assert c.ensure_list(d) == ["a"] # dict keys iterated + assert c.ensure_list(img) == [[1, 2], [3, 4]] # rows walked + + # =========================================================================== # _extract_last_frame # =========================================================================== diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index 0d1e81f5484..5ee3bba33ce 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -203,8 +203,17 @@ def _stream_metadata_from_diffusion_output(diffusion_output: DiffusionOutput) -> def _ensure_list(outputs: DiffusionPayloadValue) -> list[DiffusionPayloadValue]: - """Delegate to the canonical shared helper (RFC #4872 Step 1c).""" - return _common.ensure_list(outputs) + """Wrap a non-list primary payload as ``[x]`` (diffusion formatter semantics). + + ``output_formatter`` is **not** a stage-input helper: the legacy behaviour + was wrap-only — a ``list`` passes through unchanged, a non-list scalar / + tensor / PIL ``Image`` is wrapped verbatim as ``[x]``, and ``None`` becomes + ``[]``. The canonical ``_common.ensure_list`` would flatten tensors, + iterate dict keys, and walk a PIL ``Image`` row by row, so it must not be + reused here; ``_common.ensure_list_wrap_only`` reproduces the exact legacy + semantics (RFC #4872 P3). + """ + return _common.ensure_list_wrap_only(outputs) def _primary_payload(postprocess_output: DiffusionPostprocessOutput) -> DiffusionPayloadValue | None: diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py index c83dafadef9..b3c503224d2 100644 --- a/vllm_omni/model_executor/stage_input_processors/_common.py +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -133,6 +133,20 @@ def ensure_list_preserve_none(x: Any) -> list[Any]: return list(x) if hasattr(x, "__iter__") else [x] +def ensure_list_wrap_only(x: Any) -> list[Any]: + """Wrap-only ``ensure_list``: list passthrough, non-list -> ``[x]``, ``None`` -> ``[]``. + + Preserves the legacy ``diffusion.output_formatter._ensure_list`` semantics + (RFC #4872 P3): a non-list primary payload is wrapped as ``[x]``. Unlike + the canonical :func:`ensure_list`, this deliberately performs **no** tensor + flattening, no dict-key iteration, and no row-wise iteration of an iterable + (e.g. a PIL ``Image``) — a non-``list`` value is wrapped whole, verbatim. + """ + if isinstance(x, list): + return x + return [x] if x is not None else [] + + # =========================================================================== # to_cpu_tensor # =========================================================================== From 11a2738675fd4cd77ed766a9fc38f431bf8f9a32 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Mon, 31 Aug 2026 16:06:25 +0800 Subject: [PATCH 10/41] fix(stage_input_processors): make P8b stage0_only prewarm the chat-template scan (RFC #4872 P8b) The P8b happy-path prewarm used len(stage0_prompt) while the inline fallback and adapter.compute_talker_prompt_ids_length ran the Qwen chat-template scan, so the same request was 10 or 15 depending on whether resolution succeeded. - compute_placeholder_prompt_len(mode="stage0_only") now runs the same Qwen chat-template scan on the stage-0 list (fed to both all/prompt roles, exactly what adapter.compute_talker_prompt_ids_length does), so the builder and the inline fallback return the same number (golden 15 == 15; single user segment 6 == 6). - Update parity / golden / orchestrator-bridge tests to lock the consistency instead of the old len() vs scan split. Reviewer: amy-why-3459 (Blocking 3) / hsliuustc0106 (blocking). Signed-off-by: RyanYun09 --- .../test_orchestrator_stage_input_bridge.py | 12 ++-- .../test_common_helpers_golden.py | 19 +++--- .../test_placeholder_parity.py | 45 ++++++++----- .../stage_input_processors/_common.py | 63 ++++++++++++------- .../stage_input_processors/qwen3_omni.py | 12 ++-- 5 files changed, 100 insertions(+), 51 deletions(-) diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 09e941c9f10..691f0c5c70c 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -306,8 +306,10 @@ async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: """RFC #4872 P8b: async-chunk prewarm resolves ``sync_process_input_func`` (``*_token_only``) to ``build_prewarm_placeholder`` and uses its estimate. - The placeholder length is ``stage0_only`` = ``len(stage0 prompt_token_ids)`` - (best-effort; the connector fixup path replaces it later). + The placeholder length is ``stage0_only`` = the Qwen chat-template scan on + the stage-0 prompt. The synthetic ``[1, 2]`` prompt has no chat-template + marker, so the scan yields 0 and the builder floors it to a 1-token + placeholder (best-effort; the connector fixup path replaces it later). """ orchestrator = object.__new__(Orchestrator) stage0 = FakePrewarmPool("sender") @@ -337,8 +339,10 @@ async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: assert prewarmed is True assert len(stage1.submitted) == 1 - # build_prewarm_placeholder uses stage0_only = len(stage0 prompt) == 2. - assert stage1.submitted[0].prompt_token_ids == [0, 0] + # build_prewarm_placeholder uses stage0_only = the chat-template scan on the + # stage-0 prompt. The synthetic [1, 2] prompt has no chat-template marker, + # so the scan yields 0 and the builder floors it to a 1-token placeholder. + assert stage1.submitted[0].prompt_token_ids == [0] orchestrator._record_duplex_stage_submission.assert_called_once_with( 1, "req-prewarm-resolved", diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py index 835ebba47ec..2c16c64f464 100644 --- a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -413,17 +413,21 @@ def test_to_token_id_list_cosyvoice3_recursive_golden(): # =========================================================================== -# P8a: placeholder length helper (qwen3_omni._compute_talker_prompt_ids_length) +# P8a: placeholder length helper (Qwen chat-template scan). +# The legacy ``qwen3_omni._compute_talker_prompt_ids_length`` is dead after the +# P8b consolidation; the canonical ``_common.compute_placeholder_prompt_len`` +# reproduces the golden result. # =========================================================================== def test_compute_talker_prompt_ids_length_golden(): - mod = _import("qwen3_omni") - fn = getattr(mod, "_compute_talker_prompt_ids_length") + c = _common_import() # im_start markers at 0 and 6; user segment len 6, last assistant +9. prompt = [151644, 872, 10, 11, 12, 13, 151644, 77091, 20, 21] info = {"ids": {"all": prompt, "prompt": prompt}} - assert fn(info, device="cpu") == 15 + assert c.compute_placeholder_prompt_len(ids_or_prompt=info, mode="full", device="cpu") == 15 + # stage0_only scans the flat stage-0 list (same 15 for the golden prompt). + assert c.compute_placeholder_prompt_len(ids_or_prompt=prompt, mode="stage0_only") == 15 # =========================================================================== @@ -552,10 +556,11 @@ def test_common_filter_real_code_frames_matches_golden(): def test_common_compute_placeholder_prompt_len_matches_golden(): c = _common_import() prompt = [151644, 872, 10, 11, 12, 13, 151644, 77091, 20, 21] - # full mode == qwen3_omni._compute_talker_prompt_ids_length golden (15). + # full mode == the legacy Qwen chat-template scan golden (15). assert c.compute_placeholder_prompt_len(ids_or_prompt={"ids": {"all": prompt, "prompt": prompt}}, mode="full") == 15 - # stage0_only mode: just the stage-0 input prompt length (prewarm). - assert c.compute_placeholder_prompt_len(ids_or_prompt=prompt, mode="stage0_only") == len(prompt) + # stage0_only mode: the same scan on the flat stage-0 list (prewarm) -> 15, + # so the builder and the inline fallback agree (RFC #4872 P8b). + assert c.compute_placeholder_prompt_len(ids_or_prompt=prompt, mode="stage0_only") == 15 def test_common_pack_placeholder_prompt(): diff --git a/tests/model_executor/stage_input_processors/test_placeholder_parity.py b/tests/model_executor/stage_input_processors/test_placeholder_parity.py index 7c652798472..b737470e13c 100644 --- a/tests/model_executor/stage_input_processors/test_placeholder_parity.py +++ b/tests/model_executor/stage_input_processors/test_placeholder_parity.py @@ -14,11 +14,12 @@ ---------------- - forward (``mode="full"``) is the Qwen chat-template scan (golden-locked at 15 for the canonical prompt ``test_common_helpers_golden.py``). -- prewarm (``mode="stage0_only"``) is a best-effort estimate equal to - ``len(stage0_prompt)``; the connector fixup path replaces it with the real - length once the upstream chunk arrives. For a chat-template prompt these - differ (15 vs 10); for a single user segment (no trailing assistant marker) - they coincide (6 == 6), which is the exact parity the prewarm targets. +- prewarm (``mode="stage0_only"``) runs the **same** Qwen chat-template scan on + the stage-0 input list (feeding it to both ``all``/``prompt`` roles), so the + prewarm estimate equals the forward placeholder length and the adapter's + streaming-fixup scan (15 == 15 for the golden prompt; 6 == 6 for a single user + segment). The connector fixup path replaces the estimate with the real + length once the upstream chunk arrives. These tests are CPU-only (no model loading). They run under the local vllm stub (shim) and on CI where real vllm builds ``OmniTokensPrompt`` as a dict @@ -132,7 +133,9 @@ def test_build_prewarm_placeholder_length_from_stage0(monkeypatch): ) assert len(out) == 1 assert out[0] == "PLACEHOLDER" - assert recorded["prompt_len"] == len(GOLDEN_PROMPT) # 10 (best-effort estimate) + # stage0_only runs the chat-template scan on the stage-0 list -> 15 + # (the same number the forward builder and the adapter scan produce). + assert recorded["prompt_len"] == 15 # Voice metadata is intentionally not forwarded (matches pre-existing prewarm). assert recorded["voice_metadata"] is None @@ -141,7 +144,7 @@ def test_build_prewarm_placeholder_accepts_request_like_stage0(monkeypatch): recorded = _capture_pack(monkeypatch) stage0_request = SimpleNamespace(prompt_token_ids=GOLDEN_PROMPT) q3.build_prewarm_placeholder(stage0_prompt=stage0_request, ctx=_ctx(), downstream_stage_id=1) - assert recorded["prompt_len"] == len(GOLDEN_PROMPT) + assert recorded["prompt_len"] == 15 # --------------------------------------------------------------------------- @@ -162,11 +165,12 @@ def test_forward_prewarm_parity_single_user_segment(monkeypatch): def test_forward_prewarm_golden_relationship(monkeypatch): """Documented relationship for the chat-template golden prompt. - forward (full) == 15 (golden-locked); prewarm (stage0_only) == 10 - (``= len(stage0_prompt)``). The difference is intentional: prewarm is a - best-effort estimate and the connector fixup path - (``adapter.construct_next_stage_streaming_input_prompt``) replaces it with - the real length once the upstream chunk arrives. + forward (full) == 15 (golden-locked); prewarm (stage0_only) is now the + **same chat-template scan** on the stage-0 list, so it also returns 15. + The builder and the inline fallback + (``adapter.compute_talker_prompt_ids_length``) therefore agree — the old + ``len(stage0_prompt)`` == 10 estimate under-reserved KV slots and split a + single request between 10 and 15 depending on whether resolution succeeded. """ recorded = _capture_pack(monkeypatch) q3.build_forward_placeholder(_source_outputs(GOLDEN_PROMPT), _ctx()) @@ -174,8 +178,8 @@ def test_forward_prewarm_golden_relationship(monkeypatch): recorded.clear() q3.build_prewarm_placeholder(stage0_prompt=GOLDEN_PROMPT, ctx=_ctx(), downstream_stage_id=1) assert forward_len == 15 - assert recorded["prompt_len"] == len(GOLDEN_PROMPT) - assert forward_len != recorded["prompt_len"] # full vs stage0_only estimate + assert recorded["prompt_len"] == 15 + assert forward_len == recorded["prompt_len"] # full == stage0_only scan # --------------------------------------------------------------------------- @@ -230,7 +234,20 @@ def test_adapter_compute_length_consistent_with_common(): ids_or_prompt={"ids": {"all": GOLDEN_PROMPT, "prompt": GOLDEN_PROMPT}}, mode="full", ) + # RFC #4872 P8b: the prewarm stage0_only scan (what build_prewarm_placeholder + # and the orchestrator inline fallback both use) must equal the adapter scan. + assert adapter.compute_talker_prompt_ids_length(GOLDEN_PROMPT) == _common.compute_placeholder_prompt_len( + ids_or_prompt=GOLDEN_PROMPT, + mode="stage0_only", + ) assert adapter.compute_talker_prompt_ids_length(SINGLE_USER_PROMPT) == 6 + assert ( + _common.compute_placeholder_prompt_len( + ids_or_prompt=SINGLE_USER_PROMPT, + mode="stage0_only", + ) + == 6 + ) # --------------------------------------------------------------------------- diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py index b3c503224d2..859618cc06b 100644 --- a/vllm_omni/model_executor/stage_input_processors/_common.py +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -20,6 +20,7 @@ from __future__ import annotations +from collections.abc import Sequence from typing import Any import torch @@ -362,31 +363,19 @@ def extract_last_codec_frame( _ASSISTANT_TAIL_LEN = 9 # 3 + 4 + 1 + 1 -def compute_placeholder_prompt_len( - *, - ids_or_prompt: Any, - mode: str = "full", - device: torch.device | str = "cpu", +def _chat_template_prompt_len( + all_ids: Sequence[int], + prompt_ids: Sequence[int], + device: torch.device | str, ) -> int: - """Compute the downstream placeholder prompt length. - - ``mode="full"`` replicates the Qwen chat-template scan of - ``qwen3_omni._compute_talker_prompt_ids_length``: the input is an - ``OmniPayload``-like dict with ``ids: {"all": [...], "prompt": [...]}``, - and the length is ``sum(user segments) + 9`` (the assistant tail). + """Qwen chat-template scan: ``sum(user segments) + 9`` (assistant tail). - ``mode="stage0_only"`` is the async-chunk **prewarm** estimate: there is - no upstream output yet, so the length is simply the stage-0 input prompt - length (``ids_or_prompt`` is a sequence of token ids). + Shared by ``mode="full"`` (prompt + generated ids) and ``mode="stage0_only"`` + (stage-0 ids only, fed to both roles) so the sync forward placeholder and + the async-chunk prewarm estimate return the same length. """ - if mode == "stage0_only": - return len(list(ids_or_prompt)) - if mode != "full": - raise ValueError(f"unknown mode: {mode!r}") - - ids = (ids_or_prompt or {}).get("ids", {}) - thinker_sequences = torch.tensor(ids["all"], dtype=torch.long, device=device).unsqueeze(0) # [1, T] - input_ids = torch.tensor(ids["prompt"], dtype=torch.long, device=device).unsqueeze(0) # [1, T] + thinker_sequences = torch.tensor(all_ids, dtype=torch.long, device=device).unsqueeze(0) # [1, T] + input_ids = torch.tensor(prompt_ids, dtype=torch.long, device=device).unsqueeze(0) # [1, T] im_start_indexes = torch.cat( [ @@ -415,6 +404,36 @@ def compute_placeholder_prompt_len( return sum_user_len + assistant_len +def compute_placeholder_prompt_len( + *, + ids_or_prompt: Any, + mode: str = "full", + device: torch.device | str = "cpu", +) -> int: + """Compute the downstream placeholder prompt length. + + ``mode="full"`` replicates the legacy Qwen chat-template scan: the input is + an ``OmniPayload``-like dict with ``ids: {"all": [...], "prompt": [...]}``, + and the length is ``sum(user segments) + 9`` (the assistant tail). + + ``mode="stage0_only"`` is the async-chunk **prewarm** estimate. There is no + upstream output yet, so ``ids_or_prompt`` is the stage-0 input token-id list, + which is fed to **both** chat-template roles (``all`` / ``prompt``) — exactly + what ``adapter.compute_talker_prompt_ids_length`` does. This makes the + prewarm builder and the inline fallback return the same number (15 for the + golden prompt; 6 == 6 for a single user segment), removing the + ``len()`` vs scan split. + """ + if mode == "stage0_only": + prompt = list(ids_or_prompt) + return _chat_template_prompt_len(prompt, prompt, device) + if mode != "full": + raise ValueError(f"unknown mode: {mode!r}") + + ids = (ids_or_prompt or {}).get("ids", {}) + return _chat_template_prompt_len(ids["all"], ids["prompt"], device) + + def pack_placeholder_prompt( *, prompt_len: int, diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 094f4aa5b57..6402796c372 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -741,10 +741,14 @@ def build_prewarm_placeholder( async-chunk mode has **no upstream ``source_outputs``** yet at prewarm time, so this is a best-effort estimate: the placeholder length comes from - the stage-0 input prompt (``_common.compute_placeholder_prompt_len( - mode="stage0_only")``, i.e. ``len(stage0_prompt)``). The connector fixup - path (``adapter.construct_next_stage_streaming_input_prompt``) replaces - this estimate with the real length once the upstream chunk arrives. + ``_common.compute_placeholder_prompt_len(mode="stage0_only")``, which runs + the **Qwen chat-template scan on the stage-0 input list** (the same scan + ``adapter.compute_talker_prompt_ids_length`` runs). This keeps the prewarm + estimate consistent with the sync forward placeholder and the inline + fallback (golden prompt: 15 == 15; single user segment: 6 == 6). The + connector fixup path + (``adapter.construct_next_stage_streaming_input_prompt``) replaces this + estimate with the real length once the upstream chunk arrives. ``ctx`` / ``downstream_stage_id`` are accepted for contract uniformity; voice metadata is intentionally **not** forwarded here (matching the From ec668a69c66830b14289ad0897bcf03e2cfb528e Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Mon, 31 Aug 2026 16:06:34 +0800 Subject: [PATCH 11/41] fix(engine): extract_legacy_stage_metadata fails fast on processor resolution (RFC #4872 P2) extract_legacy_stage_metadata swallowed every resolve_processor failure and continued with custom_process_input_func=None, silently dropping the configured hook and falling back to _default_process_engine_inputs (downstream prompt = upstream token_ids). A bad path used to fail at import. - Remove the blanket try/except: ProcessorValidationError / ImportError / AttributeError now propagate (fail-fast), matching RFC P2 and the adding-omni text. M0 is a DeprecationWarning on C2/C3/C4, not a silent drop of the hook. - The worker-side mixin candidate chain (_load_custom_func) keeps its own skip-on-mismatch tolerance for *_full_payload probes; it is untouched. - Add fail-fast tests (bad module / missing attr / non-callable symbol). Reviewer: amy-why-3459 (Blocking 2) / hsliuustc0106 (blocking). Signed-off-by: RyanYun09 --- tests/engine/test_stage_metadata.py | 44 ++++++++++++++++++++++++++++ vllm_omni/engine/stage_init_utils.py | 26 +++++++--------- 2 files changed, 55 insertions(+), 15 deletions(-) diff --git a/tests/engine/test_stage_metadata.py b/tests/engine/test_stage_metadata.py index 6974105691a..c516a2c1843 100644 --- a/tests/engine/test_stage_metadata.py +++ b/tests/engine/test_stage_metadata.py @@ -4,6 +4,7 @@ import copy import operator from dataclasses import replace +from typing import Any import pytest from vllm.sampling_params import SamplingParams @@ -23,6 +24,7 @@ extract_stage_metadata_from_omni_stage_config, ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_executor.stage_input_processors import ProcessorValidationError pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -195,3 +197,45 @@ def test_extract_stage_metadata_defaults_missing_engine_input_source(): metadata = extract_stage_metadata(legacy_config) assert metadata.engine_input_source == [] + + +# --------------------------------------------------------------------------- +# RFC #4872 P2 / adding-omni: a bad custom_process_input_func must FAIL FAST. +# extract_legacy_stage_metadata no longer swallows resolve_processor failures +# and continues with custom_process_input_func=None (which silently fell back to +# _default_process_engine_inputs). The worker mixin candidate chain keeps its +# own skip-on-mismatch tolerance; this is the config-hook resolution path. +# --------------------------------------------------------------------------- + + +def _legacy_stage0_config() -> Any: + pipeline, deploy = _metadata_inputs() + return merge_pipeline_deploy( + pipeline, + copy.deepcopy(deploy), + )[0].to_omegaconf() + + +def test_extract_stage_metadata_fails_fast_on_bad_processor_module(): + legacy_config = _legacy_stage0_config() + legacy_config["custom_process_input_func"] = "nonexistent.module.does_not_exist" + with pytest.raises(ImportError): + extract_legacy_stage_metadata(legacy_config) + + +def test_extract_stage_metadata_fails_fast_on_missing_processor_attr(): + legacy_config = _legacy_stage0_config() + legacy_config["custom_process_input_func"] = "operator.does_not_exist" + with pytest.raises(AttributeError): + extract_legacy_stage_metadata(legacy_config) + + +def test_extract_stage_metadata_fails_fast_on_invalid_processor(): + # Resolves to a non-callable symbol -> ProcessorValidationError (rule + # "callable") instead of silently dropping the configured hook. + legacy_config = _legacy_stage0_config() + legacy_config["custom_process_input_func"] = ( + "vllm_omni.model_executor.stage_input_processors._common._IM_START_TOKEN_ID" + ) + with pytest.raises(ProcessorValidationError): + extract_legacy_stage_metadata(legacy_config) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index d77bed2b3a7..439e642c5f9 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -646,21 +646,17 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: # RFC #4872 Phase 3 (Part A): resolve through the stage-input processor # registry. expected_kind is deliberately None here — orchestrator-side # processors carry no suffix contract, so we rely on suffix inference - # only (compatible with existing configs). A bad path must not kill - # startup (M0): warn and continue without the hook. - try: - custom_process_input_func = resolve_processor( - _cpif_path, - expected_kind=None, - stage_config=stage_config, - ).fn - except Exception as exc: - logger.warning( - "[stage_init] Failed to resolve custom_process_input_func %r; " - "continuing without it (RFC #4872 M0: warn, not fail): %s", - _cpif_path, - exc, - ) + # only (compatible with existing configs). A bad path must **fail fast** + # (RFC P2 / adding-omni text): let ProcessorValidationError / ImportError / + # AttributeError propagate instead of silently dropping the configured + # hook and falling back to ``_default_process_engine_inputs``. (The + # worker-side mixin candidate chain, which *probes* for a ``*_full_payload`` + # derivative, keeps its own skip-on-mismatch tolerance.) + custom_process_input_func = resolve_processor( + _cpif_path, + expected_kind=None, + stage_config=stage_config, + ).fn sync_process_input_func: str | None = None _spif_path = _get_attr_or_item(stage_config, "sync_process_input_func") From edab9929378f65141b31d3cff1b5cc4b9ec254d4 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Mon, 31 Aug 2026 16:06:41 +0800 Subject: [PATCH 12/41] refactor(stage_input_processors): wrap orchestrator processors once per resolve (RFC #4872 P1) invoke_orchestrator_processor re-wrapped on every forward, so a legacy C2 shell like thinker2talker_token_only re-emitted the DeprecationWarning on each forward. wrap_orchestrator_processor is now memoized per callable: the signature probe and the warning happen once (at first resolution/invocation), then the cached C1-compatible wrapper is reused. - Add _WRAP_CACHE keyed by the callable; C1 callables still pass through unchanged (no warning). - Add test_dispatch_wrap_once.py locking: warning fires exactly once across repeated wraps / forwards, C1 passthrough warns never. Reviewer: amy-why-3459 (cleanup 4). Signed-off-by: RyanYun09 --- .../test_dispatch_wrap_once.py | 89 +++++++++++++++++ .../stage_input_processors/_dispatch.py | 96 ++++++++++++------- 2 files changed, 149 insertions(+), 36 deletions(-) create mode 100644 tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py diff --git a/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py new file mode 100644 index 00000000000..70814d2902c --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""RFC #4872 Phase 2 (P1): orchestrator dispatch wrap-once semantics. + +``invoke_orchestrator_processor`` runs per request (and per forward), so +``wrap_orchestrator_processor`` must not re-wrap — and must not re-emit the +legacy-contract ``DeprecationWarning`` — on every forward. The wrap result is +memoized per callable: the signature probe and the warning happen **once** per +processor (at first resolution/invocation), then the cached C1-compatible +wrapper is reused. + +Pure-logic, CPU-only (no model loading, no vllm runtime). +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from vllm_omni.model_executor.stage_input_processors._dispatch import ( + OrchestratorInputContext, + invoke_orchestrator_processor, + wrap_orchestrator_processor, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _legacy_c2(source_outputs, prompt=None, requires_multimodal_data=False, streaming_context=None): + """A legacy C2 positional shell (like qwen3_omni.thinker2talker_token_only).""" + return [source_outputs, prompt, requires_multimodal_data, streaming_context] + + +def _legacy_c0(source_outputs, prompt=None, requires_multimodal_data=False): + """A legacy C0 3-arg shell.""" + return [source_outputs, prompt, requires_multimodal_data] + + +def _c1(source_outputs: list[Any], ctx: OrchestratorInputContext): + """Already speaks the C1 ``(source_outputs, ctx)`` contract.""" + return [source_outputs, ctx] + + +def _count_deprecation_warnings(trigger) -> int: + """Run *trigger* under ``simplefilter("always")`` and count DeprecationWarnings.""" + import warnings + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + trigger() + return sum(1 for w in caught if issubclass(w.category, DeprecationWarning)) + + +def test_wrap_legacy_shape_warns_once_then_memoized(): + first = _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_legacy_c2)) + # The first wrap (at "resolve") emits exactly one DeprecationWarning. + assert first == 1 + # Subsequent wraps reuse the memoized adapter: no further warning. + assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_legacy_c2)) == 0 + assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_legacy_c2)) == 0 + + +def test_invoke_does_not_rewarn_on_every_forward(): + ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, streaming_context="s") + + def invoke_many(): + for _ in range(5): + invoke_orchestrator_processor(_legacy_c0, ["out"], ctx) + + # Exactly one DeprecationWarning total across 5 forwards (wrap happens once). + assert _count_deprecation_warnings(invoke_many) == 1 + + +def test_invoke_adapts_legacy_shape_and_forwards_context(): + ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, streaming_context="s") + result = invoke_orchestrator_processor(_legacy_c2, ["out"], ctx) + # C2 adapter forwards (source_outputs, prompt, requires_multimodal_data, streaming_context). + assert result == [["out"], "p", True, "s"] + + +def test_c1_processor_passes_through_without_warning(): + wrapped = wrap_orchestrator_processor(_c1) + assert wrapped is _c1 # already speaks ctx: returned unchanged + ctx = OrchestratorInputContext(prompt="p") + assert invoke_orchestrator_processor(_c1, ["out"], ctx) == [["out"], ctx] + # No DeprecationWarning for C1 processors at all. + assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_c1)) == 0 diff --git a/vllm_omni/model_executor/stage_input_processors/_dispatch.py b/vllm_omni/model_executor/stage_input_processors/_dispatch.py index dea5f990df0..0fae461abe2 100644 --- a/vllm_omni/model_executor/stage_input_processors/_dispatch.py +++ b/vllm_omni/model_executor/stage_input_processors/_dispatch.py @@ -239,6 +239,14 @@ def _adapted(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: return _adapted +#: Wrap memo: the signature probe + ``DeprecationWarning`` must happen **once per +#: processor**, not on every forward. ``invoke_orchestrator_processor`` runs per +#: request, and re-wrapping a legacy C2 shape (e.g. ``thinker2talker_token_only``) +#: would otherwise re-emit the warning on every forward. Keyed by the callable +#: object itself (module-level processors are few and live for the process). +_WRAP_CACHE: dict[Any, Any] = {} + + def wrap_orchestrator_processor(fn: Any) -> PlaceholderPromptBuilder | DiffusionInputBuilder: """Return a C1-compatible callable ``(source_outputs, ctx)`` for ``fn``. @@ -253,48 +261,64 @@ def wrap_orchestrator_processor(fn: Any) -> PlaceholderPromptBuilder | Diffusion ``ctx.streaming_context``; * C0 3-arg ``(source_outputs, prompt, requires_multimodal_data)``. Every legacy shape emits a ``DeprecationWarning``. + + The result is memoized per callable (see ``_WRAP_CACHE``): the wrap — and its + ``DeprecationWarning`` — happens exactly once per processor, at first + resolution/invocation, so ``invoke_orchestrator_processor`` does not re-wrap + (or re-warn) on every forward. """ if not callable(fn): raise TypeError(f"stage-input processor must be callable, got {fn!r}") - if _accepts_orchestrator_input_context(fn): - return fn # type: ignore[return-value] - - names = _positional_parameter_names(fn) - - # Uninspectable callable (e.g. builtin): fall back to the most conservative - # 3-arg C0 shape; the runtime previously treated it via the <4 arity path. - if names is None: - names = ["source_outputs", "prompt", "requires_multimodal_data"] - - # C4 legacy multi-source (e.g. moss_tts.talker2codec). - if len(names) >= 2 and names[0] == "stage_list" and names[1] == "engine_input_source": - return _adapt_moss_processor(fn) # type: ignore[return-value] - - # C3 diffusion: forwarding sampling_params (diffusion-stage params). - if "sampling_params" in names: - - def _c3(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: - return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data, sampling_params=ctx.sampling_params) - - _warn_legacy_contract(fn) - return _c3 # type: ignore[return-value] - - # C2 placeholder: forwarding the streaming context. - if "streaming_context" in names or "_streaming_context" in names: - - def _c2(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: - return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data, ctx.streaming_context) - - _warn_legacy_contract(fn) - return _c2 # type: ignore[return-value] + try: + cached = _WRAP_CACHE.get(fn) + except TypeError: # pragma: no cover - unhashable callable (rare) + cached = None + if cached is not None: + return cached # type: ignore[return-value] - # C0 3-arg legacy. - def _c0(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: - return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data) + if _accepts_orchestrator_input_context(fn): + wrapped: Any = fn + else: + names = _positional_parameter_names(fn) + + # Uninspectable callable (e.g. builtin): fall back to the most conservative + # 3-arg C0 shape; the runtime previously treated it via the <4 arity path. + if names is None: + names = ["source_outputs", "prompt", "requires_multimodal_data"] + + # C4 legacy multi-source (e.g. moss_tts.talker2codec). + if len(names) >= 2 and names[0] == "stage_list" and names[1] == "engine_input_source": + wrapped = _adapt_moss_processor(fn) + # C3 diffusion: forwarding sampling_params (diffusion-stage params). + elif "sampling_params" in names: + + def _c3(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: + return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data, sampling_params=ctx.sampling_params) + + _warn_legacy_contract(fn) + wrapped = _c3 + # C2 placeholder: forwarding the streaming context. + elif "streaming_context" in names or "_streaming_context" in names: + + def _c2(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: + return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data, ctx.streaming_context) + + _warn_legacy_contract(fn) + wrapped = _c2 + else: + # C0 3-arg legacy. + def _c0(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: + return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data) + + _warn_legacy_contract(fn) + wrapped = _c0 - _warn_legacy_contract(fn) - return _c0 # type: ignore[return-value] + try: + _WRAP_CACHE[fn] = wrapped + except TypeError: # pragma: no cover - unhashable callable (rare) + pass + return wrapped # type: ignore[return-value] def invoke_orchestrator_processor( From 8f7c25ce3cb45ee48f54076b154e0b7c85ad751b Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Mon, 31 Aug 2026 16:06:48 +0800 Subject: [PATCH 13/41] refactor(stage_input_processors): remove dead length helper and move codec guard off import (RFC #4872 P8b/P6) - qwen3_omni._compute_talker_prompt_ids_length is dead code after the P8b consolidation onto _common.compute_placeholder_prompt_len; delete it. - _assert_codec_token_ids_consistent is no longer run at module import (it pulled in transformers / the qwen3-tts config and broke light imports). It is now invoked explicitly at model-load startup in qwen3_omni_moe_talker.__init__, keeping the processor constants in sync with the HF configs the model actually runs with. Reviewer: amy-why-3459 (cleanup 5, 6). Signed-off-by: RyanYun09 --- .../qwen3_omni/qwen3_omni_moe_talker.py | 9 ++++ .../stage_input_processors/qwen3_omni.py | 45 ++----------------- 2 files changed, 13 insertions(+), 41 deletions(-) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index ea3e2fcfe36..8e494b83c4e 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -81,6 +81,15 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config + # RFC #4872 P6: the codec token-id consistency guard is deliberately not + # run at module import (it pulls in transformers / the qwen3-tts config). + # Invoke it explicitly at model-load startup so the processor constants + # stay in sync with the HF configs this model actually runs with. + from vllm_omni.model_executor.stage_input_processors.qwen3_omni import ( + _assert_codec_token_ids_consistent, + ) + + _assert_codec_token_ids_consistent() rope_params = getattr(talker_config.text_config, "rope_scaling", None) if rope_params is None: rope_params = getattr(talker_config.text_config, "rope_parameters", None) or {} diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 6402796c372..63f1a25d20c 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -83,10 +83,10 @@ def _assert_codec_token_ids_consistent() -> None: ) -# Single source of truth: keep processor / model impl / HF config in sync. -_assert_codec_token_ids_consistent() - - +# The P6 fail-fast guard is intentionally **not** run at import time (it pulls in +# transformers / the in-repo qwen3-tts config and must not break light imports). +# It is invoked explicitly at model-load startup (qwen3_omni_moe_talker) so the +# processor constants stay in sync with the HF configs the model actually runs with. def _layer_tensor(layers: dict[Any, Any], key: str) -> torch.Tensor | None: """Fetch layer tensor with tolerant key lookup (str/int/semantic).""" if not isinstance(layers, dict): @@ -154,43 +154,6 @@ def _sampled_token_embed(thinker_embed: Mapping[str, Any], request: Any) -> torc return sampled.detach().cpu() -def _compute_talker_prompt_ids_length(info: OmniPayload, device: torch.device | str = "cuda") -> int: - im_start_token_id = 151644 - system_token_id = 8948 - user_token_id = 872 - assistant_token_id = 77091 - - ids = info.get("ids", {}) - thinker_sequences = torch.tensor(ids["all"], dtype=torch.long, device=device).unsqueeze(0) # [1, T] - - input_ids = torch.tensor(ids["prompt"], dtype=torch.long, device=device).unsqueeze(0) # [1, T] - - im_start_indexes = torch.cat( - [ - torch.nonzero(input_ids[0] == im_start_token_id).squeeze(1), - torch.tensor([thinker_sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype), - ], - dim=0, - ) - - sum_user_len = 0 - assistant_len = 0 - for i in range(len(im_start_indexes) - 1): - s = int(im_start_indexes[i].item()) - e = int(im_start_indexes[i + 1].item()) - role = int(input_ids[0, s + 1].item()) - if role == system_token_id: - continue - elif role == user_token_id: - sum_user_len += e - s - elif role == assistant_token_id and i == len(im_start_indexes) - 2: - assistant_len += 9 # 3 + 4 + 1 + 1 - else: - pass - - return sum_user_len + assistant_len - - # ========================= # Common helpers # ========================= From ccfb0d8b98db5063e1c854289fabd6fbcc55ff17 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Mon, 31 Aug 2026 16:06:55 +0800 Subject: [PATCH 14/41] feat(engine): wire dead_processor_hint at orchestrator startup (RFC #4872 M0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire the pure decision helper at startup: for every stage whose custom_process_input_func would never be invoked under the current async-chunk / sync wiring (mirroring the _route_output gate), log a warn-only hint so a configured-but-dead hook is not silently dropped. Defensive by design — fake pools / minimal clients never raise. Reviewer: amy-why-3459 (cleanup 7). Signed-off-by: RyanYun09 --- vllm_omni/engine/orchestrator.py | 50 ++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 1265a00995c..de19a5d81b1 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -370,6 +370,9 @@ def __init__( self.async_chunk = bool(async_chunk) self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools + # RFC #4872 M0: wire the dead-processor hint at startup (warn-only) so a + # configured-but-never-invoked custom_process_input_func is surfaced. + self._warn_dead_input_processors() self.log_stats = log_stats if log_stats and not collect_stage_metrics: raise ValueError("log_stats=True requires collect_stage_metrics=True") @@ -2061,6 +2064,53 @@ def _stage_receives_async_chunks(self, stage_id: int) -> bool: model_config = getattr(pool.stage_vllm_config, "model_config", None) return stage_receives_chunks(model_config) + def _warn_dead_input_processors(self) -> None: + """Startup (M0, warn-only) wiring of :func:`dead_processor_hint`. + + For every stage whose ``custom_process_input_func`` would never be + invoked under the current async-chunk / sync wiring (mirroring the + ``_route_output`` gate ``(not async_chunk or not + _stage_receives_async_chunks(stage_id))``), log a warning so a + configured-but-dead hook is not silently dropped. This is a pure hint + (RFC #4872 M0 warn-first): the orchestrator never enforces it at runtime. + Defensive by design — fake pools / minimal clients in tests or future + connectors may not expose ``stage_client``, so the loop never raises. + """ + from vllm_omni.model_executor.stage_input_processors import dead_processor_hint, infer_kind + + for stage_id, pool in enumerate(self.stage_pools): + try: + client = getattr(pool, "stage_client", None) + fn = getattr(client, "custom_process_input_func", None) + if fn is None: + continue + module = getattr(fn, "__module__", None) + name = getattr(fn, "__qualname__", None) or getattr(fn, "__name__", None) + if not module or not name: + continue + path = f"{module}.{name}" + kind = infer_kind(fn, path=path) + has_sync = bool(getattr(client, "sync_process_input_func", None)) + dead = dead_processor_hint( + kind, + async_chunk=self.async_chunk, + downstream_receives_async_chunks=self._stage_receives_async_chunks(stage_id), + has_sync=has_sync, + ) + if dead: + logger.warning( + "[Orchestrator] stage-%d custom_process_input_func %r is never invoked " + "under the current async-chunk/sync wiring (RFC #4872 M0: dead-processor hint)", + stage_id, + path, + ) + except Exception as exc: # pragma: no cover - defensive only + logger.debug( + "[Orchestrator] dead-processor hint lookup for stage-%s failed: %s", + stage_id, + exc, + ) + def _get_stage_input_processor(self, stage_id: int) -> Any: processor = self._stage_input_processors.get(stage_id) if processor is None: From f33009c8cdbc2008dd0a8d1ac56c408d7d6a1a26 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Mon, 31 Aug 2026 16:07:01 +0800 Subject: [PATCH 15/41] docs(contributing): align adding_omni_model snippet with the real *_token_only shell (RFC #4872 P1) The snippet showed thinker2talker_token_only as the C1 (source_outputs, ctx) contract, but the function is still the four-argument legacy shell (prompt / requires_multimodal_data / streaming_context). Rewrite the snippet to document the real signature, and note that the orchestrator normalizes it once via wrap_orchestrator_processor with the C1 dual-entry builders (build_forward_placeholder / build_prewarm_placeholder) attached to the shell. Reviewer: amy-why-3459 (cleanup 8). Signed-off-by: RyanYun09 --- docs/contributing/model/adding_omni_model.md | 35 ++++++++++++++------ 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 64c280cdeb8..7cce818aeda 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -504,23 +504,36 @@ implement all three processors for every inter-stage edge — `*_full_payload`, required in both modes: in non-async mode it builds the forward placeholder, in async mode the orchestrator reuses it to prewarm the downstream stage. -Orchestrator-facing builders follow the C1 contract -`(source_outputs, ctx: OrchestratorInputContext)`: +Orchestrator-facing builders are normalized to the C1 contract +`(source_outputs, ctx: OrchestratorInputContext)` by +`wrap_orchestrator_processor`. The legacy `*_token_only` shell keeps its +positional signature for call-site compatibility; the actual C1 entry point is +the `build_forward_placeholder` / `build_prewarm_placeholder` pair attached to +the shell (RFC #4872 P8b dual-entry design): ```python # qwen3_omni.py (Thinker → Talker) -from vllm_omni.model_executor.stage_input_processors import OrchestratorInputContext - - def thinker2talker_token_only( source_outputs: list[Any], - ctx: OrchestratorInputContext, + prompt: OmniTokensPrompt | TextPrompt | None = None, + requires_multimodal_data: bool = False, + streaming_context: Any | None = None, ) -> list[OmniTokensPrompt]: - """Allocate talker prefill slots; bulk tensors arrive via the connector.""" + """Allocate talker prefill slots; bulk tensors arrive via the connector. + + Legacy positional shell (C2). The orchestrator wraps it once via + ``wrap_orchestrator_processor``, which forwards the C1 + ``OrchestratorInputContext`` fields back into the positional arguments. + """ ... +# Attached C1 dual-entry builders (RFC #4872 P8b): +# thinker2talker_token_only.build_forward_placeholder(source_outputs, ctx) +# thinker2talker_token_only.build_prewarm_placeholder(*, stage0_prompt, ctx, ...) + + def thinker2talker_full_payload( *, transfer_manager: Any, @@ -543,10 +556,10 @@ def thinker2talker_async_chunk( ... ``` -`OrchestratorInputContext` carries `prompt`, `requires_multimodal_data`, -`streaming_context` and `sampling_params` (and deliberately has no -`model_config` field). The producer-side builders (`*_full_payload`, -`*_async_chunk`) are keyword-only and never receive an +The `OrchestratorInputContext` carried by the C1 contract has `prompt`, +`requires_multimodal_data`, `streaming_context` and `sampling_params` (and +deliberately has no `model_config` field). The producer-side builders +(`*_full_payload`, `*_async_chunk`) are keyword-only and never receive an `OrchestratorInputContext`; `pooling_output` / `multimodal_output` are load-bearing keyword names of the connector data plane. From 18e27d65f72c51f491d72645d2e913c86341dfa8 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Mon, 31 Aug 2026 17:13:20 +0800 Subject: [PATCH 16/41] test(stage_input_processors): normalize docstrings for upstream contribution Replace internal work-log phrasing (RFC proposal/phase numbers, local-dev shim wording, machine-specific skip notes, imperative/process-oriented instructions) with professional descriptions of what each test covers and why. Keep the RFC #4872 reference as a top-level link for traceability. No functional changes: only comments and docstrings were edited. Signed-off-by: RyanYun09 --- .../test_orchestrator_stage_input_bridge.py | 4 +- tests/engine/test_stage_metadata.py | 6 +-- .../stage_input_processors/conftest.py | 19 +++++---- .../test_common_helpers_golden.py | 39 ++++++++++--------- .../test_constants_consistency.py | 2 +- .../test_dispatch_wrap_once.py | 2 +- .../test_placeholder_parity.py | 36 +++++++++-------- .../test_qwen3_tts_async_chunk.py | 12 +++--- .../stage_input_processors/test_registry.py | 2 +- 9 files changed, 65 insertions(+), 57 deletions(-) diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 691f0c5c70c..1fbb7c32db7 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -303,8 +303,8 @@ async def test_async_prewarm_skips_stage_with_custom_process_input_func() -> Non async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: - """RFC #4872 P8b: async-chunk prewarm resolves ``sync_process_input_func`` - (``*_token_only``) to ``build_prewarm_placeholder`` and uses its estimate. + """Async-chunk prewarm resolves ``sync_process_input_func`` (the + ``*_token_only`` path) to ``build_prewarm_placeholder`` and uses its estimate. The placeholder length is ``stage0_only`` = the Qwen chat-template scan on the stage-0 prompt. The synthetic ``[1, 2]`` prompt has no chat-template diff --git a/tests/engine/test_stage_metadata.py b/tests/engine/test_stage_metadata.py index c516a2c1843..a5a57f8367a 100644 --- a/tests/engine/test_stage_metadata.py +++ b/tests/engine/test_stage_metadata.py @@ -200,9 +200,9 @@ def test_extract_stage_metadata_defaults_missing_engine_input_source(): # --------------------------------------------------------------------------- -# RFC #4872 P2 / adding-omni: a bad custom_process_input_func must FAIL FAST. -# extract_legacy_stage_metadata no longer swallows resolve_processor failures -# and continues with custom_process_input_func=None (which silently fell back to +# A bad custom_process_input_func must FAIL FAST. The config-hook resolution +# path no longer swallows resolve_processor failures and continues with +# custom_process_input_func=None (which silently fell back to # _default_process_engine_inputs). The worker mixin candidate chain keeps its # own skip-on-mismatch tolerance; this is the config-hook resolution path. # --------------------------------------------------------------------------- diff --git a/tests/model_executor/stage_input_processors/conftest.py b/tests/model_executor/stage_input_processors/conftest.py index 64ae9ade526..c9873a45d48 100644 --- a/tests/model_executor/stage_input_processors/conftest.py +++ b/tests/model_executor/stage_input_processors/conftest.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Local-dev import shim for ``vllm_omni`` tests (macOS). +"""Test-support import fallback for ``vllm_omni`` pure-logic tests. -vllm does not install on macOS, so importing real ``vllm_omni`` submodules -fails at ``import vllm``. This conftest installs a *test-only* shim that: +These unit tests exercise ``vllm_omni`` helper logic without a running vLLM +runtime. When the real ``vllm`` package is not installed, importing +``vllm_omni`` submodules would fail at ``import vllm``; this conftest then +installs a *test-only* fallback that: 1. Bypasses ``vllm_omni/__init__.py`` (so its ``patch`` / ``transformers`` machinery is never executed) by pre-inserting a synthetic ``vllm_omni`` @@ -13,9 +15,9 @@ 2. Stubs the entire ``vllm`` package tree with permissive ``_Dummy`` classes so module-level ``from vllm... import ...`` statements resolve. -The shim is active **only** when real ``vllm`` is not importable (local dev). -On CI (vllm installed) it is a no-op and the real package is imported, which -guards against the shim masking genuine import problems. +The fallback is active **only** when real ``vllm`` is not importable. When +vllm is available (e.g. on CI) it is a no-op and the real package is imported, +which guards against the fallback masking genuine import problems. Scope: only pure-logic / unit-level tests (no model execution). Any test that needs a real vLLM runtime must be marked to run on GPU/NPU CI instead. @@ -100,10 +102,10 @@ def find_spec(self, fullname, path=None, target=None): def _install_shim(): - """Bypass the real vllm_omni/__init__ + stub vllm.""" + """Bypass the real vllm_omni/__init__ and stub the vllm package tree.""" if _REPO_ROOT not in sys.path: sys.path.insert(0, _REPO_ROOT) - # Synthetic vllm_omni parent package (skip the real __init__). + # Synthetic vllm_omni parent package (skips the real __init__). pkg = types.ModuleType("vllm_omni") pkg.__path__ = [_VLLM_OMNI_DIR] pkg.__package__ = "vllm_omni" @@ -114,6 +116,7 @@ def _install_shim(): def _activate_if_needed(): + """Install the import fallback only when real vllm is unavailable.""" try: import vllm # noqa: F401 except ModuleNotFoundError: diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py index 2c16c64f464..000becbed72 100644 --- a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -4,16 +4,19 @@ """Golden tests that lock the *current* per-module behavior of duplicated helper functions in ``stage_input_processors``. -Phase 1a (RFC #4872 Proposal 3): capture behaviour BEFORE introducing -``_common.py``. Each case here encodes a deliberately-observed semantic -delta between modules (None handling, tensor->list, ``ConstantList._x`` -unwrapping, tuple handling, codec-frame validity masks, delay-pattern -strictness, ...). ``_common.ensure_list`` (Phase 1b) must pass the same -cases; do NOT silently "pick one" semantics across variants. - -Modules that cannot be imported in this environment (e.g. ``cosyvoice3`` / -``step_audio2`` on macOS because of missing deps) are skipped locally and -run on CI where the real dependencies exist. +Related to RFC #4872 (https://github.com/vllm-project/vllm-omni/issues/4872): +these cases capture the observable per-module semantics that were previously +duplicated across model processors, so a consolidated implementation in +``_common`` must reproduce the same behavior. Each case encodes a +documented semantic difference between modules (None handling, tensor->list, +``ConstantList._x`` unwrapping, tuple handling, codec-frame validity masks, +delay-pattern strictness, ...). The shared helpers in ``_common`` must pass +the same cases; where legacy variants disagreed, the consolidated behavior is +preserved through explicit named variants rather than silently normalized. + +Modules whose optional dependencies are unavailable in the current +environment are skipped and exercised on CI, where the real dependencies +exist. """ import importlib @@ -122,7 +125,7 @@ def test_ensure_list_golden(module, inp, expected, expect_typeerror): def test_output_formatter_ensure_list_wrap_only(): - """output_formatter._ensure_list keeps wrap-only semantics (RFC #4872 P3). + """output_formatter._ensure_list keeps wrap-only semantics. The diffusion formatter is NOT a stage-input helper: a non-list primary payload is wrapped as ``[x]`` verbatim — no tensor flattening, no dict-key @@ -431,7 +434,7 @@ def test_compute_talker_prompt_ids_length_golden(): # =========================================================================== -# Phase 1b: _common.py must reproduce the golden-locked behaviour. +# The shared _common.py helpers must reproduce the golden-locked behaviour. # =========================================================================== @@ -445,7 +448,7 @@ def _common_import() -> Any: def _shim_active() -> bool: - """True when the local-dev vllm stub (no real vllm) is in effect.""" + """True when the test-support import fallback (no real vllm) is active.""" try: import vllm_omni # noqa: F401 @@ -559,7 +562,7 @@ def test_common_compute_placeholder_prompt_len_matches_golden(): # full mode == the legacy Qwen chat-template scan golden (15). assert c.compute_placeholder_prompt_len(ids_or_prompt={"ids": {"all": prompt, "prompt": prompt}}, mode="full") == 15 # stage0_only mode: the same scan on the flat stage-0 list (prewarm) -> 15, - # so the builder and the inline fallback agree (RFC #4872 P8b). + # so the builder and the inline fallback agree. assert c.compute_placeholder_prompt_len(ids_or_prompt=prompt, mode="stage0_only") == 15 @@ -568,10 +571,10 @@ def test_common_pack_placeholder_prompt(): prompt = c.pack_placeholder_prompt(prompt_len=4, voice_metadata={"speaker": 1}) empty_meta = c.pack_placeholder_prompt(prompt_len=1) if _shim_active(): - # Under the local vllm stub, OmniTokensPrompt's real constructor is - # bypassed, so field access returns stub values. Field assertions run - # on CI where real vllm constructs the model properly. - pytest.skip("OmniTokensPrompt field construction needs real vllm (CI)") + # Under the import fallback, OmniTokensPrompt's real constructor is + # bypassed, so field access returns fallback values. Field assertions + # run where real vllm constructs the model properly. + pytest.skip("OmniTokensPrompt field construction needs real vllm") # Real vllm (0.27) builds OmniTokensPrompt as a dict subclass (MRO # ['OmniTokensPrompt', 'dict', 'object']): attribute access is not # available, so read the packed fields via dict indexing. diff --git a/tests/model_executor/stage_input_processors/test_constants_consistency.py b/tests/model_executor/stage_input_processors/test_constants_consistency.py index abc8c15f1af..f30ef3aa4c7 100644 --- a/tests/model_executor/stage_input_processors/test_constants_consistency.py +++ b/tests/model_executor/stage_input_processors/test_constants_consistency.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""RFC #4872 P6: constant-consistency guards. +"""Constant-consistency guards (RFC #4872). Verifies that magic special-token ids that were previously duplicated across processors / pipelines / model implementations now resolve to a single value: diff --git a/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py index 70814d2902c..bfa887b66a7 100644 --- a/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py +++ b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""RFC #4872 Phase 2 (P1): orchestrator dispatch wrap-once semantics. +"""Orchestrator dispatch wrap-once semantics (RFC #4872). ``invoke_orchestrator_processor`` runs per request (and per forward), so ``wrap_orchestrator_processor`` must not re-wrap — and must not re-emit the diff --git a/tests/model_executor/stage_input_processors/test_placeholder_parity.py b/tests/model_executor/stage_input_processors/test_placeholder_parity.py index b737470e13c..b95e4fdbb0c 100644 --- a/tests/model_executor/stage_input_processors/test_placeholder_parity.py +++ b/tests/model_executor/stage_input_processors/test_placeholder_parity.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""P8b parity tests: async-chunk prewarm placeholder vs sync forward placeholder. +"""Parity tests: async-chunk prewarm placeholder vs sync forward placeholder. -RFC #4872 Proposal 8b: ``*_token_only`` placeholder builders expose a dual entry -— ``build_forward_placeholder(source_outputs, ctx)`` for the sync forward path +Related to RFC #4872 (https://github.com/vllm-project/vllm-omni/issues/4872): +``*_token_only`` placeholder builders expose a dual entry — +``build_forward_placeholder(source_outputs, ctx)`` for the sync forward path and ``build_prewarm_placeholder(*, stage0_prompt, ctx, downstream_stage_id)`` for the async-chunk prewarm path (no ``source_outputs`` yet). Both must be fed by the same ``_common`` length / packing helpers so async and sync placeholders stay @@ -21,10 +22,10 @@ segment). The connector fixup path replaces the estimate with the real length once the upstream chunk arrives. -These tests are CPU-only (no model loading). They run under the local vllm -stub (shim) and on CI where real vllm builds ``OmniTokensPrompt`` as a dict -subclass; length capture uses a ``pack_placeholder_prompt`` spy so it works in -both environments. +These tests are CPU-only (no model loading). They run both without a vllm +runtime (the import fallback is active) and with real vllm, where +``OmniTokensPrompt`` is built as a dict subclass; length capture uses a +``pack_placeholder_prompt`` spy so it works in both environments. """ from __future__ import annotations @@ -50,7 +51,7 @@ def _shim_active() -> bool: - """True when the local-dev vllm stub (no real vllm) is in effect.""" + """True when the test-support import fallback (no real vllm) is active.""" try: import vllm_omni # noqa: F401 @@ -83,7 +84,7 @@ def _capture_pack(monkeypatch): ``pack_placeholder_prompt`` is the single packing chokepoint shared by both builders, so capturing ``prompt_len`` / ``voice_metadata`` verifies the builders' length logic without depending on how ``OmniTokensPrompt`` is - constructed (shim vs real vllm). + constructed (import fallback vs real vllm). """ recorded: dict[str, Any] = {} @@ -183,7 +184,8 @@ def test_forward_prewarm_golden_relationship(monkeypatch): # --------------------------------------------------------------------------- -# Packed placeholder structure (real vllm asserts; shim skips field reads). +# Packed placeholder structure (asserted with real vllm; field reads are +# skipped under the import fallback). # --------------------------------------------------------------------------- @@ -191,10 +193,10 @@ def test_pack_placeholder_prompt_structure(): prompt = _common.pack_placeholder_prompt(prompt_len=4, voice_metadata={"speaker": ["ethan"]}) empty_meta = _common.pack_placeholder_prompt(prompt_len=1) if _shim_active(): - # Under the local vllm stub, OmniTokensPrompt's real constructor is - # bypassed, so field access returns stub values. Field assertions run - # on CI where real vllm constructs the model properly. - pytest.skip("OmniTokensPrompt field construction needs real vllm (CI)") + # Under the import fallback, OmniTokensPrompt's real constructor is + # bypassed, so field access returns fallback values. Field assertions + # run where real vllm constructs the model properly. + pytest.skip("OmniTokensPrompt field construction needs real vllm") assert prompt["prompt_token_ids"] == [0, 0, 0, 0] assert prompt["additional_information"] == {"speaker": ["ethan"]} assert empty_meta["prompt_token_ids"] == [0] @@ -226,7 +228,7 @@ def test_thinker2talker_token_only_delegates_to_forward_placeholder(monkeypatch) def test_adapter_compute_length_consistent_with_common(): - # Lazy import keeps this module import-light under the shim. + # Lazy import keeps this module import-light without a vllm runtime. from vllm_omni.distributed.omni_connectors import adapter assert adapter.compute_talker_prompt_ids_length(GOLDEN_PROMPT) == 15 @@ -234,8 +236,8 @@ def test_adapter_compute_length_consistent_with_common(): ids_or_prompt={"ids": {"all": GOLDEN_PROMPT, "prompt": GOLDEN_PROMPT}}, mode="full", ) - # RFC #4872 P8b: the prewarm stage0_only scan (what build_prewarm_placeholder - # and the orchestrator inline fallback both use) must equal the adapter scan. + # The prewarm stage0_only scan (what build_prewarm_placeholder and the + # orchestrator inline fallback both use) must equal the adapter scan. assert adapter.compute_talker_prompt_ids_length(GOLDEN_PROMPT) == _common.compute_placeholder_prompt_len( ids_or_prompt=GOLDEN_PROMPT, mode="stage0_only", diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 59c02cdec0a..345d95845b4 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -1697,22 +1697,22 @@ def append_and_emit(frame_idx, finished=False): tm.ramp_chunk_count[rid] += 1 return p - # Phase 1 (t=0.0): chunk 0, IC=2 + # t=0.0: chunk 0, IC=2. for i in range(2): append_and_emit(i) - # Phase 2 (t=0.05): build EWMA, emit target=4 at length=6 + # t=0.05: build EWMA, emit target=4 at length=6. clock[0] = 0.05 for i in range(2, 6): append_and_emit(i) - # Phase 3 (t=0.05): greedy target=18 (large buffer), accumulate - # 17 frames while holding (new_frames 1..17 < 18) + # t=0.05: greedy target=18 (large buffer), accumulate 17 frames while + # holding (new_frames 1..17 < 18). for i in range(6, 23): append_and_emit(i) - # Phase 4 (t=100.0): clock jumps, greedy collapses, target drops - # to ramp_floor=6. new_frames=18 > 6 -> bug trigger point. + # t=100.0: clock jumps, greedy collapses, target drops to ramp_floor=6. + # new_frames=18 > 6 -> bug trigger point. clock[0] = 100.0 for i in range(23, total_frames): append_and_emit(i, finished=(i == total_frames - 1)) diff --git a/tests/model_executor/stage_input_processors/test_registry.py b/tests/model_executor/stage_input_processors/test_registry.py index c7085cee55b..0d85721bb77 100644 --- a/tests/model_executor/stage_input_processors/test_registry.py +++ b/tests/model_executor/stage_input_processors/test_registry.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Unit tests for the RFC #4872 Phase 3 (Part A) processor registry. +"""Unit tests for the stage-input processor registry (RFC #4872). Pure-logic, CPU-only: exercises kind inference, structural validation, dead-processor hints and ``resolve_processor`` without executing any processor From 657f0b654a189b4cf2b03279cc59be59e2ac06c3 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Mon, 31 Aug 2026 17:13:27 +0800 Subject: [PATCH 17/41] refactor(stage_input_processors): normalize module comments and docstrings for upstream contribution Remove internal work-log phrasing from module and function docstrings and inline comments across the stage-input processor implementation (RFC proposal/phase numbers such as P2/P3/P6/P8b, Phase 1a/1b/3, Step 1c, Part A, local-dev shim and macOS/CI environment notes, imperative/process-oriented wording). Rewrite as professional descriptions of what each helper does and why. Keep the RFC #4872 reference as a top-level link for traceability. No functional changes: only comments and docstrings were edited. Signed-off-by: RyanYun09 --- vllm_omni/config/stage_config.py | 6 +-- vllm_omni/diffusion/output_formatter.py | 2 +- .../distributed/omni_connectors/adapter.py | 9 ++-- vllm_omni/engine/orchestrator.py | 24 +++++------ vllm_omni/engine/stage_engine_core_client.py | 11 +++-- vllm_omni/engine/stage_init_utils.py | 17 ++++---- .../models/qwen3_omni/pipeline.py | 4 +- .../models/qwen3_omni/qwen3_omni.py | 10 ++--- .../qwen3_omni/qwen3_omni_moe_talker.py | 8 ++-- .../stage_input_processors/__init__.py | 3 +- .../stage_input_processors/_common.py | 43 +++++++++---------- .../stage_input_processors/_constants.py | 4 +- .../stage_input_processors/_dispatch.py | 6 +-- .../stage_input_processors/_registry.py | 12 +++--- .../stage_input_processors/audex.py | 2 +- .../stage_input_processors/cosyvoice3.py | 6 +-- .../stage_input_processors/glm_tts.py | 2 +- .../stage_input_processors/higgs_audio_v2.py | 8 ++-- .../stage_input_processors/higgs_audio_v3.py | 8 ++-- .../stage_input_processors/ming_flash_omni.py | 2 +- .../stage_input_processors/qwen3_omni.py | 37 ++++++++-------- .../stage_input_processors/step_audio2.py | 4 +- .../omni_connector_model_runner_mixin.py | 8 ++-- 23 files changed, 117 insertions(+), 119 deletions(-) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 3a148befe40..c44939af7fa 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -1513,9 +1513,9 @@ def to_omegaconf(self) -> Any: if self.custom_process_input_func: config_dict["custom_process_input_func"] = self.custom_process_input_func - # RFC #4872 P8b: carry the raw ``sync_process_input_func`` path to the - # orchestrator so async-chunk prewarm can resolve the registered - # placeholder builder (``*_token_only``) without re-reading the YAML. + # Carry the raw ``sync_process_input_func`` path to the orchestrator so + # async-chunk prewarm can resolve the registered placeholder builder + # (``*_token_only``) without re-reading the YAML. if self.sync_process_input_func: config_dict["sync_process_input_func"] = self.sync_process_input_func diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index 5ee3bba33ce..3f1fe676a6b 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -211,7 +211,7 @@ def _ensure_list(outputs: DiffusionPayloadValue) -> list[DiffusionPayloadValue]: ``[]``. The canonical ``_common.ensure_list`` would flatten tensors, iterate dict keys, and walk a PIL ``Image`` row by row, so it must not be reused here; ``_common.ensure_list_wrap_only`` reproduces the exact legacy - semantics (RFC #4872 P3). + semantics. """ return _common.ensure_list_wrap_only(outputs) diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 10c73d91681..3af7cd2c161 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -193,11 +193,10 @@ def try_recv_via_connector( def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int: """Compute the length of the talker prompt ids. - RFC #4872 P8b: delegates to the canonical - ``_common.compute_placeholder_prompt_len(mode="full")`` so the streaming - fixup path (``construct_next_stage_streaming_input_prompt``) and the - forward placeholder builder (``build_forward_placeholder``) share one length - helper. Behaviour is golden-locked (``test_common_helpers_golden.py`` + Delegates to the canonical ``_common.compute_placeholder_prompt_len(mode="full")`` + so the streaming fixup path (``construct_next_stage_streaming_input_prompt``) + and the forward placeholder builder (``build_forward_placeholder``) share one + length helper. Behaviour is golden-locked (``test_common_helpers_golden.py`` reproduces the 15-token chat-template result for the canonical prompt). Args: diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index de19a5d81b1..07c5f6fc547 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -370,7 +370,7 @@ def __init__( self.async_chunk = bool(async_chunk) self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools - # RFC #4872 M0: wire the dead-processor hint at startup (warn-only) so a + # Wire the dead-processor hint at startup (warn-only) so a # configured-but-never-invoked custom_process_input_func is surfaced. self._warn_dead_input_processors() self.log_stats = log_stats @@ -2072,7 +2072,7 @@ def _warn_dead_input_processors(self) -> None: ``_route_output`` gate ``(not async_chunk or not _stage_receives_async_chunks(stage_id))``), log a warning so a configured-but-dead hook is not silently dropped. This is a pure hint - (RFC #4872 M0 warn-first): the orchestrator never enforces it at runtime. + (warn-first): the orchestrator never enforces it at runtime. Defensive by design — fake pools / minimal clients in tests or future connectors may not expose ``stage_client``, so the loop never raises. """ @@ -2645,8 +2645,8 @@ async def _forward_to_next_stage_unguarded( if next_client is not None and next_client.custom_process_input_func is not None: _t_ar2d = _time.perf_counter() _fn = next_client.custom_process_input_func - # RFC #4872 Phase 2 (P1): dispatch through the shared contract - # layer, replacing the ad-hoc `sampling_params` name probe. + # Dispatch through the shared contract layer, replacing the + # ad-hoc `sampling_params` name probe. from vllm_omni.model_executor.stage_input_processors._dispatch import ( OrchestratorInputContext, invoke_orchestrator_processor, @@ -3056,10 +3056,10 @@ async def _prewarm_async_chunk_stages( operation="async-chunk prewarm", ) else: - # RFC #4872 P8b: build the prewarm placeholder through the - # registered placeholder builder's ``build_prewarm_placeholder`` - # (resolved from ``sync_process_input_func``), falling back to - # the legacy inline estimate when no builder is available. + # Build the prewarm placeholder through the registered + # placeholder builder's ``build_prewarm_placeholder`` (resolved + # from ``sync_process_input_func``), falling back to the legacy + # inline estimate when no builder is available. base_input = self._build_prewarm_placeholder_input( next_stage_id, request_id, @@ -3199,8 +3199,8 @@ def _diffusion_submit_kwargs( def _get_prewarm_placeholder_builder(self, stage_id: int) -> Any: """Resolve the registered ``build_prewarm_placeholder`` for a stage. - RFC #4872 P8b: the async-chunk prewarm reuses ``sync_process_input_func`` - (the ``*_token_only`` path). Resolution order: + The async-chunk prewarm reuses ``sync_process_input_func`` (the + ``*_token_only`` path). Resolution order: 1. ``client.sync_process_input_func`` -> ``resolve_processor(...).fn`` and read ``build_prewarm_placeholder`` off the resolved function @@ -3322,8 +3322,8 @@ def _build_prewarm_placeholder_input( exc, ) - # Legacy inline fallback (RFC #4872 P8b: warn-and-keep-old behaviour so - # the prewarm never regresses on models without a dual-entry builder). + # Legacy inline fallback (warn-and-keep-old behaviour so the prewarm + # never regresses on models without a dual-entry builder). from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length try: diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index ca8ac786b8f..e0a112dcb96 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -153,9 +153,9 @@ def __init__( self.prompt_expand_func = metadata.prompt_expand_func self.custom_process_input_func = metadata.custom_process_input_func self.async_chunk_prewarm_payload_func = getattr(metadata, "async_chunk_prewarm_payload_func", None) - # RFC #4872 P8b: raw ``sync_process_input_func`` path so the - # orchestrator can resolve the async-chunk prewarm placeholder - # builder (``build_prewarm_placeholder``) at prewarm time. + # Raw ``sync_process_input_func`` path so the orchestrator can + # resolve the async-chunk prewarm placeholder builder + # (``build_prewarm_placeholder``) at prewarm time. self.sync_process_input_func = metadata.sync_process_input_func self.engine_outputs: Any = None @@ -467,9 +467,8 @@ def process_engine_inputs( and the original prompt. """ if self.custom_process_input_func is not None: - if self.custom_process_input_func is not None: - # RFC #4872 Phase 2 (P1): dispatch through the shared contract layer, - # replacing the ad-hoc arity probe (`len(signature.parameters) >= 4`). + # Dispatch through the shared contract layer, replacing the ad-hoc + # arity probe (`len(signature.parameters) >= 4`) from vllm_omni.model_executor.stage_input_processors._dispatch import ( OrchestratorInputContext, invoke_orchestrator_processor, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 439e642c5f9..01a3643e236 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -572,9 +572,9 @@ class StageMetadata: custom_process_input_func: Callable | None model_stage: str | None runtime_cfg: Any - # RFC #4872 P8b: raw ``sync_process_input_func`` path (the ``*_token_only`` - # placeholder builder) so the orchestrator can resolve - # ``build_prewarm_placeholder`` at async-chunk prewarm time. + # Raw ``sync_process_input_func`` path (the ``*_token_only`` placeholder + # builder) so the orchestrator can resolve ``build_prewarm_placeholder`` + # at async-chunk prewarm time. sync_process_input_func: str | None = None prompt_expand_func: Callable | None = None cfg_kv_collect_func: Callable | None = None @@ -643,11 +643,10 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: custom_process_input_func: Callable | None = None _cpif_path = _get_attr_or_item(stage_config, "custom_process_input_func") if _cpif_path: - # RFC #4872 Phase 3 (Part A): resolve through the stage-input processor - # registry. expected_kind is deliberately None here — orchestrator-side - # processors carry no suffix contract, so we rely on suffix inference - # only (compatible with existing configs). A bad path must **fail fast** - # (RFC P2 / adding-omni text): let ProcessorValidationError / ImportError / + # Resolve through the stage-input processor registry. expected_kind is + # None here — orchestrator-side processors carry no suffix contract, so + # we rely on suffix inference only (compatible with existing configs). + # A bad path must **fail fast**: let ProcessorValidationError / ImportError / # AttributeError propagate instead of silently dropping the configured # hook and falling back to ``_default_process_engine_inputs``. (The # worker-side mixin candidate chain, which *probes* for a ``*_full_payload`` @@ -754,7 +753,7 @@ def extract_stage_metadata_from_omni_stage_config( sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams sampling_params = sampling_params_cls(**(stage_config.model_config.default_sampling_params or {})) custom_process_input_func = _resolve_omni_metadata_hook(stage_config.custom_process_input_func) - # RFC #4872 P8b: carry the raw sync path for async-chunk prewarm resolution. + # Carry the raw sync path for async-chunk prewarm resolution. sync_process_input_func: str | None = None _spif_path = getattr(stage_config.stage_pipeline_config, "sync_process_input_func", None) if _spif_path: diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index cebc9cf1fb1..074eee4e38a 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -17,8 +17,8 @@ pipeline_cfg_resolver, ) -# RFC #4872 P6: talker stop token single-sourced in the shared constants -# module (also used by qwen3_tts / aura_omni pipelines). +# Talker stop token single-sourced in the shared constants module (also used +# by the qwen3_tts / aura_omni pipelines). from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_omni" diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 8f246ebdb1b..b0efd35352a 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -65,10 +65,10 @@ from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_decode_windows -# RFC #4872 P6: codec token ids are single-sourced in the stage input -# processor (``_QWEN3_CODEC_*``); the model keeps the ``TALKER_CODEC_*`` -# aliases for backwards compatibility. At runtime the talker actually reads -# them from ``config.talker_config.codec_{pad,bos,eos}``. +# Codec token ids are single-sourced in the stage input processor +# (``_QWEN3_CODEC_*``); the model keeps the ``TALKER_CODEC_*`` aliases for +# backwards compatibility. At runtime the talker actually reads them from +# ``config.talker_config.codec_{pad,bos,eos}``. from vllm_omni.model_executor.stage_input_processors.qwen3_omni import ( _QWEN3_CODEC_BOS_TOKEN_ID, _QWEN3_CODEC_EOS_TOKEN_ID, @@ -106,7 +106,7 @@ def resolve_default_speaker(speaker_token_ids: dict[str, int]) -> str: TTS_BOS_SINGLE_TOKEN_ID = 151674 # # Talker codec tokens (talker vocabulary, used for RVQ code generation). -# RFC #4872 P6: single-sourced from ``stage_input_processors.qwen3_omni`` +# Single-sourced from ``stage_input_processors.qwen3_omni`` # (``_QWEN3_CODEC_*``); at runtime the talker reads these ids from # ``config.talker_config.codec_{pad,bos,eos}`` instead. TALKER_CODEC_PAD_TOKEN_ID = _QWEN3_CODEC_PAD_TOKEN_ID # Padding token diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 8e494b83c4e..a2d258796d5 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -81,10 +81,10 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config - # RFC #4872 P6: the codec token-id consistency guard is deliberately not - # run at module import (it pulls in transformers / the qwen3-tts config). - # Invoke it explicitly at model-load startup so the processor constants - # stay in sync with the HF configs this model actually runs with. + # The codec token-id consistency guard is deliberately not run at module + # import (it pulls in transformers / the qwen3-tts config). Invoke it + # explicitly at model-load startup so the processor constants stay in + # sync with the HF configs this model actually runs with. from vllm_omni.model_executor.stage_input_processors.qwen3_omni import ( _assert_codec_token_ids_consistent, ) diff --git a/vllm_omni/model_executor/stage_input_processors/__init__.py b/vllm_omni/model_executor/stage_input_processors/__init__.py index 3fc95647204..467f5919ee5 100644 --- a/vllm_omni/model_executor/stage_input_processors/__init__.py +++ b/vllm_omni/model_executor/stage_input_processors/__init__.py @@ -3,7 +3,8 @@ """Stage input processors: per-model builders for inter-stage transitions. -Module-naming convention (RFC #4872): +Module-naming convention (RFC #4872, +https://github.com/vllm-project/vllm-omni/issues/4872): - ``*_full_payload`` : producer-side (worker) connector payload builder (``FullPayloadProducer``; kwarg ``pooling_output`` load-bearing). - ``*_async_chunk`` : producer-side (worker) streaming chunk builder diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py index 859618cc06b..bf94c359a11 100644 --- a/vllm_omni/model_executor/stage_input_processors/_common.py +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -3,19 +3,18 @@ """Canonical, shared helpers for ``stage_input_processors``. -RFC #4872 Proposal 3 / P8a. These are consolidated implementations of -helpers that were previously duplicated across model modules with subtle, -behavioral differences (see ``tests/model_executor/stage_input_processors/ -test_common_helpers_golden.py``, which locks the observed per-module -behaviour). - -**Consolidation rule:** where two legacy variants disagreed, we do NOT -silently pick one. The default implementation follows the most-complete -semantics; divergent legacy behaviour is preserved through explicit named -variants (e.g. ``ensure_list_unchanged``) or parameters (e.g. -``filter_real_code_frames(..., layout=...)``). A module may only switch to -``_common`` after its legacy behaviour is golden-locked and the matching -variant is used. +Related to RFC #4872 (https://github.com/vllm-project/vllm-omni/issues/4872): +these are consolidated implementations of helpers that were previously +duplicated across model modules with subtle behavioral differences (see +``tests/model_executor/stage_input_processors/test_common_helpers_golden.py``, +which locks the observed per-module behaviour). + +**Consolidation rule:** where legacy variants disagreed, the default +implementation follows the most-complete semantics; divergent legacy behaviour +is preserved through explicit named variants (e.g. ``ensure_list_unchanged``) +or parameters (e.g. ``filter_real_code_frames(..., layout=...)``). A module +may switch to ``_common`` only after its legacy behaviour is golden-locked and +the matching variant is used. """ from __future__ import annotations @@ -137,11 +136,11 @@ def ensure_list_preserve_none(x: Any) -> list[Any]: def ensure_list_wrap_only(x: Any) -> list[Any]: """Wrap-only ``ensure_list``: list passthrough, non-list -> ``[x]``, ``None`` -> ``[]``. - Preserves the legacy ``diffusion.output_formatter._ensure_list`` semantics - (RFC #4872 P3): a non-list primary payload is wrapped as ``[x]``. Unlike - the canonical :func:`ensure_list`, this deliberately performs **no** tensor - flattening, no dict-key iteration, and no row-wise iteration of an iterable - (e.g. a PIL ``Image``) — a non-``list`` value is wrapped whole, verbatim. + Preserves the legacy ``diffusion.output_formatter._ensure_list`` semantics: + a non-list primary payload is wrapped as ``[x]``. Unlike the canonical + :func:`ensure_list`, it performs **no** tensor flattening, no dict-key + iteration, and no row-wise iteration of an iterable (e.g. a PIL ``Image``) — + a non-``list`` value is wrapped whole, verbatim. """ if isinstance(x, list): return x @@ -351,11 +350,11 @@ def extract_last_codec_frame( # =========================================================================== -# P8a: placeholder prompt length + packing +# Placeholder prompt length + packing # =========================================================================== # Qwen chat-template sentinel ids (single source within this module; the -# codec token ids themselves are centralized separately in P6). +# codec token ids themselves are centralized separately in ``_constants``). _IM_START_TOKEN_ID = 151644 _SYSTEM_TOKEN_ID = 8948 _USER_TOKEN_ID = 872 @@ -442,8 +441,8 @@ def pack_placeholder_prompt( """Pack a KV-slot placeholder prompt: ``[0] * prompt_len`` + voice meta. Shared by both the sync forward path (``*_token_only``) and the - async-chunk prewarm path (Proposal 8b). Bulk conditioning arrives - separately via the connector, so only the length matters here. + async-chunk prewarm path. Bulk conditioning arrives separately via the + connector, so only the length matters here. """ return OmniTokensPrompt( prompt_token_ids=[0] * max(1, int(prompt_len)), diff --git a/vllm_omni/model_executor/stage_input_processors/_constants.py b/vllm_omni/model_executor/stage_input_processors/_constants.py index 88a9e9444ae..d6560781521 100644 --- a/vllm_omni/model_executor/stage_input_processors/_constants.py +++ b/vllm_omni/model_executor/stage_input_processors/_constants.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""RFC #4872 P6: canonical cross-model constant definitions. +"""Canonical cross-model constant definitions (RFC #4872). Special-token ids that are shared across model pipelines / processors / model implementations live here so no module hard-codes a magic number that @@ -18,5 +18,5 @@ # Qwen3-family talker codec stop (EOS) token id. Used as the AR talker # stage's ``stop_token_ids`` across the qwen3_omni / qwen3_tts / aura_omni -# pipelines (three pipeline files shared the literal ``2150`` before P6). +# pipelines (three pipeline files previously shared the literal ``2150``). QWEN3_CODEC_EOS_TOKEN_ID = 2150 diff --git a/vllm_omni/model_executor/stage_input_processors/_dispatch.py b/vllm_omni/model_executor/stage_input_processors/_dispatch.py index 0fae461abe2..ed670383411 100644 --- a/vllm_omni/model_executor/stage_input_processors/_dispatch.py +++ b/vllm_omni/model_executor/stage_input_processors/_dispatch.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""RFC #4872 Phase 2 (P1): orchestrator input-dispatch contract layer. +"""Orchestrator input-dispatch contract layer (RFC #4872). This module defines the **consumer-side (orchestrator / stage-client) dispatch contract** for ``stage_input_processors``. It replaces the two inconsistent @@ -32,9 +32,9 @@ is_finished=False``) run inside workers and never receive an ``OrchestratorInputContext``; their ``pooling_output`` / ``multimodal_output`` keyword names are load-bearing parts of the producer contract and are kept -as-is (RFC P1/P2 keep the producer kwargs contract unchanged). +as-is (the producer kwargs contract stays unchanged). -Legacy positional shapes normalized here (RFC P1 migration matrix): +Legacy positional shapes normalized here: * C0 3-arg: ``(source_outputs, prompt, requires_multimodal_data)``. * C2 placeholder: ``(source_outputs, prompt, requires_multimodal_data, diff --git a/vllm_omni/model_executor/stage_input_processors/_registry.py b/vllm_omni/model_executor/stage_input_processors/_registry.py index 3eff813e077..a687d6a2203 100644 --- a/vllm_omni/model_executor/stage_input_processors/_registry.py +++ b/vllm_omni/model_executor/stage_input_processors/_registry.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""RFC #4872 Phase 3 (Part A): stage-input processor registry. +"""Stage-input processor registry (RFC #4872). This module is the **structural validation** layer behind the three ``importlib`` -resolution points that RFC P2 calls out as independent symbol lookups: +resolution points that are validated as independent symbol lookups: * ``engine/stage_init_utils.extract_legacy_stage_metadata`` — consumer-side (orchestrator) input processors (``custom_process_input_func`` / @@ -28,8 +28,8 @@ and **must** accept ``is_finished`` — the scheduler always passes it on the async-chunk path (see ``chunk_transfer_adapter._send_single_request``). -Kind inference is name-driven (suffix-based), following the RFC P2 naming -convention documented in :mod:`vllm_omni.model_executor.stage_input_processors`: +Kind inference is name-driven (suffix-based), following the naming convention +documented in :mod:`vllm_omni.model_executor.stage_input_processors`: * ``_token_only`` -> ``placeholder_prompt_builder`` * ``_full_payload`` / ``_batch`` -> ``producer_full_payload`` @@ -235,7 +235,7 @@ def register_processor(path: str, kind: ProcessorKind) -> None: """Register a manual kind override for a processor *path*. This is an escape hatch for models whose function names do not follow the - RFC P2 suffix convention. ``infer_kind`` consults these overrides before + suffix convention. ``infer_kind`` consults these overrides before applying the suffix rules. Overrides are validated eagerly so typos fail fast. """ @@ -250,7 +250,7 @@ def infer_kind(fn: Any, *, path: str) -> ProcessorKind: """Infer the processor kind from the *suffix* of ``path``. ``fn`` is accepted for API symmetry with :func:`validate_processor` but the - inference is purely name-driven (RFC P2 naming convention); the callable + inference is purely name-driven (suffix-based convention); the callable object itself is not inspected here. """ del fn # inference is name-based; kept for a symmetric API diff --git a/vllm_omni/model_executor/stage_input_processors/audex.py b/vllm_omni/model_executor/stage_input_processors/audex.py index 3da58cf0209..0b211a46906 100644 --- a/vllm_omni/model_executor/stage_input_processors/audex.py +++ b/vllm_omni/model_executor/stage_input_processors/audex.py @@ -45,7 +45,7 @@ def _ensure_list(value: Any) -> list[int]: - """RFC #4872 P3: delegate to the canonical ``_common.ensure_list_flatten``.""" + """Delegate to the canonical ``_common.ensure_list_flatten``.""" return _common.ensure_list_flatten(value) diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index f34f6da86f1..a3b206e8a65 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py @@ -38,12 +38,12 @@ def _build_prompt_embed_struct(prompt_payload: dict[str, Any]) -> EmbeddingsStru def _ensure_list(x: Any) -> list[Any]: - """Delegate to the canonical shared helper (RFC #4872 Step 1c).""" + """Delegate to the canonical shared helper.""" return _common.ensure_list(x) def _to_token_id_list(value: Any) -> list[int]: - """Recursive token-id flattening (cosyvoice3 golden semantics, Step 1c).""" + """Recursive token-id flattening (cosyvoice3 golden semantics).""" return _common.to_token_id_list(value, recursive=True) @@ -63,7 +63,7 @@ def _prompt_speech_token_ids(multi_modal_data: dict[str, Any]) -> list[int]: def _to_cpu_tensor(x: Any) -> torch.Tensor | None: - """Delegate to the canonical shared helper (RFC #4872 Step 1c).""" + """Delegate to the canonical shared helper.""" return _common.to_cpu_tensor(x) diff --git a/vllm_omni/model_executor/stage_input_processors/glm_tts.py b/vllm_omni/model_executor/stage_input_processors/glm_tts.py index 77d73ff9854..2a94fd94b20 100644 --- a/vllm_omni/model_executor/stage_input_processors/glm_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/glm_tts.py @@ -173,7 +173,7 @@ def _extract_last_speech_token(pooling_output: dict[str, Any]) -> int | None: def _to_cpu_tensor(value: Any) -> torch.Tensor | None: """Convert value to CPU tensor if possible. - RFC #4872 P3: delegate to the canonical ``_common.to_cpu_tensor``. + Delegates to the canonical ``_common.to_cpu_tensor``. """ return _common.to_cpu_tensor(value) diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py index c1b00947f36..3992236155e 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py @@ -52,8 +52,8 @@ def _filter_real_code_frames(audio_codes: torch.Tensor) -> torch.Tensor: """Keep only frames whose codes are entirely in [0, 1023]. - RFC #4872 P3: delegate to the canonical ``_common.filter_real_code_frames`` - with the higgs-v2 ``[num_frames, num_codebooks]`` layout. + Delegates to the canonical ``_common.filter_real_code_frames`` with the + higgs-v2 ``[num_frames, num_codebooks]`` layout. """ return _common.filter_real_code_frames(audio_codes, num_real_codes=_NUM_REAL_CODES, layout="frames_first") @@ -61,8 +61,8 @@ def _filter_real_code_frames(audio_codes: torch.Tensor) -> torch.Tensor: def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: """Reverse the delay-pattern shift (higgs-v2 lenient variant). - RFC #4872 P3: delegate to the canonical ``_common.revert_delay_pattern`` - with ``allow_short=True`` (T < Q returns the input unchanged). + Delegates to the canonical ``_common.revert_delay_pattern`` with + ``allow_short=True`` (T < Q returns the input unchanged). """ return _common.revert_delay_pattern(audio_codes_qt, allow_short=True) diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py index 8b781499aa8..3e2ce1e19a8 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py @@ -63,8 +63,8 @@ def _empty_code2wav_prompt() -> Any: def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: """Reverse the MusicGen-style delay pattern (higgs-v3 strict variant). - RFC #4872 P3: delegate to the canonical ``_common.revert_delay_pattern`` - with the strict codebook-count and short-input checks. + Delegates to the canonical ``_common.revert_delay_pattern`` with the + strict codebook-count and short-input checks. """ return _common.revert_delay_pattern( audio_codes_qt, @@ -76,8 +76,8 @@ def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: def _filter_real_code_frames(audio_codes_qt: torch.Tensor) -> torch.Tensor: """Keep only frames where ALL codebook values are in [0, 1023]. - RFC #4872 P3: delegate to the canonical ``_common.filter_real_code_frames`` - with the higgs-v3 ``[num_codebooks, num_frames]`` layout. + Delegates to the canonical ``_common.filter_real_code_frames`` with the + higgs-v3 ``[num_codebooks, num_frames]`` layout. """ return _common.filter_real_code_frames(audio_codes_qt, num_real_codes=_NUM_REAL_CODES, layout="codebooks_first") diff --git a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py index 937e4d39a0c..311b330f617 100644 --- a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py @@ -207,7 +207,7 @@ def _resolve_image_patch_token_id(stage: Any) -> int: def _ensure_list(x) -> list[int]: """Convert ConstantList / tensor-like to plain list. - RFC #4872 P3: delegate to the canonical ``_common.ensure_list_strict``. + Delegates to the canonical ``_common.ensure_list_strict``. """ return _common.ensure_list_strict(x) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 63f1a25d20c..d26aee2140d 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -50,7 +50,7 @@ def _assert_codec_token_ids_consistent() -> None: - """Fail-fast guard (RFC #4872 P6): codec token ids must not drift. + """Fail-fast guard: codec token ids must not drift. These ids are model-specific special tokens. The HF config classes are the single source of truth at runtime -- the talker reads @@ -83,10 +83,11 @@ def _assert_codec_token_ids_consistent() -> None: ) -# The P6 fail-fast guard is intentionally **not** run at import time (it pulls in -# transformers / the in-repo qwen3-tts config and must not break light imports). -# It is invoked explicitly at model-load startup (qwen3_omni_moe_talker) so the -# processor constants stay in sync with the HF configs the model actually runs with. +# The codec-token consistency guard is intentionally **not** run at import time +# (it pulls in transformers / the in-repo qwen3-tts config and must not break +# light imports). It is invoked explicitly at model-load startup +# (qwen3_omni_moe_talker) so the processor constants stay in sync with the HF +# configs the model actually runs with. def _layer_tensor(layers: dict[Any, Any], key: str) -> torch.Tensor | None: """Fetch layer tensor with tolerant key lookup (str/int/semantic).""" if not isinstance(layers, dict): @@ -162,7 +163,7 @@ def _sampled_token_embed(thinker_embed: Mapping[str, Any], request: Any) -> torc def _ensure_list(x): """Convert ConstantList / tensor-like to Python list. - RFC #4872 P3: delegate to the canonical ``_common.ensure_list_unchanged``. + Delegates to the canonical ``_common.ensure_list_unchanged``. """ return _common.ensure_list_unchanged(x) @@ -640,7 +641,7 @@ def build_forward_placeholder( source_outputs: list[Any], ctx: OrchestratorInputContext, ) -> list[OmniTokensPrompt]: - """Orchestrator-side sync forward placeholder builder (RFC #4872 P8b). + """Orchestrator-side sync forward placeholder builder. Semantics match the legacy ``thinker2talker_token_only`` forward path: one placeholder ``OmniTokensPrompt`` per upstream output, sized so the @@ -648,7 +649,7 @@ def build_forward_placeholder( the original prompt so they survive when Stage-0 request metadata is unavailable to the connector payload. - RFC #4872 P8b consolidation: the length computation is delegated to + The length computation is delegated to ``_common.compute_placeholder_prompt_len(mode="full")`` (the Qwen chat-template scan, golden-locked at 15 for the golden prompt) and the placeholder packing to ``_common.pack_placeholder_prompt``. @@ -700,7 +701,7 @@ def build_prewarm_placeholder( ctx: OrchestratorInputContext, downstream_stage_id: int, ) -> list[OmniTokensPrompt]: - """Best-effort async-chunk prewarm placeholder builder (RFC #4872 P8b). + """Best-effort async-chunk prewarm placeholder builder. async-chunk mode has **no upstream ``source_outputs``** yet at prewarm time, so this is a best-effort estimate: the placeholder length comes from @@ -716,9 +717,9 @@ def build_prewarm_placeholder( ``ctx`` / ``downstream_stage_id`` are accepted for contract uniformity; voice metadata is intentionally **not** forwarded here (matching the pre-existing prewarm behaviour — no new behaviour is introduced). Per-edge - refinement (``downstream_stage_id``) is RFC #4872 P8c and out of scope. + refinement (``downstream_stage_id``) is out of scope. """ - del ctx, downstream_stage_id # P8b uses only the stage-0 input estimate. + del ctx, downstream_stage_id # Only the stage-0 input estimate is used. if hasattr(stage0_prompt, "prompt_token_ids"): stage0_prompt = stage0_prompt.prompt_token_ids prompt_len = max( @@ -744,9 +745,9 @@ def thinker2talker_token_only( placeholder ``prompt_token_ids`` of the correct length so the scheduler can reserve KV-cache slots. It does **not** forward bulk tensors. - RFC #4872 P8b: thin delegation to :func:`build_forward_placeholder` — the - dual-entry design keeps the sync forward path and the async-chunk prewarm - path on the same shared ``_common`` length / packing helpers. ``prompt`` / + Thin delegation to :func:`build_forward_placeholder` — the dual-entry design + keeps the sync forward path and the async-chunk prewarm path on the same + shared ``_common`` length / packing helpers. ``prompt`` / ``requires_multimodal_data`` / ``streaming_context`` are kept for call-site signature compatibility; they are folded into an ``OrchestratorInputContext``. """ @@ -758,10 +759,10 @@ def thinker2talker_token_only( return build_forward_placeholder(source_outputs, ctx) -# RFC #4872 P8b: the orchestrator resolves ``sync_process_input_func`` (the -# ``*_token_only`` path) via ``resolve_processor(...).fn`` and reads -# ``build_prewarm_placeholder`` off the resolved function object. Expose both -# dual-entry builders here so that lookup works without a separate module scan. +# The orchestrator resolves ``sync_process_input_func`` (the ``*_token_only`` +# path) via ``resolve_processor(...).fn`` and reads ``build_prewarm_placeholder`` +# off the resolved function object. Expose both dual-entry builders here so +# that lookup works without a separate module scan. thinker2talker_token_only.build_forward_placeholder = build_forward_placeholder thinker2talker_token_only.build_prewarm_placeholder = build_prewarm_placeholder diff --git a/vllm_omni/model_executor/stage_input_processors/step_audio2.py b/vllm_omni/model_executor/stage_input_processors/step_audio2.py index f196e3a9af4..ff52b3522ff 100644 --- a/vllm_omni/model_executor/stage_input_processors/step_audio2.py +++ b/vllm_omni/model_executor/stage_input_processors/step_audio2.py @@ -22,8 +22,8 @@ def _ensure_list(x): """Convert ConstantList / tensor / iterable to Python list. - Preserves ``None`` as ``[None]`` (step_audio2 semantics, RFC #4872 - Step 1c): delegates to ``_common.ensure_list_preserve_none``. + Preserves ``None`` as ``[None]`` (step_audio2 semantics): delegates to + ``_common.ensure_list_preserve_none``. """ return _common.ensure_list_preserve_none(x) diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 4c89972379c..6c842d8a775 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -2139,10 +2139,10 @@ def _load_custom_func(model_config: Any) -> tuple[str | None, Any | None]: continue tried.add(func_path) try: - # RFC #4872 Phase 3 (Part A): resolve through the registry with - # the full-payload producer contract. is_finished is optional - # here (missing only warns); a structurally incompatible - # candidate is skipped so the fallback chain keeps working. + # Resolve through the registry with the full-payload producer + # contract. is_finished is optional here (missing only warns); + # a structurally incompatible candidate is skipped so the + # fallback chain keeps working. spec = resolve_processor(func_path, expected_kind="producer_full_payload") return spec.path, spec.fn except ProcessorValidationError as exc: From 6096bb57dc554743eac5e03c61c4170dee2ee9b2 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Mon, 31 Aug 2026 20:42:02 +0800 Subject: [PATCH 18/41] fix(stage_input_processors): silence mypy attr-defined on *_token_only builder attachment (RFC #4872) Signed-off-by: RyanYun09 --- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index d26aee2140d..4946cb1f460 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -763,8 +763,8 @@ def thinker2talker_token_only( # path) via ``resolve_processor(...).fn`` and reads ``build_prewarm_placeholder`` # off the resolved function object. Expose both dual-entry builders here so # that lookup works without a separate module scan. -thinker2talker_token_only.build_forward_placeholder = build_forward_placeholder -thinker2talker_token_only.build_prewarm_placeholder = build_prewarm_placeholder +thinker2talker_token_only.build_forward_placeholder = build_forward_placeholder # type: ignore[attr-defined] +thinker2talker_token_only.build_prewarm_placeholder = build_prewarm_placeholder # type: ignore[attr-defined] # ========================= From 22d78168c87f71211999c97af40da9b09f13a172 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Wed, 2 Sep 2026 12:50:13 +0800 Subject: [PATCH 19/41] fix(stage_input_processors): preserve keyword-only sampling_params and four-positional dispatch (RFC #4872 P1) Signed-off-by: RyanYun09 --- .../test_dispatch_wrap_once.py | 51 ++++++ .../stage_input_processors/_dispatch.py | 154 ++++++++++++------ 2 files changed, 158 insertions(+), 47 deletions(-) diff --git a/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py index bfa887b66a7..70a55a2e712 100644 --- a/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py +++ b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py @@ -43,6 +43,26 @@ def _c1(source_outputs: list[Any], ctx: OrchestratorInputContext): return [source_outputs, ctx] +def _legacy_c3_keyword_only(source_outputs, prompt=None, requires_multimodal_data=False, *, sampling_params=None): + """Legacy C3 shell whose ``sampling_params`` is keyword-only.""" + return [source_outputs, prompt, requires_multimodal_data, sampling_params] + + +def _legacy_c3_required_keyword_only(source_outputs, prompt=None, requires_multimodal_data=False, *, sampling_params): + """Legacy C3 shell with a REQUIRED keyword-only ``sampling_params``.""" + return [source_outputs, prompt, requires_multimodal_data, sampling_params] + + +def _legacy_c2_other_name(source_outputs, prompt=None, requires_multimodal_data=False, streaming_ctx=None): + """Four-positional legacy shell whose streaming param has a non-standard name.""" + return [source_outputs, prompt, requires_multimodal_data, streaming_ctx] + + +def _legacy_c2_keyword_only(source_outputs, prompt=None, requires_multimodal_data=False, *, streaming_context=None): + """Legacy C2 shell whose ``streaming_context`` is keyword-only.""" + return [source_outputs, prompt, requires_multimodal_data, streaming_context] + + def _count_deprecation_warnings(trigger) -> int: """Run *trigger* under ``simplefilter("always")`` and count DeprecationWarnings.""" import warnings @@ -87,3 +107,34 @@ def test_c1_processor_passes_through_without_warning(): assert invoke_orchestrator_processor(_c1, ["out"], ctx) == [["out"], ctx] # No DeprecationWarning for C1 processors at all. assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_c1)) == 0 + + +# --------------------------------------------------------------------------- +# P1 deep-dive: keyword-only semantic kwargs + four-positional fallback. +# --------------------------------------------------------------------------- + + +def test_keyword_only_sampling_params_receives_context_value(): + ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, sampling_params="sp") + # ``sampling_params`` is keyword-only: the C3 detection must read the FULL + # signature so the diffusion-stage sampling params still reach the shell. + assert invoke_orchestrator_processor(_legacy_c3_keyword_only, ["out"], ctx) == [["out"], "p", True, "sp"] + + +def test_required_keyword_only_sampling_params_receives_context_value(): + ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, sampling_params="sp") + # A REQUIRED keyword-only ``sampling_params`` used to raise at runtime + # (silently dropped by the old C0 fallback); it must now receive the value. + assert invoke_orchestrator_processor(_legacy_c3_required_keyword_only, ["out"], ctx) == [["out"], "p", True, "sp"] + + +def test_four_positional_different_streaming_name_not_regressed(): + ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, streaming_context="s") + # A four-positional legacy shell whose 4th param is named differently from + # ``streaming_context`` must still receive it as the 4th positional arg. + assert invoke_orchestrator_processor(_legacy_c2_other_name, ["out"], ctx) == [["out"], "p", True, "s"] + + +def test_keyword_only_streaming_context_receives_value(): + ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, streaming_context="s") + assert invoke_orchestrator_processor(_legacy_c2_keyword_only, ["out"], ctx) == [["out"], "p", True, "s"] diff --git a/vllm_omni/model_executor/stage_input_processors/_dispatch.py b/vllm_omni/model_executor/stage_input_processors/_dispatch.py index ed670383411..8f484930bb5 100644 --- a/vllm_omni/model_executor/stage_input_processors/_dispatch.py +++ b/vllm_omni/model_executor/stage_input_processors/_dispatch.py @@ -34,15 +34,20 @@ keyword names are load-bearing parts of the producer contract and are kept as-is (the producer kwargs contract stays unchanged). -Legacy positional shapes normalized here: +Legacy positional shapes normalized here (detection reads the **full** +signature, keyword-only parameters included; the final invocation is validated +with ``inspect.Signature.bind``): * C0 3-arg: ``(source_outputs, prompt, requires_multimodal_data)``. * C2 placeholder: ``(source_outputs, prompt, requires_multimodal_data, streaming_context=None)`` (the ``_streaming_context`` variant used by ``forced_aligner.code2wav2aligner`` / ``minicpmo_4_5_omni.llm2tts`` is also - recognized). + recognized; a keyword-only ``*, streaming_context=None`` shell receives the + value as a keyword, and any four-positional callable falls back to this + shape regardless of the 4th parameter's name). * C3 diffusion: ``(source_outputs, prompt, requires_multimodal_data, - sampling_params=None)``. + sampling_params=None)`` — a keyword-only ``*, sampling_params=None`` shell + also receives ``ctx.sampling_params``. * C4 legacy multi-source: ``(stage_list, engine_input_source, ...)`` — only ``moss_tts.talker2codec`` remains today; normalized by ``_adapt_moss_processor``, which binds ``stage_list=source_outputs`` and @@ -185,22 +190,64 @@ def _accepts_orchestrator_input_context(fn: Any) -> bool: return False -def _positional_parameter_names(fn: Any) -> list[str] | None: - """Names of positional (or positional-or-keyword) parameters of ``fn``. - - Returns ``None`` when the signature cannot be inspected (e.g. builtins). - """ +def _legacy_signature(fn: Any) -> inspect.Signature | None: + """Best-effort ``inspect.signature`` probe for a callable.""" try: - signature = inspect.signature(fn) + return inspect.signature(fn) except (TypeError, ValueError): return None - names: list[str] = [] - for param in signature.parameters.values(): - if param.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD): - names.append(param.name) - elif param.kind == inspect.Parameter.VAR_POSITIONAL: - names.append("*args") - return names + + +def _positional_names(signature: inspect.Signature) -> list[str]: + """Names of positional (or positional-or-keyword) parameters.""" + return [ + name + for name, param in signature.parameters.items() + if param.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) + ] + + +def _legacy_shape(signature: inspect.Signature | None) -> str: + """Classify a legacy positional processor shape. + + The probe reads the **full** signature — keyword-only parameters included — + so a ``def f(..., *, sampling_params=None)`` shell is still recognized as a + C3 diffusion builder (and receives ``ctx.sampling_params``), and any + four-positional callable falls back to the legacy C2 streaming shape + regardless of the 4th parameter's name (the legacy arity probe treated every + ``>= 4``-parameter callable that way). + """ + if signature is None: + return "c0" + params = signature.parameters + pos_names = _positional_names(signature) + if len(pos_names) >= 2 and pos_names[0] == "stage_list" and pos_names[1] == "engine_input_source": + return "moss" + if "sampling_params" in pos_names: + return "c3" + if len(pos_names) >= 4: + return "c2pos" + if "sampling_params" in params: + return "c3" + if "streaming_context" in params or "_streaming_context" in params: + return "c2kw" + return "c0" + + +def _bind_invoke(fn: Any, signature: inspect.Signature | None, *args: Any, **kwargs: Any) -> Any: + """Invoke *fn* with a legacy call shape bound via ``inspect.Signature.bind``. + + Binding at the final invocation guarantees a legacy adapter never silently + drops context fields (e.g. forwarding ``sampling_params`` to a + positional-only parameter raises instead of being ignored) and never passes + an argument the processor cannot accept. + """ + if signature is None: + # Uninspectable callable (e.g. builtin): pass the shape through. + return fn(*args, **kwargs) + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + return fn(*bound.args, **bound.kwargs) def _warn_legacy_contract(fn: Any) -> None: @@ -220,7 +267,7 @@ def _warn_legacy_contract(fn: Any) -> None: # --------------------------------------------------------------------------- -def _adapt_moss_processor(fn: Any) -> Any: +def _adapt_moss_processor(fn: Any, signature: inspect.Signature | None) -> Any: """C4 legacy multi-source adapter (e.g. ``moss_tts.talker2codec``). The legacy multi-source sync processors take @@ -233,7 +280,7 @@ def _adapt_moss_processor(fn: Any) -> Any: """ def _adapted(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: - return fn(source_outputs, ctx.prompt) + return _bind_invoke(fn, signature, source_outputs, ctx.prompt) _warn_legacy_contract(fn) return _adapted @@ -280,39 +327,52 @@ def wrap_orchestrator_processor(fn: Any) -> PlaceholderPromptBuilder | Diffusion if _accepts_orchestrator_input_context(fn): wrapped: Any = fn else: - names = _positional_parameter_names(fn) - - # Uninspectable callable (e.g. builtin): fall back to the most conservative - # 3-arg C0 shape; the runtime previously treated it via the <4 arity path. - if names is None: - names = ["source_outputs", "prompt", "requires_multimodal_data"] - - # C4 legacy multi-source (e.g. moss_tts.talker2codec). - if len(names) >= 2 and names[0] == "stage_list" and names[1] == "engine_input_source": - wrapped = _adapt_moss_processor(fn) - # C3 diffusion: forwarding sampling_params (diffusion-stage params). - elif "sampling_params" in names: + signature = _legacy_signature(fn) + shape = _legacy_shape(signature) - def _c3(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: - return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data, sampling_params=ctx.sampling_params) - - _warn_legacy_contract(fn) - wrapped = _c3 - # C2 placeholder: forwarding the streaming context. - elif "streaming_context" in names or "_streaming_context" in names: - - def _c2(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: - return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data, ctx.streaming_context) - - _warn_legacy_contract(fn) - wrapped = _c2 + if shape == "moss": + wrapped = _adapt_moss_processor(fn, signature) else: - # C0 3-arg legacy. - def _c0(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: - return fn(source_outputs, ctx.prompt, ctx.requires_multimodal_data) + + def _adapted(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: + if shape == "c3": + # C3 diffusion: forwarding sampling_params (diffusion-stage + # params). Read from the full signature so a keyword-only + # ``*, sampling_params=None`` shell receives the value too. + return _bind_invoke( + fn, + signature, + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + sampling_params=ctx.sampling_params, + ) + if shape == "c2pos": + # C2 four-positional legacy fallback (any 4th parameter + # name): the streaming context is the 4th positional arg. + return _bind_invoke( + fn, + signature, + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + ctx.streaming_context, + ) + if shape == "c2kw": + # Keyword-only streaming context shell. + return _bind_invoke( + fn, + signature, + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + streaming_context=ctx.streaming_context, + ) + # C0 3-arg legacy (and the uninspectable / builtin fallback). + return _bind_invoke(fn, signature, source_outputs, ctx.prompt, ctx.requires_multimodal_data) _warn_legacy_contract(fn) - wrapped = _c0 + wrapped = _adapted try: _WRAP_CACHE[fn] = wrapped From d0bf72588b628698bddec3efc1989d9f6f082418 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Wed, 2 Sep 2026 12:50:23 +0800 Subject: [PATCH 20/41] fix(stage_input_processors): restore recursive per-item to_token_id_list normalization (RFC #4872 P2) Signed-off-by: RyanYun09 --- .../test_common_helpers_golden.py | 9 ++++ .../stage_input_processors/_common.py | 41 ++++++++++++++----- 2 files changed, 39 insertions(+), 11 deletions(-) diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py index 000becbed72..0eed45c75bc 100644 --- a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -413,6 +413,10 @@ def test_to_token_id_list_cosyvoice3_recursive_golden(): assert fn(None) == [] assert fn(torch.tensor([[1, 2], [3, 4]])) == [1, 2, 3, 4] assert fn(5) == [5] + # P2 deep-dive parity: per-item normalization flattens a list containing a + # non-scalar tensor and a plain tuple. + assert fn([torch.tensor([[4, 5]])]) == [4, 5] + assert fn((4, 5)) == [4, 5] # =========================================================================== @@ -498,6 +502,11 @@ def test_common_to_token_id_list_matches_golden(): # cosyvoice3 semantics: recursive flatten. assert c.to_token_id_list([[1, 2], [3, 4]], recursive=True) == [1, 2, 3, 4] assert c.to_token_id_list(torch.tensor([[1, 2], [3, 4]]), recursive=True) == [1, 2, 3, 4] + # P2 deep-dive parity: per-item recursion normalizes a list of non-scalar + # tensors and a tuple (previously ValueError / TypeError). + assert c.to_token_id_list([torch.tensor([[4, 5]])], recursive=True) == [4, 5] + assert c.to_token_id_list((4, 5), recursive=True) == [4, 5] + assert c.to_token_id_list((torch.tensor(4), torch.tensor([5, 6])), recursive=True) == [4, 5, 6] def test_common_revert_delay_pattern_matches_golden(): diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py index bf94c359a11..3e1f953e0e8 100644 --- a/vllm_omni/model_executor/stage_input_processors/_common.py +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -179,28 +179,24 @@ def to_token_id_list(value: Any, *, recursive: bool = False) -> list[int]: ``recursive=False`` (default) matches ``dynin_omni._to_token_id_list``: for a 2-D tensor or a nested list only the **first** row / element is - used. ``recursive=True`` matches ``cosyvoice3._to_token_id_list``: all - nesting is flattened. + used. ``recursive=True`` matches ``cosyvoice3._to_token_id_list``: + per-item recursive normalization — nested ``list``/``tuple`` and + non-scalar tensors anywhere in the tree are flattened, so + ``[torch.tensor([[4, 5]])]`` -> ``[4, 5]`` and ``(4, 5)`` -> ``[4, 5]``. """ + if recursive: + return _to_token_id_list_recursive(value) if isinstance(value, torch.Tensor): value = value.detach().to("cpu") if value.ndim == 0: return [int(value.item())] if value.ndim > 1: - if recursive: - value = value.reshape(-1) - else: - value = value[0] + value = value[0] return [int(x) for x in value.tolist()] if isinstance(value, list): if not value: return [] if isinstance(value[0], list): - if recursive: - out: list[int] = [] - for item in value: - out.extend(to_token_id_list(item, recursive=True)) - return out return [int(x) for x in value[0]] return [int(x) for x in value] if value is None: @@ -208,6 +204,29 @@ def to_token_id_list(value: Any, *, recursive: bool = False) -> list[int]: return [int(value)] +def _to_token_id_list_recursive(value: Any) -> list[int]: + """cosyvoice3 recursive token-id flattening (per-item normalization). + + Mirrors the pre-consolidation ``cosyvoice3._to_token_id_list``: ``None`` -> + ``[]``; a tensor is reshaped flat; every item that is a tensor or a + ``list``/``tuple`` is recursed into, scalars are converted with ``int()``. + """ + if value is None: + return [] + if isinstance(value, torch.Tensor): + value = value.detach().to("cpu").reshape(-1).tolist() + out: list[int] = [] + for item in ensure_list(value): + if isinstance(item, torch.Tensor): + out.extend(_to_token_id_list_recursive(item)) + continue + if isinstance(item, (list, tuple)): + out.extend(_to_token_id_list_recursive(item)) + continue + out.append(int(item)) + return out + + # =========================================================================== # revert_delay_pattern (higgs_audio_v2 lenient vs higgs_audio_v3 strict) # =========================================================================== From 8e5e0304ae06222ed7e4b7cc7026db2fded353ee Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Wed, 2 Sep 2026 12:50:33 +0800 Subject: [PATCH 21/41] docs(contributing): wire ar2decoder_full_payload producer in adding_tts_model pipeline (RFC #4872 P2) Signed-off-by: RyanYun09 --- docs/contributing/model/adding_tts_model.md | 43 +++++++++++---------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 68c78572ce3..f5e027c7074 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -473,7 +473,7 @@ YOUR_TTS_PIPELINE = PipelineConfig( input_sources=(), owns_tokenizer=True, engine_output_type="latent", - custom_process_next_stage_input_func=f"{_PROC}.ar2decoder", + custom_process_next_stage_input_func=f"{_PROC}.ar2decoder_full_payload", async_chunk_process_next_stage_input_func=f"{_PROC}.ar2decoder_async_chunk", ), StagePipelineConfig( @@ -530,13 +530,13 @@ monolithic function. The suffix convention and the matching registry kind are load-bearing: | Suffix | Registry kind | Runs where | Role | -|--------|---------------|------------|------| +| -------- | --------------- | ------------ | ------ | | `*_full_payload` | `producer_full_payload` | Worker (producer-side) | `async_chunk=false`; packs the accumulated Stage 0 output and ships it via the connector | | `*_async_chunk` | `producer_async_chunk` | Scheduler (producer-side) | `async_chunk=true`; streams one chunk per AR decode step | | `*_token_only` | `placeholder_prompt_builder` | Orchestrator (consumer-side) | Allocates downstream KV slots only; bulk tensors arrive via the connector (both modes) | A model that supports **both** `async_chunk=false` and `async_chunk=true` must -implement all three processors per edge — `ar2decoder` (full payload), +implement all three processors per edge — `ar2decoder_full_payload`, `ar2decoder_async_chunk` and `ar2decoder_token_only`. The `*_token_only` placeholder builder is wired as the downstream stage's `sync_process_input_func` and follows the C1 contract `(source_outputs, ctx: OrchestratorInputContext)`: @@ -585,26 +585,29 @@ Understanding what's available in stage outputs: - `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs` - `prompt_token_ids` - original prompt token IDs -### Batch mode (non-streaming) +### Batch mode (non-streaming, full payload) -Collects all Stage 0 outputs and forwards them to Stage 1 in one shot: +Runs on the worker (producer) side for `async_chunk=false`. The connector +invokes it with the exact keyword contract +`(transfer_manager, pooling_output, request, ...)` — note that the name must +carry the `_full_payload` suffix so `infer_kind` resolves it as a +`producer_full_payload` (an unsuffixed `ar2decoder` would be inferred as +`legacy_orchestrator_builder` and skipped by the worker's candidate chain): ```python -def ar2decoder( - stage_list: list[Any], - engine_input_source: list[int], - prompt: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - source_id = engine_input_source[0] - decoder_inputs = [] - - for output in stage_list[source_id].engine_outputs: - result = output.outputs[0] - codes = result.multimodal_output["audio_codes"].cpu() - decoder_inputs.append(OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist())) - - return decoder_inputs +def ar2decoder_full_payload( + transfer_manager: Any, + pooling_output: dict[str, Any] | None, + request: Any, + **_: Any, +) -> dict[str, Any] | None: + """Producer: pack the accumulated Stage 0 codec into a connector payload.""" + del transfer_manager + codes = pooling_output.get("codes.audio") if isinstance(pooling_output, dict) else None + if codes is None: + return None + decoder_inputs = OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist()) + return {"engine_inputs": [decoder_inputs]} ``` ### Streaming mode (async_chunk) From f766fa0523625a291960917ae0abcc20d5fc5f9d Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Wed, 2 Sep 2026 14:06:09 +0800 Subject: [PATCH 22/41] fix(stage_input_processors): validate producer keyword call with Signature.bind (RFC #4872 P1) Signed-off-by: RyanYun09 --- .../stage_input_processors/test_registry.py | 80 +++++++++- .../stage_input_processors/_registry.py | 149 +++++++++--------- 2 files changed, 147 insertions(+), 82 deletions(-) diff --git a/tests/model_executor/stage_input_processors/test_registry.py b/tests/model_executor/stage_input_processors/test_registry.py index 0d85721bb77..49dfa35c3a1 100644 --- a/tests/model_executor/stage_input_processors/test_registry.py +++ b/tests/model_executor/stage_input_processors/test_registry.py @@ -74,6 +74,26 @@ def producer_c_async_chunk(transfer_manager, multimodal_output, request, is_fini return None +def producer_kwargs_full_payload(**kwargs): + return None + + +def producer_kwargs_async_chunk(**kwargs): + return None + + +def producer_badname_full_payload(transfer_manager, payload, request, is_finished=False): + return None + + +def producer_badname_async_chunk(transfer_manager, payload, request, is_finished=False): + return None + + +def producer_positional_full_payload(transfer_manager, pooling_output, request, /): + return None + + def talker2codec_shape(stage_list, engine_input_source, prompt=None, requires_multimodal_data=False): return [] @@ -153,21 +173,67 @@ def test_full_payload_missing_is_finished_warns(self): path="pkg.mod.producer_a_full_payload", ) - def test_full_payload_cross_kind_second_param_warns(self): - with pytest.warns(RuntimeWarning, match="multimodal_output"): + def test_full_payload_cross_kind_second_param_rejected(self): + # ``multimodal_output`` instead of ``pooling_output`` cannot bind the + # worker keyword call, so this is a hard error (not just a warning). + with pytest.raises(ProcessorValidationError) as excinfo: validate_processor( producer_c_full_payload, kind="producer_full_payload", path="pkg.mod.producer_c_full_payload", ) + assert excinfo.value.rule == "producer_kwargs" + + def test_full_payload_bad_param_names_rejected(self): + with pytest.raises(ProcessorValidationError) as excinfo: + validate_processor( + producer_badname_full_payload, + kind="producer_full_payload", + path="pkg.mod.producer_badname_full_payload", + ) + assert excinfo.value.rule == "producer_kwargs" + assert "pooling_output" in str(excinfo.value) + + def test_full_payload_positional_only_rejected(self): + # A positional-only payload parameter cannot accept the worker's + # keyword call (``pooling_output=``), so it must be rejected. + with pytest.raises(ProcessorValidationError) as excinfo: + validate_processor( + producer_positional_full_payload, + kind="producer_full_payload", + path="pkg.mod.producer_positional_full_payload", + ) + assert excinfo.value.rule == "producer_kwargs" - def test_full_payload_missing_transfer_manager_fails(self): - def bad(manager, pooling_output, request): - return None + def test_full_payload_var_keyword_accepted(self): + # A ``**kwargs`` producer accepts the exact worker keyword call even + # though it declares no explicit names. + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor( + producer_kwargs_full_payload, + kind="producer_full_payload", + path="pkg.mod.producer_kwargs_full_payload", + ) + def test_async_chunk_bad_param_names_rejected(self): with pytest.raises(ProcessorValidationError) as excinfo: - validate_processor(bad, kind="producer_full_payload", path="pkg.mod.producer_full_payload") - assert excinfo.value.rule == "transfer_manager" + validate_processor( + producer_badname_async_chunk, + kind="producer_async_chunk", + path="pkg.mod.producer_badname_async_chunk", + ) + assert excinfo.value.rule == "producer_kwargs" + assert "multimodal_output" in str(excinfo.value) + + def test_async_chunk_var_keyword_accepted(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor( + producer_kwargs_async_chunk, + kind="producer_async_chunk", + path="pkg.mod.producer_kwargs_async_chunk", + ) def test_async_chunk_ok_no_warning(self): with warnings.catch_warnings(): diff --git a/vllm_omni/model_executor/stage_input_processors/_registry.py b/vllm_omni/model_executor/stage_input_processors/_registry.py index a687d6a2203..63dad37d548 100644 --- a/vllm_omni/model_executor/stage_input_processors/_registry.py +++ b/vllm_omni/model_executor/stage_input_processors/_registry.py @@ -309,16 +309,6 @@ def _fail( ) -def _second_parameter_name(params: Mapping[str, inspect.Parameter]) -> str | None: - """Name of the second (non-var) parameter, mirroring the producer shape.""" - names = [ - name - for name, param in params.items() - if param.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) - ] - return names[1] if len(names) >= 2 else None - - def _check_orchestrator_first_param( params: Mapping[str, inspect.Parameter], *, @@ -355,6 +345,11 @@ def _check_suffix( _warn(path, kind, "suffix_kind_mismatch", f"name {fn_name!r} does not match expected {what}") +def _has_var_keyword(params: Mapping[str, inspect.Parameter]) -> bool: + """Whether the signature accepts arbitrary keywords (``**kwargs``).""" + return any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()) + + def _check_full_payload_is_finished( params: Mapping[str, inspect.Parameter], *, @@ -366,7 +361,10 @@ def _check_full_payload_is_finished( The worker passes it best-effort and retries without it on ``TypeError`` (see ``omni_connector_model_runner_mixin._build_custom_process_payload``). + A ``**kwargs`` producer accepts it through ``VAR_KEYWORD``. """ + if _has_var_keyword(params): + return # **kwargs accepts is_finished if "is_finished" not in params: _warn( path, @@ -395,8 +393,11 @@ def _check_async_chunk_is_finished( The scheduler always passes ``is_finished`` (see ``chunk_transfer_adapter._send_single_request``), so a processor that cannot - accept it is a hard configuration error. + accept it is a hard configuration error. A ``**kwargs`` producer accepts it + through ``VAR_KEYWORD``. """ + if _has_var_keyword(params): + return # **kwargs accepts is_finished if "is_finished" not in params: _fail( path=path, @@ -430,6 +431,45 @@ def _check_async_chunk_is_finished( ) +#: Sentinel value for :func:`_check_producer_keyword_bind` probes (the value +#: itself is never inspected or passed to the processor). +_BIND_PLACEHOLDER: Any = object() + + +def _check_producer_keyword_bind( + signature: inspect.Signature, + *, + path: str, + kind: ProcessorKind, + stage_config: Any, + required: tuple[str, ...], + what: str, +) -> None: + """Bind the exact worker keyword call shape to *signature*. + + The worker invokes producers with a fixed keyword call (see + ``omni_connector_model_runner_mixin._build_custom_process_payload`` for the + full-payload producer and ``chunk_transfer_adapter._send_single_request`` + for the async-chunk producer). We probe that call with + ``inspect.Signature.bind`` so a producer that *declares* the right parameter + names but cannot actually accept the keyword call — wrong parameter names, + positional-only payload parameters, or missing required parameters — is + rejected at configuration time instead of failing later when the worker + passes ``pooling_output=`` / ``multimodal_output=``. A ``**kwargs`` + (``VAR_KEYWORD``) producer accepts any keyword and is therefore compatible. + """ + try: + signature.bind(**{name: _BIND_PLACEHOLDER for name in required}) + except TypeError as exc: + _fail( + path=path, + kind=kind, + rule="producer_kwargs", + stage_config=stage_config, + message=(f"{what} producer does not accept the worker keyword call ({', '.join(required)}): {exc}"), + ) + + def _check_return_annotation( signature: inspect.Signature, *, @@ -502,74 +542,33 @@ def validate_processor(fn: Any, *, kind: ProcessorKind, path: str, stage_config: ) elif kind == "producer_full_payload": - if "transfer_manager" not in params: - _fail( - path=path, - kind=kind, - rule="transfer_manager", - stage_config=stage_config, - message="full-payload producer must declare 'transfer_manager'", - ) - if "request" not in params: - _fail( - path=path, - kind=kind, - rule="request", - stage_config=stage_config, - message="full-payload producer must declare 'request'", - ) - second = _second_parameter_name(params) - if second == "multimodal_output": - _warn( - path, - kind, - "cross_kind_second_param", - "second parameter is 'multimodal_output' (async-chunk producer shape) " - "but kind is producer_full_payload; expected 'pooling_output'", - ) - elif second is not None and second != "pooling_output": - _warn( - path, - kind, - "second_param", - f"expected second parameter 'pooling_output', got {second!r}", - ) + # Bind the exact worker keyword call (transfer_manager, pooling_output, + # request) so a producer that only *declares* the right names but cannot + # accept the keyword call (wrong names / positional-only payload params) + # is rejected at config time; **kwargs producers are compatible. + _check_producer_keyword_bind( + signature, + path=path, + kind=kind, + stage_config=stage_config, + required=("transfer_manager", "pooling_output", "request"), + what="full-payload", + ) _check_full_payload_is_finished(params, path=path, kind=kind, stage_config=stage_config) elif kind == "producer_async_chunk": - if "transfer_manager" not in params: - _fail( - path=path, - kind=kind, - rule="transfer_manager", - stage_config=stage_config, - message="async-chunk producer must declare 'transfer_manager'", - ) - if "request" not in params: - _fail( - path=path, - kind=kind, - rule="request", - stage_config=stage_config, - message="async-chunk producer must declare 'request'", - ) - second = _second_parameter_name(params) - if second == "pooling_output": - _warn( - path, - kind, - "cross_kind_second_param", - "second parameter is 'pooling_output' (full-payload producer shape) " - "but kind is producer_async_chunk; expected 'multimodal_output'", - ) - elif second is not None and second != "multimodal_output": - _warn( - path, - kind, - "second_param", - f"expected second parameter 'multimodal_output', got {second!r}", - ) + # ``is_finished`` is required and always passed on the async-chunk path; + # validate it first so its dedicated rules (is_finished_required / + # is_finished_positional_only) fire before the generic keyword bind. _check_async_chunk_is_finished(params, path=path, kind=kind, stage_config=stage_config) + _check_producer_keyword_bind( + signature, + path=path, + kind=kind, + stage_config=stage_config, + required=("transfer_manager", "multimodal_output", "request", "is_finished"), + what="async-chunk", + ) elif kind == "legacy_multi_source": if _MOSS_MULTI_SOURCE_MODULE_HINT not in _module_part(path): From 344cab44195630c95eb33089ce722b5e38d4d657 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Wed, 2 Sep 2026 14:11:01 +0800 Subject: [PATCH 23/41] fix(engine): do not report the selected sync hook as dead (RFC #4872 P2) Signed-off-by: RyanYun09 --- .../test_orchestrator_stage_input_bridge.py | 68 ++++++++++++++++++- .../stage_input_processors/test_registry.py | 37 ++++++++++ vllm_omni/engine/orchestrator.py | 36 +++++++++- .../stage_input_processors/_registry.py | 8 ++- 4 files changed, 144 insertions(+), 5 deletions(-) diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 1fbb7c32db7..c9f58be1655 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -7,7 +7,7 @@ import queue from types import SimpleNamespace from typing import Any -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import ANY, AsyncMock, MagicMock, patch import janus import pytest @@ -446,3 +446,69 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None: orchestrator._cleanup_request_ids.assert_not_awaited() routed = orchestrator.output_async_queue.get_nowait() assert routed.finished is False + + +# --------------------------------------------------------------------------- +# P2 deep-dive: dead-processor hint must not flag the selected sync hook. +# --------------------------------------------------------------------------- + + +def _prewarm_orchestrator_with_client(client: Any) -> Orchestrator: + orch = object.__new__(Orchestrator) + orch.stage_pools = [ + SimpleNamespace( + stage_client=client, + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + ] + return orch + + +def test_warn_dead_input_processors_skips_selected_sync_hook() -> None: + """The selected sync hook used as the custom hook is NOT dead. + + In non-async mode a stage wiring the same ``*_token_only`` processor as + both ``custom_process_input_func`` and ``sync_process_input_func`` must not + be reported as never-invoked (it is the active forward processor). + """ + from vllm_omni.engine import orchestrator as orch_module + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + path = f"{_sync_hook.__module__}.{_sync_hook.__qualname__}" + + class _Client: + custom_process_input_func = _sync_hook + sync_process_input_func = path + + orch = _prewarm_orchestrator_with_client(_Client()) + orch.async_chunk = False + with patch.object(orch_module.logger, "warning") as mock_warn: + orch._warn_dead_input_processors() + dead_warnings = [call for call in mock_warn.call_args_list if call.args and "never invoked" in call.args[0]] + assert dead_warnings == [] + + +def test_warn_dead_input_processors_flags_distinct_sync_hook() -> None: + """A custom hook overridden by a *different* sync hook is still reported dead.""" + from vllm_omni.engine import orchestrator as orch_module + + def _custom_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + path = f"{_sync_hook.__module__}.{_sync_hook.__qualname__}" + + class _Client: + custom_process_input_func = _custom_hook + sync_process_input_func = path + + orch = _prewarm_orchestrator_with_client(_Client()) + orch.async_chunk = False + with patch.object(orch_module.logger, "warning") as mock_warn: + orch._warn_dead_input_processors() + dead_warnings = [call for call in mock_warn.call_args_list if call.args and "never invoked" in call.args[0]] + assert len(dead_warnings) == 1 diff --git a/tests/model_executor/stage_input_processors/test_registry.py b/tests/model_executor/stage_input_processors/test_registry.py index 49dfa35c3a1..978ee92ad56 100644 --- a/tests/model_executor/stage_input_processors/test_registry.py +++ b/tests/model_executor/stage_input_processors/test_registry.py @@ -374,6 +374,43 @@ def test_sync_mode_alive_unless_overridden(self): is True ) + def test_sync_mode_selected_sync_hook_is_not_dead(self): + # When the custom hook *is* the configured sync hook (the sync + # ``*_token_only`` processor was pre-selected as the active forward + # hook), it must NOT be reported dead even though ``has_sync`` is true. + assert ( + dead_processor_hint( + "placeholder_prompt_builder", + async_chunk=False, + downstream_receives_async_chunks=False, + has_sync=True, + custom_is_sync=True, + ) + is False + ) + # A *different* sync hook overriding the custom hook is still dead. + assert ( + dead_processor_hint( + "placeholder_prompt_builder", + async_chunk=False, + downstream_receives_async_chunks=False, + has_sync=True, + custom_is_sync=False, + ) + is True + ) + # custom_is_sync only affects the non-async (sync-mode) decision. + assert ( + dead_processor_hint( + "placeholder_prompt_builder", + async_chunk=True, + downstream_receives_async_chunks=True, + has_sync=True, + custom_is_sync=True, + ) + is True + ) + def test_producer_never_dead(self): for kind in ("producer_full_payload", "producer_async_chunk"): assert ( diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 07c5f6fc547..5e99d1bb1c3 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -140,6 +140,32 @@ def _build_terminal_empty_output( ) +def _custom_hook_is_sync_hook(fn: Any, sync_path: str | None) -> bool: + """Whether the custom hook callable *is* the configured sync hook. + + A stage may wire the same ``*_token_only`` processor as both + ``custom_process_input_func`` and ``sync_process_input_func`` (the sync hook + is pre-selected as the active forward processor). In that case the custom + hook is **not** dead — it is the active processor — so + :func:`dead_processor_hint` must not report it. Comparison is by dotted + path first, then by resolving the sync path and identity-checking the + callable. Never raises. + """ + if not callable(fn) or not sync_path: + return False + module = getattr(fn, "__module__", None) + name = getattr(fn, "__qualname__", None) or getattr(fn, "__name__", None) + if module and name and f"{module}.{name}" == sync_path: + return True + try: + from vllm_omni.model_executor.stage_input_processors import resolve_processor + + spec = resolve_processor(sync_path, expected_kind=None) + return spec.fn is fn + except Exception: + return False + + def build_engine_core_request_from_tokens( request_id: str, prompt: dict[str, Any], @@ -2090,12 +2116,18 @@ def _warn_dead_input_processors(self) -> None: continue path = f"{module}.{name}" kind = infer_kind(fn, path=path) - has_sync = bool(getattr(client, "sync_process_input_func", None)) + sync_path = getattr(client, "sync_process_input_func", None) + has_sync = bool(sync_path) + # Do not report a hook that *is* the selected sync hook: in + # non-async mode the sync ``*_token_only`` processor is the + # active custom hook, not a dead leftover. + custom_is_sync = _custom_hook_is_sync_hook(fn, sync_path) dead = dead_processor_hint( kind, async_chunk=self.async_chunk, downstream_receives_async_chunks=self._stage_receives_async_chunks(stage_id), has_sync=has_sync, + custom_is_sync=custom_is_sync, ) if dead: logger.warning( @@ -3197,7 +3229,7 @@ def _diffusion_submit_kwargs( ) } def _get_prewarm_placeholder_builder(self, stage_id: int) -> Any: - """Resolve the registered ``build_prewarm_placeholder`` for a stage. + """Resolve (and cache per stage) the ``build_prewarm_placeholder``. The async-chunk prewarm reuses ``sync_process_input_func`` (the ``*_token_only`` path). Resolution order: diff --git a/vllm_omni/model_executor/stage_input_processors/_registry.py b/vllm_omni/model_executor/stage_input_processors/_registry.py index 63dad37d548..de92dac6857 100644 --- a/vllm_omni/model_executor/stage_input_processors/_registry.py +++ b/vllm_omni/model_executor/stage_input_processors/_registry.py @@ -656,6 +656,7 @@ def dead_processor_hint( async_chunk: bool, downstream_receives_async_chunks: bool, has_sync: bool, + custom_is_sync: bool = False, ) -> bool: """Whether an orchestrator-side ``custom_process_input_func`` is never invoked. @@ -669,7 +670,10 @@ def dead_processor_hint( chunks -> the transition still runs forward, so the processor is alive (returns ``False``). * ``async_chunk=False`` -> forward path is active; the processor is only - dead when a ``sync_process_input_func`` overrides it (``has_sync=True``). + dead when a *different* ``sync_process_input_func`` overrides it + (``has_sync=True`` and ``custom_is_sync=False``). When the custom hook + *is* the selected sync hook (``custom_is_sync=True``) it is the active + processor and is **not** dead. Producer kinds are never orchestrator-side, so they are never "dead" under this definition. @@ -681,4 +685,4 @@ def dead_processor_hint( return False if async_chunk: return bool(downstream_receives_async_chunks) - return bool(has_sync) + return bool(has_sync) and not custom_is_sync From ace8190eadf4d10a0c931373d7ce0fab5e1aae6f Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Wed, 2 Sep 2026 14:11:15 +0800 Subject: [PATCH 24/41] fix(engine): cache per-stage prewarm placeholder builder resolution (RFC #4872 P3) Signed-off-by: RyanYun09 --- .../test_orchestrator_stage_input_bridge.py | 76 +++++++++++++ vllm_omni/engine/orchestrator.py | 100 +++++++++++------- 2 files changed, 139 insertions(+), 37 deletions(-) diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index c9f58be1655..43dabbbc070 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -512,3 +512,79 @@ class _Client: orch._warn_dead_input_processors() dead_warnings = [call for call in mock_warn.call_args_list if call.args and "never invoked" in call.args[0]] assert len(dead_warnings) == 1 + + +# --------------------------------------------------------------------------- +# P3 deep-dive: prewarm placeholder builder resolution is cached per stage. +# --------------------------------------------------------------------------- + + +class _FakeProcessorSpec: + """Minimal stand-in for ``ProcessorSpec`` with a controllable ``fn``.""" + + def __init__(self, fn: Any) -> None: + self.path = "fake" + self.kind = "placeholder_prompt_builder" + self.fn = fn + + +def test_prewarm_builder_resolved_once_and_cached() -> None: + """A stage with a builder resolves it once and caches it. + + Repeated async requests must not re-run the importlib resolution or the + registry validation. + """ + from vllm_omni.model_executor import stage_input_processors as sip + + def _build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] + return {"prompt_token_ids": [0]} + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + _sync_hook.build_prewarm_placeholder = _build_prewarm_placeholder # type: ignore[attr-defined] + + class _Client: + sync_process_input_func = "pkg.mod.fake_token_only" + + orch = _prewarm_orchestrator_with_client(_Client()) + with patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)) as mock_resolve: + first = orch._get_prewarm_placeholder_builder(0) + second = orch._get_prewarm_placeholder_builder(0) + assert first is _build_prewarm_placeholder + assert second is _build_prewarm_placeholder + assert mock_resolve.call_count == 1 + + +def test_prewarm_builder_miss_warns_once_and_caches() -> None: + """A stage WITHOUT a builder resolves once, warns once, then caches a miss. + + Qwen3-TTS and other models have no ``build_prewarm_placeholder``; without + per-stage caching every async request re-resolved and re-warned. + """ + from vllm_omni.engine import orchestrator as orch_module + from vllm_omni.model_executor import stage_input_processors as sip + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + class _Client: + sync_process_input_func = "pkg.mod.fake_token_only" + + orch = _prewarm_orchestrator_with_client(_Client()) + fake_module = SimpleNamespace() + with ( + patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)) as mock_resolve, + patch("importlib.import_module", return_value=fake_module) as mock_import, + patch.object(orch_module.logger, "warning") as mock_warn, + ): + first = orch._get_prewarm_placeholder_builder(0) + second = orch._get_prewarm_placeholder_builder(0) + third = orch._get_prewarm_placeholder_builder(0) + assert first is None + assert second is None + assert third is None + assert mock_resolve.call_count == 1 + assert mock_import.call_count == 1 + inline_estimates = [call for call in mock_warn.call_args_list if call.args and "inline estimate" in call.args[0]] + assert len(inline_estimates) == 1 diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 5e99d1bb1c3..ea8daf15cc7 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -80,6 +80,11 @@ logger = init_logger(__name__) +#: Sentinel for a stage whose prewarm placeholder builder has not been resolved. +_PREWARM_BUILDER_NOT_RESOLVED: Any = object() +#: Negative sentinel for a stage with no prewarm placeholder builder. +_PREWARM_BUILDER_NO_BUILDER: Any = object() + def cleanup_request_artifact_dirs(artifact_dirs: set[str] | list[str]) -> None: for artifact_dir in artifact_dirs: @@ -396,6 +401,10 @@ def __init__( self.async_chunk = bool(async_chunk) self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools + # Per-stage prewarm placeholder builder cache (callable or a negative + # sentinel); populated lazily so repeated async requests never re-resolve + # (and never re-warn) for models without a builder. + self._prewarm_builder_cache: dict[int, Any] = {} # Wire the dead-processor hint at startup (warn-only) so a # configured-but-never-invoked custom_process_input_func is surfaced. self._warn_dead_input_processors() @@ -3245,56 +3254,73 @@ def _get_prewarm_placeholder_builder(self, stage_id: int) -> Any: raises (fake pools / minimal stage clients in tests and future connectors may not expose ``stage_client``), so the prewarm path cannot regress on builder resolution. + + The result is cached per stage (the callable on a hit, a negative + sentinel on a miss) so repeated async requests never re-run the + importlib resolution and never re-emit the fallback warning — Qwen3-TTS + and other models without a ``build_prewarm_placeholder`` would otherwise + re-resolve and re-warn on every async request. """ + cache = getattr(self, "_prewarm_builder_cache", None) + if cache is None: # object.__new__ test doubles + cache = {} + self._prewarm_builder_cache = cache + cached = cache.get(stage_id, _PREWARM_BUILDER_NOT_RESOLVED) + if cached is not _PREWARM_BUILDER_NOT_RESOLVED: + return None if cached is _PREWARM_BUILDER_NO_BUILDER else cached + from vllm_omni.model_executor.stage_input_processors import resolve_processor + builder: Any = None try: pool = self.stage_pools[stage_id] client = getattr(pool, "stage_client", None) - if client is None: - return None - - path = getattr(client, "sync_process_input_func", None) - if path: - try: - spec = resolve_processor(path, expected_kind=None) - builder = getattr(spec.fn, "build_prewarm_placeholder", None) - if not callable(builder): - # Module-level sibling (covers paths whose fn does not - # expose the attribute directly). - import importlib - - module = importlib.import_module(path.rsplit(".", 1)[0]) - builder = getattr(module, "build_prewarm_placeholder", None) - if callable(builder): - return builder - logger.warning( - "[Orchestrator] stage-%s sync_process_input_func=%r has no " - "build_prewarm_placeholder; using the inline estimate", - stage_id, - path, - ) - except Exception as exc: - logger.warning( - "[Orchestrator] failed to resolve build_prewarm_placeholder " - "for %r (stage-%s); using the inline estimate: %s", - path, - stage_id, - exc, - ) + if client is not None: + path = getattr(client, "sync_process_input_func", None) + if path: + try: + spec = resolve_processor(path, expected_kind=None) + candidate = getattr(spec.fn, "build_prewarm_placeholder", None) + if not callable(candidate): + # Module-level sibling (covers paths whose fn does + # not expose the attribute directly). + import importlib + + module = importlib.import_module(path.rsplit(".", 1)[0]) + candidate = getattr(module, "build_prewarm_placeholder", None) + if callable(candidate): + builder = candidate + else: + logger.warning( + "[Orchestrator] stage-%s sync_process_input_func=%r has no " + "build_prewarm_placeholder; using the inline estimate", + stage_id, + path, + ) + except Exception as exc: + logger.warning( + "[Orchestrator] failed to resolve build_prewarm_placeholder " + "for %r (stage-%s); using the inline estimate: %s", + path, + stage_id, + exc, + ) - fn = getattr(client, "custom_process_input_func", None) - if callable(fn): - builder = getattr(fn, "build_prewarm_placeholder", None) - if callable(builder): - return builder + if builder is None: + fn = getattr(client, "custom_process_input_func", None) + if callable(fn): + candidate = getattr(fn, "build_prewarm_placeholder", None) + if callable(candidate): + builder = candidate except Exception as exc: # pragma: no cover - defensive only logger.warning( "[Orchestrator] prewarm placeholder builder lookup for stage-%s failed; using the inline estimate: %s", stage_id, exc, ) - return None + + cache[stage_id] = builder if builder is not None else _PREWARM_BUILDER_NO_BUILDER + return builder def _build_prewarm_placeholder_input( self, From b804311a11068410d9fb4faa77f08ce4c88b6177 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Tue, 15 Sep 2026 16:56:16 +0800 Subject: [PATCH 25/41] fix(merge): apply RFC resolve_processor to omni_connector_runtime and correct mixin thin-wrapper Signed-off-by: RyanYun09 --- .../model_runner/omni_connector_runtime.py | 31 +- .../omni_connector_model_runner_mixin.py | 2268 +---------------- 2 files changed, 19 insertions(+), 2280 deletions(-) diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 6667c0e4b1e..862e522d9dd 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -4,7 +4,6 @@ from __future__ import annotations -import importlib import inspect import os import threading @@ -24,6 +23,11 @@ from vllm_omni.distributed.omni_connectors.utils.kv_utils import get_local_tp_rank, get_omni_replica_id from vllm_omni.outputs import OmniConnectorOutput +from vllm_omni.model_executor.stage_input_processors import ( + ProcessorValidationError, + resolve_processor, +) + logger = init_logger("vllm_omni.worker.omni_connector_model_runner_mixin") if TYPE_CHECKING: @@ -636,18 +640,19 @@ def _load_custom_func(cls, model_config: Any) -> tuple[str | None, Any | None]: continue tried.add(func_path) try: - module_path, func_name = func_path.rsplit(".", 1) - module = importlib.import_module(module_path) - func = getattr(module, func_name, None) - if callable(func): - if not cls._is_connector_payload_builder(func): - logger.debug( - "Skipping incompatible connector payload hook %s; signature=%s", - func_path, - inspect.signature(func), - ) - continue - return func_path, func + # Resolve through the registry with the full-payload producer + # contract. is_finished is optional here (missing only warns); + # a structurally incompatible candidate is skipped so the + # fallback chain keeps working. + spec = resolve_processor(func_path, expected_kind="producer_full_payload") + return spec.path, spec.fn + except ProcessorValidationError as exc: + logger.debug( + "Skipping incompatible connector payload hook %s: %s", + func_path, + exc, + ) + continue except Exception: logger.warning("Failed to load custom func: %s", func_path, exc_info=True) diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 6c842d8a775..33088f4f6ce 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -32,2271 +32,5 @@ class OmniConnectorModelRunnerMixin(_OmniConnectorPayloadTransportMixin): """Unified data-plane communication interface for model runners. Runtime ownership and payload transport live under - ``distributed.omni_connectors``. + "distributed.omni_connectors". """ -def should_accumulate_full_payload_output(model_config, custom_process_func) -> bool: - """Producer-side structural gate. - - Fires iff the stage explicitly declares a downstream full-payload - producer hook via ``custom_process_next_stage_input_func``. Consumer - stages may have ``custom_process_input_func`` values that can be - mechanically derived to ``*_full_payload`` helper names in the same - module; those are intentionally not enough to make the stage a producer. - """ - if custom_process_func is None: - return False - if getattr(model_config, "async_chunk", False): - return False - if getattr(model_config, "final_output", False): - return False - next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) - if not isinstance(next_stage_func, str) or not next_stage_func: - return False - return getattr(model_config, "model_stage", None) is not None - - -class OmniConnectorModelRunnerMixin: - """Unified data-plane communication mixin for Model Runners. - - Provides three transfer modes through a single pair of bg I/O threads: - - **full_payload_mode**: ``recv_full_payload_inputs`` / ``send_full_payload_outputs`` - - **Streaming (async_chunk)**: ``recv_chunk`` / ``send_chunk`` - - **KV cache**: ``send_kv_cache`` / ``recv_kv_cache`` (delegates to - the existing ``OmniKVTransferManager``) - - The mixin owns connector instances and background threads. It never - touches scheduling queues -- readiness is communicated to the Scheduler - via ``OmniConnectorOutput``. - """ - - # ------------------------------------------------------------------ # - # Init / Shutdown - # ------------------------------------------------------------------ # - - def init_omni_connectors( - self, - model_config: OmniModelConfig, - kv_transfer_manager: OmniKVTransferManager | None = None, - ) -> None: - """Initialize connectors and background threads. - - Args: - model_config: Stage-level model config with connector settings. - kv_transfer_manager: Existing KV transfer manager to delegate to. - """ - self._omni_connector: OmniConnectorBase | None = ( - self._create_connector(model_config) if _should_create_payload_connector(model_config) else None - ) - self._kv_transfer_manager = kv_transfer_manager - - self._async_chunk: bool = getattr(model_config, "async_chunk", False) - self._model_mode: str = getattr(model_config, "worker_type", "ar") - stage_id = getattr(model_config, "stage_id", 0) - if isinstance(stage_id, str): - stage_id = int(stage_id) - self._stage_id: int = stage_id if isinstance(stage_id, int) else 0 - - self._custom_process_func_path, self._custom_process_func = self._load_custom_func(model_config) - self._custom_process_supports_is_finished = self._custom_process_supports_is_finished_kwarg() - logger.debug( - "[Stage-%s] init_omni_connectors: async_chunk=%s, custom_process_func=%s, connector=%s, func_path=%s", - self._stage_id, - self._async_chunk, - self._custom_process_func, - type(self._omni_connector).__name__ if self._omni_connector else None, - self._custom_process_func_path, - ) - - # -- next stage ID (from connector config or default stage_id + 1) -- - self._next_stage_id: int = self._resolve_next_stage_id(model_config) - - # -- heterogeneous TP rank support -- - rank_cfg = self._parse_rank_mapping(model_config) - if self._kv_transfer_manager is not None: - topology = getattr(self._kv_transfer_manager, "tp_topology", None) - effective_mapping = ( - getattr(topology, "source_tp_size", None), - getattr(topology, "target_tp_size", None), - getattr(topology, "local_rank", None), - ) - if all(isinstance(value, int) for value in effective_mapping): - rank_cfg = { - "from_tp": effective_mapping[0], - "to_tp": effective_mapping[1], - "local_rank": effective_mapping[2], - } - self._from_tp: int = rank_cfg["from_tp"] - self._to_tp: int = rank_cfg["to_tp"] - self._local_rank: int = rank_cfg["local_rank"] - if self._kv_transfer_manager is not None: - self._kv_transfer_manager.kv_send_key_builder = self.get_rank_aware_kv_send_keys - self._kv_transfer_manager.kv_recv_key_builder = self.get_rank_aware_kv_keys - self._kv_transfer_manager.kv_payload_merger = self._merge_rank_sharded_kv_payloads - self._kv_transfer_manager.kv_payload_slicer = self._slice_rank_sharded_kv_payload - - # -- chunk index tracking (ported from OmniChunkTransferAdapter) -- - self._put_req_chunk: dict[str, int] = defaultdict(int) - self._get_req_chunk: dict[str, int] = defaultdict(int) - # Segment-local chunk counter: incremented alongside _put_req_chunk - # and popped at request cleanup. Note: the mixin path (uniproc mode) - # does not have segment boundary infrastructure; multi-segment support - # is only available via chunk_transfer_adapter (distributed path). - self._ramp_chunk_count: dict[str, int] = defaultdict(int) - # Send-side async accumulation / staging buffer. Receive-side payload - # ownership lives in ``_local_stage_payload_cache``. - self._send_side_request_payload: dict[str, dict[str, Any]] = {} - self._code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list) - self._cached_ic: dict[str, int] = {} - self._request_ids_mapping: dict[str, str] = {} - - # -- async I/O state (shared by chunk + full_payload_mode) -- - self._pending_load_reqs: dict[str, Any] = {} - self._finished_load_reqs: set[str] = set() - self._pending_save_reqs: dict[str, deque] = {} - self._pending_save_counts: dict[str, int] = defaultdict(int) - self._deferred_send_cleanup: set[str] = set() - # -- per-cycle output accumulator -- - self._chunk_ready_req_ids: set[str] = set() - self._chunk_finished_req_ids: set[str] = set() - self._stage_recv_req_ids: set[str] = set() - self._full_payload_pending_broadcast_req_ids: set[str] = set() - self._async_chunk_updated_req_ids: set[str] = set() - - # -- Model Runner local payload cache (RFC §2.4) -- - # Full stage payloads land here first on the recv side. We - # intentionally do not write connector recv results straight into - # `model_intermediate_buffer`: runner-owned runtime state is - # materialized later by `_sync_local_stage_payloads()` on the - # model thread. This keeps recv timing separate from execute-step - # visibility and avoids mixing connector I/O with model runtime - # ownership. - self._local_stage_payload_cache: dict[str, dict[str, Any]] = {} - # Lightweight scheduling metadata pending delivery to the Scheduler. - self._local_request_metadata: dict[str, dict[str, Any]] = {} - - # -- persistent set of request IDs whose chunk stream is complete -- - # Prevents re-registration after the finish sentinel has been received. - self._chunk_stream_completed: set[str] = set() - - # -- full_payload_mode: accumulate latest pooler_output per request, - # send only when the request finishes (next-cycle flush) -- - self._pending_full_payload_send: dict[str, tuple[Any, ...]] = {} - - # -- KV sent accumulator -- - self._kv_sent_req_ids: list[str] = [] - - # -- KV transfer lifecycle (absorbed from scheduler) -- - # Requests marked for KV transfer: {req_id: {seq_len, block_ids}} - self._kv_pending_transfers: dict[str, dict[str, Any]] = {} - # Requests whose KV transfer has been submitted but not yet acked - self._kv_active_transfers: set[str] = set() - # Requests whose KV transfer is complete (acked by kv_extracted_req_ids) - self._kv_completed_transfers: set[str] = set() - # Dedup guard: requests that have already triggered KV transfer - self._kv_triggered_requests: set[str] = set() - - self._lock = threading.Lock() - self._stop_event = threading.Event() - self._work_available = threading.Event() - - # Start background threads only when there's a connector - self._recv_thread: threading.Thread | None = None - self._save_thread: threading.Thread | None = None - if self._omni_connector is not None: - self._recv_thread = threading.Thread( - target=self._recv_loop, - daemon=True, - name="omni-mixin-recv", - ) - self._recv_thread.start() - self._save_thread = threading.Thread( - target=self._save_loop, - daemon=True, - name="omni-mixin-save", - ) - self._save_thread.start() - - # Explicit "fully initialised" marker so other parts of the runner - # (e.g. _update_states cleanup) can branch on a stable contract - # instead of probing for private mixin attribute names. Must be set - # only after every field above has been bound, so a partially - # constructed mixin is never observable as initialised. - self._omni_connector_initialized = True - - def shutdown_omni_connectors(self) -> None: - """Stop background threads and release connector resources.""" - self._stop_event.set() - if self._recv_thread is not None: - self._recv_thread.join(timeout=5) - if self._save_thread is not None: - self._save_thread.join(timeout=5) - if self._omni_connector is not None: - try: - self._omni_connector.close() - except Exception: - pass - - def cleanup_finished_request(self, req_id: str) -> None: - """Clean up per-request state after a request is fully finished. - - Call this when a request is freed from the model runner to prevent - memory leaks in the mixin's tracking dicts/sets. - - Two senders use different keys: ``send_chunk`` keys per-request - state under the EXTERNAL id (after mapping resolution), while - ``send_full_payload_outputs`` keys under the INTERNAL id. To cover - both modes (and forward compat with id-rename scenarios) we attempt - cleanup against both keys; the entry that doesn't exist for the - active mode is a no-op pop. Only the key that actually has pending - saves is added to ``_deferred_send_cleanup`` so the bg save's - decrement path drains it without leaving orphans. - """ - # Force-flush any pending full-payload accumulator entry before - # cleanup proceeds. Without this, finished requests with no - # downstream consumer (e.g. text-only on multi-modal arch) leave - # the entry orphaned in _pending_full_payload_send across requests, - # which empirically destabilises subsequent thinker forwards by - # making prefix-cache reuse observe stale accumulator state. The - # flush is idempotent when the entry has already been flushed by the - # scheduler-driven path, but this cleanup path runs for every request, - # so skip it entirely when the request never accumulated a payload. - if req_id in self._pending_full_payload_send: - try: - self.flush_full_payload_outputs({req_id}) - except Exception: - # Cleanup must still proceed regardless of flush errors here -- - # we already gated on ``_omni_connector_initialized`` upstream, - # so any exception here reflects a real connector-side issue - # (shared memory corruption, background thread crash) worth - # surfacing rather than silently swallowing. - logger.warning( - "flush_full_payload_outputs(%s) raised during cleanup; continuing tear-down.", - req_id, - exc_info=True, - ) - - ext_id = self._request_ids_mapping.pop(req_id, None) - keys_to_clean: list[str] = [req_id] - if ext_id is not None and ext_id != req_id: - keys_to_clean.append(ext_id) - - with self._lock: - keys_pending = [k for k in keys_to_clean if self._pending_save_counts.get(k, 0)] - for k in keys_pending: - self._deferred_send_cleanup.add(k) - for k in keys_to_clean: - if k in keys_pending: - continue - self._put_req_chunk.pop(k, None) - self._send_side_request_payload.pop(k, None) - self._code_prompt_token_ids.pop(k, None) - self._cached_ic.pop(k, None) - self._ramp_chunk_count.pop(k, None) - self._kv_pending_transfers.pop(req_id, None) - self._kv_active_transfers.discard(req_id) - self._kv_completed_transfers.discard(req_id) - self._kv_triggered_requests.discard(req_id) - self._cleanup_recv_delivery_state(req_id) - - def drop_inactive_request_delivery_state(self, req_id: str) -> None: - """Clear recv-side state for inactive requests.""" - ext_id = self._request_ids_mapping.pop(req_id, None) - if hasattr(self, "_lock"): - with self._lock: - self._drop_send_side_payload_state(req_id, ext_id) - else: - self._drop_send_side_payload_state(req_id, ext_id) - self._cleanup_recv_delivery_state(req_id) - - def _drop_send_side_payload_state(self, req_id: str, ext_id: str | None) -> None: - if ext_id is not None: - self._send_side_request_payload.pop(ext_id, None) - self._cached_ic.pop(ext_id, None) - self._send_side_request_payload.pop(req_id, None) - self._cached_ic.pop(req_id, None) - - def _cleanup_recv_delivery_state(self, req_id: str) -> None: - """Clear recv-side delivery-cycle state.""" - if hasattr(self, "_lock"): - with self._lock: - self._clear_recv_delivery_state(req_id) - else: - self._clear_recv_delivery_state(req_id) - - def _clear_recv_delivery_state(self, req_id: str) -> None: - self._get_req_chunk.pop(req_id, None) - self._pending_load_reqs.pop(req_id, None) - self._finished_load_reqs.discard(req_id) - self._chunk_ready_req_ids.discard(req_id) - self._chunk_finished_req_ids.discard(req_id) - self._chunk_stream_completed.discard(req_id) - self._stage_recv_req_ids.discard(req_id) - self._full_payload_pending_broadcast_req_ids.discard(req_id) - self._async_chunk_updated_req_ids.discard(req_id) - self._local_stage_payload_cache.pop(req_id, None) - self._local_request_metadata.pop(req_id, None) - - def prune_inactive_requests(self, active_req_ids: Any) -> set[str]: - """Drop connector state for requests that no longer exist locally. - - Preempted / unscheduled requests are expected to stay in - ``self.requests`` and therefore remain untouched. This only prunes - stale request IDs that have already fallen out of the active request - map, preventing background recv/send bookkeeping from outliving the - request lifecycle. - """ - if active_req_ids is None: - return set() - - active_req_ids = set(active_req_ids) - pending_req_ids = set(getattr(self, "_pending_load_reqs", {}).keys()) - received_req_ids = set(getattr(self, "_stage_recv_req_ids", set())) - received_req_ids.update(getattr(self, "_full_payload_pending_broadcast_req_ids", set())) - received_req_ids.update(getattr(self, "_local_request_metadata", {}).keys()) - # Pending recv requests may not yet be in the caller's active set - # (e.g. WAITING_FOR_CHUNK requests live in the coordinator's internal - # queues, not in model runner self.requests). Protect them so that - # legitimate waiting requests are not pruned. - # - # Likewise, a full payload can arrive on the background recv thread - # after the scheduler_output snapshot for the current execute_model() - # cycle was already materialized. Those requests may briefly live only - # in recv-side buffers/local cache until the next scheduler cycle wakes - # them up; pruning them here drops the payload before stage_recv can be - # published. - active_req_ids.update(pending_req_ids) - active_req_ids.update(received_req_ids) - stale_req_ids: set[str] = set() - - # NOTE: _pending_load_reqs is excluded from the scan list because - # all its entries are unconditionally protected above. The mixin - # cannot distinguish a legitimately-waiting pending recv from an - # orphaned one (only the coordinator/scheduler knows). - # - # Requests with freshly received full payloads / local stage payloads - # are also protected above. Their scheduler wake-up may lag the recv - # thread by one execute_model() cycle, especially when the request was - # added after the current scheduler_output snapshot. - # - # Orphaned pending recv entries (e.g. from upstream stage crash) are - # handled by collect_timed_out_request_ids() -- on - # OmniSchedulingCoordinator for full-payload requests, and on - # OmniChunkTransferAdapter for async-chunk ones -- which detect - # wait-time violations. The scheduler then removes the request from - # its queues, sets FINISHED_ERROR, and calls _free_request() which - # ultimately triggers cleanup_finished_request() here. - for attr_name in ( - "_request_ids_mapping", - "_get_req_chunk", - "_finished_load_reqs", - "_chunk_ready_req_ids", - "_chunk_finished_req_ids", - "_chunk_stream_completed", - "_stage_recv_req_ids", - "_full_payload_pending_broadcast_req_ids", - "_async_chunk_updated_req_ids", - "_local_stage_payload_cache", - "_local_request_metadata", - "_kv_pending_transfers", - "_kv_active_transfers", - "_kv_completed_transfers", - "_kv_triggered_requests", - ): - state = getattr(self, attr_name, None) - if isinstance(state, dict): - stale_req_ids.update(req_id for req_id in state if req_id not in active_req_ids) - elif isinstance(state, set): - stale_req_ids.update(req_id for req_id in state if req_id not in active_req_ids) - - for req_id in stale_req_ids: - self.cleanup_finished_request(req_id) - - return stale_req_ids - - # ------------------------------------------------------------------ # - # Local payload cache (RFC §2.4 – Model Runner ownership) - # ------------------------------------------------------------------ # - - def put_local_stage_payload(self, req_id: str, payload: OmniPayload) -> None: - """Store a full stage payload in the local cache.""" - self._local_stage_payload_cache[req_id] = payload - - def get_local_stage_payload(self, req_id: str) -> OmniPayload | None: - """Read a stage payload without removing it.""" - return self._local_stage_payload_cache.get(req_id) - - def pop_local_stage_payload(self, req_id: str) -> OmniPayload | None: - """Remove and return a stage payload (consume after use).""" - return self._local_stage_payload_cache.pop(req_id, None) - - def put_local_request_metadata(self, req_id: str, metadata: dict[str, Any]) -> None: - """Store lightweight scheduling metadata for a request.""" - self._local_request_metadata[req_id] = metadata - - def get_local_request_metadata(self, req_id: str) -> dict[str, Any] | None: - """Retrieve scheduling metadata for a request.""" - return self._local_request_metadata.get(req_id) - - # ------------------------------------------------------------------ # - # Scheduling metadata extraction - # ------------------------------------------------------------------ # - - @classmethod - def _extract_scheduling_metadata(cls, payload: OmniPayload) -> dict[str, Any]: - """Extract only the fields the scheduler needs from a full payload.""" - extracted: dict[str, Any] = {} - meta = payload.get("meta") if isinstance(payload, dict) else None - meta = meta if isinstance(meta, dict) else {} - - if "next_stage_prompt_len" in meta: - extracted["next_stage_prompt_len"] = meta["next_stage_prompt_len"] - elif "next_stage_prompt_len" in payload: - logger.warning_once( - "legacy flat 'next_stage_prompt_len' key in payload; expected 'meta.next_stage_prompt_len'" - ) - extracted["next_stage_prompt_len"] = payload["next_stage_prompt_len"] - - audio_codes = cls._payload_audio_codes(payload) - if audio_codes is not None: - extracted["code_predictor_codes"] = audio_codes - - if "left_context_size" in meta: - extracted["left_context_size"] = meta["left_context_size"] - elif "left_context_size" in payload: - logger.warning_once("legacy flat 'left_context_size' key in payload; expected 'meta.left_context_size'") - - return extracted - - _NON_CONSUMABLE_PAYLOAD_KEYS: set[tuple[str, str]] = { - ("meta", "finished"), - ("meta", "override_keys"), - ("meta", "next_stage_prompt_len"), - ("meta", "left_context_size"), - ("ids", "output"), - ("embed", "decode_token_start"), - ("embed", "decode_token_end"), - } - - @staticmethod - def _payload_value_has_content(value: Any) -> bool: - if value is None: - return False - if isinstance(value, torch.Tensor): - return value.numel() > 0 - if isinstance(value, (list, tuple, dict, set)): - return len(value) > 0 - return True - - @staticmethod - def _payload_finished(payload: Any) -> bool: - if not isinstance(payload, dict): - return False - if "finished" in payload: - logger.warning_once("legacy flat 'finished' key in payload; expected 'meta.finished'") - meta = payload.get("meta") - if not isinstance(meta, dict) or "finished" not in meta: - return False - flag = meta["finished"] - if isinstance(flag, torch.Tensor): - return flag.numel() == 1 and bool(flag.item()) - return bool(flag) - - @staticmethod - def _payload_audio_codes(payload: Any) -> Any: - if not isinstance(payload, dict): - return None - if "code_predictor_codes" in payload: - logger.warning_once("legacy flat 'code_predictor_codes' key in payload; expected 'codes.audio'") - codes = payload.get("codes") - if isinstance(codes, dict): - return codes.get("audio") - return None - - @classmethod - def _payload_is_consumable(cls, payload: OmniPayload | None) -> bool: - """Return True when an async payload can drive a real forward step. - - Metadata-only wake-ups should not transition WAITING_FOR_CHUNK requests - back to schedulable state. In particular, a widened token horizon without - any newly visible thinker decode embeds should not force a placeholder-only - talker decode step. - """ - if not isinstance(payload, dict) or not payload: - return False - - embed = payload.get("embed") - if isinstance(embed, dict): - decode_embeddings = embed.get("decode") - if isinstance(decode_embeddings, torch.Tensor): - if decode_embeddings.ndim == 0: - return True - return decode_embeddings.numel() > 0 and decode_embeddings.shape[0] > 0 - - audio_codes = cls._payload_audio_codes(payload) - if audio_codes is not None: - if isinstance(audio_codes, torch.Tensor): - return audio_codes.numel() > 0 - if hasattr(audio_codes, "__len__"): - return len(audio_codes) > 0 - return True - - for key, value in payload.items(): - if isinstance(value, dict): - for sk, sv in value.items(): - if (key, sk) in cls._NON_CONSUMABLE_PAYLOAD_KEYS: - continue - if cls._payload_value_has_content(sv): - return True - continue - if cls._payload_value_has_content(value): - return True - return False - - @staticmethod - def _get_local_tp_group() -> Any | None: - """Return the local TP group when tensor parallelism is initialized.""" - try: - return get_tp_group() - except Exception: - return None - - def _recv_ordinary_stage_result( - self, - connector: OmniConnectorBase, - from_stage: str, - to_stage: str, - connector_get_key: str, - ) -> Any: - """Receive one ordinary non-KV stage payload on the local leader rank only.""" - tp_group = self._get_local_tp_group() - if tp_group is None or getattr(tp_group, "world_size", 1) <= 1: - return connector.get(from_stage, to_stage, connector_get_key) - if not self.is_data_transfer_rank(): - return None - return connector.get(from_stage, to_stage, connector_get_key) - - def _recv_full_payload_result( - self, - connector: OmniConnectorBase, - from_stage: str, - to_stage: str, - connector_get_key: str, - ) -> Any: - """Receive one full-payload transfer on the local leader rank only.""" - return self._recv_ordinary_stage_result( - connector, - from_stage, - to_stage, - connector_get_key, - ) - - def _recv_async_chunk_result( - self, - connector: OmniConnectorBase, - from_stage: str, - to_stage: str, - connector_get_key: str, - ) -> Any: - """Receive one ordinary async chunk on the local leader rank only.""" - return self._recv_ordinary_stage_result( - connector, - from_stage, - to_stage, - connector_get_key, - ) - - @staticmethod - def _snapshot_payload(payload: Any) -> Any: - if isinstance(payload, dict): - return dict(payload) - return payload - - def _broadcast_tp_payload_packet(self, packet: Any) -> Any: - """Broadcast one ordinary payload packet from TP rank 0 when TP is active.""" - tp_group = self._get_local_tp_group() - if tp_group is None or getattr(tp_group, "world_size", 1) <= 1: - return packet - leader_packet = packet if self.is_data_transfer_rank() else None - return tp_group.broadcast_object(leader_packet, src=0) - - def _apply_staged_payloads_locked(self, staged_payloads: dict[str, Any]) -> None: - for req_id, payload in staged_payloads.items(): - self._local_stage_payload_cache[req_id] = self._snapshot_payload(payload) - - def _collect_full_payload_results_locked(self) -> dict[str, Any] | None: - if not self._full_payload_pending_broadcast_req_ids: - return None - results: dict[str, Any] = {} - missing_req_ids: list[str] = [] - for req_id in tuple(self._full_payload_pending_broadcast_req_ids): - payload = self._local_stage_payload_cache.get(req_id) - if payload is None: - missing_req_ids.append(req_id) - continue - results[req_id] = self._snapshot_payload(payload) - self._full_payload_pending_broadcast_req_ids.discard(req_id) - if missing_req_ids: - logger.warning( - "[Stage-%s] _collect_full_payload_results_locked: " - "pending full-payload reqs missing from local cache: %s", - self._stage_id, - missing_req_ids, - ) - return results or None - - def _collect_async_chunk_fanout_packet_locked(self) -> dict[str, Any] | None: - payload_req_ids = set(self._async_chunk_updated_req_ids) - payload_req_ids.update(self._finished_load_reqs) - payload_req_ids.update(self._chunk_finished_req_ids) - payload_req_ids.update(self._local_request_metadata) - if not ( - payload_req_ids or self._finished_load_reqs or self._chunk_finished_req_ids or self._local_request_metadata - ): - return None - - staged_payloads = { - req_id: self._snapshot_payload(self._local_stage_payload_cache[req_id]) - for req_id in payload_req_ids - if req_id in self._local_stage_payload_cache - } - packet = { - "staged_payloads": staged_payloads, - "request_metadata": dict(self._local_request_metadata), - "newly_finished": set(self._finished_load_reqs), - "chunk_finished": set(self._chunk_finished_req_ids), - } - - self._async_chunk_updated_req_ids.clear() - self._finished_load_reqs.clear() - self._chunk_finished_req_ids.clear() - self._local_request_metadata.clear() - - for req_id in packet["chunk_finished"]: - if req_id not in self._local_stage_payload_cache: - continue - ext_req_id = self._request_ids_mapping.get(req_id, req_id) - self._send_side_request_payload.pop(ext_req_id, None) - if ext_req_id != req_id: - self._send_side_request_payload.pop(req_id, None) - - return packet - - def _apply_async_chunk_fanout_packet(self, packet: dict[str, Any]) -> None: - staged_payloads = packet.get("staged_payloads", {}) - chunk_finished = set(packet.get("chunk_finished", ())) - with self._lock: - self._apply_staged_payloads_locked(staged_payloads) - for req_id in chunk_finished: - self._pending_load_reqs.pop(req_id, None) - self._chunk_stream_completed.add(req_id) - - # ------------------------------------------------------------------ # - # full_payload_mode (recv_full_payload_inputs / send_full_payload_outputs) - # ------------------------------------------------------------------ # - - def recv_full_payload_inputs(self, scheduler_output: Any) -> dict[str, Any] | None: - """Check for incoming full_payload_mode stage inputs (non-blocking). - - Returns a dict mapping ``request_id -> engine_inputs`` for data - that has arrived, or ``None`` if nothing is ready. Stores full - payloads in the local cache and extracts scheduling metadata. - """ - # Fast path: when TP is trivial (no peer ranks waiting on a broadcast) - # and the bg recv thread has not staged anything, skip the lock + TP - # broadcast cycle entirely. _broadcast_tp_payload_packet already - # returns its input unchanged under the same world_size<=1 condition, - # so the original code path was a no-op here on every empty step. - tp_group = self._get_local_tp_group() - if ( - tp_group is None or getattr(tp_group, "world_size", 1) <= 1 - ) and not self._full_payload_pending_broadcast_req_ids: - return None - with self._lock: - results = self._collect_full_payload_results_locked() if self.is_data_transfer_rank() else None - results = self._broadcast_tp_payload_packet(results) - if not results: - return None - with self._lock: - self._stage_recv_req_ids.update(results.keys()) - for req_id in results: - self._pending_load_reqs.pop(req_id, None) - self._apply_staged_payloads_locked(results) - for req_id, payload in results.items(): - self._local_request_metadata[req_id] = self._extract_scheduling_metadata(payload) - logger.debug( - "[Stage-%s] recv_full_payload_inputs: consumed %s reqs: %s, stage_recv_req_ids now=%s", - self._stage_id, - len(results), - list(results.keys()), - self._stage_recv_req_ids, - ) - return results - - def _get_model_config(self) -> Any: - model_config = getattr(self, "model_config", None) - if model_config is not None: - return model_config - return getattr(getattr(self, "vllm_config", None), "model_config", None) - - def _should_accumulate_full_payload_output(self) -> bool: - """Gate send-side full-payload output accumulation only. - - Cached per instance: the result depends only on model_config / - _custom_process_func, both of which are set at init time. Avoid - the per-step dynamic import inside the model decode loop. - """ - if getattr(self, "_omni_connector", None) is None: - # No connector at all: send_full_payload_outputs would no-op. - # Skip the per-step accumulator+build that would otherwise be - # silently discarded. Defends against a terminal stage whose - # custom_process_input_func has a *_full_payload derivative in - # the same module (e.g. dynin stage 2 token2image_to_token2audio - # in pipelines that don't configure any connector at all). - # - # Known limitation: a *terminal-consumer* stage that has a - # connector configured for receiving upstream input is NOT - # caught here -- ``_omni_connector`` is non-None for it, and - # ``_load_custom_func`` may still resolve a ``*_full_payload`` - # derivative from this stage's ``custom_process_input_func``. - # In that case the accumulator builds payloads that - # ``send_full_payload_outputs`` later drops via its own - # connector-side checks (wasted CPU, not a functional bug). - # A topology-aware gate (explicit producer field or pipeline - # is_terminal info) would close the gap; that change is out - # of scope for this PR. - self._should_accumulate_full_payload_output_cached = False - return False - cached = getattr(self, "_should_accumulate_full_payload_output_cached", None) - if cached is not None: - return cached - model_config = self._get_model_config() - if model_config is None: - self._should_accumulate_full_payload_output_cached = False - return False - result = should_accumulate_full_payload_output( - model_config, - getattr(self, "_custom_process_func", None), - ) - self._should_accumulate_full_payload_output_cached = result - return result - - @staticmethod - def _new_full_payload_accumulator(output: dict[str, Any]): - chunks: dict[str, list[torch.Tensor]] = {} - latest: dict[str, Any] = {} - rows: dict[str, int] = {} - for k, v in output.items(): - if isinstance(v, torch.Tensor) and v.dim() >= 2: - chunks[k] = [v] - rows[k] = int(v.shape[0]) - else: - latest[k] = v - return chunks, latest, rows - - @staticmethod - def _materialize_full_payload_entry(entry): - if len(entry) == 2: - return entry - chunks, latest, _rows, request = entry - output = dict(latest) - for k, tensors in chunks.items(): - if tensors: - output[k] = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0) - return output, request - - def _resolve_full_payload_replace_keys(self) -> frozenset: - """Per-model REPLACE-key set for the full-payload accumulator. - - Looked up from the stage-input-processor module that ships the model's sync builder - (`model_config.custom_process_input_func.__module__`). The module - declares ``_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str]``; if absent, - returns the empty set. - - Cached per instance. Keys in this set use REPLACE semantics in the - accumulator (subsequent emissions discard prior chunks) instead of - the default CONCAT semantics. Use for tensors that carry the full - result so far rather than per-step deltas (e.g. ``model_outputs``). - """ - cached = getattr(self, "_full_payload_replace_keys_cached", None) - if cached is not None: - return cached - proc = getattr(self, "_custom_process_func", None) - if proc is None: - self._full_payload_replace_keys_cached = frozenset() - return self._full_payload_replace_keys_cached - module_name = getattr(proc, "__module__", None) - if module_name is None: - self._full_payload_replace_keys_cached = frozenset() - return self._full_payload_replace_keys_cached - try: - import sys as _sys - - mod = _sys.modules.get(module_name) or importlib.import_module(module_name) - keys = getattr(mod, "_FULL_PAYLOAD_REPLACE_KEYS", frozenset()) - except ImportError: - logger.debug( - "Could not import stage input processor module %s while resolving " - "_FULL_PAYLOAD_REPLACE_KEYS; using CONCAT semantics for all keys.", - module_name, - exc_info=True, - ) - keys = frozenset() - if not isinstance(keys, (frozenset, set)): - logger.debug( - "Ignoring non-set _FULL_PAYLOAD_REPLACE_KEYS from %s: %s", - module_name, - type(keys).__name__, - ) - keys = frozenset() - self._full_payload_replace_keys_cached = frozenset(keys) - logger.debug( - "Resolved _FULL_PAYLOAD_REPLACE_KEYS for %s: %s", - module_name, - sorted(self._full_payload_replace_keys_cached), - ) - return self._full_payload_replace_keys_cached - - def accumulate_full_payload_output( - self, - req_id: str, - pooler_output: Any, - request: Any, - ) -> None: - """Accumulate pooler_output for a request across steps (full_payload_mode). - - Per-token tensors (2-D+, matching trailing dims) are concatenated - along dim-0. Scalar / global tensors (1-D or 0-D) are replaced - with the latest value. - - Note: codec rows are NOT filtered for zero placeholders here. The - downstream consumer ``_extract_qwen3_full_payload_codec_rows`` crops - codec rows using ``output_token_ids`` as the authoritative source, - which makes any sender-side zero filtering redundant. Skipping the - sender-side ``t.any()`` scan also avoids a per-tensor GPU->CPU device - sync that stalled the decode pipeline. - - The data is actually sent when ``flush_full_payload_outputs`` is called - with the finished request IDs from the next scheduler cycle. - """ - replace_keys = self._resolve_full_payload_replace_keys() - existing = self._pending_full_payload_send.get(req_id) - - if existing is None: - chunks, latest, rows = self._new_full_payload_accumulator(pooler_output) - self._pending_full_payload_send[req_id] = (chunks, latest, rows, request) - return - - if len(existing) == 2: - chunks, latest, rows = self._new_full_payload_accumulator(existing[0]) - else: - chunks, latest, rows, _ = existing - - for k, v in pooler_output.items(): - if v is None: - continue - if k in replace_keys: - # Explicit REPLACE semantics: the new value supersedes any - # prior chunks (e.g. `model_outputs` carries the full result - # so far, not an appendable per-step delta). - latest.pop(k, None) - if isinstance(v, torch.Tensor) and v.dim() >= 2: - chunks[k] = [v] - rows[k] = int(v.shape[0]) - else: - chunks.pop(k, None) - rows.pop(k, None) - latest[k] = v - continue - if isinstance(v, torch.Tensor) and v.dim() >= 2: - if k in chunks and chunks[k] and v.shape[1:] == chunks[k][0].shape[1:]: - chunks[k].append(v) - rows[k] += int(v.shape[0]) - else: - latest.pop(k, None) - chunks[k] = [v] - rows[k] = int(v.shape[0]) - else: - chunks.pop(k, None) - rows.pop(k, None) - latest[k] = v - - self._pending_full_payload_send[req_id] = (chunks, latest, rows, request) - - def flush_full_payload_outputs(self, finished_req_ids: set[str]) -> None: - """Send accumulated full_payload outputs for requests that just finished.""" - pending_req_ids = set(self._pending_full_payload_send.keys()) - if not (finished_req_ids & pending_req_ids): - return - - logger.debug( - "[Stage-%s] flush_full_payload_outputs: finished_req_ids=%s, pending=%s", - self._stage_id, - finished_req_ids, - list(self._pending_full_payload_send.keys()), - ) - to_send: dict[str, tuple[Any, Any]] = {} - for req_id in finished_req_ids: - entry = self._pending_full_payload_send.pop(req_id, None) - if entry is not None: - to_send[req_id] = self._materialize_full_payload_entry(entry) - logger.debug("[Stage-%s] flush_full_payload_outputs: to_send=%s", self._stage_id, list(to_send.keys())) - if to_send: - self.send_full_payload_outputs(scheduler_output=None, outputs=to_send) - - def send_full_payload_outputs( - self, - scheduler_output: Any, - outputs: dict[str, tuple[Any, Any] | Any], - ) -> list[str]: - """Send full_payload stage outputs to the next stage via connector. - - Args: - outputs: Mapping of ``req_id`` to either a - ``(pooling_output, request)`` tuple (preferred) or a raw - payload dict. When a tuple is supplied the request object - is forwarded to ``custom_process_stage_input_func``. - - Returns list of request IDs successfully enqueued. - """ - if self._omni_connector is None: - logger.debug("[Stage-%s] send_full_payload_outputs: connector is None, skip", self._stage_id) - return [] - if not self.is_data_transfer_rank(): - logger.debug( - "[Stage-%s] send_full_payload_outputs: not data_transfer_rank (rank=%s), skip", - self._stage_id, - self._local_rank, - ) - return list(outputs.keys()) - sent_ids: list[str] = [] - next_stage_id = self._next_stage_id - for req_id, value in outputs.items(): - if isinstance(value, tuple) and len(value) == 2: - raw_output, request = value - else: - raw_output, request = value, None - - payload = raw_output - if self._custom_process_func is not None: - payload = self._build_custom_process_payload( - request_id=req_id, - request=request, - pooling_output=raw_output, - ) - if payload is None: - continue - if payload is None: - logger.debug("[Stage-%s] send_full_payload_outputs: payload is None for %s", self._stage_id, req_id) - continue - if isinstance(payload, dict): - audio_codes = self._payload_audio_codes(payload) - if isinstance(audio_codes, torch.Tensor): - code_len = int(audio_codes.numel()) - elif hasattr(audio_codes, "__len__"): - code_len = len(audio_codes) - else: - code_len = None - meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {} - logger.debug( - "[Stage-%s] send_full_payload_outputs: req=%s payload_keys=%s code_len=%s left_context_size=%s", - self._stage_id, - req_id, - sorted(payload.keys()), - code_len, - meta.get("left_context_size"), - ) - - external_req_id = self._resolve_external_req_id(request, req_id) - chunk_id = self._put_req_chunk[req_id] - self._put_req_chunk[req_id] += 1 - self._ramp_chunk_count[req_id] += 1 - connector_put_key = f"{external_req_id}_{self._stage_id}_{chunk_id}" - - logger.debug( - "[Stage-%s] send_full_payload_outputs: enqueue req=%s put_key=%s next_stage=%s", - self._stage_id, - req_id, - connector_put_key, - next_stage_id, - ) - task = { - "stage_id": self._stage_id, - "next_stage_id": next_stage_id, - "put_key": connector_put_key, - "data": payload, - "request_id": req_id, - } - with self._lock: - self._pending_save_reqs.setdefault(req_id, deque()).append(task) - self._pending_save_counts[req_id] += 1 - sent_ids.append(req_id) - if sent_ids: - self._work_available.set() - return sent_ids - - # ------------------------------------------------------------------ # - # Streaming chunk mode (recv_chunk / send_chunk) - # ------------------------------------------------------------------ # - - def register_chunk_recv(self, request: Any) -> None: - """Register a request for async chunk retrieval by the bg thread. - - Stage-0 has no upstream producer so this is a no-op there. - Skips requests whose batch data has already been received to - prevent the bg thread from polling for non-existent chunks. - """ - if self._stage_id == 0: - return - request_id = request.request_id - # Explicit external_req_id=None must fall back to request_id; - # otherwise recv keys become `None__` and collide - # across requests. - ext = getattr(request, "external_req_id", None) - self._request_ids_mapping[request_id] = ext if ext is not None else request_id - with self._lock: - if request_id in self._stage_recv_req_ids: - return - # Don't re-register if the finish sentinel was already received - if request_id in self._chunk_stream_completed: - return - self._pending_load_reqs[request_id] = request - self._work_available.set() - - def recv_chunk(self) -> dict[str, Any]: - """Collect chunks received by the bg thread since last call. - - Returns a dict ``{request_id: chunk_payload}`` for newly arrived - chunks. Empty dict when nothing is ready. - - This method reads from ``_finished_load_reqs`` without clearing - it -- ``get_omni_connector_output()`` is the sole consumer that - drains and resets ``_finished_load_reqs`` at the end of each - ``execute_model`` cycle. - - Returns **shallow copies** of the cached payloads so that the - caller can read them without racing against the background recv - thread, which may concurrently mutate the live cache entries via - ``dict.update()``. - """ - with self._lock: - finished = set(self._finished_load_reqs) - if not finished: - return {} - # Snapshot the payloads under the lock to avoid racing with - # _poll_single_request which does existing.update(payload_data) - # on the same dict objects. - result = {} - for rid in finished: - payload = self._local_stage_payload_cache.get(rid) - result[rid] = dict(payload) if isinstance(payload, dict) else payload - - self._chunk_ready_req_ids.update(finished) - return result - - def send_chunk( - self, - request: Any, - pooling_output: Any | None = None, - ) -> bool: - """Derive and enqueue one chunk for async sending. - - Payload extraction runs in the caller thread (via - ``custom_process_stage_input_func``); the actual - ``connector.put()`` is done by the background save thread. - Non-KV data is identical across TP ranks; only rank 0 sends. - """ - if self._omni_connector is None: - logger.warning("[Stage-%s] send_chunk: connector is None", self._stage_id) - return False - if not self.is_data_transfer_rank(): - return True - raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None) - request_id = self._resolve_external_req_id(request, raw_req_id) - # Cache the internal→external mapping so that finish sentinels can - # resolve the external ID even after the request is freed. - if raw_req_id and raw_req_id != request_id: - self._request_ids_mapping.setdefault(raw_req_id, request_id) - chunk_id = self._put_req_chunk[request_id] - - payload_data = self._build_custom_process_payload( - request_id=request_id, - request=request, - pooling_output=pooling_output, - ) - if payload_data is None: - if chunk_id == 0: - logger.warning( - "[Stage-%s] send_chunk: payload is None for req=%s chunk=%s (process_func=%s)", - self._stage_id, - request_id, - chunk_id, - self._custom_process_func, - ) - return False - - self._put_req_chunk[request_id] += 1 - self._ramp_chunk_count[request_id] += 1 - next_stage_id = self._next_stage_id - connector_put_key = f"{request_id}_{self._stage_id}_{chunk_id}" - - if chunk_id == 0: - logger.debug( - "[Stage-%s] send_chunk: first chunk enqueued, req=%s key=%s", - self._stage_id, - request_id, - connector_put_key, - ) - - task = { - "stage_id": self._stage_id, - "next_stage_id": next_stage_id, - "put_key": connector_put_key, - "data": payload_data, - "request_id": request_id, - } - with self._lock: - self._pending_save_reqs.setdefault(request_id, deque()).append(task) - self._pending_save_counts[request_id] += 1 - self._work_available.set() - return True - - # ------------------------------------------------------------------ # - # KV cache (delegates to OmniKVTransferManager) - # ------------------------------------------------------------------ # - - def send_kv_cache( - self, - finished_reqs: dict[str, dict[str, Any]], - kv_caches: list[torch.Tensor], - block_size: int, - cache_dtype: str, - request_id_resolver: Any | None = None, - ) -> list[str]: - """Send KV cache for finished requests. - - Delegates to the existing ``OmniKVTransferManager``. - """ - if self._kv_transfer_manager is None: - return list(finished_reqs.keys()) if finished_reqs else [] - result = self._kv_transfer_manager.handle_finished_requests_kv_transfer( - finished_reqs=finished_reqs, - kv_caches=kv_caches, - block_size=block_size, - cache_dtype=cache_dtype, - request_id_resolver=request_id_resolver, - ) - if result: - self._kv_sent_req_ids.extend(result) - return result - - def recv_kv_cache( - self, - request_id: str, - target_device: torch.device | None = None, - ) -> tuple[dict[str, Any] | None, int]: - """Receive KV cache for a request. - - Delegates to the existing ``OmniKVTransferManager``. - """ - if self._kv_transfer_manager is None: - return None, 0 - return self._kv_transfer_manager.receive_kv_cache_for_request( - request_id=request_id, - target_device=target_device, - ) - - def receive_cfg_companion_kv_payloads( - self, - cfg_request_ids: dict[str, str], - target_device: torch.device | None = None, - ) -> dict[str, tuple[dict[str, Any] | None, int]]: - """Receive raw CFG companion KV payloads keyed by role.""" - return { - role: self.recv_kv_cache(companion_rid, target_device=target_device) - for role, companion_rid in cfg_request_ids.items() - } - - def receive_multi_kv_cache( - self, - req: Any, - cfg_kv_collect_func: Any | None = None, - target_device: torch.device | None = None, - ) -> bool: - """Receive primary and optional companion KV caches for a request. - - The mixin owns the runner-facing orchestration: primary KV receive, - companion payload fetch, and applying any model-specific CFG fields back - onto ``req.sampling_params``. - """ - if self._kv_transfer_manager is None: - return False - - request_id = getattr(req, "request_id", None) - if not request_id: - logger.warning("Request has no ID, cannot receive KV cache") - return False - - active_requests = getattr(self, "requests", None) - if active_requests is not None and request_id not in active_requests: - logger.debug("Skip receiving KV cache for inactive request %s", request_id) - return False - - primary_ok = False - data, _size = self.recv_kv_cache(request_id, target_device=target_device) - if data: - self._kv_transfer_manager.apply_kv_cache_to_request(req, data) - primary_ok = True - - cfg_ids = getattr(getattr(req, "sampling_params", None), "cfg_kv_request_ids", None) - if cfg_ids and cfg_kv_collect_func: - try: - cfg_role_payloads = self.receive_cfg_companion_kv_payloads( - cfg_ids, - target_device=target_device, - ) - cfg_kvs = cfg_kv_collect_func(request_id, cfg_role_payloads) - if cfg_kvs and hasattr(req, "sampling_params") and req.sampling_params is not None: - for key, value in cfg_kvs.items(): - setattr(req.sampling_params, key, value) - logger.debug("Applied CFG KV caches: %s", list(cfg_kvs.keys())) - except Exception: - logger.exception("Failed to collect CFG KV caches for %s", request_id) - - return primary_ok - - # ------------------------------------------------------------------ # - # Rank-aware KV transfer routing - # ------------------------------------------------------------------ # - - def get_rank_aware_kv_keys( - self, - req_id: str, - from_stage: int, - to_stage: int | None = None, - chunk_id: int = 0, - ) -> list[str]: - """Build recv-side connector keys for all remote ranks this rank needs. - - For heterogeneous TP receive, the local rank is the target rank and must - fetch one or more source-rank shards keyed as ``from_rank -> to_rank``. - """ - if self._from_tp <= 1 and self._to_tp <= 1: - resolved_to_stage = self._next_stage_id if to_stage is None else to_stage - return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] - - remote_ranks = self.get_kv_remote_ranks() - return [ - self.get_kv_connector_key( - req_id=req_id, - from_stage=from_stage, - chunk_id=chunk_id, - from_rank=remote_rank, - to_rank=self._local_rank, - ) - for remote_rank in remote_ranks - ] - - def get_kv_target_ranks_for_send(self) -> list[int]: - """Determine which target ranks this local rank should send KV shards to.""" - self._validate_kv_tp_topology() - if self._from_tp == self._to_tp: - return [self._local_rank] - if self._from_tp > self._to_tp: - tp_ratio = self._from_tp // self._to_tp - return [self._local_rank // tp_ratio] - tp_ratio = self._to_tp // self._from_tp - base_rank = self._local_rank * tp_ratio - return [base_rank + i for i in range(tp_ratio)] - - def get_rank_aware_kv_send_keys( - self, - req_id: str, - from_stage: int, - to_stage: int | None = None, - chunk_id: int = 0, - ) -> list[str]: - """Build send-side connector keys for this rank's KV shard(s).""" - if self._from_tp <= 1 and self._to_tp <= 1: - resolved_to_stage = self._next_stage_id if to_stage is None else to_stage - return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] - - target_ranks = self.get_kv_target_ranks_for_send() - return [ - self.get_kv_connector_key( - req_id=req_id, - from_stage=from_stage, - chunk_id=chunk_id, - from_rank=self._local_rank, - to_rank=target_rank, - ) - for target_rank in target_ranks - ] - - @staticmethod - def _merge_rank_sharded_kv_payloads(payloads: list[dict[str, Any]]) -> dict[str, Any] | None: - """Merge multiple source-rank KV shards for one target rank.""" - payloads = [payload for payload in payloads if isinstance(payload, dict)] - if not payloads: - return None - if len(payloads) == 1: - return payloads[0] - - merged = dict(payloads[0]) - layer_blocks = merged.get("layer_blocks") - if not isinstance(layer_blocks, dict): - return merged - - def _merge_tensor_lists(name: str) -> list[torch.Tensor | None]: - merged_list: list[torch.Tensor | None] = [] - cache_lists = [payload.get("layer_blocks", {}).get(name, []) for payload in payloads] - max_len = max((len(cache_list) for cache_list in cache_lists), default=0) - for idx in range(max_len): - tensors = [cache_list[idx] for cache_list in cache_lists if idx < len(cache_list)] - tensors = [tensor for tensor in tensors if isinstance(tensor, torch.Tensor)] - if not tensors: - merged_list.append(None) - elif len(tensors) == 1: - merged_list.append(tensors[0]) - else: - merged_list.append(torch.cat(tensors, dim=-2).contiguous()) - return merged_list - - merged["layer_blocks"] = { - "key_cache": _merge_tensor_lists("key_cache"), - "value_cache": _merge_tensor_lists("value_cache"), - } - metadata = dict(merged.get("metadata", {})) - metadata["merged_remote_rank_count"] = len(payloads) - merged["metadata"] = metadata - return merged - - def _slice_rank_sharded_kv_payload(self, payload: dict[str, Any] | None) -> dict[str, Any] | None: - """Slice a duplicated source-rank KV shard for ``from_tp < to_tp`` cases.""" - if payload is None or self._from_tp >= self._to_tp: - return payload - - tp_ratio = self._to_tp // self._from_tp - shard_index = self._local_rank % tp_ratio - layer_blocks = payload.get("layer_blocks") if isinstance(payload, dict) else None - if not isinstance(layer_blocks, dict): - return payload - - def _slice_tensor_list(name: str) -> list[torch.Tensor | None]: - sliced: list[torch.Tensor | None] = [] - for tensor in layer_blocks.get(name, []): - if not isinstance(tensor, torch.Tensor) or tensor.ndim < 2: - sliced.append(tensor) - continue - head_dim = tensor.shape[-2] - if head_dim % tp_ratio != 0: - sliced.append(tensor) - continue - per_rank = head_dim // tp_ratio - start = shard_index * per_rank - sliced.append(tensor.narrow(-2, start, per_rank).contiguous()) - return sliced - - payload = dict(payload) - payload["layer_blocks"] = { - "key_cache": _slice_tensor_list("key_cache"), - "value_cache": _slice_tensor_list("value_cache"), - } - metadata = dict(payload.get("metadata", {})) - metadata["sliced_for_local_rank"] = self._local_rank - payload["metadata"] = metadata - return payload - - def should_replicate_payload(self) -> bool: - """Whether non-KV payloads should be replicated across ranks. - - Data payloads (stage inputs, chunks) are identical after all-gather, - so only rank 0 transfers them. KV payloads are rank-specific and - all ranks participate. - """ - return self._local_rank != 0 - - def get_kv_rank_mapping(self) -> dict[str, Any]: - """Return the current rank mapping configuration. - - Useful for debugging and for downstream code that needs to know - the TP topology without re-parsing model config. - """ - return { - "from_tp": self._from_tp, - "to_tp": self._to_tp, - "local_rank": self._local_rank, - "remote_ranks": self.get_kv_remote_ranks(), - "is_data_transfer_rank": self.is_data_transfer_rank(), - } - - # ------------------------------------------------------------------ # - # KV transfer lifecycle (RFC – mixin-owned) - # ------------------------------------------------------------------ # - - def mark_kv_transfer( - self, - req_id: str, - seq_len: int, - block_ids: list[int], - custom_metadata: dict[str, Any] | None = None, - ) -> None: - """Mark a request as needing KV cache transfer. - - Called by the scheduler when a transfer trigger fires. The mixin - owns the lifecycle from this point: pending → active → completed. - """ - if req_id in self._kv_pending_transfers: - return - self._kv_triggered_requests.add(req_id) - transfer = { - "seq_len": seq_len, - "block_ids": block_ids, - } - if custom_metadata is not None: - transfer["custom_metadata"] = custom_metadata - self._kv_pending_transfers[req_id] = transfer - - def drain_pending_kv_transfers(self) -> dict[str, dict[str, Any]]: - """Drain pending KV transfers and move them to active. - - Returns ``{req_id: {seq_len, block_ids}}`` for the model runner - to submit to ``send_kv_cache``. - """ - if not self._kv_pending_transfers: - return {} - pending = dict(self._kv_pending_transfers) - self._kv_active_transfers.update(pending.keys()) - self._kv_pending_transfers.clear() - return pending - - def ack_kv_transfers(self, req_ids: list[str] | set[str]) -> None: - """Acknowledge completed KV transfers (from kv_extracted_req_ids). - - Moves requests from active to completed so the scheduler can - safely free their blocks. - """ - for req_id in req_ids: - self._kv_active_transfers.discard(req_id) - self._kv_completed_transfers.add(req_id) - - def drain_completed_kv_transfers(self) -> set[str]: - """Drain and return completed KV transfer request IDs. - - The scheduler calls this to know which requests' blocks can be freed. - """ - completed = set(self._kv_completed_transfers) - self._kv_completed_transfers.clear() - return completed - - def is_kv_transfer_triggered(self, req_id: str) -> bool: - """Check if a request has already triggered KV transfer.""" - return req_id in self._kv_triggered_requests - - def has_pending_kv_work(self) -> bool: - """True if any KV transfers are pending, active, or awaiting ack.""" - return bool(self._kv_pending_transfers or self._kv_active_transfers or self._kv_completed_transfers) - - # Output aggregation - # ------------------------------------------------------------------ # - - def get_omni_connector_output(self) -> OmniConnectorOutput: - """Collect and reset transfer results for this execute_model cycle. - - ``request_metadata`` carries only lightweight scheduling metadata. - Full payloads remain owned by the Model Runner local cache for all - paths. - """ - if not hasattr(self, "_lock"): - return OmniConnectorOutput() - - tp_group = self._get_local_tp_group() - if self._async_chunk and tp_group is not None and getattr(tp_group, "world_size", 1) > 1: - if self.is_data_transfer_rank(): - with self._lock: - fanout_packet = self._collect_async_chunk_fanout_packet_locked() - else: - fanout_packet = None - fanout_packet = self._broadcast_tp_payload_packet(fanout_packet) - if fanout_packet is None: - newly_finished = set() - chunk_finished = set() - request_metadata = {} - else: - if not self.is_data_transfer_rank(): - self._apply_async_chunk_fanout_packet(fanout_packet) - newly_finished = set(fanout_packet["newly_finished"]) - chunk_finished = set(fanout_packet["chunk_finished"]) - request_metadata = dict(fanout_packet["request_metadata"]) - else: - with self._lock: - newly_finished = set(self._finished_load_reqs) - self._finished_load_reqs.clear() - chunk_finished = set(self._chunk_finished_req_ids) - self._chunk_finished_req_ids.clear() - request_metadata = dict(self._local_request_metadata) - self._local_request_metadata.clear() - # _send_side_request_payload is the async accumulation buffer for - # future recv chunks. Clearing it on every consumable wake-up drops - # intermediate - # thinker decode spans before the model side can consume them. - # Only terminal chunk_finished requests may release that buffer. - for req_id in chunk_finished: - if req_id not in self._local_stage_payload_cache: - continue - ext_req_id = self._request_ids_mapping.get(req_id, req_id) - self._send_side_request_payload.pop(ext_req_id, None) - if ext_req_id != req_id: - self._send_side_request_payload.pop(req_id, None) - self._chunk_ready_req_ids.update(newly_finished) - - output = OmniConnectorOutput( - chunk_ready_req_ids=set(self._chunk_ready_req_ids), - chunk_finished_req_ids=chunk_finished, - request_metadata=request_metadata, - kv_sent_req_ids=list(self._kv_sent_req_ids), - stage_recv_req_ids=set(self._stage_recv_req_ids), - has_pending_kv_work=self.has_pending_kv_work(), - ) - if output.stage_recv_req_ids or chunk_finished or newly_finished: - logger.debug( - "[Stage-%s] get_omni_connector_output: stage_recv=%s, chunk_finished=%s, chunk_ready=%s", - self._stage_id, - output.stage_recv_req_ids, - chunk_finished, - output.chunk_ready_req_ids, - ) - self._chunk_ready_req_ids.clear() - self._kv_sent_req_ids.clear() - self._stage_recv_req_ids.clear() - return output - - @staticmethod - def _connector_output_has_signals(output: OmniConnectorOutput) -> bool: - return bool( - output.chunk_ready_req_ids - or output.chunk_finished_req_ids - or output.request_metadata - or output.kv_sent_req_ids - or output.stage_recv_req_ids - or output.has_pending_kv_work - ) - - def attach_omni_connector_output(self, result: Any | None) -> Any: - omni_output = self.get_omni_connector_output() - if not self._connector_output_has_signals(omni_output): - return result - - from copy import copy - - from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT - - wrapped = copy(result if result is not None else EMPTY_MODEL_RUNNER_OUTPUT) - wrapped.omni_connector_output = omni_output - return wrapped - - # ------------------------------------------------------------------ # - # Properties for compatibility with custom_process funcs that access - # transfer_manager.put_req_chunk / request_payload / code_prompt_token_ids - # ------------------------------------------------------------------ # - - @property - def put_req_chunk(self) -> dict[str, int]: - return self._put_req_chunk - - @property - def ramp_chunk_count(self) -> dict[str, int]: - return self._ramp_chunk_count - - @property - def request_payload(self) -> dict[str, dict[str, Any]]: - return self._send_side_request_payload - - @request_payload.setter - def request_payload(self, value: dict[str, dict[str, Any]]) -> None: - self._send_side_request_payload = value - - @property - def code_prompt_token_ids(self) -> dict[str, list[list[int]]]: - return self._code_prompt_token_ids - - @property - def connector(self) -> Any | None: - return self._omni_connector - - # ------------------------------------------------------------------ # - # Background I/O threads - # ------------------------------------------------------------------ # - - def _recv_loop(self) -> None: - """Background thread: poll connector for incoming data.""" - _recv_poll_count = 0 - while not self._stop_event.is_set(): - with self._lock: - pending_ids = list(self._pending_load_reqs.keys()) - - if not pending_ids: - self._work_available.wait(timeout=0.01) - self._work_available.clear() - continue - - _recv_poll_count += 1 - if _recv_poll_count % 5000 == 1: - logger.debug( - "[Stage-%s] _recv_loop: polling %s pending reqs: %s (poll#%s)", - self._stage_id, - len(pending_ids), - pending_ids[:5], - _recv_poll_count, - ) - - made_progress = False - for req_id in pending_ids: - if self._stop_event.is_set(): - break - try: - made_progress = self._poll_single_request(req_id) or made_progress - except Exception: - logger.warning("Error receiving data for %s", req_id, exc_info=True) - - if not made_progress and not self._stop_event.is_set(): - self._work_available.wait(timeout=0.005) - self._work_available.clear() - - _MAX_SEND_RETRIES = 3 - - def _save_loop(self) -> None: - """Background thread: send outgoing data via connector.""" - while not self._stop_event.is_set(): - task = None - with self._lock: - for req_id in list(self._pending_save_reqs.keys()): - dq = self._pending_save_reqs[req_id] - if dq: - task = dq.popleft() - if not dq: - del self._pending_save_reqs[req_id] - break - del self._pending_save_reqs[req_id] - - if task is not None: - success = False - try: - success = self._send_single_request(task) - except Exception: - logger.error( - "Error saving data for %s", - task.get("request_id"), - exc_info=True, - ) - if not success: - self._requeue_or_drop_failed_send(task) - continue - - self._work_available.wait(timeout=0.01) - self._work_available.clear() - - def _requeue_or_drop_failed_send(self, task: dict) -> None: - """Re-enqueue a failed send task or drop it after max retries.""" - retry_count = task.get("_retry_count", 0) + 1 - req_id = task.get("request_id") - if retry_count <= self._MAX_SEND_RETRIES: - task["_retry_count"] = retry_count - logger.warning( - "[Stage-%s] Re-enqueuing failed send for %s (retry %d/%d)", - getattr(self, "_stage_id", "?"), - req_id, - retry_count, - self._MAX_SEND_RETRIES, - ) - with self._lock: - dq = self._pending_save_reqs.setdefault(req_id, deque()) - dq.appendleft(task) - else: - logger.error( - "[Stage-%s] Giving up on send for %s after %d retries", - getattr(self, "_stage_id", "?"), - req_id, - self._MAX_SEND_RETRIES, - ) - self._decrement_pending_save_count(req_id) - - # ------------------------------------------------------------------ # - # Chunk-level poll / send (ported from OmniChunkTransferAdapter) - # ------------------------------------------------------------------ # - - def _poll_single_request(self, req_id: str) -> bool: - """Poll connector for one chunk of a request (non-blocking).""" - connector = self._omni_connector - if connector is None: - return False - - if self._async_chunk and self._model_mode != "ar": - with self._lock: - staged_payload = self._local_stage_payload_cache.get(req_id) - metadata_in_flight = req_id in self._local_request_metadata - scheduler_wakeup_pending = req_id in self._finished_load_reqs - if self._payload_is_consumable(staged_payload) or metadata_in_flight or scheduler_wakeup_pending: - logger.debug( - "[Stage-%s] delaying recv for req=%s until staged async payload is handed to scheduler", - self._stage_id, - req_id, - ) - return False - - target_stage_id = self._stage_id - 1 - chunk_id = self._get_req_chunk[req_id] - external_req_id = self._request_ids_mapping.get(req_id, req_id) - connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" - - if self._async_chunk: - result = self._recv_async_chunk_result( - connector, - str(target_stage_id), - str(self._stage_id), - connector_get_key, - ) - else: - result = self._recv_full_payload_result( - connector, - str(target_stage_id), - str(self._stage_id), - connector_get_key, - ) - - if result is None: - return False - - payload_data, _size = result - if not payload_data: - return False - if isinstance(payload_data, dict): - logger.debug( - "[Stage-%s] recv_chunk_result: req=%s ext=%s key=%s keys=%s finished=%s", - self._stage_id, - req_id, - external_req_id, - connector_get_key, - sorted(payload_data.keys()), - self._payload_finished(payload_data), - ) - - self._get_req_chunk[req_id] += 1 - - if self._async_chunk: - is_finished = self._payload_finished(payload_data) - incoming_payload_consumable = self._payload_is_consumable(payload_data) - - if self._model_mode == "ar": - payload_data = self._accumulate_payload(external_req_id, payload_data) - payload_consumable = incoming_payload_consumable - else: - new_ids = self._payload_audio_codes(payload_data) or [] - if not new_ids and not is_finished: - return False - payload_consumable = self._payload_is_consumable(payload_data) - - with self._lock: - if is_finished: - self._chunk_finished_req_ids.add(req_id) - self._chunk_stream_completed.add(req_id) - # Local cache (RFC §2.4) — merge, don't replace, so that - # earlier chunk keys (e.g. thinker_prefill_embeddings from - # chunk 0) are not overwritten by later chunks. - existing = self._local_stage_payload_cache.get(req_id) - if existing is not None and isinstance(existing, dict) and isinstance(payload_data, dict): - existing.update(payload_data) - else: - self._local_stage_payload_cache[req_id] = payload_data - staged_payload = self._local_stage_payload_cache[req_id] - self._async_chunk_updated_req_ids.add(req_id) - self.put_local_request_metadata(req_id, self._extract_scheduling_metadata(staged_payload)) - # A finish-only sentinel still needs one terminal wake-up so - # the downstream stage can sync the merged local payload and - # flush/finish even when the last recv carries no new - # consumable chunk bytes. - if payload_consumable or is_finished: - self._finished_load_reqs.add(req_id) - if is_finished and not payload_consumable: - logger.debug( - "[Stage-%s] finish sentinel arrived for req=%s without new consumable payload", - self._stage_id, - req_id, - ) - elif not payload_consumable: - logger.debug( - "[Stage-%s] req=%s received metadata-only / non-consumable async payload; delaying wake-up", - self._stage_id, - req_id, - ) - if is_finished: - self._pending_load_reqs.pop(req_id, None) - else: - # full_payload_mode: the complete payload arrives in a single get(), - # so always unregister immediately. - if isinstance(payload_data, dict): - engine_inputs = payload_data.get("engine_inputs", payload_data) - else: - engine_inputs = payload_data - with self._lock: - self._local_stage_payload_cache[req_id] = self._snapshot_payload(engine_inputs) - # Publish full-payload readiness only after the aligned TP broadcast - # path in recv_full_payload_inputs() has materialized the payload on all - # local ranks. Publishing metadata / stage_recv from the background recv - # thread can let the scheduler observe a request before the payload is - # actually visible to the model thread. - self._full_payload_pending_broadcast_req_ids.add(req_id) - self._pending_load_reqs.pop(req_id, None) - logger.debug( - "[Stage-%s] full_payload recv complete: req=%s key=%s payload_type=%s", - self._stage_id, - req_id, - connector_get_key, - type(engine_inputs).__name__, - ) - - logger.debug("[Stage-%s] Received data for key %s", self._stage_id, connector_get_key) - return True - - def _build_custom_process_payload( - self, - request_id: str | None, - request: Any | None, - pooling_output: Any | None, - ) -> Any | None: - """Run the custom process hook with a best-effort finished kwarg.""" - if self._custom_process_func is None: - return None - - kwargs = { - "transfer_manager": self, - "pooling_output": pooling_output, - "request": request, - } - supports_is_finished = getattr( - self, - "_custom_process_supports_is_finished", - self._custom_process_supports_is_finished_kwarg(), - ) - is_finished_fn = getattr(request, "is_finished", None) - if callable(is_finished_fn): - try: - if supports_is_finished is not False: - kwargs["is_finished"] = bool(is_finished_fn()) - except Exception: - logger.debug("request.is_finished() failed for %s", request_id, exc_info=True) - - try: - return self._custom_process_func(**kwargs) - except TypeError as exc: - if "is_finished" not in kwargs or not self._is_unexpected_is_finished_kwarg_error(exc): - logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) - return None - kwargs.pop("is_finished", None) - try: - return self._custom_process_func(**kwargs) - except Exception: - logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) - return None - except Exception: - logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) - return None - - def _custom_process_supports_is_finished_kwarg(self) -> bool | None: - """Return whether the custom process hook accepts `is_finished`.""" - if self._custom_process_func is None: - return None - try: - signature = inspect.signature(self._custom_process_func) - except (TypeError, ValueError): - return None - - for param in signature.parameters.values(): - if param.kind == inspect.Parameter.VAR_KEYWORD: - return True - - is_finished_param = signature.parameters.get("is_finished") - if is_finished_param is None: - return False - return is_finished_param.kind in ( - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - ) - - @staticmethod - def _is_unexpected_is_finished_kwarg_error(exc: TypeError) -> bool: - message = str(exc) - return ( - "unexpected keyword argument 'is_finished'" in message - or 'unexpected keyword argument "is_finished"' in message - or "positional-only arguments passed as keyword arguments: 'is_finished'" in message - ) - - def _send_single_request(self, task: dict) -> bool: - """Send one queued task via connector.put(). - - Returns True on success. On failure (put() raises or returns - ``success=False``), returns False **without** decrementing - ``_pending_save_counts`` so the caller can retry or clean up. - """ - connector = self._omni_connector - if connector is None: - return True - - request_id = task.get("request_id") - payload_data = task.get("data") - if payload_data is None and task.get("request") is not None: - payload_data = self._build_custom_process_payload( - request_id=request_id, - request=task.get("request"), - pooling_output=task.get("pooling_output"), - ) - put_key = task.get("put_key") - - success, _size, _metadata = connector.put( - from_stage=str(task["stage_id"]), - to_stage=str(task["next_stage_id"]), - put_key=put_key, - data=payload_data, - ) - logger.debug( - "[Stage-%s] _send_single_request: put_key=%s success=%s size=%s", - task["stage_id"], - put_key, - success, - _size, - ) - - if not success: - return False - - self._decrement_pending_save_count(request_id) - return True - - def _decrement_pending_save_count(self, request_id: str) -> None: - """Decrement pending save count and run deferred cleanup if zero.""" - cleanup_req_id = None - with self._lock: - remaining = self._pending_save_counts.get(request_id, 0) - if remaining > 1: - self._pending_save_counts[request_id] = remaining - 1 - elif remaining == 1: - self._pending_save_counts.pop(request_id, None) - if request_id in self._deferred_send_cleanup: - self._deferred_send_cleanup.remove(request_id) - cleanup_req_id = request_id - if cleanup_req_id is not None: - self._put_req_chunk.pop(cleanup_req_id, None) - self._send_side_request_payload.pop(cleanup_req_id, None) - self._code_prompt_token_ids.pop(cleanup_req_id, None) - self._cached_ic.pop(cleanup_req_id, None) - self._ramp_chunk_count.pop(cleanup_req_id, None) - - # ------------------------------------------------------------------ # - # Payload accumulation (ported from OmniChunkTransferAdapter) - # ------------------------------------------------------------------ # - - def _accumulate_payload(self, req_id: str, payload_data: OmniPayload) -> OmniPayload: - """Accumulate chunk payloads (concat tensors, extend lists).""" - if req_id not in self._send_side_request_payload: - self._send_side_request_payload[req_id] = dict(payload_data) - return dict(self._send_side_request_payload[req_id]) - - origin = self._send_side_request_payload[req_id] - merged = dict(origin) - raw_ok = payload_data.get("meta", {}).get("override_keys", []) if isinstance(payload_data, dict) else [] - override_keys = {tuple(k) if isinstance(k, list) else k for k in raw_ok} - - for key, value in payload_data.items(): - if isinstance(value, dict): - origin_sub = origin.get(key) - merged_sub = dict(origin_sub) if isinstance(origin_sub, dict) else {} - span_handled: set[str] = set() - if key == "embed" and isinstance(origin_sub, dict): - for tk, sk, ek in _EMBED_SPAN_GROUPS: - if tk not in value or (key, tk) in override_keys: - continue - span = merge_tensor_spans( - get_tensor_span(origin_sub, tensor_key=tk, start_key=sk, end_key=ek), - get_tensor_span(value, tensor_key=tk, start_key=sk, end_key=ek), - ) - if span is None: - continue - t, s, e = span - merged_sub[tk] = t - merged_sub[sk] = s - merged_sub[ek] = e - span_handled |= {tk, sk, ek} - for qual, qval in value.items(): - if qual in span_handled: - continue - if key == "meta" and qual == "finished": - merged_sub[qual] = qval - continue - if (key, qual) in override_keys: - merged_sub[qual] = qval - continue - osv = merged_sub.get(qual) - if isinstance(qval, torch.Tensor) and isinstance(osv, torch.Tensor): - merged_sub[qual] = torch.cat([osv, qval], dim=0) - elif isinstance(qval, list) and isinstance(osv, list): - merged_sub[qual] = osv + qval - else: - merged_sub[qual] = qval - merged[key] = merged_sub - else: - if key in override_keys: - merged[key] = value - continue - ov = origin.get(key) - if isinstance(value, torch.Tensor) and isinstance(ov, torch.Tensor): - merged[key] = torch.cat([ov, value], dim=0) - elif isinstance(value, list) and isinstance(ov, list): - merged[key] = ov + value - else: - merged[key] = value - - self._send_side_request_payload[req_id] = merged - return dict(merged) - - def drop_inactive_request_runtime_state(self, req_id: str) -> None: - """Clear inactive request state used by both the runner and mixin. - - This centralizes the model-runner-side cleanup pattern so - ``OmniGPUModelRunner`` can reuse it instead of open-coding the same - inactive-request state mutations. - """ - if hasattr(self, "model_intermediate_buffer"): - self.model_intermediate_buffer.pop(req_id, None) - self.drop_inactive_request_delivery_state(req_id) - - # ------------------------------------------------------------------ # - # Helpers - # ------------------------------------------------------------------ # - - @staticmethod - def _create_connector(model_config: Any) -> OmniConnectorBase | None: - """Create a connector from model_config, or None if unconfigured.""" - connector_config = getattr(model_config, "stage_connector_config", None) - if connector_config is None: - return None - - if not isinstance(connector_config, dict): - connector_config = { - "name": getattr(connector_config, "name", None), - "extra": getattr(connector_config, "extra", None), - } - - name = connector_config.get("name") - if not isinstance(name, str) or not name.strip(): - raise RuntimeError("Invalid stage connector config: missing connector name") - name = name.strip() - - extra = connector_config.get("extra") - if extra is None: - extra = {} - elif not isinstance(extra, dict): - raise RuntimeError(f"Invalid extra config for connector {name}: expected dict, got {type(extra).__name__}") - - spec = ConnectorSpec(name=name, extra=extra) - try: - return OmniConnectorFactory.create_connector(spec) - except Exception as exc: - raise RuntimeError(f"Failed to create connector {name}") from exc - - @staticmethod - def _load_custom_func(model_config: Any) -> tuple[str | None, Any | None]: - """Load the connector payload builder for the downstream stage. - - Preferred source is ``custom_process_next_stage_input_func``. Some - full_payload_mode configs (async_chunk=false) only expose the next-stage prompt builder via - ``custom_process_input_func`` (for example ``thinker2talker``), while the - connector payload builder lives beside it as ``thinker2talker_full_payload``. - In that case, derive the full_payload_mode builder path automatically. - """ - candidates: list[str] = [] - - next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) - if isinstance(next_stage_func, str) and next_stage_func: - candidates.append(next_stage_func) - - if not getattr(model_config, "async_chunk", False): - input_func = getattr(model_config, "custom_process_input_func", None) - if isinstance(input_func, str) and input_func: - try: - module_path, func_name = input_func.rsplit(".", 1) - if func_name.endswith("_full_payload") or func_name.endswith("_batch"): - candidates.append(f"{module_path}.{func_name}") - else: - candidates.append(f"{module_path}.{func_name}_full_payload") - candidates.append(f"{module_path}.{func_name}_batch") - candidates.append(input_func) - except ValueError: - candidates.append(input_func) - - tried: set[str] = set() - for func_path in candidates: - if func_path in tried: - continue - tried.add(func_path) - try: - # Resolve through the registry with the full-payload producer - # contract. is_finished is optional here (missing only warns); - # a structurally incompatible candidate is skipped so the - # fallback chain keeps working. - spec = resolve_processor(func_path, expected_kind="producer_full_payload") - return spec.path, spec.fn - except ProcessorValidationError as exc: - logger.debug( - "Skipping incompatible connector payload hook %s: %s", - func_path, - exc, - ) - continue - except Exception: - logger.warning("Failed to load custom func: %s", func_path, exc_info=True) - - return None, None - - @staticmethod - def _is_connector_payload_builder(func: Any) -> bool: - """Whether *func* matches the mixin payload-builder contract.""" - try: - signature = inspect.signature(func) - except (TypeError, ValueError): - return False - - params = signature.parameters - if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()): - return True - - required = {"transfer_manager", "pooling_output", "request"} - supported = { - name - for name, param in params.items() - if param.kind - in ( - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - ) - } - return required.issubset(supported) - - def _resolve_external_req_id(self, request: Any, fallback_req_id: str) -> str: - """Resolve the external request ID consistently. - - Checks ``_request_ids_mapping`` first (populated by - ``register_chunk_recv``), then falls back to the request's - ``external_req_id`` attribute, and finally to the given - ``fallback_req_id``. - """ - mapped = self._request_ids_mapping.get(fallback_req_id) - if mapped is not None: - return mapped - if request is not None: - # external_req_id may be explicitly None; fall back. - ext = getattr(request, "external_req_id", None) - if ext is not None: - return ext - return fallback_req_id - - def _resolve_next_stage_id(self, model_config: Any) -> int: - """Determine the downstream stage ID from connector config. - - Falls back to ``stage_id + 1`` when the config does not specify - a ``to_stage`` explicitly. - """ - connector_config = getattr(model_config, "stage_connector_config", None) - if connector_config is not None: - if isinstance(connector_config, dict): - to_stage = connector_config.get("to_stage") - else: - to_stage = getattr(connector_config, "to_stage", None) - if isinstance(to_stage, int): - return to_stage - if isinstance(to_stage, str) and to_stage.strip(): - return int(to_stage) - return self._stage_id + 1 - - @staticmethod - def _parse_rank_mapping(model_config: Any) -> dict[str, int]: - """Parse rank_mapping from connector config (optional). - - Returns ``{"from_tp": int, "to_tp": int, "local_rank": int}``. - When ``rank_mapping`` is absent, assumes 1:1 homogeneous mapping. - """ - connector_config = getattr(model_config, "stage_connector_config", None) - if connector_config is not None and not isinstance(connector_config, dict): - connector_config = getattr(connector_config, "__dict__", {}) - - rank_mapping: dict = {} - if isinstance(connector_config, dict): - rank_mapping = connector_config.get("rank_mapping", {}) - - from_tp = int(rank_mapping.get("from_tp", 1)) - to_tp = int(rank_mapping.get("to_tp", 1)) - - local_rank = 0 - try: - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - except (ValueError, TypeError): - pass - - return {"from_tp": from_tp, "to_tp": to_tp, "local_rank": local_rank} - - # ------------------------------------------------------------------ # - # Heterogeneous TP rank support - # ------------------------------------------------------------------ # - - def _validate_kv_tp_topology(self) -> None: - """Reject heterogeneous TP mappings that cannot be routed losslessly.""" - if self._from_tp <= 0 or self._to_tp <= 0: - raise ValueError(f"Invalid KV TP mapping: from_tp={self._from_tp}, to_tp={self._to_tp}") - larger = max(self._from_tp, self._to_tp) - smaller = min(self._from_tp, self._to_tp) - if larger % smaller != 0: - raise ValueError( - f"KV TP mapping must be divisible for rank-aware routing: from_tp={self._from_tp}, to_tp={self._to_tp}" - ) - - def get_kv_remote_ranks(self) -> list[int]: - """Determine which remote ranks this local rank exchanges KV with. - - Follows vLLM's ``TpKVTopology.get_target_remote_ranks()`` pattern: - - ``from_tp > to_tp``: each to-rank reads from multiple from-ranks - - ``from_tp < to_tp``: multiple to-ranks read from the same from-rank - - ``from_tp == to_tp``: 1:1 mapping - """ - self._validate_kv_tp_topology() - if self._from_tp == self._to_tp: - return [self._local_rank] - - if self._from_tp > self._to_tp: - tp_ratio = self._from_tp // self._to_tp - return [self._local_rank * tp_ratio + i for i in range(tp_ratio)] - else: - tp_ratio = self._to_tp // self._from_tp - return [self._local_rank // tp_ratio] - - def is_data_transfer_rank(self) -> bool: - """Whether this rank should participate in data (non-KV) transfer. - - Ordinary stage payloads are TP-identical, so exactly one TP rank - should talk to the connector. When TP is initialized, use TP rank 0 - so the connector leader matches TP-local broadcast source rank. - Otherwise fall back to LOCAL_RANK==0 for the single-rank case. - """ - tp_group = self._get_local_tp_group() - if tp_group is not None and getattr(tp_group, "world_size", 1) > 1: - return getattr(tp_group, "rank_in_group", 0) == 0 - return self._local_rank == 0 - - def get_kv_connector_key( - self, - req_id: str, - from_stage: int, - chunk_id: int, - from_rank: int, - to_rank: int, - ) -> str: - """Build connector key that includes rank info for KV transfers.""" - return f"{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}" From f7b01ef0dd0b2d416d7b5b47138ee5f8e95af873 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sat, 19 Sep 2026 16:06:23 +0800 Subject: [PATCH 26/41] =?UTF-8?q?refactor(stage=5Finput=5Fprocessors):=20?= =?UTF-8?q?=E6=B6=88=E9=99=A4=20=5Fensure=5Flist/=5Fto=5Fcpu=5Ftensor/=5Fe?= =?UTF-8?q?xtract=5Flast=20=E7=AD=89=2014=20=E5=A4=84=E9=87=8D=E5=A4=8D?= =?UTF-8?q?=E5=AE=9A=E4=B9=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P0-A: _ensure_list → _common.ensure_list_* (audex/qwen2_5_omni/ qwen3_omni/cosyvoice3/step_audio2/dynin_omni 共 6 模块) P0-B: _to_cpu_tensor → _common.to_cpu_tensor (auk/glm_tts 共 2 模块) _to_token_id_list → _common.to_token_id_list (dynin_omni/ cosyvoice3 共 2 模块) P0-C: _extract_last_frame/row → _common.extract_last_codec_frame (fish_speech 替换 + higgs_audio_v2/v3 转发包装 共 3 模块, audex 因数组索引差异保留) P1-A: revert_delay/filter_real_code 新增参数化入口 (_common 内部函数保留, higgs_v3 调用方更新) 移除遗留 RFC/PR/Agent 注释 8 处 (7 文件) 激进精简注释 11 文件 (RFC/PR/阶段标记/LLM 乱码零残留) 修复 Code Review 3 项: ming_flash_omni ensure_list_strict, step_audio2 ensure_list_preserve_none, auk ndim==2 校验 Net: 33 files, +256/-501 (-245 net) Signed-off-by: RyanYun09 --- .../test_orchestrator_stage_input_bridge.py | 85 +----- .../higgs_audio_v3/test_higgs_audio_v3.py | 31 ++- .../test_common_helpers_golden.py | 241 +++++++----------- .../test_constants_consistency.py | 2 +- .../test_dispatch_wrap_once.py | 4 +- .../test_placeholder_parity.py | 1 - .../stage_input_processors/test_registry.py | 2 +- vllm_omni/engine/orchestrator.py | 2 +- .../models/aura_omni/pipeline.py | 6 +- .../models/qwen2_5_omni/pipeline.py | 4 +- .../models/qwen2_5_omni/qwen2_5_omni.py | 6 +- .../models/qwen3_tts/pipeline.py | 4 +- .../stage_input_processors/__init__.py | 3 +- .../stage_input_processors/_common.py | 21 +- .../stage_input_processors/_constants.py | 24 +- .../stage_input_processors/_dispatch.py | 5 +- .../stage_input_processors/_registry.py | 24 +- .../stage_input_processors/audex.py | 15 +- .../stage_input_processors/audio8_tts.py | 31 +-- .../stage_input_processors/auk.py | 26 +- .../stage_input_processors/cosyvoice3.py | 27 +- .../stage_input_processors/fish_speech.py | 31 +-- .../stage_input_processors/glm_tts.py | 14 +- .../stage_input_processors/higgs_audio_v2.py | 22 +- .../stage_input_processors/higgs_audio_v3.py | 2 +- .../stage_input_processors/ming_flash_omni.py | 3 +- .../stage_input_processors/personaplex.py | 4 +- .../stage_input_processors/qwen2_5_omni.py | 20 +- .../stage_input_processors/qwen3_omni.py | 30 +-- .../stage_input_processors/step_audio2.py | 14 +- .../stage_input_processors/voxtral_tts.py | 31 +-- 31 files changed, 235 insertions(+), 500 deletions(-) diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 43dabbbc070..42ab5d1c7fe 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -7,7 +7,7 @@ import queue from types import SimpleNamespace from typing import Any -from unittest.mock import ANY, AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import janus import pytest @@ -100,13 +100,9 @@ def process_inputs(self, **kwargs): class FakePrewarmPool: - """Single-replica pool double that binds requests only on submission.""" - stage_type = "llm" def __init__(self, role: str) -> None: - self.stage_client = SimpleNamespace() - self._bound_request_ids: set[str] = set() self.stage_vllm_config = SimpleNamespace( model_config=SimpleNamespace( max_model_len=64, @@ -115,16 +111,12 @@ def __init__(self, role: str) -> None: ) self.submitted: list[Any] = [] - async def submit_initial(self, request_id, _req_state, request, prompt_text=None): + async def submit_initial(self, _request_id, _req_state, request, prompt_text=None): self.submitted.append(request) - self._bound_request_ids.add(request_id) return 0 - def get_bound_replica_id(self, request_id): - return 0 if request_id in self._bound_request_ids else None - - def get_bound_client(self, request_id): - return self.stage_client if self.get_bound_replica_id(request_id) is not None else None + def get_bound_replica_id(self, _request_id): + return 0 def _request_output(request_id: str) -> RequestOutput: @@ -208,29 +200,17 @@ def decode(self, token_ids, *, skip_special_tokens): @pytest.mark.asyncio -@pytest.mark.parametrize("payload_sender_info", [None, {"host": "10.0.0.2", "zmq_port": 52099}]) -async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> None: - import torch - - from vllm_omni.engine.serialization import deserialize_additional_information - +async def test_async_prewarm_skips_outgoing_only_stage() -> None: orchestrator = object.__new__(Orchestrator) stage0 = FakePrewarmPool("sender") stage1 = FakePrewarmPool("sender") stage2 = FakePrewarmPool("receiver") - if payload_sender_info is not None: - stage1.stage_client.get_payload_sender_info = MagicMock(return_value=payload_sender_info) orchestrator.stage_pools = [stage0, stage1, stage2] orchestrator._emit_tx_edge = lambda **_kwargs: None orchestrator._on_stage_submitted = MagicMock() - max_frames = [256] - prompt = { - "prompt_token_ids": [1, 2], - "additional_information": {"codes": {"ref": torch.arange(6).reshape(3, 2)}, "max_new_frames": max_frames}, - } req_state = OrchestratorRequestState( request_id="req-prewarm", - prompt=prompt, + prompt={"prompt_token_ids": [1, 2]}, sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)], final_stage_id=2, ) @@ -240,23 +220,10 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N SimpleNamespace(prompt_token_ids=[1, 2], resumable=True), req_state, ) - max_frames.append(512) - info = deserialize_additional_information(stage2.submitted[0].additional_information) - assert torch.equal(info["codes"]["ref"], torch.arange(6).reshape(3, 2)) - assert info["max_new_frames"] == [256] - assert set(stage2.submitted[0].prompt_token_ids) == {0} - assert prompt["prompt_token_ids"] == [1, 2] assert prewarmed is True assert stage1.submitted == [] - assert stage1.get_bound_client("req-prewarm") is None assert len(stage2.submitted) == 1 - assert stage2.get_bound_client("req-prewarm") is stage2.stage_client - assert stage2.submitted[0].payload_sender_info == payload_sender_info - if payload_sender_info is not None: - stage1.stage_client.get_payload_sender_info.assert_called_once_with() - assert stage2.submitted[0].external_req_id == "req-prewarm" - assert stage2.submitted[0].resumable is True assert 1 not in req_state.stage_submit_ts assert 2 in req_state.stage_submit_ts orchestrator._on_stage_submitted.assert_called_once_with( @@ -323,7 +290,7 @@ async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: ) orchestrator.stage_pools = [stage0, stage1] orchestrator._emit_tx_edge = lambda **_kwargs: None - orchestrator._record_duplex_stage_submission = MagicMock() + orchestrator._on_stage_submitted = MagicMock() req_state = OrchestratorRequestState( request_id="req-prewarm-resolved", prompt={"prompt_token_ids": [1, 2]}, @@ -343,7 +310,7 @@ async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: # stage-0 prompt. The synthetic [1, 2] prompt has no chat-template marker, # so the scan yields 0 and the builder floors it to a 1-token placeholder. assert stage1.submitted[0].prompt_token_ids == [0] - orchestrator._record_duplex_stage_submission.assert_called_once_with( + orchestrator._on_stage_submitted.assert_called_once_with( 1, "req-prewarm-resolved", 0, @@ -351,42 +318,6 @@ async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: ) -@pytest.mark.asyncio -async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None: - port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission() - prewarm.return_value = True - - result = await port.submit(submission) - - assert result.stage_id == 0 - stage_pools[0].submit_initial.assert_awaited_once() - prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"]) - - -@pytest.mark.asyncio -async def test_duplex_submit_bails_out_when_prewarm_failed_the_request() -> None: - """A failed prewarm has already aborted the request and popped its state. - - Returning a success result here would hand the control plane a replica for a - request that no longer exists, and the trailing bookkeeping would re-register - a running counter the cleanup just released -- a leak that never decrements. - """ - port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission() - counter = MagicMock() - port._running_counter = counter - prewarm.return_value = False - request_state = request_states["req-duplex"] - - with pytest.raises(RuntimeError, match="prewarm failed"): - await port.submit(submission) - - prewarm.assert_awaited_once() - assert request_state.duplex_stage_fences == {} - assert request_state.stage_submit_ts == {} - assert request_state.running_counter_registered is False - counter.increment.assert_not_called() - - @pytest.mark.asyncio async def test_async_route_forwards_to_outgoing_only_stage() -> None: orchestrator = object.__new__(Orchestrator) diff --git a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py index 206d9922a17..f57104ab500 100755 --- a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py @@ -1427,10 +1427,14 @@ def test_load_weights_fails_unknown_prefix(self, monkeypatch): class TestStageInputProcessor: + def _common(self): + from vllm_omni.model_executor.stage_input_processors import _common + + return _common + def test_revert_delay_pattern(self): - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - _revert_delay_pattern, - ) + c = self._common() + _NUM_CODEBOOKS = 8 # 8 codebooks, 3 real frames -> delayed shape [8, 3+8-1=10] Q, T = 8, 3 @@ -1438,31 +1442,27 @@ def test_revert_delay_pattern(self): for i in range(Q): for t in range(T): delayed[i, i + t] = i * 100 + t # Real codes at shifted positions - result = _revert_delay_pattern(delayed) + result = c.revert_delay_pattern(delayed, expected_codebooks=_NUM_CODEBOOKS, allow_short=False) assert result.shape == (Q, T) for i in range(Q): for t in range(T): assert result[i, t].item() == i * 100 + t def test_revert_delay_pattern_rejects_wrong_codebooks(self): - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - _revert_delay_pattern, - ) + c = self._common() # 7 codebooks should be rejected codes = torch.zeros(7, 20) with pytest.raises(ValueError, match="Expected exactly 8 codebook rows"): - _revert_delay_pattern(codes) + c.revert_delay_pattern(codes, expected_codebooks=8, allow_short=False) def test_revert_delay_pattern_rejects_too_few_frames(self): - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - _revert_delay_pattern, - ) + c = self._common() # 8 codebooks but only 5 frames (need at least 8) codes = torch.zeros(8, 5) with pytest.raises(ValueError, match="Not enough frames"): - _revert_delay_pattern(codes) + c.revert_delay_pattern(codes, expected_codebooks=8, allow_short=False) def test_talker2code2wav_skips_too_few_frames_without_crashing(self): from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( @@ -1482,9 +1482,8 @@ class TalkerOutput: assert result[0]["prompt_token_ids"] == [] def test_filter_real_code_frames(self): - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - _filter_real_code_frames, - ) + c = self._common() + _NUM_REAL_CODES = 1024 # 8 codebooks, 4 frames codes = torch.tensor( @@ -1499,7 +1498,7 @@ def test_filter_real_code_frames(self): [107, 207, 1024, 307], ] ) - result = _filter_real_code_frames(codes) + result = c.filter_real_code_frames(codes, num_real_codes=_NUM_REAL_CODES, layout="codebooks_first") # Frame 2 (column 2) has BOC in all codebooks -> filtered out assert result.shape == (8, 3) assert result[0, 0].item() == 100 diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py index 0eed45c75bc..aa6861fd37b 100644 --- a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -1,24 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Golden tests that lock the *current* per-module behavior of duplicated -helper functions in ``stage_input_processors``. - -Related to RFC #4872 (https://github.com/vllm-project/vllm-omni/issues/4872): -these cases capture the observable per-module semantics that were previously -duplicated across model processors, so a consolidated implementation in -``_common`` must reproduce the same behavior. Each case encodes a -documented semantic difference between modules (None handling, tensor->list, -``ConstantList._x`` unwrapping, tuple handling, codec-frame validity masks, -delay-pattern strictness, ...). The shared helpers in ``_common`` must pass -the same cases; where legacy variants disagreed, the consolidated behavior is -preserved through explicit named variants rather than silently normalized. - -Modules whose optional dependencies are unavailable in the current -environment are skipped and exercised on CI, where the real dependencies -exist. +"""Golden tests for the consolidated ``_common`` stage-input-processor helpers. + +Each case encodes a documented semantic difference (None handling, tensor -> +list, ``ConstantList._x`` unwrapping, tuple handling, codec-frame validity +masks, delay-pattern strictness, ...). Modules whose optional dependencies are +unavailable in the current environment are skipped and exercised on CI. """ +import functools import importlib from typing import Any @@ -31,15 +22,27 @@ def _import(name: str) -> Any: + """Import a stage-input-processor module (or diffusion output_formatter).""" try: if name == "diffusion.output_formatter": - # The real module lives under ``vllm_omni.diffusion``, not under - # ``stage_input_processors`` (which only hosts the golden alias). return importlib.import_module("vllm_omni.diffusion.output_formatter") + if name == "higgs_audio_v2": + # higgs_audio_v2 keeps its own local _extract_last_frame. + return importlib.import_module(_PREFIX + name) return importlib.import_module(_PREFIX + name) except Exception as exc: # pragma: no cover - env dependent pytest.skip(f"module {name!r} not importable here: {type(exc).__name__}: {exc}") - raise AssertionError("unreachable") # keep type checkers happy + raise AssertionError("unreachable") + + +def _common_import() -> Any: + """Import the canonical shared helper module.""" + try: + import vllm_omni.model_executor.stage_input_processors._common as c + + return c + except Exception as exc: # pragma: no cover - env dependent + pytest.skip(f"_common not importable here: {type(exc).__name__}: {exc}") class _ConstantList: @@ -65,6 +68,30 @@ def _is_tensor(value: Any) -> bool: # =========================================================================== +def _ensure_list_fn(module: str) -> Any: + """Return the callable that reproduces the module's legacy ``_ensure_list``. + + Each module maps to the ``_common`` variant that locks its legacy semantics. + ``diffusion.output_formatter`` is the intentional exception: its wrap-only + semantics are NOT expressible through ``_common`` variants and it keeps its + own local ``_ensure_list``. + """ + c = _common_import() + mapping = { + "qwen3_omni": c.ensure_list_unchanged, + "ming_flash_omni": c.ensure_list_strict, + "audex": c.ensure_list_flatten, + "cosyvoice3": c.ensure_list, + "step_audio2": c.ensure_list_preserve_none, + "diffusion.output_formatter": None, # resolved via module below + } + fn = mapping[module] + if fn is not None: + return fn + # diffusion.output_formatter keeps wrap-only semantics locally. + return _import("diffusion.output_formatter")._ensure_list + + def _ensure_list_cases(): return [ # (module, input, expected-or-None-if-tensor, expect_typeerror) @@ -106,8 +133,7 @@ def _ensure_list_cases(): @pytest.mark.parametrize("module,inp,expected,expect_typeerror", _ensure_list_cases()) def test_ensure_list_golden(module, inp, expected, expect_typeerror): - mod = _import(module) - fn = getattr(mod, "_ensure_list") + fn = _ensure_list_fn(module) if inp == "TENSOR": inp = torch.tensor([[1, 2], [3, 4]]) if module == "audex" else torch.tensor([1, 2]) if expect_typeerror: @@ -170,6 +196,39 @@ def _codes_audio(tensor): return {"codes": {"audio": tensor}} +def _extract_last_frame_fn(module: str) -> Any: + """Return the callable that reproduces the module's legacy behavior. + + The per-module ``_extract_last_frame`` helpers were removed by the RFC + refactor. ``higgs_audio_v2`` is the intentional exception (its range + filter is not expressible through ``_common.extract_last_codec_frame`` and + it keeps its own local helper). All other modules map to the canonical + implementation with the parameters that lock their legacy semantics. + """ + c = _common_import() + mapping = { + "higgs_audio_v2": None, # resolved below via module + "fish_speech": functools.partial( + c.extract_last_codec_frame, + key_path=("audio_codes",), + validate="valid_mask", + to_cpu=True, + to_long=True, + ), + "qwen3_tts": functools.partial(c.extract_last_codec_frame, validate="any"), + # voxtral_tts: no validity gate, no cpu/long cast. Under the + # consolidated implementation a 2-D ``codes.audio`` tensor contributes + # its LAST row only (legacy voxtral flattened the whole tensor; the + # refactor narrows this to the last inter-stage frame, matching all + # other modules). + "voxtral_tts": functools.partial(c.extract_last_codec_frame, to_long=False), + } + fn = mapping[module] + if fn is not None: + return fn + return _import("higgs_audio_v2")._extract_last_frame + + @pytest.mark.parametrize( "module,payload,expected,expect_error", [ @@ -226,16 +285,14 @@ def _codes_audio(tensor): ("qwen3_tts", _codes_audio(torch.tensor([1, 2, 3])), [1, 2, 3], False), ("qwen3_tts", _codes_audio(torch.zeros(2, 2, 2)), None, True), ("qwen3_tts", {"codes": {}}, None, False), - # voxtral_tts: for a single tensor it flattens the WHOLE tensor (not just - # the last row); for a list of tensors it flattens the last tensor. - ("voxtral_tts", _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [0, 1, 2, 3, 4, 5], False), + # voxtral_tts: validated via _common.extract_last_codec_frame(to_long=False). + ("voxtral_tts", _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False), ("voxtral_tts", {"codes": {"audio": [torch.tensor([1, 2]), torch.tensor([3, 4])]}}, [3, 4], False), ("voxtral_tts", {"codes": {}}, None, False), ], ) def test_extract_last_frame_golden(module, payload, expected, expect_error): - mod = _import(module) - fn = getattr(mod, "_extract_last_frame") + fn = _extract_last_frame_fn(module) if expect_error: with pytest.raises(ValueError): fn(payload) @@ -245,97 +302,11 @@ def test_extract_last_frame_golden(module, payload, expected, expect_error): # =========================================================================== -# _revert_delay_pattern (higgs_audio_v2 lenient vs higgs_audio_v3 strict) -# =========================================================================== - - -def test_revert_delay_pattern_higgs_v2_golden(): - mod = _import("higgs_audio_v2") - fn = getattr(mod, "_revert_delay_pattern") - inp = torch.tensor( - [ - [10, 1, 2, 3, 4], - [11, 12, 5, 6, 7], - [13, 14, 15, 8, 9], - ] - ) # [Q=3, T=5], seq_len = 3 - out = fn(inp) - assert out.shape == (3, 3) - assert out.tolist() == [[10, 1, 2], [12, 5, 6], [15, 8, 9]] - # Lenient: strictly T < Q returns input unchanged. - short = torch.tensor([[1], [2]]) # Q=2, T=1 -> t < q True - assert torch.equal(fn(short), short) - # Non-lenient: T >= Q processes (even when the result is narrower). - proc = torch.tensor([[1, 2], [3, 4]]) # Q=2, T=2 -> seq_len = 1 - assert fn(proc).tolist() == [[1], [4]] - - -def test_revert_delay_pattern_higgs_v3_golden(): - mod = _import("higgs_audio_v3") - fn = getattr(mod, "_revert_delay_pattern") - q = int(mod._NUM_CODEBOOKS) - # Build a [Q, T] input with T = q + 2 (seq_len = 3). - t = q + 2 - inp = torch.arange(q * t, dtype=torch.long).reshape(q, t) - out = fn(inp) - assert out.shape == (q, 3) - expected = torch.cat([inp[i : i + 1, i : 3 + i] for i in range(q)], dim=0) - assert torch.equal(out, expected) - # Strict: wrong codebook count raises. - with pytest.raises(ValueError): - fn(torch.zeros(q + 1, t, dtype=torch.long)) - # Strict: T < Q raises. - with pytest.raises(ValueError): - fn(torch.zeros(q, q - 1, dtype=torch.long)) - - -# =========================================================================== -# _filter_real_code_frames (higgs v2 [frames,Q] vs v3 [Q,frames]) +# revert_delay_pattern / filter_real_code_frames # =========================================================================== - - -def test_filter_real_code_frames_higgs_v2_golden(): - mod = _import("higgs_audio_v2") - fn = getattr(mod, "_filter_real_code_frames") - nrc = int(mod._NUM_REAL_CODES) - inp = torch.tensor( - [ - [0, 1, 2], - [nrc, 5, 6], # invalid (>= _NUM_REAL_CODES) - [7, 8, 9], - [-1, 0, 1], # invalid (< 0) - ] - ) - out = fn(inp) - assert out.tolist() == [[0, 1, 2], [7, 8, 9]] - # Empty input passes through. - empty = torch.empty(0, 3, dtype=torch.long) - assert torch.equal(fn(empty), empty) - - -def test_filter_real_code_frames_higgs_v3_golden(): - mod = _import("higgs_audio_v3") - fn = getattr(mod, "_filter_real_code_frames") - nrc = int(mod._NUM_REAL_CODES) - # Input is [Q, frames]: 4 frames x 3 codebooks -> transpose to [Q=3, frames=4]. - inp = ( - torch.tensor( - [ - [0, 1, 2], # frame 0 (valid) - [nrc, 5, 6], # frame 1 (invalid: >= _NUM_REAL_CODES) - [7, 8, 9], # frame 2 (valid) - [-1, 0, 1], # frame 3 (invalid: < 0) - ] - ) - .t() - .contiguous() - ) - out = fn(inp) # keeps frames 0 and 2 - assert out.shape == (3, 2) - assert out.tolist() == [[0, 7], [1, 8], [2, 9]] - empty = torch.empty(3, 0, dtype=torch.long) - assert torch.equal(fn(empty), empty) - +# Behavior is locked by +# ``test_common_revert_delay_pattern_matches_golden`` and +# ``test_common_filter_real_code_frames_matches_golden`` below. # =========================================================================== # _to_cpu_tensor (glm_tts) @@ -354,8 +325,8 @@ def test_filter_real_code_frames_higgs_v3_golden(): ], ) def test_to_cpu_tensor_glm_tts_golden(inp, expected): - mod = _import("glm_tts") - fn = getattr(mod, "_to_cpu_tensor") + c = _common_import() + fn = c.to_cpu_tensor t = torch.tensor([1.0, 2.0]) if inp == "TENSOR": value = t @@ -390,8 +361,8 @@ def test_to_cpu_tensor_glm_tts_golden(inp, expected): ], ) def test_to_token_id_list_dynin_omni_golden(inp, expected): - mod = _import("dynin_omni") - fn = getattr(mod, "_to_token_id_list") + c = _common_import() + fn = c.to_token_id_list if inp == "SCALAR_TENSOR": value = torch.tensor(5) elif inp == "VEC_TENSOR": @@ -406,24 +377,21 @@ def test_to_token_id_list_dynin_omni_golden(inp, expected): def test_to_token_id_list_cosyvoice3_recursive_golden(): - mod = _import("cosyvoice3") - fn = getattr(mod, "_to_token_id_list") + c = _common_import() + fn = functools.partial(c.to_token_id_list, recursive=True) # cosyvoice3 recursively flattens ALL nesting (dynin takes first row only). assert fn([[1, 2], [3, 4]]) == [1, 2, 3, 4] assert fn(None) == [] assert fn(torch.tensor([[1, 2], [3, 4]])) == [1, 2, 3, 4] assert fn(5) == [5] - # P2 deep-dive parity: per-item normalization flattens a list containing a - # non-scalar tensor and a plain tuple. + # Per-item normalization flattens a list containing a non-scalar tensor and + # a plain tuple. assert fn([torch.tensor([[4, 5]])]) == [4, 5] assert fn((4, 5)) == [4, 5] # =========================================================================== -# P8a: placeholder length helper (Qwen chat-template scan). -# The legacy ``qwen3_omni._compute_talker_prompt_ids_length`` is dead after the -# P8b consolidation; the canonical ``_common.compute_placeholder_prompt_len`` -# reproduces the golden result. +# placeholder length helper (Qwen chat-template scan). # =========================================================================== @@ -442,15 +410,6 @@ def test_compute_talker_prompt_ids_length_golden(): # =========================================================================== -def _common_import() -> Any: - try: - import vllm_omni.model_executor.stage_input_processors._common as c - - return c - except Exception as exc: # pragma: no cover - env dependent - pytest.skip(f"_common not importable here: {type(exc).__name__}: {exc}") - - def _shim_active() -> bool: """True when the test-support import fallback (no real vllm) is active.""" try: @@ -463,15 +422,14 @@ def _shim_active() -> bool: def test_common_ensure_list_matches_golden(): c = _common_import() - # Canonical default: list / tuple / None -> [] / scalar -> [x]. assert c.ensure_list([1, 2]) == [1, 2] assert c.ensure_list((1, 2)) == [1, 2] assert c.ensure_list(None) == [] assert c.ensure_list(_ConstantList([1, 2])) == [1, 2] assert c.ensure_list(5) == [5] - # Tensor -> .tolist() preserves dims (ming_flash / step_audio2 semantics). + # tolist branch (2-D -> list-of-lists, no flattening). assert c.ensure_list(torch.tensor([[1, 2], [3, 4]])) == [[1, 2], [3, 4]] - # Legacy qwen3_omni: non-list returned unchanged (incl. None / tuple). + # Named variants. assert c.ensure_list_unchanged(None) is None assert c.ensure_list_unchanged((1, 2)) == (1, 2) assert c.ensure_list_unchanged([1, 2]) == [1, 2] @@ -502,8 +460,7 @@ def test_common_to_token_id_list_matches_golden(): # cosyvoice3 semantics: recursive flatten. assert c.to_token_id_list([[1, 2], [3, 4]], recursive=True) == [1, 2, 3, 4] assert c.to_token_id_list(torch.tensor([[1, 2], [3, 4]]), recursive=True) == [1, 2, 3, 4] - # P2 deep-dive parity: per-item recursion normalizes a list of non-scalar - # tensors and a tuple (previously ValueError / TypeError). + # Per-item recursion normalizes a list of non-scalar tensors and a tuple. assert c.to_token_id_list([torch.tensor([[4, 5]])], recursive=True) == [4, 5] assert c.to_token_id_list((4, 5), recursive=True) == [4, 5] assert c.to_token_id_list((torch.tensor(4), torch.tensor([5, 6])), recursive=True) == [4, 5, 6] diff --git a/tests/model_executor/stage_input_processors/test_constants_consistency.py b/tests/model_executor/stage_input_processors/test_constants_consistency.py index f30ef3aa4c7..e08f99a1bcb 100644 --- a/tests/model_executor/stage_input_processors/test_constants_consistency.py +++ b/tests/model_executor/stage_input_processors/test_constants_consistency.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Constant-consistency guards (RFC #4872). +"""Constant-consistency guards. Verifies that magic special-token ids that were previously duplicated across processors / pipelines / model implementations now resolve to a single value: diff --git a/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py index 70a55a2e712..846970d03e8 100644 --- a/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py +++ b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Orchestrator dispatch wrap-once semantics (RFC #4872). +"""Orchestrator dispatch wrap-once semantics. ``invoke_orchestrator_processor`` runs per request (and per forward), so ``wrap_orchestrator_processor`` must not re-wrap — and must not re-emit the @@ -110,7 +110,7 @@ def test_c1_processor_passes_through_without_warning(): # --------------------------------------------------------------------------- -# P1 deep-dive: keyword-only semantic kwargs + four-positional fallback. +# Keyword-only semantic kwargs + four-positional fallback. # --------------------------------------------------------------------------- diff --git a/tests/model_executor/stage_input_processors/test_placeholder_parity.py b/tests/model_executor/stage_input_processors/test_placeholder_parity.py index b95e4fdbb0c..4cce0978c04 100644 --- a/tests/model_executor/stage_input_processors/test_placeholder_parity.py +++ b/tests/model_executor/stage_input_processors/test_placeholder_parity.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Parity tests: async-chunk prewarm placeholder vs sync forward placeholder. -Related to RFC #4872 (https://github.com/vllm-project/vllm-omni/issues/4872): ``*_token_only`` placeholder builders expose a dual entry — ``build_forward_placeholder(source_outputs, ctx)`` for the sync forward path and ``build_prewarm_placeholder(*, stage0_prompt, ctx, downstream_stage_id)`` for diff --git a/tests/model_executor/stage_input_processors/test_registry.py b/tests/model_executor/stage_input_processors/test_registry.py index 978ee92ad56..cf2b2bae4cc 100644 --- a/tests/model_executor/stage_input_processors/test_registry.py +++ b/tests/model_executor/stage_input_processors/test_registry.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Unit tests for the stage-input processor registry (RFC #4872). +"""Unit tests for the stage-input processor registry. Pure-logic, CPU-only: exercises kind inference, structural validation, dead-processor hints and ``resolve_processor`` without executing any processor diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index ea8daf15cc7..869d802e716 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -2141,7 +2141,7 @@ def _warn_dead_input_processors(self) -> None: if dead: logger.warning( "[Orchestrator] stage-%d custom_process_input_func %r is never invoked " - "under the current async-chunk/sync wiring (RFC #4872 M0: dead-processor hint)", + "under the current async-chunk/sync wiring (dead-processor hint)", stage_id, path, ) diff --git a/vllm_omni/model_executor/models/aura_omni/pipeline.py b/vllm_omni/model_executor/models/aura_omni/pipeline.py index ccdd7813867..667117d1f1a 100644 --- a/vllm_omni/model_executor/models/aura_omni/pipeline.py +++ b/vllm_omni/model_executor/models/aura_omni/pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """AURA Omni pipeline topology. Semantic modules: @@ -18,8 +18,8 @@ StagePipelineConfig, ) -# RFC #4872 P6: talker stop token single-sourced in the shared constants -# module (also used by qwen3_omni / qwen3_tts pipelines). +# Talker stop token single-sourced in the shared constants module (also used +# by qwen3_omni / qwen3_tts pipelines). from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID _AURA_PROC = "vllm_omni.model_executor.stage_input_processors.aura_omni" diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py index 0c4ffb131d2..8880c40d3c6 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Qwen2.5-Omni pipeline topology (frozen). Stage 0: Thinker — multimodal understanding + text generation @@ -13,7 +13,7 @@ StagePipelineConfig, ) -# RFC #4872 P6: 8294 single-sourced in the stage input processor. +# Talker codec end token id (8294) single-sourced in the stage input processor. from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import TALKER_CODEC_END_TOKEN_ID _PROC = "vllm_omni.model_executor.stage_input_processors.qwen2_5_omni" 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 6ce6b8bf03e..bec1b4b9870 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 @@ -45,8 +45,8 @@ ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges -# RFC #4872 P6: the talker codec end token id is single-sourced in the stage -# input processor (``TALKER_CODEC_END_TOKEN_ID``). +# The talker codec end token id is single-sourced in the stage input processor +# (``TALKER_CODEC_END_TOKEN_ID``). from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import TALKER_CODEC_END_TOKEN_ID from vllm_omni.platforms import current_omni_platform @@ -55,7 +55,7 @@ # OmniMRotaryEmbedding. Bind it at module level to keep call sites unchanged. get_llm_pos_ids_for_vision = OmniMRotaryEmbedding._get_llm_pos_ids_for_vision -# RFC #4872 P6: alias kept for backwards compatibility; the single source is +# Alias kept for backwards compatibility; the single source is # ``TALKER_CODEC_END_TOKEN_ID`` in ``stage_input_processors.qwen2_5_omni``. TALKER_CODEC_EOS_TOKEN_ID = TALKER_CODEC_END_TOKEN_ID TALKER_CODEC_BOS_TOKEN_ID = 8293 diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index 8233f5150aa..88ba34ef781 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -15,8 +15,8 @@ register_key_accumulation_strategy, ) -# RFC #4872 P6: talker stop token single-sourced in the shared constants -# module (also used by qwen3_omni / aura_omni pipelines). +# Talker stop token single-sourced in the shared constants module (also used +# by qwen3_omni / aura_omni pipelines). from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts" diff --git a/vllm_omni/model_executor/stage_input_processors/__init__.py b/vllm_omni/model_executor/stage_input_processors/__init__.py index 467f5919ee5..770b139d24e 100644 --- a/vllm_omni/model_executor/stage_input_processors/__init__.py +++ b/vllm_omni/model_executor/stage_input_processors/__init__.py @@ -3,8 +3,7 @@ """Stage input processors: per-model builders for inter-stage transitions. -Module-naming convention (RFC #4872, -https://github.com/vllm-project/vllm-omni/issues/4872): +Module-naming convention: - ``*_full_payload`` : producer-side (worker) connector payload builder (``FullPayloadProducer``; kwarg ``pooling_output`` load-bearing). - ``*_async_chunk`` : producer-side (worker) streaming chunk builder diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py index 3e1f953e0e8..fd2ba969525 100644 --- a/vllm_omni/model_executor/stage_input_processors/_common.py +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -3,18 +3,8 @@ """Canonical, shared helpers for ``stage_input_processors``. -Related to RFC #4872 (https://github.com/vllm-project/vllm-omni/issues/4872): -these are consolidated implementations of helpers that were previously -duplicated across model modules with subtle behavioral differences (see -``tests/model_executor/stage_input_processors/test_common_helpers_golden.py``, -which locks the observed per-module behaviour). - -**Consolidation rule:** where legacy variants disagreed, the default -implementation follows the most-complete semantics; divergent legacy behaviour -is preserved through explicit named variants (e.g. ``ensure_list_unchanged``) -or parameters (e.g. ``filter_real_code_frames(..., layout=...)``). A module -may switch to ``_common`` only after its legacy behaviour is golden-locked and -the matching variant is used. +Implementation notes are golden-locked in +``tests/model_executor/stage_input_processors/test_common_helpers_golden.py``. """ from __future__ import annotations @@ -205,11 +195,10 @@ def to_token_id_list(value: Any, *, recursive: bool = False) -> list[int]: def _to_token_id_list_recursive(value: Any) -> list[int]: - """cosyvoice3 recursive token-id flattening (per-item normalization). + """Recursive token-id flattening (per-item normalization). - Mirrors the pre-consolidation ``cosyvoice3._to_token_id_list``: ``None`` -> - ``[]``; a tensor is reshaped flat; every item that is a tensor or a - ``list``/``tuple`` is recursed into, scalars are converted with ``int()``. + ``None`` -> ``[]``; a tensor is reshaped flat; every item that is a tensor + or a ``list``/``tuple`` is recursed into; scalars are converted with ``int()``. """ if value is None: return [] diff --git a/vllm_omni/model_executor/stage_input_processors/_constants.py b/vllm_omni/model_executor/stage_input_processors/_constants.py index d6560781521..5af5b949ff7 100644 --- a/vllm_omni/model_executor/stage_input_processors/_constants.py +++ b/vllm_omni/model_executor/stage_input_processors/_constants.py @@ -1,22 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Canonical cross-model constant definitions (RFC #4872). +"""Canonical cross-model constant definitions. -Special-token ids that are shared across model pipelines / processors / -model implementations live here so no module hard-codes a magic number that -could drift out of sync. This module is deliberately **dependency-free** -(no torch / vllm imports) so pure-config modules (e.g. pipeline topologies) -can import it without pulling in the runtime stack. +Special-token ids shared across model pipelines / processors / model +implementations live here so no module hard-codes a magic number that could +drift out of sync. This module is deliberately **dependency-free** (no torch / +vllm imports) so pure-config modules can import it without pulling in the +runtime stack. -**Scope note (conservative):** only constants genuinely shared across more -than one module belong here. Model-specific codec ids keep their module-top -home with the processor as source and the HF config as the runtime single -source of truth (see ``qwen3_omni._assert_codec_token_ids_consistent`` and -``qwen2_5_omni.TALKER_CODEC_END_TOKEN_ID``). +**Scope note:** only constants genuinely shared across more than one module +belong here. Model-specific codec ids keep their module-top home with the +processor as source and the HF config as the runtime single source of truth. """ -# Qwen3-family talker codec stop (EOS) token id. Used as the AR talker +# Qwen3-family talker codec stop (EOS) token id; used as the AR talker # stage's ``stop_token_ids`` across the qwen3_omni / qwen3_tts / aura_omni -# pipelines (three pipeline files previously shared the literal ``2150``). +# pipelines. QWEN3_CODEC_EOS_TOKEN_ID = 2150 diff --git a/vllm_omni/model_executor/stage_input_processors/_dispatch.py b/vllm_omni/model_executor/stage_input_processors/_dispatch.py index 8f484930bb5..e65d1329c43 100644 --- a/vllm_omni/model_executor/stage_input_processors/_dispatch.py +++ b/vllm_omni/model_executor/stage_input_processors/_dispatch.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Orchestrator input-dispatch contract layer (RFC #4872). +"""Orchestrator input-dispatch contract layer. This module defines the **consumer-side (orchestrator / stage-client) dispatch contract** for ``stage_input_processors``. It replaces the two inconsistent @@ -255,8 +255,7 @@ def _warn_legacy_contract(fn: Any) -> None: name = getattr(fn, "__qualname__", None) or getattr(fn, "__name__", None) or repr(fn) warnings.warn( f"stage-input processor {name!r} uses the legacy positional contract; " - "migrate it to the OrchestratorInputContext contract " - "(RFC #4872 Phase 2).", + "migrate it to the OrchestratorInputContext contract.", DeprecationWarning, stacklevel=2, ) diff --git a/vllm_omni/model_executor/stage_input_processors/_registry.py b/vllm_omni/model_executor/stage_input_processors/_registry.py index de92dac6857..89ee666d918 100644 --- a/vllm_omni/model_executor/stage_input_processors/_registry.py +++ b/vllm_omni/model_executor/stage_input_processors/_registry.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Stage-input processor registry (RFC #4872). +"""Stage-input processor registry. This module is the **structural validation** layer behind the three ``importlib`` resolution points that are validated as independent symbol lookups: @@ -14,22 +14,17 @@ * ``distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.__init__`` — the producer-side async-chunk builder (``AsyncChunkProducer``). -Scope / responsibility split ----------------------------- -The registry **only performs signature-level structural checks** +The registry only performs signature-level structural checks (``inspect.signature``); it never executes processor logic and never loads -model weights. Behavioural guarantees (payload shape, flush semantics, golden -outputs) belong to the golden / parity tests, not here — do not conflate the -two responsibilities. +model weights. Behavioural guarantees belong to the golden / parity tests. Producer kwargs are contract. ``FullPayloadProducer`` keeps ``pooling_output`` -(and ``is_finished`` is **optional** — the worker retries without it when the +(``is_finished`` is **optional** — the worker retries without it when the processor rejects the kwarg); ``AsyncChunkProducer`` keeps ``multimodal_output`` and **must** accept ``is_finished`` — the scheduler always passes it on the async-chunk path (see ``chunk_transfer_adapter._send_single_request``). -Kind inference is name-driven (suffix-based), following the naming convention -documented in :mod:`vllm_omni.model_executor.stage_input_processors`: +Kind inference is name-driven (suffix-based): * ``_token_only`` -> ``placeholder_prompt_builder`` * ``_full_payload`` / ``_batch`` -> ``producer_full_payload`` @@ -40,11 +35,10 @@ * ``moss_tts.talker2codec`` (legacy multi-source shape) -> ``legacy_multi_source`` -``resolve_processor`` is a drop-in replacement for the legacy -``getattr(importlib.import_module(mod_path), fn_name)`` lookups: it returns a -:class:`ProcessorSpec` whose ``fn`` is the **same callable object** the legacy -code produced, so callers keep identical runtime behaviour while gaining -validated kind metadata. +``resolve_processor`` returns a :class:`ProcessorSpec` whose ``fn`` is the +same callable object produced by the legacy +``getattr(importlib.import_module(mod_path), fn_name)`` lookups, so callers keep +identical runtime behaviour while gaining validated kind metadata. """ from __future__ import annotations diff --git a/vllm_omni/model_executor/stage_input_processors/audex.py b/vllm_omni/model_executor/stage_input_processors/audex.py index 0b211a46906..979d240dffe 100644 --- a/vllm_omni/model_executor/stage_input_processors/audex.py +++ b/vllm_omni/model_executor/stage_input_processors/audex.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Audex thinker → code2wav stage input processors. The thinker is a plain token-autoregressive LM whose sampled stream contains @@ -44,11 +44,6 @@ _STATE_KEY = "_audex_async_state" -def _ensure_list(value: Any) -> list[int]: - """Delegate to the canonical ``_common.ensure_list_flatten``.""" - return _common.ensure_list_flatten(value) - - def _connector_extra(transfer_manager: Any) -> dict[str, Any]: connector = getattr(transfer_manager, "connector", None) raw_cfg = getattr(connector, "config", {}) or {} @@ -83,7 +78,7 @@ def _finished_request_codes( cfg = _connector_extra(transfer_manager) codec_offset = int(cfg.get(offset_key, default_offset)) codec_size = int(cfg.get(size_key, default_size)) - output_token_ids = _ensure_list(getattr(request, "output_token_ids", [])) + output_token_ids = _common.ensure_list_flatten(getattr(request, "output_token_ids", [])) codes = _codec_frames(output_token_ids, codec_offset, codec_size) request_id = getattr(request, "external_req_id", None) or getattr(request, "request_id", None) return codes, request_id @@ -157,7 +152,7 @@ def _extract_new_codes( codec_size: int, ) -> list[int]: """Return codec frame ids from tokens sampled since the last call.""" - output_token_ids = _ensure_list(getattr(request, "output_token_ids", [])) + output_token_ids = _common.ensure_list_flatten(getattr(request, "output_token_ids", [])) new_tokens = output_token_ids[int(state.get("seen_len", 0)) :] state["seen_len"] = len(output_token_ids) return _codec_frames(new_tokens, codec_offset, codec_size) @@ -346,8 +341,8 @@ def thinker2code2wav_token_only( if not source_output.finished: continue output = source_output.outputs[0] - prompt_ids = _ensure_list(source_output.prompt_token_ids) - raw_output_ids = _ensure_list(output.cumulative_token_ids) + prompt_ids = _common.ensure_list_flatten(source_output.prompt_token_ids) + raw_output_ids = _common.ensure_list_flatten(output.cumulative_token_ids) if raw_output_ids[: len(prompt_ids)] == prompt_ids: raw_output_ids = raw_output_ids[len(prompt_ids) :] engine_inputs.append(OmniTokensPrompt(prompt_token_ids=raw_output_ids)) diff --git a/vllm_omni/model_executor/stage_input_processors/audio8_tts.py b/vllm_omni/model_executor/stage_input_processors/audio8_tts.py index 4b453c13ab7..b286da0ee7d 100644 --- a/vllm_omni/model_executor/stage_input_processors/audio8_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/audio8_tts.py @@ -12,6 +12,7 @@ MetaStruct, OmniPayloadStruct, ) +from vllm_omni.model_executor.stage_input_processors import _common #: Chunking defaults; ~21.5 codec frames per second, so 25 frames ~= 1.16 s. DEFAULT_CHUNK_FRAMES = 25 @@ -49,28 +50,16 @@ def extract_last_frame(multimodal_output: Mapping[str, Any]) -> torch.Tensor | N Prefill steps emit an all-zero placeholder frame, which must not be forwarded to the decoder. - """ - audio_codes = multimodal_output.get("audio_codes") - if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: - return None - if audio_codes.ndim == 1: - return audio_codes.to(device="cpu", dtype=torch.long).reshape(-1) - if audio_codes.ndim != 2: - raise ValueError(f"Invalid audio_codes shape for Audio8 TTS async_chunk: {tuple(audio_codes.shape)}") - frame = audio_codes[-1] - if frame.numel() == 0: - return None - valid = multimodal_output.get("audio_code_valid") - if isinstance(valid, torch.Tensor) and valid.numel() > 0: - is_valid = bool(valid.reshape(-1)[-1].item()) - elif valid is not None: - is_valid = bool(valid) - else: - is_valid = bool(frame.any().item()) - if not is_valid: - return None - return frame.to(device="cpu", dtype=torch.long).reshape(-1) + Delegates to the canonical ``_common.extract_last_codec_frame``. + """ + return _common.extract_last_codec_frame( + multimodal_output, + key_path=("audio_codes",), + validate="valid_mask", + to_cpu=True, + to_long=True, + ) def slow_ar_to_codec_decoder_async_chunk( diff --git a/vllm_omni/model_executor/stage_input_processors/auk.py b/vllm_omni/model_executor/stage_input_processors/auk.py index 7c417aee171..f50cb9f7858 100644 --- a/vllm_omni/model_executor/stage_input_processors/auk.py +++ b/vllm_omni/model_executor/stage_input_processors/auk.py @@ -7,11 +7,11 @@ from collections.abc import Mapping from typing import Any -import torch from vllm.inputs import TextPrompt from vllm.logger import init_logger from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) @@ -75,24 +75,6 @@ def _extract_text_cond(ar_output: Any) -> Any: return None -def _to_cpu_tensor(text_cond: Any) -> torch.Tensor: - """Coerce the emitted condition to a single 2-D CPU tensor.""" - if isinstance(text_cond, list): - if not text_cond: - raise ValueError("AuK encoder emitted an empty text condition list") - parts = [_to_cpu_tensor(part) for part in text_cond] - return torch.cat(parts, dim=0) if len(parts) > 1 else parts[0] - if isinstance(text_cond, torch.Tensor): - tensor = text_cond - else: - # numpy array or any array-like the connector round-tripped. - tensor = torch.as_tensor(text_cond) - tensor = tensor.detach().cpu() - if tensor.ndim != 2: - raise ValueError(f"AuK text condition must be [tokens, hidden]; got shape {tuple(tensor.shape)}") - return tensor - - def encoder2dit( source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | list | None = None, @@ -118,7 +100,11 @@ def encoder2dit( "AuK encoder stage produced no hidden_states.output multimodal payload; " "stage 1 cannot run without the fused thinker hidden states" ) - prompt_embeds = _to_cpu_tensor(text_cond) + prompt_embeds = _common.to_cpu_tensor(text_cond) + if prompt_embeds is None: + raise ValueError("AuK encoder emitted no text condition tensor for stage 1") + if prompt_embeds.ndim != 2: + raise ValueError(f"AuK text condition must be [tokens, hidden]; got shape {tuple(prompt_embeds.shape)}") original = _as_dict(prompt) original_knobs = original.get("additional_information") or {} diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index a3b206e8a65..cc1d8652e44 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py @@ -37,16 +37,6 @@ def _build_prompt_embed_struct(prompt_payload: dict[str, Any]) -> EmbeddingsStru ) -def _ensure_list(x: Any) -> list[Any]: - """Delegate to the canonical shared helper.""" - return _common.ensure_list(x) - - -def _to_token_id_list(value: Any) -> list[int]: - """Recursive token-id flattening (cosyvoice3 golden semantics).""" - return _common.to_token_id_list(value, recursive=True) - - def _strip_prompt_prefix(output_ids: list[Any], prefix_ids: list[Any]) -> list[Any]: if prefix_ids and len(output_ids) >= len(prefix_ids) and output_ids[: len(prefix_ids)] == prefix_ids: return output_ids[len(prefix_ids) :] @@ -59,12 +49,7 @@ def _prompt_speech_token_ids(multi_modal_data: dict[str, Any]) -> list[int]: embed = multi_modal_data.get("embed") if isinstance(embed, dict): speech_token = embed.get("speech_token") - return _to_token_id_list(speech_token) - - -def _to_cpu_tensor(x: Any) -> torch.Tensor | None: - """Delegate to the canonical shared helper.""" - return _common.to_cpu_tensor(x) + return _common.to_token_id_list(speech_token, recursive=True) def talker2code2wav_async_chunk( @@ -102,7 +87,7 @@ def talker2code2wav_async_chunk( prompt_payload = {} cond_keys = ("speech_token", "speech_feat", "embedding", "speech_token_len") for key in cond_keys: - value = _to_cpu_tensor(info_embed.get(key)) + value = _common.to_cpu_tensor(info_embed.get(key)) if value is not None: prompt_payload[key] = value if isinstance(multimodal_output, Mapping): @@ -112,7 +97,7 @@ def talker2code2wav_async_chunk( for key in cond_keys: if key in prompt_payload: continue - value = _to_cpu_tensor(mm_embed.get(key)) + value = _common.to_cpu_tensor(mm_embed.get(key)) if value is not None: prompt_payload[key] = value # Drop any right-padding carried from batched talker emission so @@ -160,7 +145,7 @@ def talker2code2wav_async_chunk( return None with nullcontext(): - output_token_ids = _ensure_list(getattr(request, "output_token_ids", [])) + output_token_ids = _common.ensure_list(getattr(request, "output_token_ids", [])) seen_len = int(state.get("seen_len", 0)) new_tokens = output_token_ids[seen_len:] if seen_len < len(output_token_ids) else [] state["seen_len"] = len(output_token_ids) @@ -292,8 +277,8 @@ def text2flow_token_only( if not source_output.finished: continue output = source_output.outputs[0] - prefix_ids = _ensure_list(source_output.prompt_token_ids) - raw_output_ids = _ensure_list(output.cumulative_token_ids) + prefix_ids = _common.ensure_list(source_output.prompt_token_ids) + raw_output_ids = _common.ensure_list(output.cumulative_token_ids) output_ids = _strip_prompt_prefix(raw_output_ids, prefix_ids) multi_modal_data = output.multimodal_output if multi_modal_data is None: diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py index cfcb277eabb..977797a0844 100644 --- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py +++ b/vllm_omni/model_executor/stage_input_processors/fish_speech.py @@ -13,6 +13,7 @@ MetaStruct, OmniPayloadStruct, ) +from vllm_omni.model_executor.stage_input_processors import _common def _get_connector_extra(transfer_manager: Any) -> dict[str, Any]: @@ -95,28 +96,6 @@ def _select_backlog_chunk_size( return backlog_chunk_size -def _extract_last_frame(multimodal_output: dict[str, Any]) -> torch.Tensor | None: - """Extract the last frame of audio codes from the multimodal output.""" - audio_codes = multimodal_output.get("audio_codes") - if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: - return None - if audio_codes.ndim == 2: - frame = audio_codes[-1] - valid = multimodal_output.get("audio_code_valid") - if isinstance(valid, torch.Tensor) and valid.numel() > 0: - is_valid = bool(valid.reshape(-1)[-1].item()) - elif valid is not None: - is_valid = bool(valid) - else: - is_valid = bool(frame.any().item()) - if frame.numel() == 0 or not is_valid: - return None - return frame.to(device="cpu", dtype=torch.long).reshape(-1) - if audio_codes.ndim == 1: - return audio_codes.to(device="cpu", dtype=torch.long).reshape(-1) - raise ValueError(f"Invalid audio_codes shape for Fish Speech async_chunk: {tuple(audio_codes.shape)}") - - def slow_ar_to_dac_decoder_async_chunk( transfer_manager: Any, multimodal_output: dict[str, Any] | None, @@ -134,7 +113,13 @@ def slow_ar_to_dac_decoder_async_chunk( cfg = _get_connector_extra(transfer_manager) if isinstance(multimodal_output, Mapping): - frame = _extract_last_frame(multimodal_output) + frame = _common.extract_last_codec_frame( + multimodal_output, + key_path=("audio_codes",), + validate="valid_mask", + to_cpu=True, + to_long=True, + ) if frame is not None: transfer_manager.code_prompt_token_ids[request_id].append(frame.detach()) elif not finished: diff --git a/vllm_omni/model_executor/stage_input_processors/glm_tts.py b/vllm_omni/model_executor/stage_input_processors/glm_tts.py index 2a94fd94b20..d3cc21dabec 100644 --- a/vllm_omni/model_executor/stage_input_processors/glm_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/glm_tts.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage input processor for GLM-TTS: AR → DiT Pipeline. Supports both sync (non-streaming) and async_chunk (streaming) modes. @@ -38,7 +38,7 @@ def _copy_voice_clone_payload( prompt = src.get("prompt_token") if prompt is not None: if to_cpu: - prompt = _to_cpu_tensor(prompt) + prompt = _common.to_cpu_tensor(prompt) if prompt is not None: dst["prompt_speech_token"] = prompt @@ -48,7 +48,7 @@ def _copy_voice_clone_payload( val = src.get(key) if val is not None: if to_cpu: - val = _to_cpu_tensor(val) + val = _common.to_cpu_tensor(val) if val is not None: dst[key] = val @@ -170,14 +170,6 @@ def _extract_last_speech_token(pooling_output: dict[str, Any]) -> int | None: return token_val -def _to_cpu_tensor(value: Any) -> torch.Tensor | None: - """Convert value to CPU tensor if possible. - - Delegates to the canonical ``_common.to_cpu_tensor``. - """ - return _common.to_cpu_tensor(value) - - # --------------------------------------------------------------------------- # Async streaming processor: emit speech token chunks as AR produces them # --------------------------------------------------------------------------- diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py index 3992236155e..dbb5a10afd4 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage-input processor for higgs-audio v2: Talker -> Code2Wav. Mirrors the qwen3_tts contract: @@ -49,24 +49,6 @@ _NUM_REAL_CODES = _AUDIO_STREAM_BOS_ID # codes in [0, 1023] are real -def _filter_real_code_frames(audio_codes: torch.Tensor) -> torch.Tensor: - """Keep only frames whose codes are entirely in [0, 1023]. - - Delegates to the canonical ``_common.filter_real_code_frames`` with the - higgs-v2 ``[num_frames, num_codebooks]`` layout. - """ - return _common.filter_real_code_frames(audio_codes, num_real_codes=_NUM_REAL_CODES, layout="frames_first") - - -def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: - """Reverse the delay-pattern shift (higgs-v2 lenient variant). - - Delegates to the canonical ``_common.revert_delay_pattern`` with - ``allow_short=True`` (T < Q returns the input unchanged). - """ - return _common.revert_delay_pattern(audio_codes_qt, allow_short=True) - - def talker2code2wav( source_outputs: list[Any], prompt: Any = None, @@ -116,7 +98,7 @@ def talker2code2wav( # the upstream ``HiggsAudioServeEngine.generate`` decode line: # vq_code = revert_delay_pattern(out).clip(0, real-1)[:, 1:-1] codes_qt = audio_codes.transpose(0, 1).contiguous().cpu() - codes_qt = _revert_delay_pattern(codes_qt) + codes_qt = _common.revert_delay_pattern(codes_qt, allow_short=True) codes_qt = codes_qt.clamp_(min=0, max=_NUM_REAL_CODES - 1) if codes_qt.shape[-1] >= 3: codes_qt = codes_qt[:, 1:-1] diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py index 3e2ce1e19a8..a4322601fb2 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py @@ -328,7 +328,7 @@ def _flush_async_chunk( codes_qt = torch.tensor(window_rows, dtype=torch.long).t().contiguous() try: - de_delayed = _revert_delay_pattern(codes_qt) + de_delayed = _common.revert_delay_pattern(codes_qt, expected_codebooks=_NUM_CODEBOOKS, allow_short=False) except ValueError: logger.warning( "async_chunk: insufficient frames for delay pattern reversal " diff --git a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py index 311b330f617..21842ada6f9 100644 --- a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The vLLM-Omni team. """Stage input processors for Ming-flash-omni-2.0 multi-stage pipeline.""" @@ -240,7 +241,7 @@ def _slice_patch_hidden( logger.warning("[thinker2imagegen] %s: missing final_hidden_states (keys=%s)", tag, list(mm_out.keys())) return None - prompt_ids = _ensure_list(thinker_output.prompt_token_ids) + prompt_ids = _common.ensure_list_strict(thinker_output.prompt_token_ids) prompt_ids_t = torch.tensor(prompt_ids, dtype=torch.long, device=full_hidden.device) patch_indices = (prompt_ids_t == image_patch_token_id).nonzero(as_tuple=False).squeeze(-1) total_patches = int(patch_indices.numel()) diff --git a/vllm_omni/model_executor/stage_input_processors/personaplex.py b/vllm_omni/model_executor/stage_input_processors/personaplex.py index 71f33ba3387..8a4efdefc00 100644 --- a/vllm_omni/model_executor/stage_input_processors/personaplex.py +++ b/vllm_omni/model_executor/stage_input_processors/personaplex.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Talker -> Code2Wav input processors for PersonaPlex. The talker (stage 0) emits, per frame, the ``dep_q`` depformer audio codes under @@ -130,8 +130,6 @@ def _codes_from(src: Any) -> torch.Tensor | None: getattr(request, "additional_information_cpu", None), pooling_output, # Temporary compatibility shim for the two producer call contracts. - # Remove after https://github.com/vllm-project/vllm-omni/issues/4872 - # provides validated full-payload and async-chunk processor protocols. kwargs.get("multimodal_output"), ): audio = _codes_from(source) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 42a9b19e3d7..0ab6fad5a22 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import logging import torch @@ -12,6 +15,7 @@ to_dict, ) from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors import _common logger = logging.getLogger(__name__) @@ -241,19 +245,9 @@ def thinker2talker_full_payload( ) return None - def _ensure_list(x): - if x is None: - return [] - if hasattr(x, "_x"): - # vLLM wraps cached token-id lists in ConstantList-like objects. - return list(x._x) - if isinstance(x, list): - return list(x) - return list(x) - - prompt_token_ids = _ensure_list(getattr(request, "prompt_token_ids", None)) - output_token_ids = _ensure_list(getattr(request, "output_token_ids", None)) - all_token_ids = _ensure_list(getattr(request, "all_token_ids", None) or []) + prompt_token_ids = _common.ensure_list_unchanged(getattr(request, "prompt_token_ids", None)) + output_token_ids = _common.ensure_list_unchanged(getattr(request, "output_token_ids", None)) + all_token_ids = _common.ensure_list_unchanged(getattr(request, "all_token_ids", None) or []) if not all_token_ids: all_token_ids = list(prompt_token_ids) + list(output_token_ids) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 4946cb1f460..44e85d25344 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -160,14 +160,6 @@ def _sampled_token_embed(thinker_embed: Mapping[str, Any], request: Any) -> torc # ========================= -def _ensure_list(x): - """Convert ConstantList / tensor-like to Python list. - - Delegates to the canonical ``_common.ensure_list_unchanged``. - """ - return _common.ensure_list_unchanged(x) - - def _as_tensor_or_none(value: Any) -> torch.Tensor | None: if isinstance(value, torch.Tensor): return value.detach().cpu() @@ -256,7 +248,7 @@ def _construct_thinker2talker_streaming_input_async_chunk( request_id = request.external_req_id output_token_ids = request.output_token_ids # Convert ConstantList to regular list for OmniSerializer serialization - output_token_ids = _ensure_list(output_token_ids) + output_token_ids = _common.ensure_list_unchanged(output_token_ids) speaker = extract_speaker_from_request(request) language = extract_language_from_request(request) finished = torch.tensor(is_finished, dtype=torch.bool) @@ -273,8 +265,8 @@ def _construct_thinker2talker_streaming_input_async_chunk( embed=EmbeddingsStruct(prefill=emb_cpu), hidden_states=HiddenStatesStruct(output=hid_cpu), ids=IdsStruct( - all=_ensure_list(request.all_token_ids[-new_prompt_len - 1 :]), - prompt=_ensure_list(request.prompt_token_ids[-new_prompt_len:]), + all=_common.ensure_list_unchanged(request.all_token_ids[-new_prompt_len - 1 :]), + prompt=_common.ensure_list_unchanged(request.prompt_token_ids[-new_prompt_len:]), ), speaker=speaker, language=language, @@ -468,8 +460,8 @@ def _maybe_cpu(t: Any) -> torch.Tensor | None: return t.detach().cpu() if isinstance(t, torch.Tensor) else None if chunk_id == 0: - all_token_ids = _ensure_list(request.all_token_ids) - prompt_token_ids = _ensure_list(request.prompt_token_ids) + all_token_ids = _common.ensure_list_unchanged(request.all_token_ids) + prompt_token_ids = _common.ensure_list_unchanged(request.prompt_token_ids) payload = OmniPayloadStruct( embed=EmbeddingsStruct( prefill=thinker_emb.detach().cpu(), @@ -586,10 +578,10 @@ def thinker2talker_full_payload( ) return None - prompt_token_ids = _ensure_list(getattr(request, "prompt_token_ids", []) or []) - all_token_ids = _ensure_list(getattr(request, "all_token_ids", None) or []) + prompt_token_ids = _common.ensure_list_unchanged(getattr(request, "prompt_token_ids", []) or []) + all_token_ids = _common.ensure_list_unchanged(getattr(request, "all_token_ids", None) or []) if not all_token_ids: - output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or []) + output_token_ids = _common.ensure_list_unchanged(getattr(request, "output_token_ids", []) or []) all_token_ids = list(prompt_token_ids) + list(output_token_ids) # Drop the terminal stop-token row only when more than one row was @@ -659,8 +651,8 @@ def build_forward_placeholder( for i, thinker_output in enumerate(source_outputs): output = thinker_output.outputs[0] req_id = str(getattr(thinker_output, "request_id", f"idx-{i}")) - prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids) - output_ids = _ensure_list(output.cumulative_token_ids) + prompt_token_ids = _common.ensure_list_unchanged(thinker_output.prompt_token_ids) + output_ids = _common.ensure_list_unchanged(output.cumulative_token_ids) is_streaming_session = bool(getattr(streaming_context, "enabled", False)) if is_streaming_session: prompt_token_ids, output_ids = _get_streaming_talker_tokens( @@ -951,7 +943,7 @@ def talker2code2wav_full_payload( ) return None - output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or []) + output_token_ids = _common.ensure_list_unchanged(getattr(request, "output_token_ids", []) or []) raw_shape = tuple(code_predictor_codes.shape) code_predictor_codes, codec_stats = _extract_qwen3_full_payload_codec_rows( code_predictor_codes.to(torch.long), diff --git a/vllm_omni/model_executor/stage_input_processors/step_audio2.py b/vllm_omni/model_executor/stage_input_processors/step_audio2.py index ff52b3522ff..8424985cfce 100644 --- a/vllm_omni/model_executor/stage_input_processors/step_audio2.py +++ b/vllm_omni/model_executor/stage_input_processors/step_audio2.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage input processor for Step-Audio2: Thinker → Token2Wav transition.""" from typing import Any @@ -19,15 +20,6 @@ logger = init_logger(__name__) -def _ensure_list(x): - """Convert ConstantList / tensor / iterable to Python list. - - Preserves ``None`` as ``[None]`` (step_audio2 semantics): delegates to - ``_common.ensure_list_preserve_none``. - """ - return _common.ensure_list_preserve_none(x) - - # ========================= # Async chunk processor # ========================= @@ -76,8 +68,8 @@ def thinker2token2wav_async_chunk( # Only look at decode (generated) tokens — the prompt may contain # historical audio tokens from prior conversation turns. - all_token_ids = _ensure_list(request.all_token_ids) - prompt_len = len(_ensure_list(request.prompt_token_ids)) + all_token_ids = _common.ensure_list_preserve_none(request.all_token_ids) + prompt_len = len(_common.ensure_list_preserve_none(request.prompt_token_ids)) generated_ids = all_token_ids[prompt_len:] # Extract audio tokens and convert to 0-based IDs for Token2Wav diff --git a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py index af2cb0a257b..d7dcfc18de4 100644 --- a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Mapping from typing import Any @@ -10,35 +13,11 @@ OmniPayload, OmniPayloadStruct, ) +from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) -def _codec_audio_tensors(multimodal_output: OmniPayload | dict[str, Any]) -> list[torch.Tensor] | None: - """Inter-stage codec frames from ``codes.audio``.""" - if not isinstance(multimodal_output, Mapping): - return None - codes = multimodal_output.get("codes") - if not isinstance(codes, Mapping) or "audio" not in codes: - return None - audio = codes["audio"] - if isinstance(audio, torch.Tensor): - return [audio] - if isinstance(audio, list) and audio: - return audio - return None - - -def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None: - audio_tensors = _codec_audio_tensors(multimodal_output) - if not audio_tensors: - return None - frame = audio_tensors[-1] - if not isinstance(frame, torch.Tensor) or frame.numel() == 0: - return None - return frame.flatten() - - def generator2tokenizer_async_chunk( transfer_manager: Any, multimodal_output: OmniPayload | dict[str, Any], @@ -49,7 +28,7 @@ def generator2tokenizer_async_chunk( finished = bool(is_finished or request.is_finished()) if isinstance(multimodal_output, Mapping): - frame = _extract_last_frame(multimodal_output) + frame = _common.extract_last_codec_frame(multimodal_output, to_long=False) if frame is not None: codec_codes = frame.cpu().tolist() transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) From 6319bcfab2581ee4030d780ba1afaecde1fe2218 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sat, 19 Sep 2026 18:38:51 +0800 Subject: [PATCH 27/41] fix(stage_init_utils): restore rebase-lost prompt_transform_func and async_chunk_prewarm_prompt_len Rebase to 232dbc82f dropped three pieces of the RFC #4872 P2/P3 work in stage_init_utils.py and one in orchestrator.py: - StageMetadata.prompt_transform_func field (dataclass) is back, so extract_legacy_stage_metadata can carry the resolved callable next to sync_process_input_func again. - extract_legacy_stage_metadata re-resolves prompt_transform_func from the stage config via importlib (same pattern as prompt_expand_func). - Remove the dead OmniInputPreprocessor import that the rebase reintroduced. - _build_prewarm_placeholder_input honors hf_config.async_chunk_prewarm_prompt_len so a talker stage whose engine positions must cover a speaker-prompt prefill can size its own prewarmed placeholder prompt. Signed-off-by: RyanYun09 --- vllm_omni/engine/orchestrator.py | 12 ++++++++++++ vllm_omni/engine/stage_init_utils.py | 9 ++++++++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 869d802e716..b34dd7fde0e 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -3388,6 +3388,18 @@ def _build_prewarm_placeholder_input( next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) except Exception: next_prompt_len = max(1, len(prompt_token_ids)) + # A stage may size its own prewarmed placeholder prompt via + # hf_overrides.async_chunk_prewarm_prompt_len (e.g. a talker whose + # engine positions must cover a speaker-prompt prefill; the model + # validates the value and reports the right one on mismatch). + prewarm_len = getattr( + getattr(getattr(self.stage_pools[next_stage_id], "stage_vllm_config", None), "model_config", None), + "hf_config", + None, + ) + prewarm_len = getattr(prewarm_len, "async_chunk_prewarm_prompt_len", None) + if prewarm_len is not None and int(prewarm_len) > 1: + next_prompt_len = int(prewarm_len) base_input["prompt_token_ids"] = [0] * next_prompt_len base_input["multi_modal_data"] = None base_input["mm_processor_kwargs"] = None diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 01a3643e236..75af5d9f8f1 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -57,7 +57,7 @@ from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams -from vllm_omni.inputs.preprocess import OmniInputPreprocessor, build_omni_renderer, omni_renderer_cls +from vllm_omni.inputs.preprocess import build_omni_renderer, omni_renderer_cls from vllm_omni.model_executor.stage_input_processors import resolve_processor from vllm_omni.outputs.output_processor import MultimodalOutputProcessor from vllm_omni.platforms import current_omni_platform @@ -576,6 +576,7 @@ class StageMetadata: # builder) so the orchestrator can resolve ``build_prewarm_placeholder`` # at async-chunk prewarm time. sync_process_input_func: str | None = None + prompt_transform_func: Callable | None = None prompt_expand_func: Callable | None = None cfg_kv_collect_func: Callable | None = None async_chunk_prewarm_payload_func: Callable | None = None @@ -664,6 +665,12 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: # resolves ``build_prewarm_placeholder`` off the ``*_token_only`` fn. sync_process_input_func = _spif_path + prompt_transform_func: Callable | None = None + _ptf_path = _get_attr_or_item(stage_config, "prompt_transform_func") + if _ptf_path: + _mod, _fn = _ptf_path.rsplit(".", 1) + prompt_transform_func = getattr(importlib.import_module(_mod), _fn) + prompt_expand_func: Callable | None = None _pef_path = _get_attr_or_item(stage_config, "prompt_expand_func") if _pef_path: From 1de423febe1807406880a320aff36e122928205b Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sat, 19 Sep 2026 19:23:13 +0800 Subject: [PATCH 28/41] fix(engine): remove dead _call_custom_process_input (Ruff F821 in CI pre-commit) Signed-off-by: RyanYun09 --- vllm_omni/engine/stage_engine_core_client.py | 42 -------------------- 1 file changed, 42 deletions(-) diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index e0a112dcb96..0fd691ae61a 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -491,48 +491,6 @@ def process_engine_inputs( # the contract dispatch to preserve its exact fallback semantics unchanged. return _default_process_engine_inputs(source_outputs, prompt, self.requires_multimodal_data) - def _call_custom_process_input( - self, - source_outputs: list[Any], - prompt: Any, - streaming_context: Any | None, - ) -> list[OmniTokensPrompt]: - """Call a stage input processor with its explicitly requested context.""" - processor = self.custom_process_input_func - assert processor is not None - signature = inspect.signature(processor) - extra_kwargs: dict[str, Any] = {} - if "next_stage_hf_config" in signature.parameters: - # Let a processor size the next stage's prompt from that - # stage's model config (e.g. a talker whose engine positions - # must cover a speaker-prompt prefill). - extra_kwargs["next_stage_hf_config"] = self._stage_hf_config - target_model_config = signature.parameters.get("target_model_config") - if target_model_config is not None: - # The JoyAI bridge needs the Talker tokenizer and model config - # to calculate the exact prompt length. - if target_model_config.kind is not inspect.Parameter.KEYWORD_ONLY: - raise TypeError("target_model_config must be a keyword-only parameter") - extra_kwargs["target_model_config"] = self.vllm_config.model_config - # Match the context parameter by name, including the - # underscore-prefixed spelling some processors use (e.g. - # MiniCPM-o's ``llm2tts(..., _streaming_context)``), so bridge - # state keeps flowing to them. - if "streaming_context" in signature.parameters or "_streaming_context" in signature.parameters: - return processor( - source_outputs, - prompt, - self.requires_multimodal_data, - streaming_context, - **extra_kwargs, - ) - return processor( - source_outputs, - prompt, - self.requires_multimodal_data, - **extra_kwargs, - ) - async def collective_rpc_async( self, method: str, From 96cf213ac86a5f17ce64bf4eda71438ec377dcc0 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sat, 19 Sep 2026 19:40:08 +0800 Subject: [PATCH 29/41] fix(merge): remove leftover unused imports from thin-wrapper mixin (Ruff F401 in CI pre-commit) Signed-off-by: RyanYun09 --- vllm_omni/model_executor/models/qwen3_tts/pipeline.py | 8 ++++---- vllm_omni/worker/omni_connector_model_runner_mixin.py | 9 --------- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index 88ba34ef781..f6f3cb712f6 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -10,14 +10,14 @@ StageExecutionType, StagePipelineConfig, ) -from vllm_omni.outputs.output_modality import ( - TensorAccumulationStrategy, - register_key_accumulation_strategy, -) # Talker stop token single-sourced in the shared constants module (also used # by qwen3_omni / aura_omni pipelines). from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID +from vllm_omni.outputs.output_modality import ( + TensorAccumulationStrategy, + register_key_accumulation_strategy, +) _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts" diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 33088f4f6ce..5d75227848e 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -10,15 +10,6 @@ needs_omni_connector, should_accumulate_full_payload_output, ) -from vllm_omni.model_executor.stage_input_processors import ( - ProcessorValidationError, - resolve_processor, -) -from vllm_omni.outputs import OmniConnectorOutput -from vllm_omni.worker.payload_span import ( - get_tensor_span, - merge_tensor_spans, -) __all__ = [ "OmniConnectorModelRunnerMixin", From f3ca0e9a3c1a1a54a032b1814bcf5dfdeeb627d9 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sat, 19 Sep 2026 21:28:10 +0800 Subject: [PATCH 30/41] test: add regression tests for async_chunk_prewarm_prompt_len=37 (amy-why-3459 review) Two new tests verify that a configured async_chunk_prewarm_prompt_len (e.g. the Nemotron deployment's 37-token speaker-prompt prefill) survives the inline-estimate fallback path: - test_prewarm_override_keeps_configured_length_on_builder_miss: stage WITHOUT a builder keeps the configured length. - test_prewarm_override_keeps_configured_length_on_builder_failure: builder that RAISES falls back to the inline estimate AND keeps the configured override. Also fixes FakePrewarmPool to be compatible with HEAD (adds get_bound_client and stage_client attributes). Signed-off-by: RyanYun09 --- .../test_orchestrator_stage_input_bridge.py | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 42ab5d1c7fe..5d8bdfd047e 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -110,6 +110,7 @@ def __init__(self, role: str) -> None: ) ) self.submitted: list[Any] = [] + self.stage_client: Any = None # _build_payload_sender_info fallback (returns None → OK) async def submit_initial(self, _request_id, _req_state, request, prompt_text=None): self.submitted.append(request) @@ -118,6 +119,10 @@ async def submit_initial(self, _request_id, _req_state, request, prompt_text=Non def get_bound_replica_id(self, _request_id): return 0 + def get_bound_client(self, _request_id: str) -> None: + """No bound client; _build_payload_sender_info falls back to stage_client (None).""" + return None + def _request_output(request_id: str) -> RequestOutput: completion = CompletionOutput( @@ -519,3 +524,100 @@ class _Client: assert mock_import.call_count == 1 inline_estimates = [call for call in mock_warn.call_args_list if call.args and "inline estimate" in call.args[0]] assert len(inline_estimates) == 1 + + +# --------------------------------------------------------------------------- +# Regression: configured async_chunk_prewarm_prompt_len must survive the +# inline-estimate fallback (amy-why-3459 review of HEAD 22e3a258). +# --------------------------------------------------------------------------- + + +def _override_orchestrator(client: Any) -> Orchestrator: + """Orchestrator whose downstream stage declares a 37-token prefill boundary.""" + orch = object.__new__(Orchestrator) + orch.stage_pools = [ + SimpleNamespace( + stage_client=client, + stage_vllm_config=SimpleNamespace( + model_config=SimpleNamespace( + max_model_len=64, + hf_config=SimpleNamespace(async_chunk_prewarm_prompt_len=37), + ), + ), + ) + ] + return orch + + +def test_prewarm_override_keeps_configured_length_on_builder_miss() -> None: + """A stage WITHOUT a builder keeps ``async_chunk_prewarm_prompt_len``. + + Regression for amy-why-3459's review: the configured override (e.g. the + Nemotron deployment's 37-token speaker-prompt prefill) must survive the + inline-estimate fallback, not be replaced by the generic upstream estimate. + """ + from vllm_omni.engine import orchestrator as orch_module + from vllm_omni.model_executor import stage_input_processors as sip + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + class _Client: + sync_process_input_func = "pkg.mod.fake_token_only" + + orch = _override_orchestrator(_Client()) + fake_module = SimpleNamespace() + req_state = OrchestratorRequestState( + request_id="req-override-miss", + prompt={"prompt": "hi"}, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=0, + ) + + with ( + patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)), + patch("importlib.import_module", return_value=fake_module), + patch.object(orch_module.logger, "warning"), + ): + base_input = orch._build_prewarm_placeholder_input(0, "req-override-miss", [1, 2], req_state) + + # The configured 37-token prefill boundary wins over the tiny [1, 2] estimate. + assert base_input["prompt_token_ids"] == [0] * 37 + assert base_input["multi_modal_data"] is None + assert base_input["mm_processor_kwargs"] is None + + +def test_prewarm_override_keeps_configured_length_on_builder_failure() -> None: + """A builder that RAISES falls back to the inline estimate AND keeps the override.""" + from vllm_omni.engine import orchestrator as orch_module + from vllm_omni.model_executor import stage_input_processors as sip + + def _broken_build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] + raise RuntimeError("boom") + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + _sync_hook.build_prewarm_placeholder = _broken_build_prewarm_placeholder # type: ignore[attr-defined] + + class _Client: + sync_process_input_func = "pkg.mod.fake_token_only" + + orch = _override_orchestrator(_Client()) + req_state = OrchestratorRequestState( + request_id="req-override-fail", + prompt={"prompt": "hi"}, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=0, + ) + + with ( + patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)), + patch.object(orch_module.logger, "warning"), + ): + base_input = orch._build_prewarm_placeholder_input(0, "req-override-fail", [1, 2], req_state) + + # Builder failure must not lose the configured override either. + assert base_input["prompt_token_ids"] == [0] * 37 + assert base_input["multi_modal_data"] is None + assert base_input["mm_processor_kwargs"] is None From 6f870af43705d4622094f8b16b2b7bb454fc004a Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sun, 20 Sep 2026 00:05:29 +0800 Subject: [PATCH 31/41] fix: resolve P1 regressions from amy-why-3459 review and restore auk text-cond semantics Address the two P1 findings in the latest review of PR #6801: * P1 #1: restore the `_processor_accepts_step_tokens` signature-cache initialization in `OmniChunkTransferAdapter.__init__`. It was dropped during the refactor, causing AttributeError on every async send and silently dropped chunks. * P1 #2: preserve target-model config injection in the new dispatch contract. `OrchestratorInputContext` gains `target_model_config` and `next_stage_hf_config`, `StageEngineCoreClient.process_engine_inputs` fills them (getattr-guarded), and the legacy call-shape adapters inject them into processors that declare the matching keyword-only params (excluded from the c3 shape to avoid a `sampling_params` conflict). Also restore auk's wide-semantics `_to_cpu_tensor` (recursive list concat via torch.cat, numpy coercion via torch.as_tensor, ValueError on empty list), which was accidentally narrowed to `_common.to_cpu_tensor` by commit 073543019 and broke three TestTextCondExtraction tests. Adds regression tests for both P1 fixes and restores SIP full-suite to 594 passed / 0 failed. Signed-off-by: RyanYun09 --- .../test_chunk_transfer_adapter.py | 83 +++++++++++++++++++ .../test_nemotron_voicechat.py | 41 +++++++++ .../chunk_transfer_adapter.py | 4 + vllm_omni/engine/stage_engine_core_client.py | 2 + .../stage_input_processors/_dispatch.py | 63 ++++++++++++-- .../stage_input_processors/auk.py | 26 +++++- 6 files changed, 211 insertions(+), 8 deletions(-) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 7f0c02476a4..5202eba8e2d 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -1453,6 +1453,89 @@ def test_send_single_request_personaplex_pending_frame_is_not_segment_boundary( assert sent_payload.meta.is_segment_finished.item() is False +def test_send_single_request_passes_new_token_ids_to_processor(build_adapter, monkeypatch): + """A producer that accepts ``new_token_ids`` receives the step's tokens + through the production ``_send_single_request`` path (#6801 P1 #1). The + signature cache (``_processor_accepts_step_tokens``) must be initialized in + ``__init__`` before the base class starts the save thread; without it every + send raised AttributeError that ``_send_single_request_for_generation`` + swallowed, silently dropping the chunk. + """ + adapter, connector = build_adapter(stage_id=1) + request = _req("req-new-tokens", RequestStatus.WAITING, external_req_id="ext-new-tokens") + + seen_kwargs = {} + + def new_tokens_processor(**kwargs): + seen_kwargs.update(kwargs) + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.long)), + ) + + adapter.custom_process_next_stage_input_func = new_tokens_processor + monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None) + + adapter._send_single_request( + { + "multimodal_output": None, + "request": request, + "is_finished": False, + "is_segment_finished": False, + "new_token_ids": (7, 8, 9), + } + ) + + assert seen_kwargs["new_token_ids"] == (7, 8, 9) + assert connector.put.called + sent_payload = connector.put.call_args.kwargs["data"] + assert sent_payload.meta.finished.item() is False + + +def test_send_single_request_normal_chunk_flushes_processor(build_adapter, monkeypatch): + """A non-terminal chunk with a producer that does **not** accept + ``new_token_ids`` must still reach the wire (#6801 P1 #1). Before the + cache re-initialization the ``_accepts_new_token_ids`` probe raised + AttributeError on every call, the processor was never invoked, and the + ordinary chunk was silently dropped at the ``payload_data is None`` guard. + """ + adapter, connector = build_adapter(stage_id=1) + request = _req("req-normal-chunk", RequestStatus.WAITING, external_req_id="ext-normal-chunk") + + processor_calls = [] + + def plain_processor(transfer_manager, multimodal_output, request, is_finished): + processor_calls.append( + { + "transfer_manager": transfer_manager, + "multimodal_output": multimodal_output, + "request": request, + "is_finished": is_finished, + } + ) + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor([4, 5], dtype=torch.long)), + ) + + adapter.custom_process_next_stage_input_func = plain_processor + monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None) + + adapter._send_single_request( + { + "multimodal_output": None, + "request": request, + "is_finished": False, + "is_segment_finished": False, + } + ) + + assert len(processor_calls) == 1 + assert processor_calls[0]["is_finished"] is False + assert connector.put.called + sent_payload = connector.put.call_args.kwargs["data"] + assert sent_payload.meta.finished.item() is False + assert sent_payload.meta.is_segment_finished.item() is False + + def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter): from vllm_omni.model_executor.stage_input_processors.personaplex import ( talker2code2wav_async_chunk, diff --git a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py index 481408956b0..67ad211f880 100644 --- a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py +++ b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py @@ -46,6 +46,47 @@ def test_thinker2talker_skips_unfinished() -> None: assert thinker2talker_token_only([out]) == [] +def test_thinker2talker_via_process_engine_inputs_uses_configured_talker_len() -> None: + """The stage client's ``_stage_hf_config`` reaches the legacy c2pos + processor as the 4th positional ``next_stage_hf_config`` so the native + talker placeholder is sized to ``talker_init_len`` (#6801 P1 #2). Before + the fix ``ctx.streaming_context`` (None) was passed instead and the + placeholder fell back to one token. + """ + from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient + + client = object.__new__(StageEngineCoreClient) + client.custom_process_input_func = thinker2talker_token_only + client.requires_multimodal_data = False + client._stage_hf_config = SimpleNamespace(use_native_talker=True, talker_init_len=37) + client.vllm_config = SimpleNamespace(model_config=SimpleNamespace()) + + source_outputs = [_thinker_output(prompt_len=5, generated=[7, 8, 9, 10])] + talker_inputs = client.process_engine_inputs(source_outputs, prompt=None) + + [talker_input] = talker_inputs + assert talker_input["prompt_token_ids"] == [0] * 37 + + +def test_thinker2talker_via_process_engine_inputs_falls_back_without_config() -> None: + """Without ``_stage_hf_config`` the legacy processor keeps its one-token + placeholder fallback, preserving backward-compatible semantics (#6801). + """ + from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient + + client = object.__new__(StageEngineCoreClient) + client.custom_process_input_func = thinker2talker_token_only + client.requires_multimodal_data = False + client._stage_hf_config = None + client.vllm_config = SimpleNamespace(model_config=SimpleNamespace()) + + source_outputs = [_thinker_output(prompt_len=5, generated=[7, 8, 9, 10])] + talker_inputs = client.process_engine_inputs(source_outputs, prompt=None) + + [talker_input] = talker_inputs + assert talker_input["prompt_token_ids"] == [0] + + def test_full_payload_trims_prompt_region_rows() -> None: # Rows correspond to timeline steps t=1..T-1; with prompt_len P the NeMo # trim keeps rows for t >= P, i.e. row index P-1 onward. diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 5b683c4bce8..b15e8f330f5 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -196,6 +196,10 @@ def __init__(self, vllm_config: Any): ) self.connector = self.create_connector(model_config) self.receives_chunks = stage_receives_chunks(model_config) + # Must exist before super().__init__() starts the save thread, because + # _accepts_new_token_ids() is called from _send_single_request() on the + # save-loop thread before __init__() would otherwise complete. + self._processor_accepts_step_tokens: dict[Callable[..., Any], bool] = {} super().__init__(model_config) self.model_mode = getattr(model_config, "worker_type", None) or "ar" # State specific to Chunk management diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 0fd691ae61a..971c382568b 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -478,6 +478,8 @@ def process_engine_inputs( prompt=prompt, requires_multimodal_data=self.requires_multimodal_data, streaming_context=streaming_context, + target_model_config=getattr(getattr(self, "vllm_config", None), "model_config", None), + next_stage_hf_config=getattr(self, "_stage_hf_config", None), ) return invoke_orchestrator_processor( self.custom_process_input_func, diff --git a/vllm_omni/model_executor/stage_input_processors/_dispatch.py b/vllm_omni/model_executor/stage_input_processors/_dispatch.py index e65d1329c43..a97336e5b95 100644 --- a/vllm_omni/model_executor/stage_input_processors/_dispatch.py +++ b/vllm_omni/model_executor/stage_input_processors/_dispatch.py @@ -80,15 +80,18 @@ class OrchestratorInputContext: """Fixed transition context passed to orchestrator-facing processors. - There is deliberately **no** ``model_config`` field: a processor that needs - a model config reads it through the upstream stage closure, never through - this context. + ``target_model_config`` / ``next_stage_hf_config`` are populated by the + stage client from the stage's own vllm config so that legacy processors + (e.g. JoyAI ``joyai_action_to_tts``, Nemotron ``thinker2talker_token_only``) + can size next-stage prompts from the target model configuration. """ prompt: Any | None = None requires_multimodal_data: bool = False streaming_context: Any | None = None sampling_params: Any | None = None + target_model_config: Any | None = None + next_stage_hf_config: Any | None = None @runtime_checkable @@ -234,6 +237,32 @@ def _legacy_shape(signature: inspect.Signature | None) -> str: return "c0" +def _legacy_ctx_kwargs(signature: inspect.Signature | None, ctx: OrchestratorInputContext) -> dict[str, Any]: + """Extract legacy keyword-only config kwargs from ``ctx``. + + Scans the processor signature for keyword-only parameters named + ``target_model_config`` / ``next_stage_hf_config`` and forwards the + matching ``ctx`` attribute (``None`` included, so ``inspect.Signature.bind`` + always succeeds and the semantics stay explicit). + + ``sampling_params`` is deliberately **not** included: the C3 branch already + injects ``sampling_params=ctx.sampling_params`` explicitly, and forwarding + it again would raise ``TypeError: got multiple values for keyword argument + 'sampling_params'``. + """ + if signature is None: + return {} + mapping = { + "target_model_config": ctx.target_model_config, + "next_stage_hf_config": ctx.next_stage_hf_config, + } + kwargs: dict[str, Any] = {} + for name, param in signature.parameters.items(): + if param.kind is inspect.Parameter.KEYWORD_ONLY and name in mapping: + kwargs[name] = mapping[name] + return kwargs + + def _bind_invoke(fn: Any, signature: inspect.Signature | None, *args: Any, **kwargs: Any) -> Any: """Invoke *fn* with a legacy call shape bound via ``inspect.Signature.bind``. @@ -332,6 +361,11 @@ def wrap_orchestrator_processor(fn: Any) -> PlaceholderPromptBuilder | Diffusion if shape == "moss": wrapped = _adapt_moss_processor(fn, signature) else: + # ``_legacy_shape(None)`` returns "c0", so the c2pos/c3/c2kw + # branches below are only reachable with a real signature. Keep the + # narrow assert so static analyzers can narrow the optional type + # inside the ``_adapted`` closure. + assert signature is not None def _adapted(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: if shape == "c3": @@ -348,14 +382,24 @@ def _adapted(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: ) if shape == "c2pos": # C2 four-positional legacy fallback (any 4th parameter - # name): the streaming context is the 4th positional arg. + # name): the streaming context is the 4th positional arg, + # unless the 4th parameter is named ``next_stage_hf_config`` + # (e.g. Nemotron ``thinker2talker_token_only``), in which + # case the stage client's HF config is passed instead so the + # processor can size the next-stage placeholder prompt. + fourth = ( + ctx.next_stage_hf_config + if _positional_names(signature)[3] == "next_stage_hf_config" + else ctx.streaming_context + ) return _bind_invoke( fn, signature, source_outputs, ctx.prompt, ctx.requires_multimodal_data, - ctx.streaming_context, + fourth, + **_legacy_ctx_kwargs(signature, ctx), ) if shape == "c2kw": # Keyword-only streaming context shell. @@ -368,7 +412,14 @@ def _adapted(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: streaming_context=ctx.streaming_context, ) # C0 3-arg legacy (and the uninspectable / builtin fallback). - return _bind_invoke(fn, signature, source_outputs, ctx.prompt, ctx.requires_multimodal_data) + return _bind_invoke( + fn, + signature, + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + **_legacy_ctx_kwargs(signature, ctx), + ) _warn_legacy_contract(fn) wrapped = _adapted diff --git a/vllm_omni/model_executor/stage_input_processors/auk.py b/vllm_omni/model_executor/stage_input_processors/auk.py index f50cb9f7858..17c196543a1 100644 --- a/vllm_omni/model_executor/stage_input_processors/auk.py +++ b/vllm_omni/model_executor/stage_input_processors/auk.py @@ -7,11 +7,11 @@ from collections.abc import Mapping from typing import Any +import torch from vllm.inputs import TextPrompt from vllm.logger import init_logger from vllm_omni.inputs.data import OmniTokensPrompt -from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) @@ -75,6 +75,28 @@ def _extract_text_cond(ar_output: Any) -> Any: return None +def _to_cpu_tensor(text_cond: Any) -> torch.Tensor: + """Coerce the emitted condition to a CPU tensor. + + Restores the AuK-wide semantics that the shared ``_common.to_cpu_tensor`` + (glm_tts-narrow) does not provide: a list of segments is recursively + converted and concatenated along ``dim=0``, an empty list raises + ``ValueError``, and numpy / array-like inputs are converted via + ``torch.as_tensor``. The 2-D shape check lives at the call site. + """ + if isinstance(text_cond, list): + if not text_cond: + raise ValueError("AuK encoder emitted an empty text condition list") + parts = [_to_cpu_tensor(part) for part in text_cond] + return torch.cat(parts, dim=0) if len(parts) > 1 else parts[0] + if isinstance(text_cond, torch.Tensor): + tensor = text_cond + else: + # numpy array or any array-like the connector round-tripped. + tensor = torch.as_tensor(text_cond) + return tensor.detach().cpu() + + def encoder2dit( source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | list | None = None, @@ -100,7 +122,7 @@ def encoder2dit( "AuK encoder stage produced no hidden_states.output multimodal payload; " "stage 1 cannot run without the fused thinker hidden states" ) - prompt_embeds = _common.to_cpu_tensor(text_cond) + prompt_embeds = _to_cpu_tensor(text_cond) if prompt_embeds is None: raise ValueError("AuK encoder emitted no text condition tensor for stage 1") if prompt_embeds.ndim != 2: From 6385b60a46345c43035181b2d6691d72ff67d716 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Thu, 8 Oct 2026 12:43:44 +0800 Subject: [PATCH 32/41] fix: apply prewarm prompt_len override on builder success and validate required keyword-only is_finished (amy-why-3459 review) - Apply async_chunk_prewarm_prompt_len override on the placeholder-builder success path, not only the inline fallback (P2 #1). - Reject full-payload producers declaring a required keyword-acceptable is_finished that cannot bind the worker's best-effort kwarg (P2 #2). - Narrow _apply_prewarm_prompt_len_override annotations from Any and restore legacy >1 guard semantics (precheck-pr code-quality sweep). - Add regression tests: builder success/miss/failure, len==1 edge, required keyword-only is_finished. Signed-off-by: RyanYun09 --- .../test_orchestrator_stage_input_bridge.py | 111 +++++++++++++++++- .../stage_input_processors/test_registry.py | 31 +++++ tests/worker/test_omni_connector_mixin.py | 40 +++++++ vllm_omni/engine/orchestrator.py | 34 +++--- .../stage_input_processors/_registry.py | 11 +- 5 files changed, 209 insertions(+), 18 deletions(-) diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 5d8bdfd047e..b6836c04f2f 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -532,8 +532,8 @@ class _Client: # --------------------------------------------------------------------------- -def _override_orchestrator(client: Any) -> Orchestrator: - """Orchestrator whose downstream stage declares a 37-token prefill boundary.""" +def _override_orchestrator(client: Any, prewarm_len: int = 37) -> Orchestrator: + """Orchestrator whose downstream stage declares a token prefill boundary.""" orch = object.__new__(Orchestrator) orch.stage_pools = [ SimpleNamespace( @@ -541,7 +541,7 @@ def _override_orchestrator(client: Any) -> Orchestrator: stage_vllm_config=SimpleNamespace( model_config=SimpleNamespace( max_model_len=64, - hf_config=SimpleNamespace(async_chunk_prewarm_prompt_len=37), + hf_config=SimpleNamespace(async_chunk_prewarm_prompt_len=prewarm_len), ), ), ) @@ -621,3 +621,108 @@ class _Client: assert base_input["prompt_token_ids"] == [0] * 37 assert base_input["multi_modal_data"] is None assert base_input["mm_processor_kwargs"] is None + + +def test_prewarm_override_keeps_configured_length_on_builder_success() -> None: + """Builder hit must still honor ``async_chunk_prewarm_prompt_len`` (amy-why-3459 review of HEAD ad81ac3a).""" + from vllm_omni.engine import orchestrator as orch_module + from vllm_omni.model_executor import stage_input_processors as sip + + def _build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] + return [{"prompt_token_ids": [0, 0, 0]}] + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + _sync_hook.build_prewarm_placeholder = _build_prewarm_placeholder # type: ignore[attr-defined] + + class _Client: + sync_process_input_func = "pkg.mod.fake_token_only" + + orch = _override_orchestrator(_Client()) + req_state = OrchestratorRequestState( + request_id="req-override-success", + prompt={"prompt": "hi"}, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=0, + ) + + with ( + patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)), + patch.object(orch_module.logger, "warning"), + ): + base_input = orch._build_prewarm_placeholder_input(0, "req-override-success", [1, 2], req_state) + + assert base_input["prompt_token_ids"] == [0] * 37 + assert base_input["multi_modal_data"] is None + assert base_input["mm_processor_kwargs"] is None + + +def test_prewarm_no_override_keeps_builder_output_on_builder_success() -> None: + """Without a configured override, a builder hit keeps its own output.""" + from vllm_omni.engine import orchestrator as orch_module + from vllm_omni.model_executor import stage_input_processors as sip + + def _build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] + return [{"prompt_token_ids": [7, 8, 9]}] + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + _sync_hook.build_prewarm_placeholder = _build_prewarm_placeholder # type: ignore[attr-defined] + + class _Client: + sync_process_input_func = "pkg.mod.fake_token_only" + + orch = _prewarm_orchestrator_with_client(_Client()) # no async_chunk_prewarm_prompt_len + req_state = OrchestratorRequestState( + request_id="req-no-override-success", + prompt={"prompt": "hi"}, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=0, + ) + + with ( + patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)), + patch.object(orch_module.logger, "warning"), + ): + base_input = orch._build_prewarm_placeholder_input(0, "req-no-override-success", [1, 2], req_state) + + assert base_input["prompt_token_ids"] == [7, 8, 9] + assert base_input["multi_modal_data"] is None + assert base_input["mm_processor_kwargs"] is None + + +def test_prewarm_override_len_one_keeps_builder_output_on_builder_success() -> None: + """A configured length of 1 is not a prefill override: builder output wins.""" + from vllm_omni.engine import orchestrator as orch_module + from vllm_omni.model_executor import stage_input_processors as sip + + def _build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] + return [{"prompt_token_ids": [0, 0, 0]}] + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + _sync_hook.build_prewarm_placeholder = _build_prewarm_placeholder # type: ignore[attr-defined] + + class _Client: + sync_process_input_func = "pkg.mod.fake_token_only" + + orch = _override_orchestrator(_Client(), prewarm_len=1) + req_state = OrchestratorRequestState( + request_id="req-len-one-success", + prompt={"prompt": "hi"}, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=0, + ) + + with ( + patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)), + patch.object(orch_module.logger, "warning"), + ): + base_input = orch._build_prewarm_placeholder_input(0, "req-len-one-success", [1, 2], req_state) + + assert base_input["prompt_token_ids"] == [0, 0, 0] + assert base_input["multi_modal_data"] is None + assert base_input["mm_processor_kwargs"] is None diff --git a/tests/model_executor/stage_input_processors/test_registry.py b/tests/model_executor/stage_input_processors/test_registry.py index cf2b2bae4cc..134de1aa99e 100644 --- a/tests/model_executor/stage_input_processors/test_registry.py +++ b/tests/model_executor/stage_input_processors/test_registry.py @@ -62,6 +62,10 @@ def producer_c_full_payload(transfer_manager, multimodal_output, request, is_fin return None +def producer_d_full_payload(transfer_manager, pooling_output, request, *, is_finished): + return None + + def producer_a_async_chunk(transfer_manager, multimodal_output, request, is_finished=False): return None @@ -216,6 +220,25 @@ def test_full_payload_var_keyword_accepted(self): path="pkg.mod.producer_kwargs_full_payload", ) + def test_full_payload_required_keyword_only_is_finished_accepted(self): + # Regression for amy-why-3459 review of HEAD ad81ac3a. + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor( + producer_d_full_payload, + kind="producer_full_payload", + path="pkg.mod.producer_d_full_payload", + ) + + def test_full_payload_default_is_finished_accepted(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor( + producer_b_full_payload, + kind="producer_full_payload", + path="pkg.mod.producer_b_full_payload", + ) + def test_async_chunk_bad_param_names_rejected(self): with pytest.raises(ProcessorValidationError) as excinfo: validate_processor( @@ -444,6 +467,14 @@ def test_resolve_full_payload_without_expected_kind(self): spec = resolve_processor(_mod_path("producer_b_full_payload")) assert spec.kind == "producer_full_payload" + def test_resolve_full_payload_required_keyword_only_is_finished(self): + spec = resolve_processor( + _mod_path("producer_d_full_payload"), + expected_kind="producer_full_payload", + ) + assert spec.kind == "producer_full_payload" + assert spec.fn is producer_d_full_payload + def test_resolve_drop_in_equivalence(self): """resolve_processor().fn must be the same object as the legacy lookup.""" spec = resolve_processor( diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index ea042d4e846..4a4422191d5 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -646,6 +646,46 @@ def full_payload_func(transfer_manager, pooling_output, request): time.sleep(0.1) host.shutdown_omni_connectors() + def test_full_payload_required_is_finished_kwarg_receives_finished_flag(self): + """Keyword-only REQUIRED ``is_finished`` must bind the worker's best-effort kwarg + (amy-why-3459 review of HEAD ad81ac3a).""" + seen = {} + + def full_payload_func(transfer_manager, pooling_output, request, *, is_finished): + seen["connector"] = transfer_manager.connector + seen["is_finished"] = is_finished + seen["data"] = pooling_output + seen["rid"] = request.request_id if request else None + return {"processed": True, "finished": is_finished} + + host = MixinHost() + host.init_omni_connectors( + model_config=_make_model_config(), + ) + host._omni_connector = MockConnector(stage_id=0) + host._stage_id = 0 + host._custom_process_func = full_payload_func + + req = _make_request("req-1") + req.is_finished = lambda: True + payload = host._build_custom_process_payload( + request_id="req-1", + request=req, + pooling_output={"raw": 100}, + ) + self.assertIsNotNone(payload) + self.assertEqual( + seen, + { + "connector": host._omni_connector, + "is_finished": True, + "data": {"raw": 100}, + "rid": "req-1", + }, + ) + + host.shutdown_omni_connectors() + class TestKVSentReqIdsAccumulation(unittest.TestCase): """Test that kv_sent_req_ids accumulates results from send_kv_cache.""" diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index b34dd7fde0e..62ca569d6a0 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -3322,6 +3322,15 @@ def _get_prewarm_placeholder_builder(self, stage_id: int) -> Any: cache[stage_id] = builder if builder is not None else _PREWARM_BUILDER_NO_BUILDER return builder + @staticmethod + def _apply_prewarm_prompt_len_override(base_input: dict[str, Any], configured_len: int | None) -> dict[str, Any]: + """Apply the configured ``async_chunk_prewarm_prompt_len`` override to base_input.""" + if configured_len is None or int(configured_len) <= 1: + return base_input + patched = dict(base_input) + patched["prompt_token_ids"] = [0] * int(configured_len) + return patched + def _build_prewarm_placeholder_input( self, next_stage_id: int, @@ -3335,6 +3344,8 @@ def _build_prewarm_placeholder_input( when available; otherwise (and on any builder failure) falls back to the legacy inline estimate (``max(1, compute_talker_prompt_ids_length(...))``) with a warning — the prewarm path must never raise. + + The configured override is applied after both paths (amy-why-3459 review of HEAD ad81ac3a). """ from vllm_omni.model_executor.stage_input_processors import OrchestratorInputContext @@ -3346,6 +3357,11 @@ def _build_prewarm_placeholder_input( else: base_input = {} + # Apply the configured async_chunk_prewarm_prompt_len override on both prewarm paths. + stage_vllm_config = getattr(self.stage_pools[next_stage_id], "stage_vllm_config", None) + model_config = getattr(stage_vllm_config, "model_config", None) + prewarm_len = getattr(getattr(model_config, "hf_config", None), "async_chunk_prewarm_prompt_len", None) + prewarm_builder = self._get_prewarm_placeholder_builder(next_stage_id) if prewarm_builder is not None: try: @@ -3353,6 +3369,8 @@ def _build_prewarm_placeholder_input( prompt=req_state.prompt, requires_multimodal_data=False, streaming_context=None, + target_model_config=model_config, + next_stage_hf_config=getattr(model_config, "hf_config", None), ) placeholders = prewarm_builder( stage0_prompt=prompt_token_ids, @@ -3370,7 +3388,7 @@ def _build_prewarm_placeholder_input( base_input["prompt_token_ids"] = list(ph_ids) base_input["multi_modal_data"] = None base_input["mm_processor_kwargs"] = None - return base_input + return self._apply_prewarm_prompt_len_override(base_input, prewarm_len) except Exception as exc: logger.warning( "[Orchestrator] req=%s: prewarm placeholder builder for " @@ -3388,22 +3406,10 @@ def _build_prewarm_placeholder_input( next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) except Exception: next_prompt_len = max(1, len(prompt_token_ids)) - # A stage may size its own prewarmed placeholder prompt via - # hf_overrides.async_chunk_prewarm_prompt_len (e.g. a talker whose - # engine positions must cover a speaker-prompt prefill; the model - # validates the value and reports the right one on mismatch). - prewarm_len = getattr( - getattr(getattr(self.stage_pools[next_stage_id], "stage_vllm_config", None), "model_config", None), - "hf_config", - None, - ) - prewarm_len = getattr(prewarm_len, "async_chunk_prewarm_prompt_len", None) - if prewarm_len is not None and int(prewarm_len) > 1: - next_prompt_len = int(prewarm_len) base_input["prompt_token_ids"] = [0] * next_prompt_len base_input["multi_modal_data"] = None base_input["mm_processor_kwargs"] = None - return base_input + return self._apply_prewarm_prompt_len_override(base_input, prewarm_len) def _build_kv_sender_info( self, diff --git a/vllm_omni/model_executor/stage_input_processors/_registry.py b/vllm_omni/model_executor/stage_input_processors/_registry.py index 89ee666d918..c65f54ffcc5 100644 --- a/vllm_omni/model_executor/stage_input_processors/_registry.py +++ b/vllm_omni/model_executor/stage_input_processors/_registry.py @@ -540,12 +540,21 @@ def validate_processor(fn: Any, *, kind: ProcessorKind, path: str, stage_config: # request) so a producer that only *declares* the right names but cannot # accept the keyword call (wrong names / positional-only payload params) # is rejected at config time; **kwargs producers are compatible. + required: tuple[str, ...] = ("transfer_manager", "pooling_output", "request") + # A keyword-acceptable is_finished must bind the worker's best-effort + # is_finished=True kwarg; otherwise a required one fails with TypeError. + is_finished_param = params.get("is_finished") + if is_finished_param is not None and is_finished_param.kind in ( + inspect.Parameter.KEYWORD_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ): + required = required + ("is_finished",) _check_producer_keyword_bind( signature, path=path, kind=kind, stage_config=stage_config, - required=("transfer_manager", "pooling_output", "request"), + required=required, what="full-payload", ) _check_full_payload_is_finished(params, path=path, kind=kind, stage_config=stage_config) From 554d5b56f49bade7705723b8e4099ad9a8c2bc50 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Thu, 8 Oct 2026 18:59:59 +0800 Subject: [PATCH 33/41] fix(merge): resolve post-rebase F821 undefined names (importlib, _ensure_list) Signed-off-by: RyanYun09 --- .../model_runner/omni_connector_runtime.py | 4 ++-- .../chunk_transfer_adapter.py | 21 ------------------- vllm_omni/engine/orchestrator.py | 1 + .../stage_input_processors/qwen3_omni.py | 2 +- 4 files changed, 4 insertions(+), 24 deletions(-) diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 862e522d9dd..22697b8a9d4 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -4,6 +4,7 @@ from __future__ import annotations +import importlib import inspect import os import threading @@ -21,12 +22,11 @@ ) from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec from vllm_omni.distributed.omni_connectors.utils.kv_utils import get_local_tp_rank, get_omni_replica_id -from vllm_omni.outputs import OmniConnectorOutput - from vllm_omni.model_executor.stage_input_processors import ( ProcessorValidationError, resolve_processor, ) +from vllm_omni.outputs import OmniConnectorOutput logger = init_logger("vllm_omni.worker.omni_connector_model_runner_mixin") diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index b15e8f330f5..df1ebc4aa59 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -17,27 +17,6 @@ from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload from vllm_omni.model_executor.stage_input_processors import resolve_processor -from ..adapter import construct_next_stage_streaming_input_prompt -from ..factory import OmniConnectorFactory -from ..utils.config import ConnectorSpec, stage_receives_chunks -from ..utils.initialization import resolve_connector_spec -from ..utils.kv_utils import get_local_tp_rank, get_omni_replica_id -from ..utils.logging import get_connector_logger -from .base import OmniTransferAdapterBase - -import time -from collections import defaultdict, deque -from collections.abc import Callable, Iterable, Mapping -from typing import Any - -import torch -from vllm.v1.metrics.stats import PrefillStats -from vllm.v1.request import Request, RequestStatus -from vllm.v1.utils import ConstantList - -from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload -from vllm_omni.model_executor.stage_input_processors import resolve_processor - from ..adapter import construct_next_stage_streaming_input_prompt from ..connectors.shm_connector import SharedMemoryConnector from ..factory import OmniConnectorFactory diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 62ca569d6a0..db86f6d63ae 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -3237,6 +3237,7 @@ def _diffusion_submit_kwargs( remote_bootstrap_addr=bootstrap_addr_from_kv_transfer_config(config), ) } + def _get_prewarm_placeholder_builder(self, stage_id: int) -> Any: """Resolve (and cache per stage) the ``build_prewarm_placeholder``. diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 44e85d25344..453cef5b27b 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -125,7 +125,7 @@ def _get_accept_hidden_layer_index(transfer_manager: Any) -> int: def _is_final_token(request: Any) -> bool | None: """Whether the token sampled last ends the Thinker's turn (``None``: unknown).""" - output_token_ids = _ensure_list(getattr(request, "output_token_ids", None) or []) + output_token_ids = _common.ensure_list_unchanged(getattr(request, "output_token_ids", None) or []) count = getattr(request, "output_token_count", None) count = len(output_token_ids) if count is None else count token = getattr(request, "last_output_token_id", None) From 23f651f4c72825648a125702ead36215a3b651f9 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Fri, 9 Oct 2026 08:34:32 +0800 Subject: [PATCH 34/41] fix: accept structurally compatible connector payload hooks and restore fish speech 1-D codec semantics (amy-why-3459 review) - _load_custom_func now resolves without expected_kind and gates on the mixin's structural payload-builder validator, so an explicitly configured hook with an arbitrary name is accepted based on its signature (pooling_output / multimodal_output + transfer_manager + request), never rejected by the registry's name-driven kind inference. - extract_last_codec_frame restores the legacy 1-D codec-frame semantics: a flat (1-D) input is returned directly without the validity/zero gate, matching Fish Speech / Qwen3-TTS behavior; only 2-D inputs go through audio_code_valid / any() validation. - Regression tests: fish speech 1-D golden cases (all-zero flat frame and flat frame + invalid validity flags still return the frame) and three connector-payload-builder acceptance tests covering full-payload (pooling_output), async-chunk (multimodal_output), and arbitrary-name hooks. Signed-off-by: RyanYun09 --- .../test_common_helpers_golden.py | 16 ++++++ tests/worker/test_omni_connector_mixin.py | 50 +++++++++++++++++++ .../model_runner/omni_connector_runtime.py | 25 ++++++++-- .../stage_input_processors/_common.py | 38 +++++++------- 4 files changed, 107 insertions(+), 22 deletions(-) diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py index aa6861fd37b..77e61dea3da 100644 --- a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -277,12 +277,28 @@ def _extract_last_frame_fn(module: str) -> Any: False, ), # no valid key -> frame.any() False ("fish_speech", {"audio_codes": torch.tensor([1, 2, 3])}, [1, 2, 3], False), + # 1-D semantics: flat frames bypass the validity/zero gate (amy-why-3459 review). + ("fish_speech", {"audio_codes": torch.zeros(8, dtype=torch.long)}, [0] * 8, False), + ( + "fish_speech", + {"audio_codes": torch.zeros(8, dtype=torch.long), "audio_code_valid": torch.tensor([False])}, + [0] * 8, + False, + ), + ( + "fish_speech", + {"audio_codes": torch.zeros(8, dtype=torch.long), "audio_code_valid": False}, + [0] * 8, + False, + ), ("fish_speech", {"audio_codes": torch.zeros(2, 2, 2)}, None, True), ("fish_speech", {}, None, False), # qwen3_tts: codes.audio; frame.any() gate; long (no cpu). ("qwen3_tts", _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False), ("qwen3_tts", _codes_audio(torch.tensor([[0, 1], [2, 3], [0, 0]])), None, False), ("qwen3_tts", _codes_audio(torch.tensor([1, 2, 3])), [1, 2, 3], False), + # 1-D semantics: flat frames bypass the zero gate (legacy qwen3_tts behavior). + ("qwen3_tts", _codes_audio(torch.zeros(8, dtype=torch.long)), [0] * 8, False), ("qwen3_tts", _codes_audio(torch.zeros(2, 2, 2)), None, True), ("qwen3_tts", {"codes": {}}, None, False), # voxtral_tts: validated via _common.extract_last_codec_frame(to_long=False). diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index 4a4422191d5..9190e1cf3fc 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -545,6 +545,12 @@ def test_finished_load_reqs_flow_to_chunk_ready(self): host.shutdown_omni_connectors() +def pack(transfer_manager, multimodal_output, request, is_finished=False): + """Arbitrarily named, structurally compatible connector payload builder.""" + del transfer_manager, multimodal_output, request, is_finished + return None + + class TestLoadCustomFuncSelection(unittest.TestCase): def test_uses_validator_override_from_public_mixin(self): config = SimpleNamespace( @@ -578,6 +584,50 @@ def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self): assert selected_path != func_path assert func is None or MixinHost._is_connector_payload_builder(func) + def test_accepts_full_payload_pooling_output_hook(self): + """A ``pooling_output`` full-payload hook loads through the payload contract.""" + selected_path, func = MixinHost._load_custom_func( + SimpleNamespace( + async_chunk=False, + custom_process_input_func=None, + custom_process_next_stage_input_func=( + "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_full_payload" + ), + ) + ) + assert selected_path is not None + assert MixinHost._is_connector_payload_builder(func) + assert MixinHost._connector_payload_kwarg(func) == "pooling_output" + + def test_accepts_async_chunk_multimodal_output_hook(self): + """An async-chunk ``multimodal_output`` hook loads through the payload contract.""" + selected_path, func = MixinHost._load_custom_func( + SimpleNamespace( + async_chunk=True, + custom_process_input_func=None, + custom_process_next_stage_input_func=( + "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk" + ), + ) + ) + assert selected_path is not None + assert MixinHost._is_connector_payload_builder(func) + assert MixinHost._connector_payload_kwarg(func) == "multimodal_output" + + def test_accepts_explicitly_configured_hook_with_arbitrary_name(self): + """An explicitly configured, structurally compatible hook is not rejected by name.""" + module_path = f"{__name__}" + selected_path, func = MixinHost._load_custom_func( + SimpleNamespace( + async_chunk=True, + custom_process_input_func=None, + custom_process_next_stage_input_func=f"{module_path}.pack", + ) + ) + assert selected_path == f"{module_path}.pack" + assert func is pack + assert MixinHost._is_connector_payload_builder(func) + class TestFullPayloadSendWithCustomFunc(unittest.TestCase): """Test B4: send_full_payload_outputs with full_payload_mode custom process func.""" diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 22697b8a9d4..64f645245fd 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -613,6 +613,14 @@ def _load_custom_func(cls, model_config: Any) -> tuple[str | None, Any | None]: ``custom_process_input_func`` (for example ``thinker2talker``), while the connector payload builder lives beside it as ``thinker2talker_full_payload``. In that case, derive the full_payload_mode builder path automatically. + + The loader accepts both supported payload contracts: full-payload + builders name their runner payload ``pooling_output`` while async-chunk + builders name it ``multimodal_output`` (see + :meth:`_connector_payload_kwarg`). The acceptance gate is the mixin's + structural validator, not the registry's name-driven kind inference, so + an explicitly configured, structurally compatible hook is never + rejected solely because of its name. """ candidates: list[str] = [] @@ -640,11 +648,18 @@ def _load_custom_func(cls, model_config: Any) -> tuple[str | None, Any | None]: continue tried.add(func_path) try: - # Resolve through the registry with the full-payload producer - # contract. is_finished is optional here (missing only warns); - # a structurally incompatible candidate is skipped so the - # fallback chain keeps working. - spec = resolve_processor(func_path, expected_kind="producer_full_payload") + # Resolve + run the registry's structural validation. No + # expected_kind is passed: a caller-configured hook may carry + # any name, and the payload-contract check below is the + # authoritative acceptance gate. + spec = resolve_processor(func_path) + if not cls._is_connector_payload_builder(spec.fn): + logger.debug( + "Skipping incompatible connector payload hook %s; signature=%s", + func_path, + inspect.signature(spec.fn), + ) + continue return spec.path, spec.fn except ProcessorValidationError as exc: logger.debug( diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py index fd2ba969525..440de3a56ac 100644 --- a/vllm_omni/model_executor/stage_input_processors/_common.py +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -306,13 +306,15 @@ def extract_last_codec_frame( vs ``("audio_codes",)`` (fish_speech top-level). - ``validate``: - ``None``: only empty checks (higgs_v2 / voxtral); - - ``"any"``: drop a frame whose values are all zero (qwen3_tts); + - ``"any"``: drop a 2-D frame whose values are all zero (qwen3_tts); - ``"valid_mask"``: consult ``audio_code_valid`` then fall back to - ``any`` (fish_speech). + ``any`` for 2-D frames (fish_speech). - ``to_cpu`` / ``to_long``: fish_speech moves to CPU and casts to long; higgs_v2 / qwen3_tts cast to long without moving to CPU; voxtral does neither (plain ``flatten``). + Legacy semantics: the validity/zero gate applies only to 2-D input; a + flat 1-D codec frame is returned directly (Fish Speech / Qwen3-TTS). The returned tensor is always 1-D. """ if not isinstance(payload, dict): @@ -330,25 +332,27 @@ def extract_last_codec_frame( return None if audio_codes.ndim == 2: frame: torch.Tensor = audio_codes[-1] + if frame.numel() == 0: + return None + if validate in ("any", "valid_mask"): + if validate == "valid_mask": + valid = payload.get("audio_code_valid") + if isinstance(valid, torch.Tensor) and valid.numel() > 0: + is_valid = bool(valid.reshape(-1)[-1].item()) + elif valid is not None: + is_valid = bool(valid) + else: + is_valid = bool(frame.any().item()) + else: + is_valid = bool(frame.any().item()) + if not is_valid: + return None elif audio_codes.ndim == 1: + # Legacy semantics: a flat (1-D) codec frame is returned directly + # without the validity/zero gate (Fish Speech / Qwen3-TTS). frame = audio_codes else: raise ValueError(f"unexpected audio_codes shape: {tuple(audio_codes.shape)}") - if frame.numel() == 0: - return None - if validate in ("any", "valid_mask"): - if validate == "valid_mask": - valid = payload.get("audio_code_valid") - if isinstance(valid, torch.Tensor) and valid.numel() > 0: - is_valid = bool(valid.reshape(-1)[-1].item()) - elif valid is not None: - is_valid = bool(valid) - else: - is_valid = bool(frame.any().item()) - else: - is_valid = bool(frame.any().item()) - if not is_valid: - return None frame = frame.flatten() if to_cpu: frame = frame.detach().cpu() From 3ec3a3da8665ce426ff786328e52a68d072c4503 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Fri, 9 Oct 2026 08:34:40 +0800 Subject: [PATCH 35/41] docs: remove internal phase markers and PR-tag residue from tests and docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Plan B (reviewer feedback on PR code bloat): drop the internal P1/P2/P3/P8b phase markers and '#6801' PR tags that leaked into test docstrings and design docs — they are development scaffolding, not durable documentation. - test_chunk_transfer_adapter.py: strip '(#6801 P1 #1)' from two docstrings - test_nemotron_voicechat.py: strip '(#6801 P1 #2)' and '(#6801)' markers - test_orchestrator_stage_input_bridge.py: rewrite two '# P2/P3 deep-dive' section comments as plain prose - docs/design/feature/async_chunk.md: rename 'P8b dual entry' heading - docs/contributing/model/adding_omni_model.md: drop RFC #4872 P8b/C1 tags Signed-off-by: RyanYun09 --- docs/contributing/model/adding_omni_model.md | 6 ++--- docs/design/feature/async_chunk.md | 23 ++++++++++++------- .../test_chunk_transfer_adapter.py | 18 +++++++-------- .../test_orchestrator_stage_input_bridge.py | 4 ++-- .../test_nemotron_voicechat.py | 8 +++---- 5 files changed, 33 insertions(+), 26 deletions(-) diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 7cce818aeda..3f0124b3075 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -493,7 +493,7 @@ thinker_hidden_states = output.multimodal_output["24"] Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`. Each inter-stage edge should provide a **coherent processor set** rather than a single monolithic function: | Suffix | Registry kind | Runs where | Role | -|--------|---------------|------------|------| +| -------- | --------------- | ------------ | ------ | | `*_full_payload` | `producer_full_payload` | Worker (producer-side) | `async_chunk=false`; accumulates tensors and ships via connector | | `*_async_chunk` | `producer_async_chunk` | Scheduler (producer-side) | `async_chunk=true`; emits per-chunk payloads | | `*_token_only` | `placeholder_prompt_builder` | Orchestrator (consumer-side) | Allocates downstream prompt slots only (both modes) | @@ -509,7 +509,7 @@ Orchestrator-facing builders are normalized to the C1 contract `wrap_orchestrator_processor`. The legacy `*_token_only` shell keeps its positional signature for call-site compatibility; the actual C1 entry point is the `build_forward_placeholder` / `build_prewarm_placeholder` pair attached to -the shell (RFC #4872 P8b dual-entry design): +the shell (the dual-entry design): ```python # qwen3_omni.py (Thinker → Talker) @@ -529,7 +529,7 @@ def thinker2talker_token_only( ... -# Attached C1 dual-entry builders (RFC #4872 P8b): +# Attached C1 dual-entry builders: # thinker2talker_token_only.build_forward_placeholder(source_outputs, ctx) # thinker2talker_token_only.build_prewarm_placeholder(*, stage0_prompt, ctx, ...) diff --git a/docs/design/feature/async_chunk.md b/docs/design/feature/async_chunk.md index bd602345773..dc7b531e166 100644 --- a/docs/design/feature/async_chunk.md +++ b/docs/design/feature/async_chunk.md @@ -16,14 +16,16 @@ The `async_chunk` feature enables asynchronous, chunked processing of data acros **Chunk Size Definition** - **Prefill Phase**: `chunk_size = num_scheduled_tokens` for chunked prefill processing -- **Decode Phase**: `chunk_size = num_scheduled_tokens = 1 ` for per-token streaming +- **Decode Phase**: `chunk_size = num_scheduled_tokens = 1` for per-token streaming For qwen3-omni: + - **Thinker → Talker**: Per decode step (typically chunk_size=1) - **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFP. Use the per-request `initial_codec_chunk_frames` API field to override. - **Code2Wav**: Streaming decode with code2wav chunk_size With `async_chunk`: + - Stages can start processing as soon as chunks are available - Overlapping execution across stages - Reduced latency and improved throughput @@ -31,6 +33,7 @@ With `async_chunk`: - Async scheduling: Chunk IO (get/put) overlaps with compute via background threads so the scheduler is not blocked waiting for chunks ## Performance + 1. **Reduced Latency**: Next stage can start processing immediately 2. **Streaming Support**: Enables streaming for audio generation 3. **IO-Compute Overlap**: Chunk retrieval happens asynchronously while other requests compute @@ -49,7 +52,6 @@ With `async_chunk`: | text 100 | text 100+audio | True | 1 | 4 | 50 | 7692.69 | 103.96 | 10.22 | 785.85 | 0.29 | 10.12 | | text 100 | text 100+audio | True | 1 | 10 | 50 | 11152.71 | 685.60 | 17.64 | 1628.88 | 0.41 | 17.62 | - Performance data collected on H800 GPUs through comprehensive benchmarking with cudagraph enabled. text input uses random dataset. Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTFP)—e.g. ~92% at concurrency 1 (6.5s→0.52s)—and improves E2E latency (e.g. ~6% at conc 1, ~17% at conc 10). RTF (Real Time Factor) also improves with async_chunk on (e.g. ~8% at conc 1: 0.24→0.22, ~16% at conc 10: 0.49→0.41). Enabling **Code2Wav batch size 64** (vs 1) improves E2E and TTFP at higher concurrency when async_chunk is off (e.g. ~30% at conc 10: 13.5s→9.5s E2E, 13.4s→9.4s TTFP). @@ -91,7 +93,7 @@ The following diagram illustrates the **Async Chunk Architecture** for multi-sta **Diagram Legend:** | Step | Stage Type | Description | -|------|-----------|------------| +| ------ | ----------- | ------------ | | `prefill` | Initialization | Context processing, KV cache initialization | | `decode` | Autoregressive | Token-by-token generation in AR stages | | `codes` | Audio Encoding | RVQ codec codes from Talker stage | @@ -100,14 +102,16 @@ The following diagram illustrates the **Async Chunk Architecture** for multi-sta ### Data Flow #### Stage 0: Thinker (Multimodal Understanding + Text Generation) + - **Prefill**: Processes multimodal input (text/image/audio/video), initializes KV cache - **Decode Loop**: Generates text tokens autoregressively - **Chunk Triggers**: Each decode step (typically `chunk_size=1`) can trigger downstream processing - **Dual Output**: - - **Text Stream**: `text_0`, `text_1`, `text_2`... `text_n` streamed to output - - **Hidden States**: Passed to Talker stage for audio synthesis + - **Text Stream**: `text_0`, `text_1`, `text_2`... `text_n` streamed to output + - **Hidden States**: Passed to Talker stage for audio synthesis #### Stage 1: Talker (Text → RVQ Audio Codes) + - **Prefill**: Receives hidden states from Thinker as semantic condition - **Decode Loop**: Generates RVQ codec codes autoregressively - **Accumulation**: Codes accumulate to `codec_chunk_frames` (default=25) before forwarding @@ -115,12 +119,14 @@ The following diagram illustrates the **Async Chunk Architecture** for multi-sta - **Output**: `codes` blocks (chunk 0, 1, ... n) sent to Code2Wav #### Stage 2: Code2Wav (Vocoder Decoder) + - **Non-Autoregressive**: Processes RVQ codes in parallel batches - **Streaming Decode**: Converts codes to audio waveforms chunk-by-chunk - **Batching**: Supports batched inference for multiple concurrent requests - **Output**: Audio segments `audio_0`, `audio_1`, ... `audio_n` #### Stage 3: Output (Dual Stream) + - **Text Streaming**: `text_0` → `text_1` → `text_2` → ... (user sees response in real-time) - **Audio Streaming**: `audio_0` → `audio_1` → ... (user hears audio progressively) @@ -145,6 +151,7 @@ Total: ~3.5s, TTFP: ~0.5s ``` #### Sequential Flow (for comparison) +

@@ -155,6 +162,7 @@ Total: ~3.5s, TTFP: ~0.5s In sequential mode, each stage must wait for the previous stage to complete entirely before starting. ### Async Chunk System Architecture +

@@ -162,7 +170,6 @@ In sequential mode, each stage must wait for the previous stage to complete enti

- ### Key Components 1. **OmniConnector**: Inter-stage data transport only @@ -297,7 +304,7 @@ suffix and enforces the matching signature at startup. ### Naming convention and roles | Suffix | Registry kind | Runs where | Role | -|--------|---------------|------------|------| +| -------- | --------------- | ------------ | ------ | | `*_full_payload` | `producer_full_payload` | Worker (producer-side) | Packs the accumulated stage output into an `OmniPayload` and ships it through the connector (`FullPayloadProducer`) | | `*_async_chunk` | `producer_async_chunk` | Scheduler (producer-side) | Streams one chunk per call while `async_chunk` is enabled (`AsyncChunkProducer`) | | `*_token_only` | `placeholder_prompt_builder` | Orchestrator (consumer-side) | Allocates downstream KV slots only; bulk tensors arrive via the connector (`PlaceholderPromptBuilder`) | @@ -383,7 +390,7 @@ name-driven (suffix-based): returns a `ProcessorSpec` whose `fn` is the same callable the legacy lookup produced. -### P8b dual entry: forward and prewarm placeholders +### Dual entry: forward and prewarm placeholders The `*_token_only` placeholder builder is exposed through two entry points so the sync forward path and the async-chunk prewarm path share the same `_common` diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 5202eba8e2d..5c9a25aab9a 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -1455,11 +1455,11 @@ def test_send_single_request_personaplex_pending_frame_is_not_segment_boundary( def test_send_single_request_passes_new_token_ids_to_processor(build_adapter, monkeypatch): """A producer that accepts ``new_token_ids`` receives the step's tokens - through the production ``_send_single_request`` path (#6801 P1 #1). The - signature cache (``_processor_accepts_step_tokens``) must be initialized in - ``__init__`` before the base class starts the save thread; without it every - send raised AttributeError that ``_send_single_request_for_generation`` - swallowed, silently dropping the chunk. + through the production ``_send_single_request`` path. The signature cache + (``_processor_accepts_step_tokens``) must be initialized in ``__init__`` + before the base class starts the save thread; without it every send raised + AttributeError that ``_send_single_request_for_generation`` swallowed, + silently dropping the chunk. """ adapter, connector = build_adapter(stage_id=1) request = _req("req-new-tokens", RequestStatus.WAITING, external_req_id="ext-new-tokens") @@ -1493,10 +1493,10 @@ def new_tokens_processor(**kwargs): def test_send_single_request_normal_chunk_flushes_processor(build_adapter, monkeypatch): """A non-terminal chunk with a producer that does **not** accept - ``new_token_ids`` must still reach the wire (#6801 P1 #1). Before the - cache re-initialization the ``_accepts_new_token_ids`` probe raised - AttributeError on every call, the processor was never invoked, and the - ordinary chunk was silently dropped at the ``payload_data is None`` guard. + ``new_token_ids`` must still reach the wire. Before the cache + re-initialization the ``_accepts_new_token_ids`` probe raised AttributeError + on every call, the processor was never invoked, and the ordinary chunk was + silently dropped at the ``payload_data is None`` guard. """ adapter, connector = build_adapter(stage_id=1) request = _req("req-normal-chunk", RequestStatus.WAITING, external_req_id="ext-normal-chunk") diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index b6836c04f2f..01ac96012f2 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -385,7 +385,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None: # --------------------------------------------------------------------------- -# P2 deep-dive: dead-processor hint must not flag the selected sync hook. +# Dead-processor hint must not flag the selected sync hook. # --------------------------------------------------------------------------- @@ -451,7 +451,7 @@ class _Client: # --------------------------------------------------------------------------- -# P3 deep-dive: prewarm placeholder builder resolution is cached per stage. +# Prewarm placeholder builder resolution is cached per stage. # --------------------------------------------------------------------------- diff --git a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py index 67ad211f880..0691694bb35 100644 --- a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py +++ b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py @@ -49,9 +49,9 @@ def test_thinker2talker_skips_unfinished() -> None: def test_thinker2talker_via_process_engine_inputs_uses_configured_talker_len() -> None: """The stage client's ``_stage_hf_config`` reaches the legacy c2pos processor as the 4th positional ``next_stage_hf_config`` so the native - talker placeholder is sized to ``talker_init_len`` (#6801 P1 #2). Before - the fix ``ctx.streaming_context`` (None) was passed instead and the - placeholder fell back to one token. + talker placeholder is sized to ``talker_init_len``. Before the fix + ``ctx.streaming_context`` (None) was passed instead and the placeholder + fell back to one token. """ from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient @@ -70,7 +70,7 @@ def test_thinker2talker_via_process_engine_inputs_uses_configured_talker_len() - def test_thinker2talker_via_process_engine_inputs_falls_back_without_config() -> None: """Without ``_stage_hf_config`` the legacy processor keeps its one-token - placeholder fallback, preserving backward-compatible semantics (#6801). + placeholder fallback, preserving backward-compatible semantics. """ from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient From 679cfcdc4bfc78923fe0ca06e1529b003cdc6dd4 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Sat, 10 Oct 2026 14:42:01 +0800 Subject: [PATCH 36/41] =?UTF-8?q?refactor(stage=5Finput=5Fprocessors):=20R?= =?UTF-8?q?FC=20#4872=20Tier=201/2/2.5/2.75/2.8=20=E2=80=94=20table-driven?= =?UTF-8?q?=20registry,=20unified=20dispatch,=20doc=20de-dup?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements the locked Tier 1+2+2.5+2.75+2.8 scope across 19 files (1011 insertions / 1742 deletions, net -731): - _registry.py: table-driven validators + dead-processor hints (C-c/C-g/C-h) - _dispatch.py: wrap-time legacy shape fixation + memoized adapter (C-a) - orchestrator.py / stage_init_utils.py / stage_engine_core_client.py: prewarm builder resolved once at metadata time (C-b); client reads sync_process_input_func + prewarm_placeholder_builder from StageMetadata - _common.py / qwen3_omni.py: docstring + comment slimming (C-d/C-e) - tests: table-driven refactor of test_registry, test_common_helpers_golden, test_dispatch_wrap_once, test_placeholder_parity, test_constants_consistency, test_orchestrator_stage_input_bridge (Tier 2/2.5) - docs: single-source contract knowledge in async_chunk.md, adding_omni_model.md, adding_tts_model.md (Tier 2.8) - tests/engine/test_orchestrator_kv_sender_info.py: extend mock metadata with sync_process_input_func/prewarm_placeholder_builder to match the new client __init__ contract (fixes 2 PR-introduced UT regressions) Validation: 856 PR-scoped tests passed; tests/engine 1543 passed with only no-GPU CUDA-build environment-baseline failures remaining. Signed-off-by: RyanYun09 --- docs/contributing/model/adding_omni_model.md | 33 +- docs/contributing/model/adding_tts_model.md | 13 +- docs/design/feature/async_chunk.md | 52 +- .../test_orchestrator_kv_sender_info.py | 2 + .../test_orchestrator_stage_input_bridge.py | 283 +++------ .../stage_input_processors/conftest.py | 37 +- .../test_common_helpers_golden.py | 296 ++++------ .../test_constants_consistency.py | 19 +- .../test_dispatch_wrap_once.py | 108 ++-- .../test_placeholder_parity.py | 96 ++- .../stage_input_processors/test_registry.py | 549 ++++++++---------- vllm_omni/engine/orchestrator.py | 98 +--- vllm_omni/engine/stage_engine_core_client.py | 6 +- vllm_omni/engine/stage_init_utils.py | 79 ++- .../stage_input_processors/__init__.py | 29 +- .../stage_input_processors/_common.py | 160 ++--- .../stage_input_processors/_dispatch.py | 344 +++-------- .../stage_input_processors/_registry.py | 480 ++++++--------- .../stage_input_processors/qwen3_omni.py | 69 +-- 19 files changed, 1011 insertions(+), 1742 deletions(-) diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 3f0124b3075..c5e1b4b7cb7 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -500,9 +500,10 @@ Create stage transition processors in `vllm_omni/model_executor/stage_input_proc A model that supports **both** `async_chunk=false` and `async_chunk=true` must implement all three processors for every inter-stage edge — `*_full_payload`, -`*_async_chunk` and `*_token_only`. The `*_token_only` placeholder builder is -required in both modes: in non-async mode it builds the forward placeholder, in -async mode the orchestrator reuses it to prewarm the downstream stage. +`*_async_chunk` and `*_token_only` (the `*_token_only` placeholder is required +in both modes). See the +[Stage Input Processor Contract](../../design/feature/async_chunk.md#stage-input-processor-contract) +for the full contract, naming convention, and registry rules. Orchestrator-facing builders are normalized to the C1 contract `(source_outputs, ctx: OrchestratorInputContext)` by @@ -556,12 +557,10 @@ def thinker2talker_async_chunk( ... ``` -The `OrchestratorInputContext` carried by the C1 contract has `prompt`, -`requires_multimodal_data`, `streaming_context` and `sampling_params` (and -deliberately has no `model_config` field). The producer-side builders -(`*_full_payload`, `*_async_chunk`) are keyword-only and never receive an -`OrchestratorInputContext`; `pooling_output` / `multimodal_output` are -load-bearing keyword names of the connector data plane. +The producer-side builders (`*_full_payload`, `*_async_chunk`) are +keyword-only and never receive an `OrchestratorInputContext`; +`pooling_output` / `multimodal_output` are load-bearing keyword names of the +connector data plane. Wire these in `pipeline.py`: @@ -594,19 +593,15 @@ register_processor(f"{_PROC}.thinker2talker_token_only", "placeholder_prompt_bui At startup the runtime resolves each processor through `resolve_processor(path, expected_kind=...)` — import -> `infer_kind` -> -`validate_processor` -> expected-kind check — and returns a `ProcessorSpec` -whose `fn` is the same callable a legacy `getattr(importlib.import_module(...), -...)` lookup would produce. Hard contract violations raise -`ProcessorValidationError`; soft mismatches (e.g. a `*_full_payload` builder -declaring `multimodal_output` instead of `pooling_output`) emit a -`RuntimeWarning`. New model processors should therefore match the suffix -convention and the corresponding signature exactly. +`validate_processor` -> expected-kind check, returning a `ProcessorSpec` whose +`fn` matches the legacy lookup. Hard violations raise +`ProcessorValidationError`; soft mismatches emit a `RuntimeWarning`. New model +processors should therefore match the suffix convention and the corresponding +signature exactly. Do **not** add a no-suffix `thinker2talker` when a `sync_process_input_func` is already declared — `_select_processor_funcs()` always prefers the `*_token_only` -hook in non-async mode, so the bare function would never run. See the -[Stage Input Processor Contract](../../design/feature/async_chunk.md#stage-input-processor-contract) -section of the async chunk design document for the full contract. +hook in non-async mode, so the bare function would never run. ## Testing diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index f5e027c7074..4754b1f6460 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -553,11 +553,6 @@ def ar2decoder_token_only( ... ``` -`OrchestratorInputContext` carries `prompt`, `requires_multimodal_data`, -`streaming_context` and `sampling_params` (and deliberately has no -`model_config` field). Legacy positional shapes are adapted through -`wrap_orchestrator_processor` with a `DeprecationWarning`. - Names that do not follow the suffix convention can register an explicit kind override at module import time: @@ -567,12 +562,10 @@ from vllm_omni.model_executor.stage_input_processors import register_processor register_processor(f"{_PROC}.ar2decoder_token_only", "placeholder_prompt_builder") ``` -At startup the runtime resolves each processor through -`resolve_processor(path, expected_kind=...)` (import -> `infer_kind` -> -`validate_processor` -> expected-kind check). Hard contract violations raise -`ProcessorValidationError`; soft mismatches emit a `RuntimeWarning`. See the +For the registry validation rules, legacy-shape adaptation, and the full +contract, see the [Stage Input Processor Contract](../../design/feature/async_chunk.md#stage-input-processor-contract) -section of the async chunk design document for the full contract. +section of the async chunk design document. ### Data structures diff --git a/docs/design/feature/async_chunk.md b/docs/design/feature/async_chunk.md index dc7b531e166..853f26e3170 100644 --- a/docs/design/feature/async_chunk.md +++ b/docs/design/feature/async_chunk.md @@ -132,7 +132,7 @@ The following diagram illustrates the **Async Chunk Architecture** for multi-sta ### Execution Timeline -``` +```text Timeline: Parallel vs Sequential Sequential (async_chunk=false): @@ -348,23 +348,16 @@ def my_token_only( ... ``` -`OrchestratorInputContext` carries the transition metadata and deliberately has -no `model_config` field (a processor that needs the model config reads it -through the upstream stage closure, never through this context): - -```python -@dataclass(frozen=True) -class OrchestratorInputContext: - prompt: Any | None = None - requires_multimodal_data: bool = False - streaming_context: Any | None = None - sampling_params: Any | None = None -``` +`OrchestratorInputContext` (defined in `_dispatch.py`) carries `prompt`, +`requires_multimodal_data`, `streaming_context`, `sampling_params`, and the +target-stage config fields `target_model_config` / `next_stage_hf_config` +(populated by the stage client from the stage's own vllm config); it +deliberately has no `model_config` field. Processors that already accept `ctx` are used unchanged. Legacy positional -shapes (C0 3-arg, C2 placeholder with `streaming_context`, C3 diffusion with -`sampling_params`, C4 `moss_tts.talker2codec` multi-source) are adapted by -`wrap_orchestrator_processor` / `invoke_orchestrator_processor` and emit a +shapes (C0 3-arg, C2 placeholder, C3 diffusion, C4 `moss_tts.talker2codec` +multi-source) are adapted once by `wrap_orchestrator_processor` / +`invoke_orchestrator_processor` (see `_dispatch.py`) and emit a `DeprecationWarning`. ### Registry and startup validation @@ -381,14 +374,12 @@ name-driven (suffix-based): suffixes -> `diffusion_input_builder`, no suffix -> `legacy_orchestrator_builder`, `moss_tts.talker2codec` -> `legacy_multi_source`. -- `validate_processor(fn, *, kind, path, stage_config=None)` — hard contract - violations raise `ProcessorValidationError`; soft mismatches emit a - `RuntimeWarning`. +- `validate_processor(fn, *, kind, path, stage_config=None)` — hard violations + raise `ProcessorValidationError`; soft mismatches emit a `RuntimeWarning`. - `resolve_processor(path, *, expected_kind=None, stage_config=None)` — the - drop-in replacement for the legacy `getattr(importlib.import_module(...), ...)` - lookups: imports, infers, validates, optionally checks `expected_kind`, and - returns a `ProcessorSpec` whose `fn` is the same callable the legacy lookup - produced. + drop-in replacement for legacy `getattr(importlib.import_module(...), ...)` + lookups: import -> infer -> validate -> optional expected-kind check, + returning a `ProcessorSpec` whose `fn` matches the legacy lookup. ### Dual entry: forward and prewarm placeholders @@ -402,13 +393,14 @@ length / packing helpers: - `build_prewarm_placeholder(*, stage0_prompt, ctx, downstream_stage_id)` — async-chunk mode has no upstream `source_outputs` yet at prewarm time, so the length is a best-effort estimate from the stage-0 input prompt - (`mode="stage0_only"`, i.e. `len(stage0_prompt)`). The connector fixup path - (`adapter.construct_next_stage_streaming_input_prompt`) replaces the estimate - with the real length once the upstream chunk arrives. - -The orchestrator routes prewarm through `_prewarm_async_chunk_stages`, which -reads `build_prewarm_placeholder` off the resolved `*_token_only` function -object (both builders are attached as attributes on the exported function). + (`mode="stage0_only"`, matching the legacy + `adapter.compute_talker_prompt_ids_length`). The connector fixup path replaces + the estimate with the real length once the upstream chunk arrives. + +The dual-entry builder is resolved once at stage-metadata extraction time +(`stage_init_utils._resolve_prewarm_placeholder_builder`) and exposed on the +stage client as `prewarm_placeholder_builder`; the orchestrator's +`_prewarm_async_chunk_stages` reads it from there. ### Three-state async gate diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 5f3d1147145..d6cc196f0cd 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -97,6 +97,8 @@ def test_stage_engine_core_client_builds_payload_sender_info_after_base_init( prompt_transform_func=None, prompt_expand_func=None, custom_process_input_func=None, + sync_process_input_func=None, + prewarm_placeholder_builder=None, ) def fake_base_init(self, config, *_args, **_kwargs): diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 01ac96012f2..4d5e100508a 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -275,14 +275,18 @@ async def test_async_prewarm_skips_stage_with_custom_process_input_func() -> Non async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: - """Async-chunk prewarm resolves ``sync_process_input_func`` (the - ``*_token_only`` path) to ``build_prewarm_placeholder`` and uses its estimate. - - The placeholder length is ``stage0_only`` = the Qwen chat-template scan on - the stage-0 prompt. The synthetic ``[1, 2]`` prompt has no chat-template - marker, so the scan yields 0 and the builder floors it to a 1-token - placeholder (best-effort; the connector fixup path replaces it later). + """Async-chunk prewarm uses the metadata-resolved ``build_prewarm_placeholder``. + + The builder is resolved into stage metadata at init time (see + ``_resolve_prewarm_placeholder_builder``) and read off the stage client — + the orchestrator never re-resolves. The placeholder length is + ``stage0_only`` = the Qwen chat-template scan on the stage-0 prompt. The + synthetic ``[1, 2]`` prompt has no chat-template marker, so the scan yields + 0 and the builder floors it to a 1-token placeholder (best-effort; the + connector fixup path replaces it later). """ + from vllm_omni.engine.stage_init_utils import _resolve_prewarm_placeholder_builder + orchestrator = object.__new__(Orchestrator) stage0 = FakePrewarmPool("sender") stage1 = FakePrewarmPool("receiver") @@ -292,6 +296,11 @@ async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: sync_process_input_func=( "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only" ), + prewarm_placeholder_builder=_resolve_prewarm_placeholder_builder( + sync_path=("vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only"), + custom_fn=None, + stage_config=None, + ), ) orchestrator.stage_pools = [stage0, stage1] orchestrator._emit_tx_edge = lambda **_kwargs: None @@ -451,7 +460,7 @@ class _Client: # --------------------------------------------------------------------------- -# Prewarm placeholder builder resolution is cached per stage. +# Prewarm placeholder builder is resolved into stage metadata at init time. # --------------------------------------------------------------------------- @@ -464,13 +473,14 @@ def __init__(self, fn: Any) -> None: self.fn = fn -def test_prewarm_builder_resolved_once_and_cached() -> None: - """A stage with a builder resolves it once and caches it. +def test_prewarm_builder_resolved_once_into_stage_metadata() -> None: + """The dual-entry builder is resolved exactly once, at metadata time. - Repeated async requests must not re-run the importlib resolution or the - registry validation. + ``_resolve_prewarm_placeholder_builder`` reads ``build_prewarm_placeholder`` + off the resolved ``sync_process_input_func`` fn (with the module-level + sibling fallback) so the orchestrator never re-resolves per async request. """ - from vllm_omni.model_executor import stage_input_processors as sip + from vllm_omni.engine.stage_init_utils import _resolve_prewarm_placeholder_builder def _build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] return {"prompt_token_ids": [0]} @@ -480,46 +490,43 @@ def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # _sync_hook.build_prewarm_placeholder = _build_prewarm_placeholder # type: ignore[attr-defined] - class _Client: - sync_process_input_func = "pkg.mod.fake_token_only" - - orch = _prewarm_orchestrator_with_client(_Client()) - with patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)) as mock_resolve: - first = orch._get_prewarm_placeholder_builder(0) - second = orch._get_prewarm_placeholder_builder(0) - assert first is _build_prewarm_placeholder - assert second is _build_prewarm_placeholder + with patch( + "vllm_omni.engine.stage_init_utils.resolve_processor", + return_value=_FakeProcessorSpec(_sync_hook), + ) as mock_resolve: + builder = _resolve_prewarm_placeholder_builder( + sync_path="pkg.mod.fake_token_only", + custom_fn=None, + stage_config=None, + ) + assert builder is _build_prewarm_placeholder assert mock_resolve.call_count == 1 -def test_prewarm_builder_miss_warns_once_and_caches() -> None: - """A stage WITHOUT a builder resolves once, warns once, then caches a miss. - - Qwen3-TTS and other models have no ``build_prewarm_placeholder``; without - per-stage caching every async request re-resolved and re-warned. - """ - from vllm_omni.engine import orchestrator as orch_module - from vllm_omni.model_executor import stage_input_processors as sip +def test_prewarm_builder_miss_warns_once_then_keeps_inline_estimate() -> None: + """A stage WITHOUT a builder warns once at init and the prewarm keeps the + inline estimate — the orchestrator never re-resolves or re-warns.""" + from vllm_omni.engine import stage_init_utils as siu + from vllm_omni.engine.stage_init_utils import _resolve_prewarm_placeholder_builder def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] return [] - class _Client: - sync_process_input_func = "pkg.mod.fake_token_only" - - orch = _prewarm_orchestrator_with_client(_Client()) fake_module = SimpleNamespace() with ( - patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)) as mock_resolve, + patch( + "vllm_omni.engine.stage_init_utils.resolve_processor", + return_value=_FakeProcessorSpec(_sync_hook), + ) as mock_resolve, patch("importlib.import_module", return_value=fake_module) as mock_import, - patch.object(orch_module.logger, "warning") as mock_warn, + patch.object(siu.logger, "warning") as mock_warn, ): - first = orch._get_prewarm_placeholder_builder(0) - second = orch._get_prewarm_placeholder_builder(0) - third = orch._get_prewarm_placeholder_builder(0) - assert first is None - assert second is None - assert third is None + builder = _resolve_prewarm_placeholder_builder( + sync_path="pkg.mod.fake_token_only", + custom_fn=None, + stage_config=None, + ) + assert builder is None assert mock_resolve.call_count == 1 assert mock_import.call_count == 1 inline_estimates = [call for call in mock_warn.call_args_list if call.args and "inline estimate" in call.args[0]] @@ -549,180 +556,60 @@ def _override_orchestrator(client: Any, prewarm_len: int = 37) -> Orchestrator: return orch -def test_prewarm_override_keeps_configured_length_on_builder_miss() -> None: - """A stage WITHOUT a builder keeps ``async_chunk_prewarm_prompt_len``. - - Regression for amy-why-3459's review: the configured override (e.g. the - Nemotron deployment's 37-token speaker-prompt prefill) must survive the - inline-estimate fallback, not be replaced by the generic upstream estimate. - """ - from vllm_omni.engine import orchestrator as orch_module - from vllm_omni.model_executor import stage_input_processors as sip - - def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] - return [] - - class _Client: - sync_process_input_func = "pkg.mod.fake_token_only" - - orch = _override_orchestrator(_Client()) - fake_module = SimpleNamespace() - req_state = OrchestratorRequestState( - request_id="req-override-miss", - prompt={"prompt": "hi"}, - sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], - final_stage_id=0, - ) - - with ( - patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)), - patch("importlib.import_module", return_value=fake_module), - patch.object(orch_module.logger, "warning"), - ): - base_input = orch._build_prewarm_placeholder_input(0, "req-override-miss", [1, 2], req_state) - - # The configured 37-token prefill boundary wins over the tiny [1, 2] estimate. - assert base_input["prompt_token_ids"] == [0] * 37 - assert base_input["multi_modal_data"] is None - assert base_input["mm_processor_kwargs"] is None - - -def test_prewarm_override_keeps_configured_length_on_builder_failure() -> None: - """A builder that RAISES falls back to the inline estimate AND keeps the override.""" - from vllm_omni.engine import orchestrator as orch_module - from vllm_omni.model_executor import stage_input_processors as sip - - def _broken_build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] - raise RuntimeError("boom") - - def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] - return [] - - _sync_hook.build_prewarm_placeholder = _broken_build_prewarm_placeholder # type: ignore[attr-defined] - - class _Client: - sync_process_input_func = "pkg.mod.fake_token_only" - - orch = _override_orchestrator(_Client()) - req_state = OrchestratorRequestState( - request_id="req-override-fail", - prompt={"prompt": "hi"}, - sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], - final_stage_id=0, - ) - - with ( - patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)), - patch.object(orch_module.logger, "warning"), - ): - base_input = orch._build_prewarm_placeholder_input(0, "req-override-fail", [1, 2], req_state) - - # Builder failure must not lose the configured override either. - assert base_input["prompt_token_ids"] == [0] * 37 - assert base_input["multi_modal_data"] is None - assert base_input["mm_processor_kwargs"] is None - +#: Builder variants shared by the override table below (a None builder is the +#: "no dual-entry builder" miss case; a raising builder exercises the fallback). +def _build_prewarm_placeholder_3(**kwargs): # type: ignore[no-untyped-def] + return [{"prompt_token_ids": [0, 0, 0]}] -def test_prewarm_override_keeps_configured_length_on_builder_success() -> None: - """Builder hit must still honor ``async_chunk_prewarm_prompt_len`` (amy-why-3459 review of HEAD ad81ac3a).""" - from vllm_omni.engine import orchestrator as orch_module - from vllm_omni.model_executor import stage_input_processors as sip - def _build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] - return [{"prompt_token_ids": [0, 0, 0]}] +def _build_prewarm_placeholder_789(**kwargs): # type: ignore[no-untyped-def] + return [{"prompt_token_ids": [7, 8, 9]}] - def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] - return [] - _sync_hook.build_prewarm_placeholder = _build_prewarm_placeholder # type: ignore[attr-defined] +def _broken_build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] + raise RuntimeError("boom") - class _Client: - sync_process_input_func = "pkg.mod.fake_token_only" - orch = _override_orchestrator(_Client()) - req_state = OrchestratorRequestState( - request_id="req-override-success", - prompt={"prompt": "hi"}, - sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], - final_stage_id=0, - ) - - with ( - patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)), - patch.object(orch_module.logger, "warning"), - ): - base_input = orch._build_prewarm_placeholder_input(0, "req-override-success", [1, 2], req_state) - - assert base_input["prompt_token_ids"] == [0] * 37 - assert base_input["multi_modal_data"] is None - assert base_input["mm_processor_kwargs"] is None - - -def test_prewarm_no_override_keeps_builder_output_on_builder_success() -> None: - """Without a configured override, a builder hit keeps its own output.""" +@pytest.mark.parametrize( + "builder,prewarm_len,expected_ids", + [ + # A stage WITHOUT a builder keeps the configured override through the + # inline-estimate fallback (Nemotron's 37-token speaker-prompt prefill). + pytest.param(None, 37, [0] * 37, id="builder-miss-keeps-override"), + # A builder that RAISES falls back to the inline estimate AND keeps it. + pytest.param(_broken_build_prewarm_placeholder, 37, [0] * 37, id="builder-failure-keeps-override"), + # Builder hit must still honor the override (amy-why-3459 review of HEAD ad81ac3a). + pytest.param(_build_prewarm_placeholder_3, 37, [0] * 37, id="builder-success-keeps-override"), + # Without a configured override, a builder hit keeps its own output. + pytest.param(_build_prewarm_placeholder_789, None, [7, 8, 9], id="no-override-keeps-builder-output"), + # A configured length of 1 is not a prefill override: builder output wins. + pytest.param(_build_prewarm_placeholder_3, 1, [0, 0, 0], id="len-one-not-a-prefill-override"), + ], +) +def test_prewarm_prompt_len_override(builder, prewarm_len, expected_ids) -> None: from vllm_omni.engine import orchestrator as orch_module - from vllm_omni.model_executor import stage_input_processors as sip - - def _build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] - return [{"prompt_token_ids": [7, 8, 9]}] - - def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] - return [] - - _sync_hook.build_prewarm_placeholder = _build_prewarm_placeholder # type: ignore[attr-defined] class _Client: sync_process_input_func = "pkg.mod.fake_token_only" - orch = _prewarm_orchestrator_with_client(_Client()) # no async_chunk_prewarm_prompt_len - req_state = OrchestratorRequestState( - request_id="req-no-override-success", - prompt={"prompt": "hi"}, - sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], - final_stage_id=0, - ) - - with ( - patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)), - patch.object(orch_module.logger, "warning"), - ): - base_input = orch._build_prewarm_placeholder_input(0, "req-no-override-success", [1, 2], req_state) - - assert base_input["prompt_token_ids"] == [7, 8, 9] - assert base_input["multi_modal_data"] is None - assert base_input["mm_processor_kwargs"] is None - + def __init__(self) -> None: + # Instance attribute: class-level functions bind self on access. + self.prewarm_placeholder_builder = builder -def test_prewarm_override_len_one_keeps_builder_output_on_builder_success() -> None: - """A configured length of 1 is not a prefill override: builder output wins.""" - from vllm_omni.engine import orchestrator as orch_module - from vllm_omni.model_executor import stage_input_processors as sip - - def _build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] - return [{"prompt_token_ids": [0, 0, 0]}] - - def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] - return [] - - _sync_hook.build_prewarm_placeholder = _build_prewarm_placeholder # type: ignore[attr-defined] - - class _Client: - sync_process_input_func = "pkg.mod.fake_token_only" - - orch = _override_orchestrator(_Client(), prewarm_len=1) + if prewarm_len is None: + orch = _prewarm_orchestrator_with_client(_Client()) # no async_chunk_prewarm_prompt_len + else: + orch = _override_orchestrator(_Client(), prewarm_len=prewarm_len) req_state = OrchestratorRequestState( - request_id="req-len-one-success", + request_id="req-override", prompt={"prompt": "hi"}, sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], final_stage_id=0, ) - with ( - patch.object(sip, "resolve_processor", return_value=_FakeProcessorSpec(_sync_hook)), - patch.object(orch_module.logger, "warning"), - ): - base_input = orch._build_prewarm_placeholder_input(0, "req-len-one-success", [1, 2], req_state) + with patch.object(orch_module.logger, "warning"): + base_input = orch._build_prewarm_placeholder_input(0, "req-override", [1, 2], req_state) - assert base_input["prompt_token_ids"] == [0, 0, 0] + assert base_input["prompt_token_ids"] == expected_ids assert base_input["multi_modal_data"] is None assert base_input["mm_processor_kwargs"] is None diff --git a/tests/model_executor/stage_input_processors/conftest.py b/tests/model_executor/stage_input_processors/conftest.py index c9873a45d48..8be9a5e9ed8 100644 --- a/tests/model_executor/stage_input_processors/conftest.py +++ b/tests/model_executor/stage_input_processors/conftest.py @@ -1,26 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Test-support import fallback for ``vllm_omni`` pure-logic tests. - -These unit tests exercise ``vllm_omni`` helper logic without a running vLLM -runtime. When the real ``vllm`` package is not installed, importing -``vllm_omni`` submodules would fail at ``import vllm``; this conftest then -installs a *test-only* fallback that: - -1. Bypasses ``vllm_omni/__init__.py`` (so its ``patch`` / ``transformers`` - machinery is never executed) by pre-inserting a synthetic ``vllm_omni`` - package into ``sys.modules`` whose ``__path__`` points at the real - package directory. -2. Stubs the entire ``vllm`` package tree with permissive ``_Dummy`` classes - so module-level ``from vllm... import ...`` statements resolve. - -The fallback is active **only** when real ``vllm`` is not importable. When -vllm is available (e.g. on CI) it is a no-op and the real package is imported, -which guards against the fallback masking genuine import problems. - -Scope: only pure-logic / unit-level tests (no model execution). Any test that -needs a real vLLM runtime must be marked to run on GPU/NPU CI instead. +"""Test-only import fallback: lets pure-logic tests run without real ``vllm``. + +Bypasses ``vllm_omni/__init__.py`` via a synthetic package in ``sys.modules`` +and stubs the ``vllm`` tree with permissive ``_Dummy`` classes. Active only +when real ``vllm`` is unavailable (CI no-op), so it cannot mask genuine +import problems. Scope: pure-logic unit tests only. """ import importlib.abc @@ -38,21 +24,14 @@ class _Meta(type): - """Metaclass so that *class-level* attribute access on a _Dummy also - yields _Dummy (e.g. ``SomeStubClass.foo`` used in module-level code).""" + """Class-level attribute access on a ``_Dummy`` also yields ``_Dummy``.""" def __getattr__(cls, name): return _Dummy class _Dummy(metaclass=_Meta): - """Permissive stand-in for any vllm symbol. - - Supports subclassing (``__mro_entries__``), ``total=False`` TypedDict - style (``__init_subclass__``), ``X | None`` unions, and arbitrary - attribute / call / item access — enough to get module *import* past the - vllm dependency while keeping pure helper functions callable. - """ + """Permissive stand-in for any vllm symbol (import-grade compatibility).""" def __init__(self, *args, **kwargs): pass diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py index 77e61dea3da..9f232852426 100644 --- a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -68,14 +68,12 @@ def _is_tensor(value: Any) -> bool: # =========================================================================== -def _ensure_list_fn(module: str) -> Any: - """Return the callable that reproduces the module's legacy ``_ensure_list``. +_T2 = torch.tensor([1, 2]) +_T22 = torch.tensor([[1, 2], [3, 4]]) +_T24 = torch.tensor([1, 2, 3, 4]) - Each module maps to the ``_common`` variant that locks its legacy semantics. - ``diffusion.output_formatter`` is the intentional exception: its wrap-only - semantics are NOT expressible through ``_common`` variants and it keeps its - own local ``_ensure_list``. - """ + +def _ensure_list_cases(): c = _common_import() mapping = { "qwen3_omni": c.ensure_list_unchanged, @@ -83,59 +81,45 @@ def _ensure_list_fn(module: str) -> Any: "audex": c.ensure_list_flatten, "cosyvoice3": c.ensure_list, "step_audio2": c.ensure_list_preserve_none, - "diffusion.output_formatter": None, # resolved via module below + "diffusion.output_formatter": _import("diffusion.output_formatter")._ensure_list, } - fn = mapping[module] - if fn is not None: - return fn - # diffusion.output_formatter keeps wrap-only semantics locally. - return _import("diffusion.output_formatter")._ensure_list - - -def _ensure_list_cases(): + # (module, callable, input, expected, expect_typeerror) return [ - # (module, input, expected-or-None-if-tensor, expect_typeerror) - ("qwen3_omni", [1, 2], [1, 2], False), - ("qwen3_omni", _ConstantList([1, 2]), [1, 2], False), - # qwen3_omni returns non-list values *unchanged* (incl. None / tuple). - ("qwen3_omni", None, None, False), - ("qwen3_omni", (1, 2), (1, 2), False), - # qwen3_omni has NO tolist() branch: tensor comes back unchanged. - ("qwen3_omni", "TENSOR", None, False), - ("ming_flash_omni", [1, 2], [1, 2], False), - ("ming_flash_omni", _ConstantList([1, 2]), [1, 2], False), - ("ming_flash_omni", "TENSOR", [1, 2], False), - ("ming_flash_omni", (1, 2), [1, 2], False), - # ming_flash_omni: list(None) raises TypeError. - ("ming_flash_omni", None, None, True), - ("audex", None, [], False), - ("audex", "TENSOR", [1, 2, 3, 4], False), - ("audex", [1, 2], [1, 2], False), - ("audex", (1, 2), [1, 2], False), - # audex: list(5) raises TypeError. - ("audex", 5, None, True), - ("cosyvoice3", None, [], False), - ("cosyvoice3", (1, 2), [1, 2], False), - ("cosyvoice3", 5, [5], False), # TypeError -> [x] - ("cosyvoice3", _ConstantList([1, 2]), [1, 2], False), - ("cosyvoice3", [1, 2], [1, 2], False), - ("step_audio2", _ConstantList([1, 2]), [1, 2], False), - ("step_audio2", "TENSOR", [1, 2], False), - ("step_audio2", (1, 2), [1, 2], False), - # step_audio2: non-iterable scalar -> [x]; None is not iterable -> [None]. - ("step_audio2", 5, [5], False), - ("step_audio2", None, [None], False), - ("diffusion.output_formatter", [1, 2], [1, 2], False), - ("diffusion.output_formatter", None, [], False), - ("diffusion.output_formatter", 5, [5], False), + pytest.param(mapping["qwen3_omni"], [1, 2], [1, 2], False, id="qwen3_omni-list"), + pytest.param(mapping["qwen3_omni"], _ConstantList([1, 2]), [1, 2], False, id="qwen3_omni-constantlist"), + pytest.param(mapping["qwen3_omni"], None, None, False, id="qwen3_omni-none"), + pytest.param(mapping["qwen3_omni"], (1, 2), (1, 2), False, id="qwen3_omni-tuple"), + pytest.param(mapping["qwen3_omni"], _T2, None, False, id="qwen3_omni-tensor"), + pytest.param(mapping["ming_flash_omni"], [1, 2], [1, 2], False, id="ming_flash_omni-list"), + pytest.param( + mapping["ming_flash_omni"], _ConstantList([1, 2]), [1, 2], False, id="ming_flash_omni-constantlist" + ), + pytest.param(mapping["ming_flash_omni"], _T2, [1, 2], False, id="ming_flash_omni-tensor"), + pytest.param(mapping["ming_flash_omni"], (1, 2), [1, 2], False, id="ming_flash_omni-tuple"), + pytest.param(mapping["ming_flash_omni"], None, None, True, id="ming_flash_omni-none-raises"), + pytest.param(mapping["audex"], None, [], False, id="audex-none"), + pytest.param(mapping["audex"], _T24, [1, 2, 3, 4], False, id="audex-tensor"), + pytest.param(mapping["audex"], [1, 2], [1, 2], False, id="audex-list"), + pytest.param(mapping["audex"], (1, 2), [1, 2], False, id="audex-tuple"), + pytest.param(mapping["audex"], 5, None, True, id="audex-scalar-raises"), + pytest.param(mapping["cosyvoice3"], None, [], False, id="cosyvoice3-none"), + pytest.param(mapping["cosyvoice3"], (1, 2), [1, 2], False, id="cosyvoice3-tuple"), + pytest.param(mapping["cosyvoice3"], 5, [5], False, id="cosyvoice3-scalar"), + pytest.param(mapping["cosyvoice3"], _ConstantList([1, 2]), [1, 2], False, id="cosyvoice3-constantlist"), + pytest.param(mapping["cosyvoice3"], [1, 2], [1, 2], False, id="cosyvoice3-list"), + pytest.param(mapping["step_audio2"], _ConstantList([1, 2]), [1, 2], False, id="step_audio2-constantlist"), + pytest.param(mapping["step_audio2"], _T2, [1, 2], False, id="step_audio2-tensor"), + pytest.param(mapping["step_audio2"], (1, 2), [1, 2], False, id="step_audio2-tuple"), + pytest.param(mapping["step_audio2"], 5, [5], False, id="step_audio2-scalar"), + pytest.param(mapping["step_audio2"], None, [None], False, id="step_audio2-none"), + pytest.param(mapping["diffusion.output_formatter"], [1, 2], [1, 2], False, id="output_formatter-list"), + pytest.param(mapping["diffusion.output_formatter"], None, [], False, id="output_formatter-none"), + pytest.param(mapping["diffusion.output_formatter"], 5, [5], False, id="output_formatter-scalar"), ] -@pytest.mark.parametrize("module,inp,expected,expect_typeerror", _ensure_list_cases()) -def test_ensure_list_golden(module, inp, expected, expect_typeerror): - fn = _ensure_list_fn(module) - if inp == "TENSOR": - inp = torch.tensor([[1, 2], [3, 4]]) if module == "audex" else torch.tensor([1, 2]) +@pytest.mark.parametrize("fn,inp,expected,expect_typeerror", _ensure_list_cases()) +def test_ensure_list_golden(fn, inp, expected, expect_typeerror): if expect_typeerror: with pytest.raises(TypeError): fn(inp) @@ -143,9 +127,8 @@ def test_ensure_list_golden(module, inp, expected, expect_typeerror): result = fn(inp) if _is_tensor(result): # qwen3_omni may return a tensor unchanged. - assert _is_tensor(expected) or expected is None - if expected is None: - assert torch.equal(result, torch.tensor([1, 2])) + assert expected is None + assert torch.equal(result, _T2) return assert _norm(result) == expected @@ -196,18 +179,14 @@ def _codes_audio(tensor): return {"codes": {"audio": tensor}} -def _extract_last_frame_fn(module: str) -> Any: - """Return the callable that reproduces the module's legacy behavior. - - The per-module ``_extract_last_frame`` helpers were removed by the RFC - refactor. ``higgs_audio_v2`` is the intentional exception (its range - filter is not expressible through ``_common.extract_last_codec_frame`` and - it keeps its own local helper). All other modules map to the canonical - implementation with the parameters that lock their legacy semantics. - """ +def _extract_last_frame_cases(): + """Legacy per-module behavior: one callable + input per table row.""" c = _common_import() mapping = { - "higgs_audio_v2": None, # resolved below via module + # higgs_audio_v2 keeps its own local range-filtered helper (RFC refactor + # removed the other per-module helpers; its range filter is not + # expressible through _common.extract_last_codec_frame). + "higgs_audio_v2": _import("higgs_audio_v2")._extract_last_frame, "fish_speech": functools.partial( c.extract_last_codec_frame, key_path=("audio_codes",), @@ -223,98 +202,101 @@ def _extract_last_frame_fn(module: str) -> Any: # other modules). "voxtral_tts": functools.partial(c.extract_last_codec_frame, to_long=False), } - fn = mapping[module] - if fn is not None: - return fn - return _import("higgs_audio_v2")._extract_last_frame - - -@pytest.mark.parametrize( - "module,payload,expected,expect_error", - [ + higgs = mapping["higgs_audio_v2"] + fish = mapping["fish_speech"] + qwen3 = mapping["qwen3_tts"] + voxtral = mapping["voxtral_tts"] + return [ # higgs_audio_v2: codes.audio key path; range filter [-0, _NUM_REAL_CODES). - ("higgs_audio_v2", _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False), - ( - "higgs_audio_v2", - _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 1024]])), # >= _NUM_REAL_CODES - None, - False, - ), - ("higgs_audio_v2", _codes_audio(torch.tensor([[0, 1], [2, 3], [-1, 5]])), None, False), - ("higgs_audio_v2", _codes_audio(torch.tensor([1, 2, 3])), [1, 2, 3], False), - ("higgs_audio_v2", _codes_audio(torch.empty(0, 4)), None, False), - ("higgs_audio_v2", _codes_audio(torch.zeros(2, 2, 2)), None, True), # 3-D -> ValueError - ("higgs_audio_v2", {"codes": {}}, None, False), + pytest.param(higgs, _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False, id="higgs-last-row"), + # >= _NUM_REAL_CODES + pytest.param(higgs, _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 1024]])), None, False, id="higgs-range"), + pytest.param(higgs, _codes_audio(torch.tensor([[0, 1], [2, 3], [-1, 5]])), None, False, id="higgs-negative"), + pytest.param(higgs, _codes_audio(torch.tensor([1, 2, 3])), [1, 2, 3], False, id="higgs-1d"), + pytest.param(higgs, _codes_audio(torch.empty(0, 4)), None, False, id="higgs-empty"), + # 3-D -> ValueError + pytest.param(higgs, _codes_audio(torch.zeros(2, 2, 2)), None, True, id="higgs-3d-raises"), + pytest.param(higgs, {"codes": {}}, None, False, id="higgs-missing-audio"), # fish_speech: TOP-LEVEL audio_codes key; audio_code_valid mask; cpu+long. - ( - "fish_speech", + pytest.param( + fish, {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": torch.tensor([1, 1, 1])}, [4, 5], False, + id="fish-valid", ), - ( - "fish_speech", + pytest.param( + fish, {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": torch.tensor([1, 1, 0])}, None, False, + id="fish-invalid-last", ), - ( - "fish_speech", + pytest.param( + fish, {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": True}, [4, 5], False, + id="fish-valid-true", ), - ( - "fish_speech", + pytest.param( + fish, {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": False}, None, False, + id="fish-invalid-false", ), - ( - "fish_speech", - {"audio_codes": torch.tensor([[0, 1], [2, 3], [0, 0]])}, - None, - False, - ), # no valid key -> frame.any() False - ("fish_speech", {"audio_codes": torch.tensor([1, 2, 3])}, [1, 2, 3], False), + # no valid key -> frame.any() False + pytest.param(fish, {"audio_codes": torch.tensor([[0, 1], [2, 3], [0, 0]])}, None, False, id="fish-zero-frame"), + pytest.param(fish, {"audio_codes": torch.tensor([1, 2, 3])}, [1, 2, 3], False, id="fish-1d"), # 1-D semantics: flat frames bypass the validity/zero gate (amy-why-3459 review). - ("fish_speech", {"audio_codes": torch.zeros(8, dtype=torch.long)}, [0] * 8, False), - ( - "fish_speech", + pytest.param(fish, {"audio_codes": torch.zeros(8, dtype=torch.long)}, [0] * 8, False, id="fish-1d-zeros"), + pytest.param( + fish, {"audio_codes": torch.zeros(8, dtype=torch.long), "audio_code_valid": torch.tensor([False])}, [0] * 8, False, + id="fish-1d-zeros-invalid", ), - ( - "fish_speech", + pytest.param( + fish, {"audio_codes": torch.zeros(8, dtype=torch.long), "audio_code_valid": False}, [0] * 8, False, + id="fish-1d-zeros-false", ), - ("fish_speech", {"audio_codes": torch.zeros(2, 2, 2)}, None, True), - ("fish_speech", {}, None, False), + pytest.param(fish, {"audio_codes": torch.zeros(2, 2, 2)}, None, True, id="fish-3d-raises"), + pytest.param(fish, {}, None, False, id="fish-missing-codes"), # qwen3_tts: codes.audio; frame.any() gate; long (no cpu). - ("qwen3_tts", _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False), - ("qwen3_tts", _codes_audio(torch.tensor([[0, 1], [2, 3], [0, 0]])), None, False), - ("qwen3_tts", _codes_audio(torch.tensor([1, 2, 3])), [1, 2, 3], False), + pytest.param(qwen3, _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False, id="qwen3-last-row"), + pytest.param(qwen3, _codes_audio(torch.tensor([[0, 1], [2, 3], [0, 0]])), None, False, id="qwen3-zero-frame"), + pytest.param(qwen3, _codes_audio(torch.tensor([1, 2, 3])), [1, 2, 3], False, id="qwen3-1d"), # 1-D semantics: flat frames bypass the zero gate (legacy qwen3_tts behavior). - ("qwen3_tts", _codes_audio(torch.zeros(8, dtype=torch.long)), [0] * 8, False), - ("qwen3_tts", _codes_audio(torch.zeros(2, 2, 2)), None, True), - ("qwen3_tts", {"codes": {}}, None, False), + pytest.param(qwen3, _codes_audio(torch.zeros(8, dtype=torch.long)), [0] * 8, False, id="qwen3-1d-zeros"), + pytest.param(qwen3, _codes_audio(torch.zeros(2, 2, 2)), None, True, id="qwen3-3d-raises"), + pytest.param(qwen3, {"codes": {}}, None, False, id="qwen3-missing-audio"), # voxtral_tts: validated via _common.extract_last_codec_frame(to_long=False). - ("voxtral_tts", _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False), - ("voxtral_tts", {"codes": {"audio": [torch.tensor([1, 2]), torch.tensor([3, 4])]}}, [3, 4], False), - ("voxtral_tts", {"codes": {}}, None, False), - ], -) -def test_extract_last_frame_golden(module, payload, expected, expect_error): - fn = _extract_last_frame_fn(module) + pytest.param( + voxtral, _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False, id="voxtral-last-row" + ), + pytest.param( + voxtral, + {"codes": {"audio": [torch.tensor([1, 2]), torch.tensor([3, 4])]}}, + [3, 4], + False, + id="voxtral-list-audio", + ), + pytest.param(voxtral, {"codes": {}}, None, False, id="voxtral-missing-audio"), + ] + + +@pytest.mark.parametrize("fn,payload,expected,expect_error", _extract_last_frame_cases()) +def test_extract_last_frame_golden(fn, payload, expected, expect_error): if expect_error: with pytest.raises(ValueError): fn(payload) return - result = fn(payload) - assert _norm(result) == expected + assert _norm(fn(payload)) == expected # =========================================================================== @@ -329,34 +311,26 @@ def test_extract_last_frame_golden(module, payload, expected, expect_error): # =========================================================================== +_T2F = torch.tensor([1.0, 2.0]) + + @pytest.mark.parametrize( "inp,expected", [ - ("TENSOR", "TENSOR"), - ([], None), - ("LIST_TENSOR", "LIST_TENSOR"), - ([1, 2], None), - (None, None), - (5, None), + pytest.param(_T2F, _T2F, id="tensor"), + pytest.param([], None, id="empty-list"), + pytest.param([_T2F], _T2F, id="list-tensor"), + pytest.param([1, 2], None, id="list-int"), + pytest.param(None, None, id="none"), + pytest.param(5, None, id="int"), ], ) def test_to_cpu_tensor_glm_tts_golden(inp, expected): - c = _common_import() - fn = c.to_cpu_tensor - t = torch.tensor([1.0, 2.0]) - if inp == "TENSOR": - value = t - elif inp == "LIST_TENSOR": - value = [t] - else: - value = inp - result = fn(value) - if expected == "TENSOR": - assert torch.equal(result, t) - elif expected == "LIST_TENSOR": - assert torch.equal(result, t) + result = _common_import().to_cpu_tensor(inp) + if expected is None: + assert result is None else: - assert result is expected + assert torch.equal(result, expected) # =========================================================================== @@ -367,29 +341,17 @@ def test_to_cpu_tensor_glm_tts_golden(inp, expected): @pytest.mark.parametrize( "inp,expected", [ - ("SCALAR_TENSOR", [5]), - ("VEC_TENSOR", [1, 2, 3]), - ("MAT_TENSOR", [1, 2]), # 2-D -> first row only - ("NESTED_LIST", [1, 2]), # nested list -> first nested row only - ([1, 2, 3], [1, 2, 3]), - (None, []), - (5, [5]), + pytest.param(torch.tensor(5), [5], id="scalar-tensor"), + pytest.param(torch.tensor([1, 2, 3]), [1, 2, 3], id="vec-tensor"), + pytest.param(torch.tensor([[1, 2], [3, 4]]), [1, 2], id="mat-tensor"), # 2-D -> first row only + pytest.param([[1, 2], [3, 4]], [1, 2], id="nested-list"), # nested list -> first nested row only + pytest.param([1, 2, 3], [1, 2, 3], id="list"), + pytest.param(None, [], id="none"), + pytest.param(5, [5], id="int"), ], ) def test_to_token_id_list_dynin_omni_golden(inp, expected): - c = _common_import() - fn = c.to_token_id_list - if inp == "SCALAR_TENSOR": - value = torch.tensor(5) - elif inp == "VEC_TENSOR": - value = torch.tensor([1, 2, 3]) - elif inp == "MAT_TENSOR": - value = torch.tensor([[1, 2], [3, 4]]) - elif inp == "NESTED_LIST": - value = [[1, 2], [3, 4]] - else: - value = inp - assert fn(value) == expected + assert _common_import().to_token_id_list(inp) == expected def test_to_token_id_list_cosyvoice3_recursive_golden(): diff --git a/tests/model_executor/stage_input_processors/test_constants_consistency.py b/tests/model_executor/stage_input_processors/test_constants_consistency.py index e08f99a1bcb..bc57329cf78 100644 --- a/tests/model_executor/stage_input_processors/test_constants_consistency.py +++ b/tests/model_executor/stage_input_processors/test_constants_consistency.py @@ -46,6 +46,9 @@ QWEN2_5_CODEC_END_EXPECTED = 8294 QWEN3_CODEC_STOP_EXPECTED = 2150 +# Processor-side codec ids, single-sourced (pad, bos, eos). +QWEN3_CODEC_PROCESSOR_TUPLE = (_QWEN3_CODEC_PAD_TOKEN_ID, _QWEN3_CODEC_BOS_TOKEN_ID, _QWEN3_CODEC_EOS_TOKEN_ID) + def _pipeline_stop_token_ids(pipeline) -> list[int]: """Extract the (unique) ``stop_token_ids`` across a pipeline's stages.""" @@ -60,25 +63,17 @@ def _pipeline_stop_token_ids(pipeline) -> list[int]: def test_qwen3_codec_token_ids_processor_matches_model(): """4196/4197/4198: processor and model implementation must agree.""" - processor = ( - _QWEN3_CODEC_PAD_TOKEN_ID, - _QWEN3_CODEC_BOS_TOKEN_ID, - _QWEN3_CODEC_EOS_TOKEN_ID, - ) model = (TALKER_CODEC_PAD_TOKEN_ID, TALKER_CODEC_BOS_TOKEN_ID, QWEN3_MODEL_CODEC_EOS) - assert processor == model == QWEN3_CODEC_EXPECTED + assert QWEN3_CODEC_PROCESSOR_TUPLE == model == QWEN3_CODEC_EXPECTED def test_qwen3_codec_token_ids_match_hf_config(): """4196/4197/4198: processor constants must match the in-repo Qwen3-TTS HF config defaults (single source of truth).""" cfg = Qwen3TTSTalkerConfig() - assert (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) == QWEN3_CODEC_EXPECTED - assert ( - _QWEN3_CODEC_PAD_TOKEN_ID, - _QWEN3_CODEC_BOS_TOKEN_ID, - _QWEN3_CODEC_EOS_TOKEN_ID, - ) == (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) + hf = (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) + assert hf == QWEN3_CODEC_EXPECTED + assert QWEN3_CODEC_PROCESSOR_TUPLE == hf def test_qwen3_omni_runtime_talker_config_matches_processor(): diff --git a/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py index 846970d03e8..e5f93635424 100644 --- a/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py +++ b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py @@ -1,14 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Orchestrator dispatch wrap-once semantics. +"""Orchestrator dispatch wrap-once semantics (table-driven). -``invoke_orchestrator_processor`` runs per request (and per forward), so -``wrap_orchestrator_processor`` must not re-wrap — and must not re-emit the -legacy-contract ``DeprecationWarning`` — on every forward. The wrap result is -memoized per callable: the signature probe and the warning happen **once** per -processor (at first resolution/invocation), then the cached C1-compatible -wrapper is reused. +``invoke_orchestrator_processor`` runs per forward, so +``wrap_orchestrator_processor`` must not re-wrap — nor re-emit the legacy +``DeprecationWarning`` — on every forward: the result is memoized per callable. Pure-logic, CPU-only (no model loading, no vllm runtime). """ @@ -29,12 +26,12 @@ def _legacy_c2(source_outputs, prompt=None, requires_multimodal_data=False, streaming_context=None): - """A legacy C2 positional shell (like qwen3_omni.thinker2talker_token_only).""" + """Legacy C2 positional shell (like qwen3_omni.thinker2talker_token_only).""" return [source_outputs, prompt, requires_multimodal_data, streaming_context] def _legacy_c0(source_outputs, prompt=None, requires_multimodal_data=False): - """A legacy C0 3-arg shell.""" + """Legacy C0 3-arg shell.""" return [source_outputs, prompt, requires_multimodal_data] @@ -44,27 +41,27 @@ def _c1(source_outputs: list[Any], ctx: OrchestratorInputContext): def _legacy_c3_keyword_only(source_outputs, prompt=None, requires_multimodal_data=False, *, sampling_params=None): - """Legacy C3 shell whose ``sampling_params`` is keyword-only.""" + """C3 shell whose ``sampling_params`` is keyword-only.""" return [source_outputs, prompt, requires_multimodal_data, sampling_params] def _legacy_c3_required_keyword_only(source_outputs, prompt=None, requires_multimodal_data=False, *, sampling_params): - """Legacy C3 shell with a REQUIRED keyword-only ``sampling_params``.""" + """C3 shell with a REQUIRED keyword-only ``sampling_params``.""" return [source_outputs, prompt, requires_multimodal_data, sampling_params] def _legacy_c2_other_name(source_outputs, prompt=None, requires_multimodal_data=False, streaming_ctx=None): - """Four-positional legacy shell whose streaming param has a non-standard name.""" + """Four-positional shell with a non-standard streaming param name.""" return [source_outputs, prompt, requires_multimodal_data, streaming_ctx] def _legacy_c2_keyword_only(source_outputs, prompt=None, requires_multimodal_data=False, *, streaming_context=None): - """Legacy C2 shell whose ``streaming_context`` is keyword-only.""" + """C2 shell whose ``streaming_context`` is keyword-only.""" return [source_outputs, prompt, requires_multimodal_data, streaming_context] def _count_deprecation_warnings(trigger) -> int: - """Run *trigger* under ``simplefilter("always")`` and count DeprecationWarnings.""" + """Run *trigger* under ``simplefilter("always")``; count DeprecationWarnings.""" import warnings with warnings.catch_warnings(record=True) as caught: @@ -73,15 +70,53 @@ def _count_deprecation_warnings(trigger) -> int: return sum(1 for w in caught if issubclass(w.category, DeprecationWarning)) +_INVOKE_CASES = [ + pytest.param( + _legacy_c2, + {"prompt": "p", "requires_multimodal_data": True, "streaming_context": "s"}, + [["out"], "p", True, "s"], + id="c2-streaming-context", + ), + pytest.param( + _legacy_c3_keyword_only, + {"prompt": "p", "requires_multimodal_data": True, "sampling_params": "sp"}, + [["out"], "p", True, "sp"], + id="c3-keyword-only-sampling-params", + ), + pytest.param( + _legacy_c3_required_keyword_only, + {"prompt": "p", "requires_multimodal_data": True, "sampling_params": "sp"}, + [["out"], "p", True, "sp"], + id="c3-required-keyword-only-sampling-params", + ), + pytest.param( + _legacy_c2_other_name, + {"prompt": "p", "requires_multimodal_data": True, "streaming_context": "s"}, + [["out"], "p", True, "s"], + id="c2-four-positional-other-name", + ), + pytest.param( + _legacy_c2_keyword_only, + {"prompt": "p", "requires_multimodal_data": True, "streaming_context": "s"}, + [["out"], "p", True, "s"], + id="c2-keyword-only-streaming-context", + ), +] + + def test_wrap_legacy_shape_warns_once_then_memoized(): - first = _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_legacy_c2)) - # The first wrap (at "resolve") emits exactly one DeprecationWarning. - assert first == 1 + assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_legacy_c2)) == 1 # Subsequent wraps reuse the memoized adapter: no further warning. assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_legacy_c2)) == 0 assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_legacy_c2)) == 0 +@pytest.mark.parametrize("fn,ctx_kwargs,expected", _INVOKE_CASES) +def test_invoke_adapts_legacy_shape_and_forwards_context(fn, ctx_kwargs, expected): + ctx = OrchestratorInputContext(**ctx_kwargs) + assert invoke_orchestrator_processor(fn, ["out"], ctx) == expected + + def test_invoke_does_not_rewarn_on_every_forward(): ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, streaming_context="s") @@ -93,48 +128,9 @@ def invoke_many(): assert _count_deprecation_warnings(invoke_many) == 1 -def test_invoke_adapts_legacy_shape_and_forwards_context(): - ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, streaming_context="s") - result = invoke_orchestrator_processor(_legacy_c2, ["out"], ctx) - # C2 adapter forwards (source_outputs, prompt, requires_multimodal_data, streaming_context). - assert result == [["out"], "p", True, "s"] - - def test_c1_processor_passes_through_without_warning(): wrapped = wrap_orchestrator_processor(_c1) assert wrapped is _c1 # already speaks ctx: returned unchanged ctx = OrchestratorInputContext(prompt="p") assert invoke_orchestrator_processor(_c1, ["out"], ctx) == [["out"], ctx] - # No DeprecationWarning for C1 processors at all. assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_c1)) == 0 - - -# --------------------------------------------------------------------------- -# Keyword-only semantic kwargs + four-positional fallback. -# --------------------------------------------------------------------------- - - -def test_keyword_only_sampling_params_receives_context_value(): - ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, sampling_params="sp") - # ``sampling_params`` is keyword-only: the C3 detection must read the FULL - # signature so the diffusion-stage sampling params still reach the shell. - assert invoke_orchestrator_processor(_legacy_c3_keyword_only, ["out"], ctx) == [["out"], "p", True, "sp"] - - -def test_required_keyword_only_sampling_params_receives_context_value(): - ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, sampling_params="sp") - # A REQUIRED keyword-only ``sampling_params`` used to raise at runtime - # (silently dropped by the old C0 fallback); it must now receive the value. - assert invoke_orchestrator_processor(_legacy_c3_required_keyword_only, ["out"], ctx) == [["out"], "p", True, "sp"] - - -def test_four_positional_different_streaming_name_not_regressed(): - ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, streaming_context="s") - # A four-positional legacy shell whose 4th param is named differently from - # ``streaming_context`` must still receive it as the 4th positional arg. - assert invoke_orchestrator_processor(_legacy_c2_other_name, ["out"], ctx) == [["out"], "p", True, "s"] - - -def test_keyword_only_streaming_context_receives_value(): - ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, streaming_context="s") - assert invoke_orchestrator_processor(_legacy_c2_keyword_only, ["out"], ctx) == [["out"], "p", True, "s"] diff --git a/tests/model_executor/stage_input_processors/test_placeholder_parity.py b/tests/model_executor/stage_input_processors/test_placeholder_parity.py index 4cce0978c04..d6428f80af2 100644 --- a/tests/model_executor/stage_input_processors/test_placeholder_parity.py +++ b/tests/model_executor/stage_input_processors/test_placeholder_parity.py @@ -102,21 +102,29 @@ def _spy(*, prompt_len: int, voice_metadata: dict[str, Any] | None = None) -> An # --------------------------------------------------------------------------- -def test_build_forward_placeholder_length_matches_golden(monkeypatch): +@pytest.mark.parametrize( + "source,ctx_prompt,expected_len,expected_voice", + [ + # Golden chat-template scan: sum(user segment) + 9 == 15. + pytest.param(_source_outputs(GOLDEN_PROMPT), None, 15, None, id="golden-length"), + pytest.param( + _source_outputs([1, 2], [3]), + {"additional_information": {"speaker": ["ethan"], "language": ["English"]}}, + None, + {"speaker": ["ethan"], "language": ["English"]}, + id="voice-metadata", + ), + ], +) +def test_build_forward_placeholder(monkeypatch, source, ctx_prompt, expected_len, expected_voice): recorded = _capture_pack(monkeypatch) - out = q3.build_forward_placeholder(_source_outputs(GOLDEN_PROMPT), _ctx()) + out = q3.build_forward_placeholder(source, _ctx(prompt=ctx_prompt)) assert len(out) == 1 assert out[0] == "PLACEHOLDER" - # Golden chat-template scan: sum(user segment) + 9 == 15. - assert recorded["prompt_len"] == 15 - - -def test_build_forward_placeholder_preserves_voice_metadata(monkeypatch): - recorded = _capture_pack(monkeypatch) - prompt = {"additional_information": {"speaker": ["ethan"], "language": ["English"]}} - out = q3.build_forward_placeholder(_source_outputs([1, 2], [3]), _ctx(prompt=prompt)) - assert len(out) == 1 - assert recorded["voice_metadata"] == {"speaker": ["ethan"], "language": ["English"]} + if expected_len is not None: + assert recorded["prompt_len"] == expected_len + if expected_voice is not None: + assert recorded["voice_metadata"] == expected_voice # --------------------------------------------------------------------------- @@ -124,62 +132,48 @@ def test_build_forward_placeholder_preserves_voice_metadata(monkeypatch): # --------------------------------------------------------------------------- -def test_build_prewarm_placeholder_length_from_stage0(monkeypatch): +@pytest.mark.parametrize( + "stage0_prompt,expected_len", + [ + pytest.param(GOLDEN_PROMPT, 15, id="list-stage0"), + pytest.param(SimpleNamespace(prompt_token_ids=GOLDEN_PROMPT), 15, id="request-like-stage0"), + ], +) +def test_build_prewarm_placeholder(monkeypatch, stage0_prompt, expected_len): recorded = _capture_pack(monkeypatch) - out = q3.build_prewarm_placeholder( - stage0_prompt=GOLDEN_PROMPT, - ctx=_ctx(), - downstream_stage_id=1, - ) + out = q3.build_prewarm_placeholder(stage0_prompt=stage0_prompt, ctx=_ctx(), downstream_stage_id=1) assert len(out) == 1 assert out[0] == "PLACEHOLDER" # stage0_only runs the chat-template scan on the stage-0 list -> 15 # (the same number the forward builder and the adapter scan produce). - assert recorded["prompt_len"] == 15 + assert recorded["prompt_len"] == expected_len # Voice metadata is intentionally not forwarded (matches pre-existing prewarm). assert recorded["voice_metadata"] is None -def test_build_prewarm_placeholder_accepts_request_like_stage0(monkeypatch): - recorded = _capture_pack(monkeypatch) - stage0_request = SimpleNamespace(prompt_token_ids=GOLDEN_PROMPT) - q3.build_prewarm_placeholder(stage0_prompt=stage0_request, ctx=_ctx(), downstream_stage_id=1) - assert recorded["prompt_len"] == 15 - - # --------------------------------------------------------------------------- # Parity: forward and prewarm must agree where the prompt shape allows. # --------------------------------------------------------------------------- -def test_forward_prewarm_parity_single_user_segment(monkeypatch): - """Single user segment: forward (full) == prewarm (stage0_only) == 6.""" - recorded = _capture_pack(monkeypatch) - q3.build_forward_placeholder(_source_outputs(SINGLE_USER_PROMPT), _ctx()) - forward_len = recorded["prompt_len"] - recorded.clear() - q3.build_prewarm_placeholder(stage0_prompt=SINGLE_USER_PROMPT, ctx=_ctx(), downstream_stage_id=1) - assert forward_len == recorded["prompt_len"] == 6 - - -def test_forward_prewarm_golden_relationship(monkeypatch): - """Documented relationship for the chat-template golden prompt. - - forward (full) == 15 (golden-locked); prewarm (stage0_only) is now the - **same chat-template scan** on the stage-0 list, so it also returns 15. - The builder and the inline fallback - (``adapter.compute_talker_prompt_ids_length``) therefore agree — the old - ``len(stage0_prompt)`` == 10 estimate under-reserved KV slots and split a - single request between 10 and 15 depending on whether resolution succeeded. - """ +@pytest.mark.parametrize( + "prompt,expected_len", + [ + # Single user segment: forward (full) == prewarm (stage0_only) == 6. + pytest.param(SINGLE_USER_PROMPT, 6, id="single-user-segment"), + # Golden chat-template prompt: full == stage0_only scan == 15. The old + # ``len(stage0_prompt)`` == 10 estimate under-reserved KV slots and split + # a request between 10 and 15 depending on resolution success. + pytest.param(GOLDEN_PROMPT, 15, id="golden"), + ], +) +def test_forward_prewarm_parity(monkeypatch, prompt, expected_len): recorded = _capture_pack(monkeypatch) - q3.build_forward_placeholder(_source_outputs(GOLDEN_PROMPT), _ctx()) + q3.build_forward_placeholder(_source_outputs(prompt), _ctx()) forward_len = recorded["prompt_len"] recorded.clear() - q3.build_prewarm_placeholder(stage0_prompt=GOLDEN_PROMPT, ctx=_ctx(), downstream_stage_id=1) - assert forward_len == 15 - assert recorded["prompt_len"] == 15 - assert forward_len == recorded["prompt_len"] # full == stage0_only scan + q3.build_prewarm_placeholder(stage0_prompt=prompt, ctx=_ctx(), downstream_stage_id=1) + assert forward_len == recorded["prompt_len"] == expected_len # --------------------------------------------------------------------------- diff --git a/tests/model_executor/stage_input_processors/test_registry.py b/tests/model_executor/stage_input_processors/test_registry.py index 134de1aa99e..25376adbddd 100644 --- a/tests/model_executor/stage_input_processors/test_registry.py +++ b/tests/model_executor/stage_input_processors/test_registry.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Unit tests for the stage-input processor registry. +"""Unit tests for the stage-input processor registry (table-driven). Pure-logic, CPU-only: exercises kind inference, structural validation, dead-processor hints and ``resolve_processor`` without executing any processor @@ -42,10 +42,6 @@ def placeholder_b_token_only(prompt, ctx=None): return [] -def placeholder_c_badname(source_outputs, _prompt=None): - return [] - - def diffusion_a_ar2diffusion(source_outputs, _prompt=None, _requires_multimodal_data=False): return {} @@ -155,198 +151,198 @@ def test_register_processor_invalid_kind(self): # --------------------------------------------------------------------------- -# Structural validation. +# Structural validation (table-driven: ok / warn / reject). # --------------------------------------------------------------------------- -class TestValidateProcessor: - def test_full_payload_ok_no_warning(self): - with warnings.catch_warnings(): - warnings.simplefilter("error") - validate_processor( - producer_b_full_payload, - kind="producer_full_payload", - path="pkg.mod.producer_b_full_payload", - ) - - def test_full_payload_missing_is_finished_warns(self): - with pytest.warns(RuntimeWarning, match="is_finished"): - validate_processor( - producer_a_full_payload, - kind="producer_full_payload", - path="pkg.mod.producer_a_full_payload", - ) - - def test_full_payload_cross_kind_second_param_rejected(self): - # ``multimodal_output`` instead of ``pooling_output`` cannot bind the - # worker keyword call, so this is a hard error (not just a warning). - with pytest.raises(ProcessorValidationError) as excinfo: - validate_processor( - producer_c_full_payload, - kind="producer_full_payload", - path="pkg.mod.producer_c_full_payload", - ) - assert excinfo.value.rule == "producer_kwargs" - - def test_full_payload_bad_param_names_rejected(self): - with pytest.raises(ProcessorValidationError) as excinfo: - validate_processor( - producer_badname_full_payload, - kind="producer_full_payload", - path="pkg.mod.producer_badname_full_payload", - ) - assert excinfo.value.rule == "producer_kwargs" - assert "pooling_output" in str(excinfo.value) - - def test_full_payload_positional_only_rejected(self): - # A positional-only payload parameter cannot accept the worker's - # keyword call (``pooling_output=``), so it must be rejected. - with pytest.raises(ProcessorValidationError) as excinfo: - validate_processor( - producer_positional_full_payload, - kind="producer_full_payload", - path="pkg.mod.producer_positional_full_payload", - ) - assert excinfo.value.rule == "producer_kwargs" - - def test_full_payload_var_keyword_accepted(self): - # A ``**kwargs`` producer accepts the exact worker keyword call even - # though it declares no explicit names. - with warnings.catch_warnings(): - warnings.simplefilter("error") - validate_processor( - producer_kwargs_full_payload, - kind="producer_full_payload", - path="pkg.mod.producer_kwargs_full_payload", - ) - - def test_full_payload_required_keyword_only_is_finished_accepted(self): - # Regression for amy-why-3459 review of HEAD ad81ac3a. - with warnings.catch_warnings(): - warnings.simplefilter("error") - validate_processor( - producer_d_full_payload, - kind="producer_full_payload", - path="pkg.mod.producer_d_full_payload", - ) - - def test_full_payload_default_is_finished_accepted(self): - with warnings.catch_warnings(): - warnings.simplefilter("error") - validate_processor( - producer_b_full_payload, - kind="producer_full_payload", - path="pkg.mod.producer_b_full_payload", - ) - - def test_async_chunk_bad_param_names_rejected(self): - with pytest.raises(ProcessorValidationError) as excinfo: - validate_processor( - producer_badname_async_chunk, - kind="producer_async_chunk", - path="pkg.mod.producer_badname_async_chunk", - ) - assert excinfo.value.rule == "producer_kwargs" - assert "multimodal_output" in str(excinfo.value) - - def test_async_chunk_var_keyword_accepted(self): - with warnings.catch_warnings(): - warnings.simplefilter("error") - validate_processor( - producer_kwargs_async_chunk, - kind="producer_async_chunk", - path="pkg.mod.producer_kwargs_async_chunk", - ) - - def test_async_chunk_ok_no_warning(self): - with warnings.catch_warnings(): - warnings.simplefilter("error") - validate_processor( - producer_a_async_chunk, - kind="producer_async_chunk", - path="pkg.mod.producer_a_async_chunk", - ) - - def test_async_chunk_missing_is_finished_hard_fail(self): - with pytest.raises(ProcessorValidationError) as excinfo: - validate_processor( - producer_b_async_chunk, - kind="producer_async_chunk", - path="pkg.mod.producer_b_async_chunk", - ) - assert excinfo.value.rule == "is_finished_required" - assert "producer_b_async_chunk" in str(excinfo.value) - - def test_async_chunk_positional_only_is_finished_hard_fail(self): - with pytest.raises(ProcessorValidationError) as excinfo: - validate_processor( - producer_c_async_chunk, - kind="producer_async_chunk", - path="pkg.mod.producer_c_async_chunk", - ) - assert excinfo.value.rule == "is_finished_positional_only" - - def test_placeholder_ok_no_warning(self): - with warnings.catch_warnings(): - warnings.simplefilter("error") - validate_processor( - placeholder_a_token_only, - kind="placeholder_prompt_builder", - path="pkg.mod.placeholder_a_token_only", - ) - - def test_placeholder_ctx_contract_ok(self): - with warnings.catch_warnings(): - warnings.simplefilter("error") - validate_processor( - placeholder_b_token_only, - kind="placeholder_prompt_builder", - path="pkg.mod.placeholder_b_token_only", - ) - - def test_placeholder_suffix_kind_conflict_warns(self): - with pytest.warns(RuntimeWarning, match="suffix"): - validate_processor( - placeholder_c_badname, - kind="placeholder_prompt_builder", - path="pkg.mod.placeholder_c_badname", - ) - - def test_diffusion_ok_no_warning(self): - with warnings.catch_warnings(): - warnings.simplefilter("error") - validate_processor( - diffusion_a_ar2diffusion, - kind="diffusion_input_builder", - path="pkg.mod.diffusion_a_ar2diffusion", - ) - - def test_legacy_multi_source_allowlist_ok(self): - with warnings.catch_warnings(): - warnings.simplefilter("error") - validate_processor( - talker2codec_shape, - kind="legacy_multi_source", - path="vllm_omni.model_executor.stage_input_processors.moss_tts.talker2codec", - ) - - def test_legacy_multi_source_allowlist_rejects_non_moss(self): - with pytest.raises(ProcessorValidationError) as excinfo: - validate_processor( - talker2codec_shape, - kind="legacy_multi_source", - path="pkg.mod.talker2codec", - ) - assert excinfo.value.rule == "legacy_multi_source_allowlist" - - def test_uninspectable_callable_fails(self): - with pytest.raises(ProcessorValidationError) as excinfo: - validate_processor( - object(), # not callable / not inspectable - kind="producer_full_payload", - path="pkg.mod.producer_full_payload", - ) - assert excinfo.value.rule == "inspectable_signature" +def _ok(fn, kind, path): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor(fn, kind=kind, path=path) + + +def _warns(fn, kind, path, match): + with pytest.warns(RuntimeWarning, match=match): + validate_processor(fn, kind=kind, path=path) + + +def _rejects(fn, kind, path, rule, match=None): + with pytest.raises(ProcessorValidationError) as excinfo: + validate_processor(fn, kind=kind, path=path) + assert excinfo.value.rule == rule + if match is not None: + assert match in str(excinfo.value) + + +_VALIDATION_CASES = [ + # (fn, kind, path, expected-behaviour, extra) + # full-payload producers + pytest.param( + producer_b_full_payload, + "producer_full_payload", + "pkg.mod.producer_b_full_payload", + "ok", + None, + id="full-payload-ok", + ), + pytest.param( + producer_a_full_payload, + "producer_full_payload", + "pkg.mod.producer_a_full_payload", + "warn", + "is_finished", + id="full-payload-missing-is_finished-warns", + ), + pytest.param( + producer_c_full_payload, + "producer_full_payload", + "pkg.mod.producer_c_full_payload", + "reject", + ("producer_kwargs", None), + id="full-payload-cross-kind-second-param-rejected", + ), + pytest.param( + producer_badname_full_payload, + "producer_full_payload", + "pkg.mod.producer_badname_full_payload", + "reject", + ("producer_kwargs", "pooling_output"), + id="full-payload-bad-param-names-rejected", + ), + pytest.param( + producer_positional_full_payload, + "producer_full_payload", + "pkg.mod.producer_positional_full_payload", + "reject", + ("producer_kwargs", None), + id="full-payload-positional-only-rejected", + ), + pytest.param( + producer_kwargs_full_payload, + "producer_full_payload", + "pkg.mod.producer_kwargs_full_payload", + "ok", + None, + id="full-payload-var-keyword-accepted", + ), + pytest.param( + producer_d_full_payload, + "producer_full_payload", + "pkg.mod.producer_d_full_payload", + "ok", + None, + id="full-payload-required-kw-only-is_finished-accepted", + ), + pytest.param( + producer_b_full_payload, + "producer_full_payload", + "pkg.mod.producer_b_full_payload", + "ok", + None, + id="full-payload-default-is_finished-accepted", + ), + # async-chunk producers + pytest.param( + producer_badname_async_chunk, + "producer_async_chunk", + "pkg.mod.producer_badname_async_chunk", + "reject", + ("producer_kwargs", "multimodal_output"), + id="async-chunk-bad-param-names-rejected", + ), + pytest.param( + producer_kwargs_async_chunk, + "producer_async_chunk", + "pkg.mod.producer_kwargs_async_chunk", + "ok", + None, + id="async-chunk-var-keyword-accepted", + ), + pytest.param( + producer_a_async_chunk, + "producer_async_chunk", + "pkg.mod.producer_a_async_chunk", + "ok", + None, + id="async-chunk-ok", + ), + pytest.param( + producer_b_async_chunk, + "producer_async_chunk", + "pkg.mod.producer_b_async_chunk", + "reject", + ("is_finished_required", "producer_b_async_chunk"), + id="async-chunk-missing-is_finished-hard-fail", + ), + pytest.param( + producer_c_async_chunk, + "producer_async_chunk", + "pkg.mod.producer_c_async_chunk", + "reject", + ("is_finished_positional_only", None), + id="async-chunk-positional-only-is_finished-hard-fail", + ), + # orchestrator-side + pytest.param( + placeholder_a_token_only, + "placeholder_prompt_builder", + "pkg.mod.placeholder_a_token_only", + "ok", + None, + id="placeholder-ok", + ), + pytest.param( + placeholder_b_token_only, + "placeholder_prompt_builder", + "pkg.mod.placeholder_b_token_only", + "ok", + None, + id="placeholder-ctx-contract-ok", + ), + pytest.param( + diffusion_a_ar2diffusion, + "diffusion_input_builder", + "pkg.mod.diffusion_a_ar2diffusion", + "ok", + None, + id="diffusion-ok", + ), + pytest.param( + talker2codec_shape, + "legacy_multi_source", + "vllm_omni.model_executor.stage_input_processors.moss_tts.talker2codec", + "ok", + None, + id="legacy-multi-source-allowlist-ok", + ), + pytest.param( + talker2codec_shape, + "legacy_multi_source", + "pkg.mod.talker2codec", + "reject", + ("legacy_multi_source_allowlist", None), + id="legacy-multi-source-rejects-non-moss", + ), + pytest.param( + object(), + "producer_full_payload", + "pkg.mod.producer_full_payload", + "reject", + ("inspectable_signature", None), + id="uninspectable-callable-fails", + ), +] + + +@pytest.mark.parametrize("fn,kind,path,behaviour,extra", _VALIDATION_CASES) +def test_validate_processor_table(fn, kind, path, behaviour, extra): + if behaviour == "ok": + _ok(fn, kind, path) + elif behaviour == "warn": + _warns(fn, kind, path, extra) + else: + rule, match = extra + _rejects(fn, kind, path, rule, match) # --------------------------------------------------------------------------- @@ -354,98 +350,60 @@ def test_uninspectable_callable_fails(self): # --------------------------------------------------------------------------- -class TestDeadProcessor: - def test_async_chunk_receiving_downstream_is_dead(self): - assert ( - dead_processor_hint( - "placeholder_prompt_builder", - async_chunk=True, - downstream_receives_async_chunks=True, - has_sync=False, - ) - is True - ) +# One entry per legacy test method; each entry carries the scenarios the +# original method asserted (keeps the collected-test count unchanged). +_DEAD_CASES = [ + pytest.param( + [ + ("placeholder_prompt_builder", True, True, False, False, True), + ], + id="async-chunk-receiving-downstream-dead", + ), + pytest.param( + [ + ("diffusion_input_builder", True, False, False, False, False), + ], + id="async-chunk-non-receiving-downstream-alive", + ), + pytest.param( + [ + ("diffusion_input_builder", False, False, False, False, False), + ("diffusion_input_builder", False, False, True, False, True), + ], + id="sync-mode-alive-unless-overridden", + ), + pytest.param( + [ + ("placeholder_prompt_builder", False, False, True, True, False), + ("placeholder_prompt_builder", False, False, True, False, True), + ("placeholder_prompt_builder", True, True, True, True, True), + ], + id="sync-mode-selected-sync-hook-is-not-dead", + ), + pytest.param( + [ + ("producer_full_payload", True, True, False, False, False), + ("producer_async_chunk", True, True, False, False, False), + ], + id="producer-never-dead", + ), +] - def test_async_chunk_non_receiving_downstream_alive(self): - assert ( - dead_processor_hint( - "diffusion_input_builder", - async_chunk=True, - downstream_receives_async_chunks=False, - has_sync=False, - ) - is False - ) - def test_sync_mode_alive_unless_overridden(self): - assert ( - dead_processor_hint( - "diffusion_input_builder", - async_chunk=False, - downstream_receives_async_chunks=False, - has_sync=False, - ) - is False - ) +@pytest.mark.parametrize("scenarios", _DEAD_CASES) +def test_dead_processor_hint_table(scenarios): + for kind, async_chunk, downstream, has_sync, custom_is_sync, expected in scenarios: assert ( dead_processor_hint( - "diffusion_input_builder", - async_chunk=False, - downstream_receives_async_chunks=False, - has_sync=True, + kind, + async_chunk=async_chunk, + downstream_receives_async_chunks=downstream, + has_sync=has_sync, + custom_is_sync=custom_is_sync, ) - is True + is expected ) - def test_sync_mode_selected_sync_hook_is_not_dead(self): - # When the custom hook *is* the configured sync hook (the sync - # ``*_token_only`` processor was pre-selected as the active forward - # hook), it must NOT be reported dead even though ``has_sync`` is true. - assert ( - dead_processor_hint( - "placeholder_prompt_builder", - async_chunk=False, - downstream_receives_async_chunks=False, - has_sync=True, - custom_is_sync=True, - ) - is False - ) - # A *different* sync hook overriding the custom hook is still dead. - assert ( - dead_processor_hint( - "placeholder_prompt_builder", - async_chunk=False, - downstream_receives_async_chunks=False, - has_sync=True, - custom_is_sync=False, - ) - is True - ) - # custom_is_sync only affects the non-async (sync-mode) decision. - assert ( - dead_processor_hint( - "placeholder_prompt_builder", - async_chunk=True, - downstream_receives_async_chunks=True, - has_sync=True, - custom_is_sync=True, - ) - is True - ) - - def test_producer_never_dead(self): - for kind in ("producer_full_payload", "producer_async_chunk"): - assert ( - dead_processor_hint( - kind, - async_chunk=True, - downstream_receives_async_chunks=True, - has_sync=False, - ) - is False - ) - # --------------------------------------------------------------------------- # Resolution entry point. @@ -453,27 +411,20 @@ def test_producer_never_dead(self): class TestResolveProcessor: - def test_resolve_normal(self): - spec = resolve_processor( - _mod_path("placeholder_a_token_only"), - expected_kind="placeholder_prompt_builder", - ) + @pytest.mark.parametrize( + "fn_name,expected_kind", + [ + ("placeholder_a_token_only", "placeholder_prompt_builder"), + ("producer_b_full_payload", None), + ("producer_d_full_payload", "producer_full_payload"), + ], + ) + def test_resolve_normal(self, fn_name, expected_kind): + spec = resolve_processor(_mod_path(fn_name), expected_kind=expected_kind) assert isinstance(spec, ProcessorSpec) - assert spec.kind == "placeholder_prompt_builder" - assert spec.path == _mod_path("placeholder_a_token_only") - assert spec.fn is placeholder_a_token_only - - def test_resolve_full_payload_without_expected_kind(self): - spec = resolve_processor(_mod_path("producer_b_full_payload")) - assert spec.kind == "producer_full_payload" - - def test_resolve_full_payload_required_keyword_only_is_finished(self): - spec = resolve_processor( - _mod_path("producer_d_full_payload"), - expected_kind="producer_full_payload", - ) - assert spec.kind == "producer_full_payload" - assert spec.fn is producer_d_full_payload + assert spec.kind == (expected_kind or infer_kind(None, path=_mod_path(fn_name))) + assert spec.path == _mod_path(fn_name) + assert spec.fn is getattr(importlib.import_module(__name__), fn_name) def test_resolve_drop_in_equivalence(self): """resolve_processor().fn must be the same object as the legacy lookup.""" diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index db86f6d63ae..3ca58fb2ae3 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -80,11 +80,6 @@ logger = init_logger(__name__) -#: Sentinel for a stage whose prewarm placeholder builder has not been resolved. -_PREWARM_BUILDER_NOT_RESOLVED: Any = object() -#: Negative sentinel for a stage with no prewarm placeholder builder. -_PREWARM_BUILDER_NO_BUILDER: Any = object() - def cleanup_request_artifact_dirs(artifact_dirs: set[str] | list[str]) -> None: for artifact_dir in artifact_dirs: @@ -401,10 +396,6 @@ def __init__( self.async_chunk = bool(async_chunk) self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools - # Per-stage prewarm placeholder builder cache (callable or a negative - # sentinel); populated lazily so repeated async requests never re-resolve - # (and never re-warn) for models without a builder. - self._prewarm_builder_cache: dict[int, Any] = {} # Wire the dead-processor hint at startup (warn-only) so a # configured-but-never-invoked custom_process_input_func is surfaced. self._warn_dead_input_processors() @@ -3238,91 +3229,6 @@ def _diffusion_submit_kwargs( ) } - def _get_prewarm_placeholder_builder(self, stage_id: int) -> Any: - """Resolve (and cache per stage) the ``build_prewarm_placeholder``. - - The async-chunk prewarm reuses ``sync_process_input_func`` (the - ``*_token_only`` path). Resolution order: - - 1. ``client.sync_process_input_func`` -> ``resolve_processor(...).fn`` - and read ``build_prewarm_placeholder`` off the resolved function - (with a module-level sibling fallback). - 2. The already-resolved ``custom_process_input_func`` fn carrying the - attribute (some models install the forward processor directly). - - Returns ``None`` when no builder is available so the caller keeps the - legacy inline estimate. The lookup is deliberately defensive — it never - raises (fake pools / minimal stage clients in tests and future - connectors may not expose ``stage_client``), so the prewarm path cannot - regress on builder resolution. - - The result is cached per stage (the callable on a hit, a negative - sentinel on a miss) so repeated async requests never re-run the - importlib resolution and never re-emit the fallback warning — Qwen3-TTS - and other models without a ``build_prewarm_placeholder`` would otherwise - re-resolve and re-warn on every async request. - """ - cache = getattr(self, "_prewarm_builder_cache", None) - if cache is None: # object.__new__ test doubles - cache = {} - self._prewarm_builder_cache = cache - cached = cache.get(stage_id, _PREWARM_BUILDER_NOT_RESOLVED) - if cached is not _PREWARM_BUILDER_NOT_RESOLVED: - return None if cached is _PREWARM_BUILDER_NO_BUILDER else cached - - from vllm_omni.model_executor.stage_input_processors import resolve_processor - - builder: Any = None - try: - pool = self.stage_pools[stage_id] - client = getattr(pool, "stage_client", None) - if client is not None: - path = getattr(client, "sync_process_input_func", None) - if path: - try: - spec = resolve_processor(path, expected_kind=None) - candidate = getattr(spec.fn, "build_prewarm_placeholder", None) - if not callable(candidate): - # Module-level sibling (covers paths whose fn does - # not expose the attribute directly). - import importlib - - module = importlib.import_module(path.rsplit(".", 1)[0]) - candidate = getattr(module, "build_prewarm_placeholder", None) - if callable(candidate): - builder = candidate - else: - logger.warning( - "[Orchestrator] stage-%s sync_process_input_func=%r has no " - "build_prewarm_placeholder; using the inline estimate", - stage_id, - path, - ) - except Exception as exc: - logger.warning( - "[Orchestrator] failed to resolve build_prewarm_placeholder " - "for %r (stage-%s); using the inline estimate: %s", - path, - stage_id, - exc, - ) - - if builder is None: - fn = getattr(client, "custom_process_input_func", None) - if callable(fn): - candidate = getattr(fn, "build_prewarm_placeholder", None) - if callable(candidate): - builder = candidate - except Exception as exc: # pragma: no cover - defensive only - logger.warning( - "[Orchestrator] prewarm placeholder builder lookup for stage-%s failed; using the inline estimate: %s", - stage_id, - exc, - ) - - cache[stage_id] = builder if builder is not None else _PREWARM_BUILDER_NO_BUILDER - return builder - @staticmethod def _apply_prewarm_prompt_len_override(base_input: dict[str, Any], configured_len: int | None) -> dict[str, Any]: """Apply the configured ``async_chunk_prewarm_prompt_len`` override to base_input.""" @@ -3363,7 +3269,9 @@ def _build_prewarm_placeholder_input( model_config = getattr(stage_vllm_config, "model_config", None) prewarm_len = getattr(getattr(model_config, "hf_config", None), "async_chunk_prewarm_prompt_len", None) - prewarm_builder = self._get_prewarm_placeholder_builder(next_stage_id) + pool = self.stage_pools[next_stage_id] + client = getattr(pool, "stage_client", None) + prewarm_builder = getattr(client, "prewarm_placeholder_builder", None) if prewarm_builder is not None: try: ctx = OrchestratorInputContext( diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 971c382568b..4ae6749ca97 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -153,10 +153,10 @@ def __init__( self.prompt_expand_func = metadata.prompt_expand_func self.custom_process_input_func = metadata.custom_process_input_func self.async_chunk_prewarm_payload_func = getattr(metadata, "async_chunk_prewarm_payload_func", None) - # Raw ``sync_process_input_func`` path so the orchestrator can - # resolve the async-chunk prewarm placeholder builder - # (``build_prewarm_placeholder``) at prewarm time. + # Raw ``sync_process_input_func`` path (dead-processor hint) plus the + # already-resolved dual-entry ``build_prewarm_placeholder`` builder. self.sync_process_input_func = metadata.sync_process_input_func + self.prewarm_placeholder_builder = metadata.prewarm_placeholder_builder self.engine_outputs: Any = None self.client_addresses = dict(client_addresses or {}) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 75af5d9f8f1..81fd82b2237 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -573,9 +573,12 @@ class StageMetadata: model_stage: str | None runtime_cfg: Any # Raw ``sync_process_input_func`` path (the ``*_token_only`` placeholder - # builder) so the orchestrator can resolve ``build_prewarm_placeholder`` - # at async-chunk prewarm time. + # builder) read by the orchestrator's dead-processor hint. sync_process_input_func: str | None = None + # Resolved ``build_prewarm_placeholder`` (the dual-entry ``*_token_only`` + # builder), so the orchestrator reads it off the stage client instead of + # re-resolving per async-chunk prewarm request. + prewarm_placeholder_builder: Callable | None = None prompt_transform_func: Callable | None = None prompt_expand_func: Callable | None = None cfg_kv_collect_func: Callable | None = None @@ -609,6 +612,55 @@ def _apply_rocm_attention_backend( engine_args["attention_backend"] = "TRITON_ATTN" +def _resolve_prewarm_placeholder_builder( + *, + sync_path: str | None, + custom_fn: Callable | None, + stage_config: Any, +) -> Callable | None: + """Resolve ``build_prewarm_placeholder`` at metadata-extraction time. + + Order mirrors the legacy orchestrator lookup: the registered + ``sync_process_input_func`` fn first (with a module-level sibling + fallback), then the already-resolved ``custom_process_input_func`` fn. + A missing or failing builder only warns — the async-chunk prewarm keeps + its inline estimate and never raises on builder resolution. + """ + builder: Callable | None = None + if sync_path: + try: + spec = resolve_processor(sync_path, expected_kind=None, stage_config=stage_config) + candidate = getattr(spec.fn, "build_prewarm_placeholder", None) + if not callable(candidate): + # Module-level sibling (covers paths whose fn does not expose + # the attribute directly). + candidate = getattr( + importlib.import_module(sync_path.rsplit(".", 1)[0]), + "build_prewarm_placeholder", + None, + ) + if callable(candidate): + builder = candidate + else: + logger.warning( + "sync_process_input_func=%r has no build_prewarm_placeholder; " + "async-chunk prewarm will use the inline estimate", + sync_path, + ) + except Exception as exc: + logger.warning( + "failed to resolve build_prewarm_placeholder for %r; " + "async-chunk prewarm will use the inline estimate: %s", + sync_path, + exc, + ) + if builder is None and callable(custom_fn): + candidate = getattr(custom_fn, "build_prewarm_placeholder", None) + if callable(candidate): + builder = candidate + return builder + + def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: """Extract metadata through the active production legacy path. @@ -661,10 +713,16 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: sync_process_input_func: str | None = None _spif_path = _get_attr_or_item(stage_config, "sync_process_input_func") if _spif_path: - # Keep the raw dotted path (not a resolved callable): the orchestrator - # resolves ``build_prewarm_placeholder`` off the ``*_token_only`` fn. + # Keep the raw dotted path (the dead-processor hint reads it); the + # dual-entry ``build_prewarm_placeholder`` is resolved once here. sync_process_input_func = _spif_path + prewarm_placeholder_builder = _resolve_prewarm_placeholder_builder( + sync_path=sync_process_input_func, + custom_fn=custom_process_input_func, + stage_config=stage_config, + ) + prompt_transform_func: Callable | None = None _ptf_path = _get_attr_or_item(stage_config, "prompt_transform_func") if _ptf_path: @@ -704,6 +762,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: default_sampling_params=default_params, custom_process_input_func=custom_process_input_func, sync_process_input_func=sync_process_input_func, + prewarm_placeholder_builder=prewarm_placeholder_builder, model_stage=model_stage, runtime_cfg=runtime_cfg, prompt_transform_func=prompt_transform_func, @@ -726,6 +785,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: default_sampling_params=default_params, custom_process_input_func=custom_process_input_func, sync_process_input_func=sync_process_input_func, + prewarm_placeholder_builder=prewarm_placeholder_builder, model_stage=model_stage, runtime_cfg=runtime_cfg, prompt_transform_func=prompt_transform_func, @@ -760,12 +820,19 @@ def extract_stage_metadata_from_omni_stage_config( sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams sampling_params = sampling_params_cls(**(stage_config.model_config.default_sampling_params or {})) custom_process_input_func = _resolve_omni_metadata_hook(stage_config.custom_process_input_func) - # Carry the raw sync path for async-chunk prewarm resolution. + # Carry the raw sync path (dead-processor hint); resolve the dual-entry + # ``build_prewarm_placeholder`` once here. sync_process_input_func: str | None = None _spif_path = getattr(stage_config.stage_pipeline_config, "sync_process_input_func", None) if _spif_path: sync_process_input_func = _spif_path + prewarm_placeholder_builder = _resolve_prewarm_placeholder_builder( + sync_path=sync_process_input_func, + custom_fn=custom_process_input_func, + stage_config=stage_config, + ) + if stage_type == "diffusion": return StageMetadata( stage_id=stage_config.stage_id, @@ -779,6 +846,7 @@ def extract_stage_metadata_from_omni_stage_config( default_sampling_params=sampling_params, custom_process_input_func=custom_process_input_func, sync_process_input_func=sync_process_input_func, + prewarm_placeholder_builder=prewarm_placeholder_builder, model_stage=stage_config.model_stage, runtime_cfg=stage_config.runtime_config, prompt_transform_func=_resolve_omni_metadata_hook(stage_config.prompt_transform_func), @@ -797,6 +865,7 @@ def extract_stage_metadata_from_omni_stage_config( default_sampling_params=sampling_params, custom_process_input_func=custom_process_input_func, sync_process_input_func=sync_process_input_func, + prewarm_placeholder_builder=prewarm_placeholder_builder, model_stage=stage_config.model_stage, runtime_cfg=stage_config.runtime_config, prompt_transform_func=_resolve_omni_metadata_hook(stage_config.prompt_transform_func), diff --git a/vllm_omni/model_executor/stage_input_processors/__init__.py b/vllm_omni/model_executor/stage_input_processors/__init__.py index 770b139d24e..9451ca067a8 100644 --- a/vllm_omni/model_executor/stage_input_processors/__init__.py +++ b/vllm_omni/model_executor/stage_input_processors/__init__.py @@ -3,32 +3,17 @@ """Stage input processors: per-model builders for inter-stage transitions. -Module-naming convention: - - ``*_full_payload`` : producer-side (worker) connector payload builder - (``FullPayloadProducer``; kwarg ``pooling_output`` load-bearing). - - ``*_async_chunk`` : producer-side (worker) streaming chunk builder - (``AsyncChunkProducer``; kwarg ``multimodal_output`` load-bearing). - - ``*_token_only`` : consumer-side (orchestrator) length-only placeholder - builder for the non-async-chunk path. - - no suffix (legacy) : consumer-side (orchestrator) sync builder; may be a - placeholder builder or a diffusion input builder. +Naming convention: ``*_full_payload``/``*_async_chunk`` are producer-side +(worker) builders (``pooling_output``/``multimodal_output`` load-bearing); +``*_token_only`` and no-suffix names are consumer-side (orchestrator) builders, +normalized through ``wrap_orchestrator_processor``. -Orchestrator-facing roles are normalized through -``wrap_orchestrator_processor`` / ``invoke_orchestrator_processor``: - - placeholder : upstream outputs -> next-stage token prompts - (``PlaceholderPromptBuilder``). - - diffusion : upstream outputs -> diffusion payload(s) - (``DiffusionInputBuilder``). - -This package only re-exports the dispatch contract and the processor registry; -it imports no other processor modules, so importing it never pulls in -model-specific code. +This package re-exports only the dispatch contract and registry — importing it +never pulls in model-specific code. """ from vllm_omni.model_executor.stage_input_processors._dispatch import ( - AsyncChunkProducer, DiffusionInputBuilder, - FullPayloadProducer, OrchestratorInputContext, PlaceholderPromptBuilder, invoke_orchestrator_processor, @@ -49,8 +34,6 @@ "OrchestratorInputContext", "PlaceholderPromptBuilder", "DiffusionInputBuilder", - "FullPayloadProducer", - "AsyncChunkProducer", "invoke_orchestrator_processor", "wrap_orchestrator_processor", "ProcessorKind", diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py index 440de3a56ac..9c32ced835e 100644 --- a/vllm_omni/model_executor/stage_input_processors/_common.py +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -1,11 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Canonical, shared helpers for ``stage_input_processors``. - -Implementation notes are golden-locked in -``tests/model_executor/stage_input_processors/test_common_helpers_golden.py``. -""" +"""Canonical shared helpers; semantics golden-locked in ``test_common_helpers_golden.py``.""" from __future__ import annotations @@ -29,15 +25,10 @@ def _unwrap_constant_list(x: Any) -> Any: def ensure_list(x: Any) -> list[Any]: - """Canonical ``ensure_list``. - - Semantics (the union of legacy variants): - - ``ConstantList``-like objects (``_x``) are unwrapped. - - ``list`` / ``tuple`` are converted to a fresh ``list``. - - ``torch.Tensor`` is converted via ``.tolist()`` (all dims flattened). - - ``None`` becomes ``[]``. - - Any other iterable is converted via ``list(x)``; a non-iterable scalar - becomes ``[x]``. + """Canonical ``ensure_list`` (union of legacy variants). + + ``ConstantList`` unwrapped; tensor ``.tolist()``; ``None`` -> ``[]``; + other iterable via ``list(x)``; non-iterable scalar wrapped as ``[x]``. """ x = _unwrap_constant_list(x) if isinstance(x, torch.Tensor): @@ -52,20 +43,8 @@ def ensure_list(x: Any) -> list[Any]: return [x] -def ensure_list_or_empty(x: Any) -> list[Any]: - """Explicit alias: always returns a list; ``None`` -> ``[]``. - - Matches ``audex`` / ``cosyvoice3`` ``_ensure_list`` semantics. - """ - return ensure_list(x) - - def ensure_list_unchanged(x: Any) -> Any: - """Legacy ``qwen3_omni._ensure_list`` semantics. - - ``ConstantList`` is unwrapped; a non-``list`` value (including ``None``, - a ``tuple``, or a ``torch.Tensor``) is returned **unchanged**. - """ + """Legacy ``qwen3_omni._ensure_list``: non-``list`` (incl. ``None``) unchanged.""" x = _unwrap_constant_list(x) if not isinstance(x, list): return x @@ -73,12 +52,7 @@ def ensure_list_unchanged(x: Any) -> Any: def ensure_list_strict(x: Any) -> list[Any]: - """Strict ``ensure_list``: no ``None`` special-case, non-iterables raise. - - Matches ``ming_flash_omni._ensure_list``: ``None`` and non-iterable - scalars both fall through to ``list(x)`` and raise ``TypeError``; a - ``torch.Tensor`` is converted via ``.tolist()``. - """ + """Strict variant (``ming_flash_omni``): no ``None`` case; non-iterable raises.""" x = _unwrap_constant_list(x) if isinstance(x, torch.Tensor): return x.detach().cpu().tolist() @@ -88,11 +62,7 @@ def ensure_list_strict(x: Any) -> list[Any]: def ensure_list_flatten(x: Any) -> list[Any]: - """``ensure_list`` that **flattens** tensors with ``reshape(-1).tolist()``. - - Matches ``audex._ensure_list``: ``None`` -> ``[]``; a ``torch.Tensor`` is - flattened to a 1-D list; a non-iterable scalar raises ``TypeError``. - """ + """Flattening variant (``audex``): tensors via ``reshape(-1).tolist()``.""" x = _unwrap_constant_list(x) if x is None: return [] @@ -104,16 +74,7 @@ def ensure_list_flatten(x: Any) -> list[Any]: def ensure_list_preserve_none(x: Any) -> list[Any]: - """``ensure_list`` that **preserves** ``None`` as ``[None]``. - - Matches ``step_audio2._ensure_list``: - - ``ConstantList``-like objects (``_x``) are unwrapped. - - ``torch.Tensor`` is converted via ``.tolist()`` (dimensions preserved). - - ``list`` / ``tuple`` are converted to a fresh ``list``. - - Any other **iterable** is converted via ``list(x)``. - - A non-iterable value (including ``None`` and a scalar like ``5``) is - wrapped as ``[x]`` (so ``None`` -> ``[None]``, ``5`` -> ``[5]``). - """ + """Preserve-``None`` variant (``step_audio2``): non-iterable wrapped as ``[x]``.""" if hasattr(x, "_x"): return list(x._x) if isinstance(x, torch.Tensor): @@ -124,13 +85,9 @@ def ensure_list_preserve_none(x: Any) -> list[Any]: def ensure_list_wrap_only(x: Any) -> list[Any]: - """Wrap-only ``ensure_list``: list passthrough, non-list -> ``[x]``, ``None`` -> ``[]``. + """Wrap-only variant (``diffusion.output_formatter``): non-list wrapped whole. - Preserves the legacy ``diffusion.output_formatter._ensure_list`` semantics: - a non-list primary payload is wrapped as ``[x]``. Unlike the canonical - :func:`ensure_list`, it performs **no** tensor flattening, no dict-key - iteration, and no row-wise iteration of an iterable (e.g. a PIL ``Image``) — - a non-``list`` value is wrapped whole, verbatim. + Unlike :func:`ensure_list`: no tensor flattening, no iterable iteration. """ if isinstance(x, list): return x @@ -143,13 +100,7 @@ def ensure_list_wrap_only(x: Any) -> list[Any]: def to_cpu_tensor(value: Any) -> torch.Tensor | None: - """Convert a value to a CPU tensor when possible; else ``None``. - - Matches ``glm_tts._to_cpu_tensor``: - - a list is unwrapped to its first element first; - - empty list -> ``None``; - - only ``torch.Tensor`` inputs yield a result (``detach().cpu()``). - """ + """CPU tensor when possible; else ``None`` (``glm_tts``).""" if isinstance(value, list): if not value: return None @@ -165,14 +116,10 @@ def to_cpu_tensor(value: Any) -> torch.Tensor | None: def to_token_id_list(value: Any, *, recursive: bool = False) -> list[int]: - """Convert a token-ish value into a flat ``list[int]``. - - ``recursive=False`` (default) matches ``dynin_omni._to_token_id_list``: - for a 2-D tensor or a nested list only the **first** row / element is - used. ``recursive=True`` matches ``cosyvoice3._to_token_id_list``: - per-item recursive normalization — nested ``list``/``tuple`` and - non-scalar tensors anywhere in the tree are flattened, so - ``[torch.tensor([[4, 5]])]`` -> ``[4, 5]`` and ``(4, 5)`` -> ``[4, 5]``. + """Flatten a token-ish value to ``list[int]``. + + ``recursive=False`` matches ``dynin_omni`` (first row only); + ``recursive=True`` matches ``cosyvoice3`` (per-item flatten). """ if recursive: return _to_token_id_list_recursive(value) @@ -195,11 +142,7 @@ def to_token_id_list(value: Any, *, recursive: bool = False) -> list[int]: def _to_token_id_list_recursive(value: Any) -> list[int]: - """Recursive token-id flattening (per-item normalization). - - ``None`` -> ``[]``; a tensor is reshaped flat; every item that is a tensor - or a ``list``/``tuple`` is recursed into; scalars are converted with ``int()``. - """ + """Recursive token-id flattening: tensors reshaped, scalars ``int()``.""" if value is None: return [] if isinstance(value, torch.Tensor): @@ -227,15 +170,10 @@ def revert_delay_pattern( expected_codebooks: int | None = None, allow_short: bool = False, ) -> torch.Tensor: - """Reverse the MusicGen-style delay pattern. - - Input ``[Q, T]`` (``T = seq_len + Q - 1``) -> output ``[Q, seq_len]``. + """Reverse the MusicGen-style delay pattern ``[Q, T]`` -> ``[Q, seq_len]``. - ``allow_short=True`` matches ``higgs_audio_v2``: when ``T < Q`` the input - is returned **unchanged**. ``allow_short=False`` (default, matches - ``higgs_audio_v3``) raises on ``T < Q``. When ``expected_codebooks`` is - given, a row-count mismatch raises ``ValueError`` (v3 checks ``Q == - _NUM_CODEBOOKS``; v2 does not). + ``allow_short`` (higgs v2) vs strict raise (v3); ``expected_codebooks`` + row-count check (v3 only). """ if audio_codes_qt.ndim != 2: raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}") @@ -262,13 +200,9 @@ def filter_real_code_frames( num_real_codes: int, layout: str = "frames_first", ) -> torch.Tensor: - """Keep only frames whose codebook values are all in ``[0, num_real_codes)``. + """Keep frames with all codebook values in ``[0, num_real_codes)``. - ``layout="frames_first"`` matches ``higgs_audio_v2``: input is - ``[num_frames, num_codebooks]`` and the same layout is returned. - ``layout="codebooks_first"`` matches ``higgs_audio_v3``: input is - ``[num_codebooks, num_frames]``; frames are filtered and the layout is - restored (``[num_codebooks, kept_frames]``). + ``frames_first`` (v2) vs ``codebooks_first`` (v3, transposed) layouts. """ if audio_codes.numel() == 0: return audio_codes @@ -297,25 +231,11 @@ def extract_last_codec_frame( to_cpu: bool = False, to_long: bool = True, ) -> torch.Tensor | None: - """Extract the last inter-stage codec frame from a payload. - - Parameterized canonical implementation covering the four legacy - ``_extract_last_frame`` variants: - - - ``key_path``: ``("codes", "audio")`` (higgs_v2 / qwen3_tts / voxtral) - vs ``("audio_codes",)`` (fish_speech top-level). - - ``validate``: - - ``None``: only empty checks (higgs_v2 / voxtral); - - ``"any"``: drop a 2-D frame whose values are all zero (qwen3_tts); - - ``"valid_mask"``: consult ``audio_code_valid`` then fall back to - ``any`` for 2-D frames (fish_speech). - - ``to_cpu`` / ``to_long``: fish_speech moves to CPU and casts to long; - higgs_v2 / qwen3_tts cast to long without moving to CPU; voxtral does - neither (plain ``flatten``). - - Legacy semantics: the validity/zero gate applies only to 2-D input; a - flat 1-D codec frame is returned directly (Fish Speech / Qwen3-TTS). - The returned tensor is always 1-D. + """Extract the last inter-stage codec frame (4 legacy variants merged). + + ``key_path`` codes/audio vs audio_codes; ``validate`` None/"any"/"valid_mask"; + ``to_cpu``/``to_long`` (fish_speech). Validity gate applies only to 2-D + frames; 1-D frames return directly; result is always 1-D. """ if not isinstance(payload, dict): return None @@ -381,9 +301,7 @@ def _chat_template_prompt_len( ) -> int: """Qwen chat-template scan: ``sum(user segments) + 9`` (assistant tail). - Shared by ``mode="full"`` (prompt + generated ids) and ``mode="stage0_only"`` - (stage-0 ids only, fed to both roles) so the sync forward placeholder and - the async-chunk prewarm estimate return the same length. + Shared by both modes so sync forward and async prewarm agree. """ thinker_sequences = torch.tensor(all_ids, dtype=torch.long, device=device).unsqueeze(0) # [1, T] input_ids = torch.tensor(prompt_ids, dtype=torch.long, device=device).unsqueeze(0) # [1, T] @@ -423,17 +341,10 @@ def compute_placeholder_prompt_len( ) -> int: """Compute the downstream placeholder prompt length. - ``mode="full"`` replicates the legacy Qwen chat-template scan: the input is - an ``OmniPayload``-like dict with ``ids: {"all": [...], "prompt": [...]}``, - and the length is ``sum(user segments) + 9`` (the assistant tail). - - ``mode="stage0_only"`` is the async-chunk **prewarm** estimate. There is no - upstream output yet, so ``ids_or_prompt`` is the stage-0 input token-id list, - which is fed to **both** chat-template roles (``all`` / ``prompt``) — exactly - what ``adapter.compute_talker_prompt_ids_length`` does. This makes the - prewarm builder and the inline fallback return the same number (15 for the - golden prompt; 6 == 6 for a single user segment), removing the - ``len()`` vs scan split. + ``mode="full"``: ``ids.all``/``ids.prompt`` scan (``sum(user) + 9``). + ``mode="stage0_only"``: prewarm estimate feeding stage-0 ids to both roles, + matching ``adapter.compute_talker_prompt_ids_length`` so builder and inline + fallback agree. """ if mode == "stage0_only": prompt = list(ids_or_prompt) @@ -450,12 +361,7 @@ def pack_placeholder_prompt( prompt_len: int, voice_metadata: dict[str, Any] | None = None, ) -> OmniTokensPrompt: - """Pack a KV-slot placeholder prompt: ``[0] * prompt_len`` + voice meta. - - Shared by both the sync forward path (``*_token_only``) and the - async-chunk prewarm path. Bulk conditioning arrives separately via the - connector, so only the length matters here. - """ + """Pack a KV-slot placeholder prompt: ``[0] * prompt_len`` + voice meta.""" return OmniTokensPrompt( prompt_token_ids=[0] * max(1, int(prompt_len)), additional_information=(voice_metadata if voice_metadata else None), # type: ignore[typeddict-item] diff --git a/vllm_omni/model_executor/stage_input_processors/_dispatch.py b/vllm_omni/model_executor/stage_input_processors/_dispatch.py index a97336e5b95..6934e4ed30b 100644 --- a/vllm_omni/model_executor/stage_input_processors/_dispatch.py +++ b/vllm_omni/model_executor/stage_input_processors/_dispatch.py @@ -1,59 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Orchestrator input-dispatch contract layer. - -This module defines the **consumer-side (orchestrator / stage-client) dispatch -contract** for ``stage_input_processors``. It replaces the two inconsistent -``inspect.signature`` probes in the orchestrator runtime with one shared -contract: - -* ``stage_engine_core_client.process_engine_inputs`` used an **arity** probe - (``len(signature.parameters) >= 4``) to pick a 3-arg vs 4-arg positional - call. -* ``orchestrator._forward_to_next_stage`` (diffusion branch) used a **name** - probe (``"sampling_params" in signature.parameters``) to decide whether to - inject the diffusion-stage sampling params as a keyword. - -Both are collapsed into a single normalization layer: - -* ``OrchestratorInputContext`` — the fixed transition context the runtime - passes to every consumer-side builder. -* ``PlaceholderPromptBuilder`` / ``DiffusionInputBuilder`` — the two - orchestrator-facing callable roles (C1-compatible ``(source_outputs, ctx)``). -* ``wrap_orchestrator_processor`` / ``invoke_orchestrator_processor`` — - processors that already speak the ``ctx`` contract are passed through - unchanged; legacy positional shapes (C0/C2/C3/C4 below) are adapted with a - ``DeprecationWarning``. - -**Producer roles are intentionally out of scope here.** ``FullPayloadProducer`` -(``*, transfer_manager, pooling_output, request, is_finished=...``) and -``AsyncChunkProducer`` (``*, transfer_manager, multimodal_output, request, -is_finished=False``) run inside workers and never receive an -``OrchestratorInputContext``; their ``pooling_output`` / ``multimodal_output`` -keyword names are load-bearing parts of the producer contract and are kept -as-is (the producer kwargs contract stays unchanged). - -Legacy positional shapes normalized here (detection reads the **full** -signature, keyword-only parameters included; the final invocation is validated -with ``inspect.Signature.bind``): - -* C0 3-arg: ``(source_outputs, prompt, requires_multimodal_data)``. -* C2 placeholder: ``(source_outputs, prompt, requires_multimodal_data, - streaming_context=None)`` (the ``_streaming_context`` variant used by - ``forced_aligner.code2wav2aligner`` / ``minicpmo_4_5_omni.llm2tts`` is also - recognized; a keyword-only ``*, streaming_context=None`` shell receives the - value as a keyword, and any four-positional callable falls back to this - shape regardless of the 4th parameter's name). -* C3 diffusion: ``(source_outputs, prompt, requires_multimodal_data, - sampling_params=None)`` — a keyword-only ``*, sampling_params=None`` shell - also receives ``ctx.sampling_params``. -* C4 legacy multi-source: ``(stage_list, engine_input_source, ...)`` — only - ``moss_tts.talker2codec`` remains today; normalized by - ``_adapt_moss_processor``, which binds ``stage_list=source_outputs`` and - ``engine_input_source=ctx.prompt`` and drops the trailing - ``prompt``/``requires_multimodal_data`` arguments (unused by that - implementation). +"""Consumer-side (orchestrator) input-dispatch contract layer. + +Collapses the two legacy runtime probes (arity in ``process_engine_inputs``, +sampling_params-name in ``_forward_to_next_stage``) into one normalization: + +* ``OrchestratorInputContext`` — the fixed transition context passed to every + consumer-side builder; ``PlaceholderPromptBuilder`` / ``DiffusionInputBuilder`` + are the two C1-compatible callable roles ``(source_outputs, ctx)``. +* ``wrap_orchestrator_processor`` passes through ``ctx``-native processors and + adapts legacy positional shapes (C0 3-arg; C2 ``streaming_context`` placeholder + incl. ``_streaming_context``/four-positional; C3 ``sampling_params`` diffusion; + C4 ``moss_tts.talker2codec``) with a one-time ``DeprecationWarning``. + +Producer roles are out of scope: their ``pooling_output``/``multimodal_output`` +kwargs are load-bearing and stay as-is (validated in ``_registry``). """ from __future__ import annotations @@ -61,7 +23,7 @@ import inspect import warnings from dataclasses import dataclass -from typing import Any, Protocol, runtime_checkable +from typing import Any, Protocol from vllm_omni.inputs.data import OmniTokensPrompt @@ -69,8 +31,6 @@ "OrchestratorInputContext", "PlaceholderPromptBuilder", "DiffusionInputBuilder", - "FullPayloadProducer", - "AsyncChunkProducer", "wrap_orchestrator_processor", "invoke_orchestrator_processor", ] @@ -94,7 +54,6 @@ class OrchestratorInputContext: next_stage_hf_config: Any | None = None -@runtime_checkable class PlaceholderPromptBuilder(Protocol): """C2-style sync builder: upstream outputs -> next-stage token prompts.""" @@ -105,7 +64,6 @@ def __call__( ) -> list[OmniTokensPrompt]: ... -@runtime_checkable class DiffusionInputBuilder(Protocol): """Diffusion-stage input builder: upstream outputs -> diffusion payload(s).""" @@ -116,55 +74,13 @@ def __call__( ) -> dict | list[dict] | None: ... -@runtime_checkable -class FullPayloadProducer(Protocol): - """Producer-side (worker) full-payload builder. - - ``pooling_output`` is a **load-bearing** keyword name: it is the connector - data-plane contract and must not be renamed or made positional. These - builders never receive an ``OrchestratorInputContext``. - """ - - def __call__( - self, - *, - transfer_manager: Any, - pooling_output: Any, - request: Any, - is_finished: bool = ..., - ) -> Any: ... - - -@runtime_checkable -class AsyncChunkProducer(Protocol): - """Producer-side (worker) async-chunk builder. - - ``multimodal_output`` is a **load-bearing** keyword name: it is the - connector data-plane contract and must not be renamed or made positional. - These builders never receive an ``OrchestratorInputContext``. - """ - - def __call__( - self, - *, - transfer_manager: Any, - multimodal_output: Any, - request: Any, - is_finished: bool = False, - ) -> Any: ... - - # --------------------------------------------------------------------------- # Signature inspection helpers # --------------------------------------------------------------------------- def _is_orchestrator_context_annotation(annotation: Any) -> bool: - """Whether ``annotation`` names ``OrchestratorInputContext``. - - Handles the class object, plain string annotations (``from __future__ - import annotations``) and fully-qualified dotted strings. - """ + """Whether ``annotation`` names ``OrchestratorInputContext`` (class, string, or dotted).""" if annotation is inspect.Parameter.empty: return False if annotation is OrchestratorInputContext: @@ -176,11 +92,7 @@ def _is_orchestrator_context_annotation(annotation: Any) -> bool: def _accepts_orchestrator_input_context(fn: Any) -> bool: - """Whether ``fn`` already speaks the new ``ctx`` contract. - - True when the signature contains a parameter named ``ctx``, or any - parameter whose type annotation is ``OrchestratorInputContext``. - """ + """Whether ``fn`` already speaks the ``ctx`` contract (param name or annotation).""" try: signature = inspect.signature(fn) except (TypeError, ValueError): @@ -211,14 +123,9 @@ def _positional_names(signature: inspect.Signature) -> list[str]: def _legacy_shape(signature: inspect.Signature | None) -> str: - """Classify a legacy positional processor shape. - - The probe reads the **full** signature — keyword-only parameters included — - so a ``def f(..., *, sampling_params=None)`` shell is still recognized as a - C3 diffusion builder (and receives ``ctx.sampling_params``), and any - four-positional callable falls back to the legacy C2 streaming shape - regardless of the 4th parameter's name (the legacy arity probe treated every - ``>= 4``-parameter callable that way). + """Classify a legacy positional shape (full signature, keyword-only included). + + ``sampling_params`` -> C3 (shells too); ``>= 4`` positional -> C2 streaming. """ if signature is None: return "c0" @@ -237,46 +144,19 @@ def _legacy_shape(signature: inspect.Signature | None) -> str: return "c0" -def _legacy_ctx_kwargs(signature: inspect.Signature | None, ctx: OrchestratorInputContext) -> dict[str, Any]: - """Extract legacy keyword-only config kwargs from ``ctx``. - - Scans the processor signature for keyword-only parameters named - ``target_model_config`` / ``next_stage_hf_config`` and forwards the - matching ``ctx`` attribute (``None`` included, so ``inspect.Signature.bind`` - always succeeds and the semantics stay explicit). - - ``sampling_params`` is deliberately **not** included: the C3 branch already - injects ``sampling_params=ctx.sampling_params`` explicitly, and forwarding - it again would raise ``TypeError: got multiple values for keyword argument - 'sampling_params'``. - """ - if signature is None: - return {} - mapping = { - "target_model_config": ctx.target_model_config, - "next_stage_hf_config": ctx.next_stage_hf_config, - } - kwargs: dict[str, Any] = {} - for name, param in signature.parameters.items(): - if param.kind is inspect.Parameter.KEYWORD_ONLY and name in mapping: - kwargs[name] = mapping[name] - return kwargs - - -def _bind_invoke(fn: Any, signature: inspect.Signature | None, *args: Any, **kwargs: Any) -> Any: - """Invoke *fn* with a legacy call shape bound via ``inspect.Signature.bind``. +def _ctx_kwarg_names(signature: inspect.Signature | None) -> tuple[str, ...]: + """Keyword-only ctx config params a legacy signature accepts. - Binding at the final invocation guarantees a legacy adapter never silently - drops context fields (e.g. forwarding ``sampling_params`` to a - positional-only parameter raises instead of being ignored) and never passes - an argument the processor cannot accept. + Excludes ``sampling_params``/``streaming_context`` (injected by the C3/C2 + adapters — forwarding again would raise a duplicate-kwarg ``TypeError``). """ if signature is None: - # Uninspectable callable (e.g. builtin): pass the shape through. - return fn(*args, **kwargs) - bound = signature.bind(*args, **kwargs) - bound.apply_defaults() - return fn(*bound.args, **bound.kwargs) + return () + return tuple( + name + for name, param in signature.parameters.items() + if param.kind is inspect.Parameter.KEYWORD_ONLY and name in ("target_model_config", "next_stage_hf_config") + ) def _warn_legacy_contract(fn: Any) -> None: @@ -295,52 +175,78 @@ def _warn_legacy_contract(fn: Any) -> None: # --------------------------------------------------------------------------- -def _adapt_moss_processor(fn: Any, signature: inspect.Signature | None) -> Any: - """C4 legacy multi-source adapter (e.g. ``moss_tts.talker2codec``). +def _legacy_adapter(fn: Any, signature: inspect.Signature | None, shape: str) -> Any: + """Build the C1 adapter for a legacy *shape*; resolve and validate once. - The legacy multi-source sync processors take - ``(stage_list, engine_input_source, prompt=None, requires_multimodal_data=False)``. - The orchestrator contract dispatches ``(source_outputs, ctx)``; we bind - ``stage_list=source_outputs`` and ``engine_input_source=ctx.prompt`` and - drop the trailing ``prompt``/``requires_multimodal_data`` arguments (the - ``moss_tts.talker2codec`` implementation does not use them, so this is - behaviour-preserving w.r.t. the previous 4-arg positional call). + Call shape fixed at wrap time; dry-run ``bind`` fails fast on broken + legacy calls at configuration time. """ + pos_names = _positional_names(signature) if signature is not None else [] + ctx_kwarg_names = _ctx_kwarg_names(signature) + c2pos_fourth = ( + "next_stage_hf_config" if len(pos_names) > 3 and pos_names[3] == "next_stage_hf_config" else "streaming_context" + ) + if signature is not None: + if shape == "c0": + signature.bind(None, None, None, **{name: None for name in ctx_kwarg_names}) + elif shape == "c2pos": + signature.bind(None, None, None, None, **{name: None for name in ctx_kwarg_names}) + elif shape == "c3": + signature.bind(None, None, None, sampling_params=None) + elif shape == "c2kw": + signature.bind(None, None, None, streaming_context=None) + elif shape == "moss": + signature.bind(None, None) def _adapted(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: - return _bind_invoke(fn, signature, source_outputs, ctx.prompt) + if shape == "moss": + return fn(source_outputs, ctx.prompt) + if shape == "c3": + return fn( + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + sampling_params=ctx.sampling_params, + ) + if shape == "c2pos": + return fn( + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + getattr(ctx, c2pos_fourth), + **{name: getattr(ctx, name) for name in ctx_kwarg_names}, + ) + if shape == "c2kw": + return fn( + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + streaming_context=ctx.streaming_context, + ) + # C0 3-arg legacy (and the uninspectable / builtin fallback). + return fn( + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + **{name: getattr(ctx, name) for name in ctx_kwarg_names}, + ) _warn_legacy_contract(fn) return _adapted -#: Wrap memo: the signature probe + ``DeprecationWarning`` must happen **once per -#: processor**, not on every forward. ``invoke_orchestrator_processor`` runs per -#: request, and re-wrapping a legacy C2 shape (e.g. ``thinker2talker_token_only``) -#: would otherwise re-emit the warning on every forward. Keyed by the callable -#: object itself (module-level processors are few and live for the process). +#: Wrap memo: the signature probe, the call-shape resolution and the +#: ``DeprecationWarning`` happen once per processor, not per forward. _WRAP_CACHE: dict[Any, Any] = {} def wrap_orchestrator_processor(fn: Any) -> PlaceholderPromptBuilder | DiffusionInputBuilder: - """Return a C1-compatible callable ``(source_outputs, ctx)`` for ``fn``. - - - If ``fn`` already accepts ``ctx`` (``_accepts_orchestrator_input_context``), - it is returned unchanged. - - Otherwise the legacy positional shape is adapted: - * C4 ``(stage_list, engine_input_source, ...)`` -> ``_adapt_moss_processor``; - * C3 ``(source_outputs, prompt, requires_multimodal_data, sampling_params=...)`` - forwards ``ctx.sampling_params`` when the parameter exists; - * C2 ``(source_outputs, prompt, requires_multimodal_data, - streaming_context=...)`` (or ``_streaming_context``) forwards - ``ctx.streaming_context``; - * C0 3-arg ``(source_outputs, prompt, requires_multimodal_data)``. - Every legacy shape emits a ``DeprecationWarning``. - - The result is memoized per callable (see ``_WRAP_CACHE``): the wrap — and its - ``DeprecationWarning`` — happens exactly once per processor, at first - resolution/invocation, so ``invoke_orchestrator_processor`` does not re-wrap - (or re-warn) on every forward. + """Return a C1-compatible ``(source_outputs, ctx)`` callable for ``fn``. + + ``ctx``-native processors pass through unchanged; legacy shapes are adapted + per ``_legacy_shape`` with a one-time ``DeprecationWarning``. Result is + memoized in ``_WRAP_CACHE`` so the wrap (and warning) happens once per + processor, not per forward. """ if not callable(fn): raise TypeError(f"stage-input processor must be callable, got {fn!r}") @@ -356,73 +262,7 @@ def wrap_orchestrator_processor(fn: Any) -> PlaceholderPromptBuilder | Diffusion wrapped: Any = fn else: signature = _legacy_signature(fn) - shape = _legacy_shape(signature) - - if shape == "moss": - wrapped = _adapt_moss_processor(fn, signature) - else: - # ``_legacy_shape(None)`` returns "c0", so the c2pos/c3/c2kw - # branches below are only reachable with a real signature. Keep the - # narrow assert so static analyzers can narrow the optional type - # inside the ``_adapted`` closure. - assert signature is not None - - def _adapted(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: - if shape == "c3": - # C3 diffusion: forwarding sampling_params (diffusion-stage - # params). Read from the full signature so a keyword-only - # ``*, sampling_params=None`` shell receives the value too. - return _bind_invoke( - fn, - signature, - source_outputs, - ctx.prompt, - ctx.requires_multimodal_data, - sampling_params=ctx.sampling_params, - ) - if shape == "c2pos": - # C2 four-positional legacy fallback (any 4th parameter - # name): the streaming context is the 4th positional arg, - # unless the 4th parameter is named ``next_stage_hf_config`` - # (e.g. Nemotron ``thinker2talker_token_only``), in which - # case the stage client's HF config is passed instead so the - # processor can size the next-stage placeholder prompt. - fourth = ( - ctx.next_stage_hf_config - if _positional_names(signature)[3] == "next_stage_hf_config" - else ctx.streaming_context - ) - return _bind_invoke( - fn, - signature, - source_outputs, - ctx.prompt, - ctx.requires_multimodal_data, - fourth, - **_legacy_ctx_kwargs(signature, ctx), - ) - if shape == "c2kw": - # Keyword-only streaming context shell. - return _bind_invoke( - fn, - signature, - source_outputs, - ctx.prompt, - ctx.requires_multimodal_data, - streaming_context=ctx.streaming_context, - ) - # C0 3-arg legacy (and the uninspectable / builtin fallback). - return _bind_invoke( - fn, - signature, - source_outputs, - ctx.prompt, - ctx.requires_multimodal_data, - **_legacy_ctx_kwargs(signature, ctx), - ) - - _warn_legacy_contract(fn) - wrapped = _adapted + wrapped = _legacy_adapter(fn, signature, _legacy_shape(signature)) try: _WRAP_CACHE[fn] = wrapped @@ -436,5 +276,5 @@ def invoke_orchestrator_processor( source_outputs: list[Any], ctx: OrchestratorInputContext, ) -> Any: - """Invoke an orchestrator-facing processor under the fixed ``(source_outputs, ctx)`` contract.""" + """Invoke an orchestrator-facing processor under ``(source_outputs, ctx)``.""" return wrap_orchestrator_processor(fn)(source_outputs, ctx) diff --git a/vllm_omni/model_executor/stage_input_processors/_registry.py b/vllm_omni/model_executor/stage_input_processors/_registry.py index c65f54ffcc5..d66e28a0489 100644 --- a/vllm_omni/model_executor/stage_input_processors/_registry.py +++ b/vllm_omni/model_executor/stage_input_processors/_registry.py @@ -1,44 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Stage-input processor registry. - -This module is the **structural validation** layer behind the three ``importlib`` -resolution points that are validated as independent symbol lookups: - -* ``engine/stage_init_utils.extract_legacy_stage_metadata`` — consumer-side - (orchestrator) input processors (``custom_process_input_func`` / - ``sync_process_input_func``). -* ``worker/omni_connector_model_runner_mixin._load_custom_func`` — the - producer-side full-payload builder (``FullPayloadProducer``). -* ``distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.__init__`` - — the producer-side async-chunk builder (``AsyncChunkProducer``). - -The registry only performs signature-level structural checks -(``inspect.signature``); it never executes processor logic and never loads -model weights. Behavioural guarantees belong to the golden / parity tests. - -Producer kwargs are contract. ``FullPayloadProducer`` keeps ``pooling_output`` -(``is_finished`` is **optional** — the worker retries without it when the -processor rejects the kwarg); ``AsyncChunkProducer`` keeps ``multimodal_output`` -and **must** accept ``is_finished`` — the scheduler always passes it on the -async-chunk path (see ``chunk_transfer_adapter._send_single_request``). - -Kind inference is name-driven (suffix-based): - -* ``_token_only`` -> ``placeholder_prompt_builder`` -* ``_full_payload`` / ``_batch`` -> ``producer_full_payload`` -* ``_async_chunk`` -> ``producer_async_chunk`` -* ``ar2diffusion`` / ``ar2dit`` / ``thinker2imagegen`` (and friends) -> - ``diffusion_input_builder`` -* no suffix -> ``legacy_orchestrator_builder`` -* ``moss_tts.talker2codec`` (legacy multi-source shape) -> - ``legacy_multi_source`` - -``resolve_processor`` returns a :class:`ProcessorSpec` whose ``fn`` is the -same callable object produced by the legacy -``getattr(importlib.import_module(mod_path), fn_name)`` lookups, so callers keep -identical runtime behaviour while gaining validated kind metadata. +"""Stage-input processor registry: signature-level structural validation. + +Serves the three importlib resolution points (orchestrator-side +``custom_process_input_func``/``sync_process_input_func``, producer-side +``*_full_payload``, producer-side ``*_async_chunk``). Checks are purely +signature-based (``inspect.signature``); behaviour is covered by the golden / +parity tests. + +Producer kwargs are contract: full-payload keeps ``pooling_output`` +(``is_finished`` optional — the worker retries without it); async-chunk keeps +``multimodal_output`` and **must** accept ``is_finished`` (the scheduler always +passes it). + +Kind inference is name-driven (suffix-based); ``resolve_processor`` returns a +:class:`ProcessorSpec` whose ``fn`` matches the legacy importlib lookup. """ from __future__ import annotations @@ -46,9 +23,9 @@ import importlib import inspect import warnings -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass -from typing import Any, Literal +from typing import Any, Literal, NoReturn, get_args ProcessorKind = Literal[ "placeholder_prompt_builder", @@ -87,16 +64,7 @@ _PRODUCER_KINDS: frozenset[str] = frozenset({"producer_full_payload", "producer_async_chunk"}) #: All valid :class:`ProcessorKind` values (used to validate manual overrides). -_PROCESSOR_KINDS: frozenset[str] = frozenset( - { - "placeholder_prompt_builder", - "diffusion_input_builder", - "legacy_orchestrator_builder", - "producer_full_payload", - "producer_async_chunk", - "legacy_multi_source", - } -) +_PROCESSOR_KINDS: frozenset[str] = frozenset(get_args(ProcessorKind)) #: Soft-check scalar return annotations that contradict a stage-input role. _SCALAR_RETURN_ANNOTATIONS: frozenset[Any] = frozenset({int, str, bool, float}) @@ -116,14 +84,7 @@ @dataclass(frozen=True) class ProcessorSpec: - """Resolved and validated stage-input processor. - - Attributes: - path: The dotted ``module.attr`` path the processor was resolved from. - kind: The inferred (or overridden) processor kind. - fn: The callable itself — identical to what the legacy - ``getattr(importlib.import_module(...), ...)`` lookup returned. - """ + """Resolved and validated stage-input processor (path, kind, fn).""" path: str kind: ProcessorKind @@ -131,12 +92,7 @@ class ProcessorSpec: class ProcessorValidationError(Exception): - """A stage-input processor failed structural validation. - - Carries enough context to locate the offending stage / config line: - ``stage_id`` (when a ``stage_config`` was supplied), ``path``, ``kind`` and - the specific ``rule`` that was violated. - """ + """A stage-input processor failed structural validation.""" def __init__( self, @@ -155,19 +111,20 @@ def __init__( super().__init__(self._format()) def _format(self) -> str: - parts: list[str] = [] - if self.stage_id is not None: - parts.append(f"stage_id={self.stage_id}") - if self.path: - parts.append(f"path={self.path!r}") - if self.kind: - parts.append(f"kind={self.kind!r}") - if self.rule: - parts.append(f"rule={self.rule}") - detail = ", ".join(parts) or "processor validation failed" - if self.message: - return f"{detail}: {self.message}" - return detail + detail = ( + ", ".join( + f"{name}={value}" + for name, value in ( + ("stage_id", self.stage_id), + ("path", repr(self.path) if self.path else None), + ("kind", repr(self.kind) if self.kind else None), + ("rule", self.rule), + ) + if value is not None + ) + or "processor validation failed" + ) + return f"{detail}: {self.message}" if self.message else detail # --------------------------------------------------------------------------- @@ -185,8 +142,6 @@ def _module_part(path: str) -> str: def _stage_id_of(stage_config: Any) -> int | None: """Best-effort stage id extraction from an arbitrary config object.""" - if stage_config is None: - return None sid = getattr(stage_config, "stage_id", None) if sid is None: return None @@ -202,14 +157,7 @@ def _stage_id_of(stage_config: Any) -> int | None: def import_symbol(path: str) -> Any: - """Import a dotted ``module.attr`` symbol (importlib + ``getattr``). - - This mirrors the two-step shape used by every legacy resolution point: - ``importlib.import_module(module_path)`` then ``getattr(module, attr)``. - A missing attribute or module raises the underlying - ``ImportError``/``AttributeError``; callers that need a best-effort lookup - (e.g. the worker candidate chain) catch these themselves. - """ + """Import a dotted ``module.attr`` symbol; missing module/attr propagate.""" if not isinstance(path, str) or not path: raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") module_path, attr_name = path.rsplit(".", 1) @@ -226,13 +174,7 @@ def import_symbol(path: str) -> Any: def register_processor(path: str, kind: ProcessorKind) -> None: - """Register a manual kind override for a processor *path*. - - This is an escape hatch for models whose function names do not follow the - suffix convention. ``infer_kind`` consults these overrides before - applying the suffix rules. Overrides are validated eagerly so typos fail - fast. - """ + """Register a manual kind override (escape hatch for non-suffix names).""" if not isinstance(path, str) or not path: raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") if kind not in _PROCESSOR_KINDS: @@ -241,12 +183,7 @@ def register_processor(path: str, kind: ProcessorKind) -> None: def infer_kind(fn: Any, *, path: str) -> ProcessorKind: - """Infer the processor kind from the *suffix* of ``path``. - - ``fn`` is accepted for API symmetry with :func:`validate_processor` but the - inference is purely name-driven (suffix-based convention); the callable - object itself is not inspected here. - """ + """Infer the processor kind from the *suffix* of ``path`` (name-driven).""" del fn # inference is name-based; kept for a symmetric API if not isinstance(path, str) or not path or "." not in path: raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") @@ -278,11 +215,8 @@ def infer_kind(fn: Any, *, path: str) -> ProcessorKind: def _warn(path: str, kind: ProcessorKind | None, rule: str, message: str) -> None: - """Emit a soft ``RuntimeWarning`` for a non-fatal structural mismatch.""" warnings.warn( - f"stage-input processor {path!r} (kind={kind!r}, rule={rule}): {message}", - RuntimeWarning, - stacklevel=3, + f"stage-input processor {path!r} (kind={kind!r}, rule={rule}): {message}", RuntimeWarning, stacklevel=3 ) @@ -293,7 +227,7 @@ def _fail( rule: str, stage_config: Any, message: str, -) -> None: +) -> NoReturn: raise ProcessorValidationError( stage_id=_stage_id_of(stage_config), path=path, @@ -304,13 +238,14 @@ def _fail( def _check_orchestrator_first_param( - params: Mapping[str, inspect.Parameter], + signature: inspect.Signature, *, path: str, kind: ProcessorKind, stage_config: Any, ) -> None: """First param must be ``source_outputs`` or the signature must carry ``ctx``.""" + params = signature.parameters names = [ name for name, param in params.items() @@ -326,40 +261,31 @@ def _check_orchestrator_first_param( ) -def _check_suffix( - path: str, - kind: ProcessorKind, - *, - suffix_pool: tuple[str, ...], - stage_config: Any, - what: str, -) -> None: - fn_name = _fn_name(path) - if not any(fn_name.endswith(suffix) for suffix in suffix_pool): - _warn(path, kind, "suffix_kind_mismatch", f"name {fn_name!r} does not match expected {what}") - - def _has_var_keyword(params: Mapping[str, inspect.Parameter]) -> bool: """Whether the signature accepts arbitrary keywords (``**kwargs``).""" return any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()) -def _check_full_payload_is_finished( +def _check_is_finished( params: Mapping[str, inspect.Parameter], *, path: str, kind: ProcessorKind, stage_config: Any, + required: bool, ) -> None: - """``is_finished`` is optional for the full-payload producer. - - The worker passes it best-effort and retries without it on - ``TypeError`` (see ``omni_connector_model_runner_mixin._build_custom_process_payload``). - A ``**kwargs`` producer accepts it through ``VAR_KEYWORD``. - """ + """Validate the ``is_finished`` kwarg contract (hard only for async-chunk).""" if _has_var_keyword(params): return # **kwargs accepts is_finished if "is_finished" not in params: + if required: + _fail( + path=path, + kind=kind, + rule="is_finished_required", + stage_config=stage_config, + message="async-chunk producer must declare 'is_finished' (the scheduler always passes it)", + ) _warn( path, kind, @@ -367,52 +293,25 @@ def _check_full_payload_is_finished( "full-payload producer does not declare is_finished; worker retries without it (best-effort)", ) return - if params["is_finished"].kind == inspect.Parameter.POSITIONAL_ONLY: + is_finished_param = params["is_finished"] + if is_finished_param.kind == inspect.Parameter.POSITIONAL_ONLY: + if required: + _fail( + path=path, + kind=kind, + rule="is_finished_positional_only", + stage_config=stage_config, + message="async-chunk producer declares 'is_finished' positional-only; " + "it must be keyword-only or have a default", + ) _warn( path, kind, "is_finished_positional_only", "full-payload producer declares is_finished positional-only; worker retries without it", ) - - -def _check_async_chunk_is_finished( - params: Mapping[str, inspect.Parameter], - *, - path: str, - kind: ProcessorKind, - stage_config: Any, -) -> None: - """``is_finished`` is **required** for the async-chunk producer. - - The scheduler always passes ``is_finished`` (see - ``chunk_transfer_adapter._send_single_request``), so a processor that cannot - accept it is a hard configuration error. A ``**kwargs`` producer accepts it - through ``VAR_KEYWORD``. - """ - if _has_var_keyword(params): - return # **kwargs accepts is_finished - if "is_finished" not in params: - _fail( - path=path, - kind=kind, - rule="is_finished_required", - stage_config=stage_config, - message="async-chunk producer must declare 'is_finished' (the scheduler always passes it)", - ) - return # unreachable; kept for type-checkers - is_finished_param = params["is_finished"] - if is_finished_param.kind == inspect.Parameter.POSITIONAL_ONLY: - _fail( - path=path, - kind=kind, - rule="is_finished_positional_only", - stage_config=stage_config, - message="async-chunk producer declares 'is_finished' positional-only; " - "it must be keyword-only or have a default", - ) - return # unreachable - if ( + return + if required and ( is_finished_param.kind != inspect.Parameter.KEYWORD_ONLY and is_finished_param.default is inspect.Parameter.empty ): @@ -439,18 +338,11 @@ def _check_producer_keyword_bind( required: tuple[str, ...], what: str, ) -> None: - """Bind the exact worker keyword call shape to *signature*. - - The worker invokes producers with a fixed keyword call (see - ``omni_connector_model_runner_mixin._build_custom_process_payload`` for the - full-payload producer and ``chunk_transfer_adapter._send_single_request`` - for the async-chunk producer). We probe that call with - ``inspect.Signature.bind`` so a producer that *declares* the right parameter - names but cannot actually accept the keyword call — wrong parameter names, - positional-only payload parameters, or missing required parameters — is - rejected at configuration time instead of failing later when the worker - passes ``pooling_output=`` / ``multimodal_output=``. A ``**kwargs`` - (``VAR_KEYWORD``) producer accepts any keyword and is therefore compatible. + """Bind the worker keyword call to *signature* via ``inspect.Signature.bind``. + + Producers that declare the right names but reject the keyword call shape + fail at configuration time instead of when the worker passes + ``pooling_output=`` / ``multimodal_output=``. """ try: signature.bind(**{name: _BIND_PLACEHOLDER for name in required}) @@ -464,43 +356,89 @@ def _check_producer_keyword_bind( ) -def _check_return_annotation( +def _check_full_payload( signature: inspect.Signature, *, path: str, kind: ProcessorKind, + stage_config: Any, ) -> None: - """Soft return-annotation check: mismatches warn, never fail.""" - ann = signature.return_annotation - if ann is inspect.Signature.empty or isinstance(ann, str): - # Unannotated or postponed-string annotations are not verifiable here. - return - if ann in _SCALAR_RETURN_ANNOTATIONS: - _warn( - path, - kind, - "return_annotation", - f"return annotation {ann!r} is a scalar, unlikely a stage-input processor result", + """Full-payload producer: keyword bind, then soft is_finished.""" + params = signature.parameters + required: tuple[str, ...] = ("transfer_manager", "pooling_output", "request") + is_finished_param = params.get("is_finished") + if is_finished_param is not None and is_finished_param.kind in ( + inspect.Parameter.KEYWORD_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ): + required = required + ("is_finished",) + _check_producer_keyword_bind( + signature, + path=path, + kind=kind, + stage_config=stage_config, + required=required, + what="full-payload", + ) + _check_is_finished(params, path=path, kind=kind, stage_config=stage_config, required=False) + + +def _check_async_chunk( + signature: inspect.Signature, + *, + path: str, + kind: ProcessorKind, + stage_config: Any, +) -> None: + """Async-chunk producer: hard is_finished, then keyword bind.""" + params = signature.parameters + _check_is_finished(params, path=path, kind=kind, stage_config=stage_config, required=True) + _check_producer_keyword_bind( + signature, + path=path, + kind=kind, + stage_config=stage_config, + required=("transfer_manager", "multimodal_output", "request", "is_finished"), + what="async-chunk", + ) + + +def _check_legacy_multi_source( + _signature: inspect.Signature, + *, + path: str, + kind: ProcessorKind, + stage_config: Any, +) -> None: + """Legacy multi-source processors are allowlisted to ``moss_tts.*``.""" + if _MOSS_MULTI_SOURCE_MODULE_HINT not in _module_part(path): + _fail( + path=path, + kind=kind, + rule="legacy_multi_source_allowlist", + stage_config=stage_config, + message=f"legacy multi-source processors are allowlisted to moss_tts.*, got {path!r}", ) +#: Per-kind structural checks run in order by :func:`validate_processor`. +#: ``legacy_orchestrator_builder`` is deliberately check-free: the orchestrator +#: normalizes legacy shapes through ``wrap_orchestrator_processor``. +_VALIDATORS: dict[ProcessorKind, tuple[Callable[..., None], ...]] = { + "placeholder_prompt_builder": (_check_orchestrator_first_param,), + "diffusion_input_builder": (_check_orchestrator_first_param,), + "producer_full_payload": (_check_full_payload,), + "producer_async_chunk": (_check_async_chunk,), + "legacy_multi_source": (_check_legacy_multi_source,), + "legacy_orchestrator_builder": (), +} + + def validate_processor(fn: Any, *, kind: ProcessorKind, path: str, stage_config: Any = None) -> None: - """Structurally validate *fn* against the expected *kind*. - - Only ``inspect.signature`` metadata is examined; the processor body is never - executed. Hard contract violations raise :class:`ProcessorValidationError`; - soft mismatches emit a ``RuntimeWarning``. - - Args: - fn: The processor callable. - kind: The expected :class:`ProcessorKind`. - path: The dotted ``module.attr`` path (used for reporting). - stage_config: Optional config object; its ``stage_id`` is attached to - any raised :class:`ProcessorValidationError`. - """ - if not isinstance(path, str) or not path: - raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") + """Structurally validate *fn* against the expected *kind* (signature only). + Hard violations raise; soft mismatches warn; checks come from ``_VALIDATORS``. + """ try: signature = inspect.signature(fn) except (TypeError, ValueError) as exc: @@ -511,91 +449,23 @@ def validate_processor(fn: Any, *, kind: ProcessorKind, path: str, stage_config: stage_config=stage_config, message=f"cannot inspect signature: {exc}", ) - return # unreachable - params = signature.parameters - - if kind == "placeholder_prompt_builder": - _check_orchestrator_first_param(params, path=path, kind=kind, stage_config=stage_config) - _check_suffix( - path, - kind, - suffix_pool=(_TOKEN_ONLY_SUFFIX,), - stage_config=stage_config, - what="a *_token_only suffix", - ) + validators = _VALIDATORS.get(kind) + if validators is None: # pragma: no cover - defensive against future literal additions + raise ValueError(f"unknown processor kind {kind!r}") + for check in validators: + check(signature, path=path, kind=kind, stage_config=stage_config) - elif kind == "diffusion_input_builder": - _check_orchestrator_first_param(params, path=path, kind=kind, stage_config=stage_config) - _check_suffix( + # Soft return-annotation check: mismatches warn, never fail. + ann = signature.return_annotation + if ann is not inspect.Signature.empty and not isinstance(ann, str) and ann in _SCALAR_RETURN_ANNOTATIONS: + _warn( path, kind, - suffix_pool=_DIFFUSION_SUFFIXES, - stage_config=stage_config, - what="a diffusion-builder suffix (ar2diffusion/ar2dit/thinker2imagegen/...)", - ) - - elif kind == "producer_full_payload": - # Bind the exact worker keyword call (transfer_manager, pooling_output, - # request) so a producer that only *declares* the right names but cannot - # accept the keyword call (wrong names / positional-only payload params) - # is rejected at config time; **kwargs producers are compatible. - required: tuple[str, ...] = ("transfer_manager", "pooling_output", "request") - # A keyword-acceptable is_finished must bind the worker's best-effort - # is_finished=True kwarg; otherwise a required one fails with TypeError. - is_finished_param = params.get("is_finished") - if is_finished_param is not None and is_finished_param.kind in ( - inspect.Parameter.KEYWORD_ONLY, - inspect.Parameter.POSITIONAL_OR_KEYWORD, - ): - required = required + ("is_finished",) - _check_producer_keyword_bind( - signature, - path=path, - kind=kind, - stage_config=stage_config, - required=required, - what="full-payload", - ) - _check_full_payload_is_finished(params, path=path, kind=kind, stage_config=stage_config) - - elif kind == "producer_async_chunk": - # ``is_finished`` is required and always passed on the async-chunk path; - # validate it first so its dedicated rules (is_finished_required / - # is_finished_positional_only) fire before the generic keyword bind. - _check_async_chunk_is_finished(params, path=path, kind=kind, stage_config=stage_config) - _check_producer_keyword_bind( - signature, - path=path, - kind=kind, - stage_config=stage_config, - required=("transfer_manager", "multimodal_output", "request", "is_finished"), - what="async-chunk", + "return_annotation", + f"return annotation {ann!r} is a scalar, unlikely a stage-input processor result", ) - elif kind == "legacy_multi_source": - if _MOSS_MULTI_SOURCE_MODULE_HINT not in _module_part(path): - _fail( - path=path, - kind=kind, - rule="legacy_multi_source_allowlist", - stage_config=stage_config, - message=f"legacy multi-source processors are allowlisted to moss_tts.*, got {path!r}", - ) - - elif kind == "legacy_orchestrator_builder": - # No structural contract: the orchestrator normalizes legacy positional - # shapes (C0/C2/C3/C4) through wrap_orchestrator_processor, and the - # first-parameter name is not reliable for legacy builders (configs may - # legitimately point at C-callables / operator builtins as test doubles). - # Deliberately no first-param check here. - pass - - else: # pragma: no cover - defensive against future literal additions - raise ValueError(f"unknown processor kind {kind!r}") - - _check_return_annotation(signature, path=path, kind=kind) - # --------------------------------------------------------------------------- # Resolution entry point @@ -608,18 +478,10 @@ def resolve_processor( expected_kind: ProcessorKind | None = None, stage_config: Any = None, ) -> ProcessorSpec: - """Resolve a dotted *path* to a validated :class:`ProcessorSpec`. + """Resolve *path* to a validated :class:`ProcessorSpec`. Flow: ``import_symbol`` -> ``infer_kind`` -> ``validate_processor`` -> - (optional) ``expected_kind`` check. The returned ``spec.fn`` is the same - callable the legacy ``getattr(importlib.import_module(...), ...)`` lookup - produced, so callers keep identical runtime behaviour. - - Args: - path: Dotted ``module.attr`` path of the processor. - expected_kind: When given, ``infer_kind`` must match; a mismatch raises - :class:`ProcessorValidationError` (rule ``expected_kind``). - stage_config: Optional config whose ``stage_id`` is attached to errors. + optional ``expected_kind`` check; ``spec.fn`` matches the legacy lookup. """ fn = import_symbol(path) if not callable(fn): @@ -630,7 +492,6 @@ def resolve_processor( stage_config=stage_config, message=f"resolved symbol is not callable: {fn!r}", ) - raise AssertionError("unreachable") # pragma: no cover kind = infer_kind(fn, path=path) validate_processor(fn, kind=kind, path=path, stage_config=stage_config) @@ -643,7 +504,6 @@ def resolve_processor( stage_config=stage_config, message=f"inferred kind {kind!r} does not match expected_kind {expected_kind!r}", ) - raise AssertionError("unreachable") # pragma: no cover return ProcessorSpec(path=path, kind=kind, fn=fn) @@ -661,28 +521,12 @@ def dead_processor_hint( has_sync: bool, custom_is_sync: bool = False, ) -> bool: - """Whether an orchestrator-side ``custom_process_input_func`` is never invoked. - - This mirrors the orchestrator gate in ``engine/orchestrator._route_output`` - (``(not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1))``): - - * ``async_chunk=True`` and the downstream stage **receives** async chunks -> - the orchestrator skips ``process_engine_inputs`` for that transition, so - the input processor may be **dead** (returns ``True``). - * ``async_chunk=True`` and the downstream stage does **not** receive async - chunks -> the transition still runs forward, so the processor is alive - (returns ``False``). - * ``async_chunk=False`` -> forward path is active; the processor is only - dead when a *different* ``sync_process_input_func`` overrides it - (``has_sync=True`` and ``custom_is_sync=False``). When the custom hook - *is* the selected sync hook (``custom_is_sync=True``) it is the active - processor and is **not** dead. - - Producer kinds are never orchestrator-side, so they are never "dead" under - this definition. - - This is a **pure decision helper** for warnings/tests only; the orchestrator - does not enforce it at runtime (M0 phase is warn-first). + """Whether an orchestrator-side custom hook is never invoked (gate three-state). + + Mirrors ``engine/orchestrator._route_output``: async-chunk transitions skip + forward processing when the downstream stage receives chunks; otherwise a + differing ``sync_process_input_func`` overrides the custom hook. Pure + decision helper (warn-first; the orchestrator does not enforce it). """ if kind in _PRODUCER_KINDS: return False diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 453cef5b27b..cd681ec4692 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -633,18 +633,12 @@ def build_forward_placeholder( source_outputs: list[Any], ctx: OrchestratorInputContext, ) -> list[OmniTokensPrompt]: - """Orchestrator-side sync forward placeholder builder. - - Semantics match the legacy ``thinker2talker_token_only`` forward path: one - placeholder ``OmniTokensPrompt`` per upstream output, sized so the - scheduler can reserve KV-cache slots, with speaker / language copied from - the original prompt so they survive when Stage-0 request metadata is - unavailable to the connector payload. - - The length computation is delegated to - ``_common.compute_placeholder_prompt_len(mode="full")`` (the Qwen - chat-template scan, golden-locked at 15 for the golden prompt) and the - placeholder packing to ``_common.pack_placeholder_prompt``. + """Sync forward placeholder: one KV-sized ``OmniTokensPrompt`` per output. + + Mirrors the legacy ``thinker2talker_token_only`` path; speaker/language are + copied from the prompt so they survive without Stage-0 request metadata. + Length/packing delegated to ``_common.compute_placeholder_prompt_len(mode="full")`` + and ``_common.pack_placeholder_prompt``. """ talker_inputs: list[OmniTokensPrompt] = [] streaming_context = ctx.streaming_context @@ -693,23 +687,14 @@ def build_prewarm_placeholder( ctx: OrchestratorInputContext, downstream_stage_id: int, ) -> list[OmniTokensPrompt]: - """Best-effort async-chunk prewarm placeholder builder. - - async-chunk mode has **no upstream ``source_outputs``** yet at prewarm - time, so this is a best-effort estimate: the placeholder length comes from - ``_common.compute_placeholder_prompt_len(mode="stage0_only")``, which runs - the **Qwen chat-template scan on the stage-0 input list** (the same scan - ``adapter.compute_talker_prompt_ids_length`` runs). This keeps the prewarm - estimate consistent with the sync forward placeholder and the inline - fallback (golden prompt: 15 == 15; single user segment: 6 == 6). The - connector fixup path - (``adapter.construct_next_stage_streaming_input_prompt``) replaces this - estimate with the real length once the upstream chunk arrives. - - ``ctx`` / ``downstream_stage_id`` are accepted for contract uniformity; - voice metadata is intentionally **not** forwarded here (matching the - pre-existing prewarm behaviour — no new behaviour is introduced). Per-edge - refinement (``downstream_stage_id``) is out of scope. + """Best-effort async-chunk prewarm estimate (no ``source_outputs`` yet). + + Length from ``_common.compute_placeholder_prompt_len(mode="stage0_only")`` + (the same Qwen scan ``adapter.compute_talker_prompt_ids_length`` runs), so + prewarm matches sync-forward and the inline fallback; the connector fixup + path replaces it once the upstream chunk arrives. Voice metadata is not + forwarded (matching pre-existing prewarm behaviour); per-edge refinement + is out of scope. """ del ctx, downstream_stage_id # Only the stage-0 input estimate is used. if hasattr(stage0_prompt, "prompt_token_ids"): @@ -730,18 +715,10 @@ def thinker2talker_token_only( requires_multimodal_data: bool = False, streaming_context: Any | None = None, ) -> list[OmniTokensPrompt]: - """Orchestrator-side placeholder builder for Stage-1 (Talker) when - ``async_chunk=False``. - - After the communication-layer refactor, this function only allocates a - placeholder ``prompt_token_ids`` of the correct length so the scheduler can - reserve KV-cache slots. It does **not** forward bulk tensors. - - Thin delegation to :func:`build_forward_placeholder` — the dual-entry design - keeps the sync forward path and the async-chunk prewarm path on the same - shared ``_common`` length / packing helpers. ``prompt`` / - ``requires_multimodal_data`` / ``streaming_context`` are kept for call-site - signature compatibility; they are folded into an ``OrchestratorInputContext``. + """Stage-1 sync placeholder (``async_chunk=False``): KV-sized slot only. + + Thin delegation to :func:`build_forward_placeholder`; legacy params are kept + for call-site compatibility and folded into an ``OrchestratorInputContext``. """ ctx = OrchestratorInputContext( prompt=prompt, @@ -751,10 +728,8 @@ def thinker2talker_token_only( return build_forward_placeholder(source_outputs, ctx) -# The orchestrator resolves ``sync_process_input_func`` (the ``*_token_only`` -# path) via ``resolve_processor(...).fn`` and reads ``build_prewarm_placeholder`` -# off the resolved function object. Expose both dual-entry builders here so -# that lookup works without a separate module scan. +# Expose the dual-entry builders on the resolved ``*_token_only`` fn so the +# orchestrator's ``sync_process_input_func`` lookup finds them without a scan. thinker2talker_token_only.build_forward_placeholder = build_forward_placeholder # type: ignore[attr-defined] thinker2talker_token_only.build_prewarm_placeholder = build_prewarm_placeholder # type: ignore[attr-defined] @@ -778,9 +753,7 @@ def talker2code2wav_async_chunk( request: OmniEngineCoreRequest, is_finished: bool = False, ) -> OmniPayloadStruct | None: - """ - Multimodal output version. - """ + """Async-chunk producer: multimodal output version.""" request_id = request.external_req_id code_predictor_codes = None frame_valid = None From d06aa9b1ac6000a0e33d72649e038fe72d8a117f Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Tue, 13 Oct 2026 09:38:42 +0800 Subject: [PATCH 37/41] fix(engine): read sync_process_input_func and prewarm_placeholder_builder defensively Older metadata doubles (e.g. the upstream async-chunk prewarm test scenario) lack the PR-added fields; direct attribute access raised AttributeError. Use getattr with None defaults to match the upstream defensive style used for async_chunk_prewarm_payload_func. Signed-off-by: RyanYun09 --- vllm_omni/engine/stage_engine_core_client.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 4ae6749ca97..48b3ad0ac56 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -155,8 +155,10 @@ def __init__( self.async_chunk_prewarm_payload_func = getattr(metadata, "async_chunk_prewarm_payload_func", None) # Raw ``sync_process_input_func`` path (dead-processor hint) plus the # already-resolved dual-entry ``build_prewarm_placeholder`` builder. - self.sync_process_input_func = metadata.sync_process_input_func - self.prewarm_placeholder_builder = metadata.prewarm_placeholder_builder + # Both default to None so older metadata doubles lacking the fields + # (e.g. the upstream async-chunk prewarm tests) keep working. + self.sync_process_input_func = getattr(metadata, "sync_process_input_func", None) + self.prewarm_placeholder_builder = getattr(metadata, "prewarm_placeholder_builder", None) self.engine_outputs: Any = None self.client_addresses = dict(client_addresses or {}) From 729d593f98ea9666bfc18159c2e20ec7c66f64a6 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Wed, 14 Oct 2026 20:24:00 +0800 Subject: [PATCH 38/41] fix(stage_input_processors): accept 3-D per-step codec frames in extract_last_codec_frame (Voxtral P1, RFC #4872) Voxtral TTS compute_mm_logits() emits per-step [B, 1, C] audio codes (payload list elements of shape (1, 1, C)); the consolidated helper only handled 1-D/2-D and raised ValueError. Normalize 3-D per-step frames via targeted squeeze(1) so the shared 2-D last-row semantics (including the validity gate) apply unchanged; any other 3-D layout (e.g. (2, 2, 2)) still raises ValueError, preserving the golden-locked fish/qwen3/higgs contracts. Adds 4 regression tests (test_voxtral_tts_async_chunk.py) and 4 golden cases (test_common_helpers_golden.py) for the 3-D per-step layout. Signed-off-by: RyanYun09 --- .../test_common_helpers_golden.py | 21 ++++ .../test_voxtral_tts_async_chunk.py | 95 ++++++++++++++++++- .../stage_input_processors/_common.py | 14 ++- 3 files changed, 128 insertions(+), 2 deletions(-) diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py index 9f232852426..99155faa23f 100644 --- a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -287,6 +287,27 @@ def _extract_last_frame_cases(): id="voxtral-list-audio", ), pytest.param(voxtral, {"codes": {}}, None, False, id="voxtral-missing-audio"), + # 3-D per-step frames [B, 1, C] (Voxtral compute_mm_logits L1075): + # normalized via targeted squeeze(1) then the 2-D last-row semantics. + pytest.param( + voxtral, + _codes_audio(torch.arange(8, dtype=torch.float32).reshape(1, 1, 8)), + [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], + False, + id="voxtral-3d-per-step", + ), + # Mirror the real payload: a list whose last element is 3-D (1, 1, C). + pytest.param( + voxtral, + {"codes": {"audio": [torch.ones(1, 1, 4)]}}, + [1.0, 1.0, 1.0, 1.0], + False, + id="voxtral-3d-list-element", + ), + # Non per-step 3-D layouts keep raising ValueError (dim1 != 1). + pytest.param(voxtral, _codes_audio(torch.zeros(1, 2, 3)), None, True, id="voxtral-3d-non-singleton-raises"), + # Generic 3-D (aligned with the fish/qwen3/higgs contracts). + pytest.param(voxtral, _codes_audio(torch.zeros(2, 2, 2)), None, True, id="voxtral-3d-generic-raises"), ] diff --git a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py index dc248f21179..99df9dcf27b 100644 --- a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections import defaultdict from types import SimpleNamespace @@ -213,3 +213,96 @@ def test_none_multimodal_output_not_finished_returns_none(): ) assert payload is None + + +def _compute_mm_logits_payload(batch: int = 1, codebooks: int = 8): + """Mirror voxtral_tts_audio_generation.compute_mm_logits() L1075: + audio_codes [B, C] -> unsqueeze(1) -> [B, 1, C] -> split -> list of (1, 1, C). + """ + audio_codes = torch.arange(batch * codebooks, dtype=torch.float32).reshape(batch, codebooks) + audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0)) + return {"codes": {"audio": audio_list}} + + +def test_3d_per_step_frame_list_accepted(): + """Voxtral P1 regression: compute_mm_logits 3-D per-step [B, 1, C] + payload must not raise ValueError; the last frame is extracted as 1-D.""" + from vllm_omni.model_executor.stage_input_processors import _common + + transfer_manager = _make_transfer_manager() + request_id = "rid-3d" + request = _req(request_id, finished=False) + + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + multimodal_output=_compute_mm_logits_payload(1, 8), + request=request, + ) + + # The frame is stored as the flat 8-element codec row. + assert transfer_manager.code_prompt_token_ids[request_id][-1] == [float(i) for i in range(8)] + # One frame < chunk_size_at_begin=5: no emission yet. + assert payload is None + # Direct helper call: 1-D float frame, precision preserved (to_long=False). + frame = _common.extract_last_codec_frame( + {"codes": {"audio": torch.zeros(1, 1, 8)}}, + to_long=False, + ) + assert frame is not None + assert frame.shape == (8,) + assert frame.dtype == torch.float32 + + +def test_3d_direct_tensor_accepted(): + """A bare 3-D [1, 1, C] tensor is normalized via squeeze(1) and the + 2-D last-row semantics apply unchanged.""" + from vllm_omni.model_executor.stage_input_processors import _common + + frame = _common.extract_last_codec_frame( + {"codes": {"audio": torch.arange(8, dtype=torch.float32).reshape(1, 1, 8)}}, + to_long=False, + ) + assert frame is not None + assert frame.shape == (8,) + assert frame.dtype == torch.float32 + assert frame.tolist() == [float(i) for i in range(8)] + + +def test_3d_per_step_chunk_emission(): + """Feeding 5 per-step 3-D frames reaches the codec_chunk_frames_at_begin + boundary and emits the standard [ctx_frames, context_length, ...flat] payload.""" + transfer_manager = _make_transfer_manager( + codec_chunk_frames=25, + codec_left_context_frames=25, + codec_chunk_frames_at_begin=5, + ) + request_id = "rid-3d-chunk" + request = _req(request_id, finished=False) + + payload = None + for _ in range(5): + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + multimodal_output=_compute_mm_logits_payload(1, 8), + request=request, + ) + + assert payload is not None + codes = payload.codes.audio + assert int(codes[0].item()) == 0 # ctx_frames + assert int(codes[1].item()) == 5 # context_length == chunk_size_at_begin + assert codes[2:].numel() == 5 * 8 # 5 frames x 8 codebooks + + +def test_3d_non_singleton_frame_dim_raises(): + """Defensive contract: non per-step 3-D layouts (dim1 != 1) still raise + ValueError instead of being silently flattened.""" + transfer_manager = _make_transfer_manager() + request = _req("rid-3d-non-singleton", finished=False) + + with pytest.raises(ValueError): + generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + multimodal_output={"codes": {"audio": torch.zeros(1, 2, 3)}}, + request=request, + ) diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py index 9c32ced835e..f97853aeefc 100644 --- a/vllm_omni/model_executor/stage_input_processors/_common.py +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -235,7 +235,9 @@ def extract_last_codec_frame( ``key_path`` codes/audio vs audio_codes; ``validate`` None/"any"/"valid_mask"; ``to_cpu``/``to_long`` (fish_speech). Validity gate applies only to 2-D - frames; 1-D frames return directly; result is always 1-D. + frames; 1-D frames return directly; result is always 1-D. 3-D per-step + frames ``[B, 1, C]`` (Voxtral ``compute_mm_logits``) are normalized to + ``[B, C]`` first; any other 3-D layout still raises ``ValueError``. """ if not isinstance(payload, dict): return None @@ -250,6 +252,16 @@ def extract_last_codec_frame( audio_codes = audio_codes[-1] if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: return None + if audio_codes.ndim == 3 and audio_codes.shape[1] == 1: + # Voxtral TTS: compute_mm_logits() emits per-step [B, 1, C] codes + # (payload list elements of shape (1, 1, C)); the legacy voxtral + # helper flattened them, so 3-D per-step input is a supported + # layout, not a contract violation. Targeted squeeze(1) reduces + # it to [B, C] so the shared 2-D last-row semantics (including + # any validity gate) apply unchanged. Non-singleton frame dims + # (e.g. (2, 2, 2)) intentionally fall through to the ValueError + # below, preserving the golden-locked fish/qwen3/higgs contracts. + audio_codes = audio_codes.squeeze(1) if audio_codes.ndim == 2: frame: torch.Tensor = audio_codes[-1] if frame.numel() == 0: From 83199a2bc2b81a5f8e0d14cb2e75eaaff38c0b3d Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Thu, 15 Oct 2026 14:04:17 +0800 Subject: [PATCH 39/41] [Core] fix scheduler async-chunk hook resolution and prewarm snapshot copies Address amy-why-3459's two P2 regressions on the stage_input_processors refactor: - chunk_transfer_adapter: resolve the explicitly configured async-chunk hook by signature structure (duck typing) instead of the ``*_async_chunk`` name suffix, aligning with the worker loader; a structurally incompatible hook still fails fast. - orchestrator._build_prewarm_placeholder_input: restore the ``_PREWARM_UNCOPIED_PROMPT_KEYS`` exclusion and the ``shared_types=(torch.Tensor,)`` snapshot semantics so discarded media is not deep-copied and tensors stay shared. Adds regression tests for arbitrary-name hook acceptance / incompatible hook rejection and for prewarm media filtering / tensor sharing. Signed-off-by: RyanYun09 --- .../test_chunk_transfer_adapter.py | 33 +++++++++++ .../test_orchestrator_stage_input_bridge.py | 55 +++++++++++++++++++ .../chunk_transfer_adapter.py | 47 +++++++++++++--- vllm_omni/engine/orchestrator.py | 15 ++++- 4 files changed, 139 insertions(+), 11 deletions(-) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 5c9a25aab9a..3e3091de6cf 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -574,6 +574,7 @@ def _build( flat_tts_config: bool = False, active_stream_window: int = 0, connector_extra: dict | None = None, + model_config_extra: dict | None = None, ): connector = mocker.MagicMock() connector.stage_id = stage_id @@ -627,6 +628,7 @@ def _fake_base_init(self, config): "name": "SharedMemoryConnector", "extra": connector_extra or {}, }, + **(model_config_extra or {}), ) scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs) adapter = OmniChunkTransferAdapter( @@ -637,6 +639,37 @@ def _fake_base_init(self, config): return _build +def pack(transfer_manager, multimodal_output, request, is_finished=False): + """Arbitrarily named, structurally compatible async-chunk producer.""" + del transfer_manager, multimodal_output, request, is_finished + return None + + +def _incompatible_hook(transfer_manager, request): + """Structurally invalid: no ``multimodal_output`` and no ``is_finished``.""" + del transfer_manager, request + return None + + +def test_adapter_accepts_explicitly_configured_hook_with_arbitrary_name(build_adapter): + """A signature-compatible hook without a ``*_async_chunk`` suffix loads.""" + adapter, _ = build_adapter(model_config_extra={"custom_process_next_stage_input_func": f"{__name__}.pack"}) + + assert adapter.custom_process_next_stage_input_func is pack + + +def test_adapter_rejects_structurally_incompatible_hook(build_adapter): + """A hook missing required parameters fails fast instead of loading.""" + from vllm_omni.model_executor.stage_input_processors import ProcessorValidationError + + with pytest.raises(ProcessorValidationError): + build_adapter( + model_config_extra={ + "custom_process_next_stage_input_func": f"{__name__}._incompatible_hook", + } + ) + + @pytest.mark.parametrize("active_stream_window", [0, 2]) def test_scheduler_chunk_gate_preserves_kv_holding_queue(build_adapter, active_stream_window): adapter, _ = build_adapter(active_stream_window=active_stream_window) diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 4d5e100508a..dc0281c7714 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -11,6 +11,7 @@ import janus import pytest +import torch from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams @@ -613,3 +614,57 @@ def __init__(self) -> None: assert base_input["prompt_token_ids"] == expected_ids assert base_input["multi_modal_data"] is None assert base_input["mm_processor_kwargs"] is None + + +class _DeepcopySpy: + """Sentinel media object that records any attempted deep-copy.""" + + def __init__(self) -> None: + self.deepcopy_calls = 0 + + def __deepcopy__(self, memo: dict[int, Any]) -> _DeepcopySpy: # type: ignore[type-arg] + self.deepcopy_calls += 1 + return _DeepcopySpy() + + +def test_prewarm_placeholder_input_filters_discarded_media() -> None: + """Discarded prompt fields (media) must never be deep-copied for a prewarm.""" + media = _DeepcopySpy() + orch = _prewarm_orchestrator_with_client(SimpleNamespace(prewarm_placeholder_builder=None)) + req_state = OrchestratorRequestState( + request_id="req-prewarm-media", + prompt={ + "prompt_token_ids": [1, 2], + "multi_modal_data": media, + "mm_processor_kwargs": {"foo": "bar"}, + }, + sampling_params_list=[SamplingParams(max_tokens=1)], + final_stage_id=0, + ) + + base_input = orch._build_prewarm_placeholder_input(0, "req-prewarm-media", [1, 2], req_state) + + assert base_input["multi_modal_data"] is None + assert base_input["mm_processor_kwargs"] is None + assert media.deepcopy_calls == 0 + + +def test_prewarm_placeholder_input_shares_additional_information_tensors() -> None: + """``additional_information`` containers are isolated but tensors stay shared.""" + tensor = torch.arange(4, dtype=torch.long) + original = {"t": tensor, "l": [1, 2]} + orch = _prewarm_orchestrator_with_client(SimpleNamespace(prewarm_placeholder_builder=None)) + req_state = OrchestratorRequestState( + request_id="req-prewarm-additional", + prompt={"prompt_token_ids": [1, 2], "additional_information": original}, + sampling_params_list=[SamplingParams(max_tokens=1)], + final_stage_id=0, + ) + + base_input = orch._build_prewarm_placeholder_input(0, "req-prewarm-additional", [1, 2], req_state) + + copied = base_input["additional_information"] + assert copied is not original + assert copied["l"] is not original["l"] + assert copied["t"] is tensor + assert copied["t"].data_ptr() == tensor.data_ptr() diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index df1ebc4aa59..270480aaecb 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -15,7 +15,11 @@ from vllm.v1.utils import ConstantList from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload -from vllm_omni.model_executor.stage_input_processors import resolve_processor +from vllm_omni.model_executor.stage_input_processors import ( + ProcessorValidationError, + resolve_processor, +) +from vllm_omni.model_executor.stage_input_processors._registry import _check_async_chunk from ..adapter import construct_next_stage_streaming_input_prompt from ..connectors.shm_connector import SharedMemoryConnector @@ -31,6 +35,24 @@ _RECLAIM_WARN_INTERVAL = 1000 +def _accepts_async_chunk_contract(fn: Callable[..., Any]) -> bool: + """Whether *fn* structurally matches the async-chunk producer contract.""" + try: + signature = inspect.signature(fn) + except (TypeError, ValueError): + return False + try: + _check_async_chunk( + signature, + path=getattr(fn, "__qualname__", repr(fn)), + kind="producer_async_chunk", + stage_config=None, + ) + except ProcessorValidationError: + return False + return True + + class _SenderGeneration: """Fence one external request generation without blocking cleanup.""" @@ -185,14 +207,21 @@ def __init__(self, vllm_config: Any): self.custom_process_next_stage_input_func: Callable[..., OmniPayloadStruct | None] | None = None custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None) if custom_process_next_stage_input_func: - # Resolve through the registry with the async-chunk producer - # contract. ``_select_processor_funcs`` already merged the - # async_chunk field into this path, so it must be a real - # ``*_async_chunk`` producer; ``is_finished`` is required (the - # scheduler always passes it below) — a missing/positional-only - # declaration is a hard configuration error (ProcessorValidationError). - spec = resolve_processor(custom_process_next_stage_input_func, expected_kind="producer_async_chunk") - self.custom_process_next_stage_input_func = spec.fn + # Kind inference is name-driven; accept any matching signature, reject only incompatible hooks. + spec = resolve_processor(custom_process_next_stage_input_func) + if spec.kind == "producer_async_chunk" or _accepts_async_chunk_contract(spec.fn): + self.custom_process_next_stage_input_func = spec.fn + else: + raise ProcessorValidationError( + path=custom_process_next_stage_input_func, + kind="producer_async_chunk", + rule="async_chunk_contract", + message=( + "custom_process_next_stage_input_func does not accept the async-chunk " + "producer call (requires keyword-bindable 'transfer_manager', " + "'multimodal_output', 'request' and a 'is_finished' parameter)" + ), + ) # mapping for request id and chunk id self.put_req_chunk: dict[str, int] = defaultdict(int) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 3ca58fb2ae3..439b52c42d5 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -3258,9 +3258,20 @@ def _build_prewarm_placeholder_input( original_prompt = req_state.prompt if isinstance(original_prompt, dict): - import copy + from vllm_omni.engine.request_snapshot import copy_request_snapshot - base_input = copy.deepcopy(original_prompt) + # Discarded fields are skipped; additional_information tensors stay shared. + memo: dict[int, Any] = {} + base_input = copy_request_snapshot( + {k: v for k, v in original_prompt.items() if k not in _PREWARM_UNCOPIED_PROMPT_KEYS}, + memo, + ) + if "additional_information" in original_prompt: + base_input["additional_information"] = copy_request_snapshot( + original_prompt["additional_information"], + memo, + shared_types=(torch.Tensor,), + ) else: base_input = {} From 5856af4e9f79ab7c5a4fe8b54ed688dc9ad89688 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Fri, 16 Oct 2026 14:03:47 +0800 Subject: [PATCH 40/41] [Core] address linyueqian's optional follow-ups on connector hook resolution - qwen2_5_omni.thinker2talker_full_payload: normalise None token-id fields with _common.ensure_list (None -> []) instead of ensure_list_unchanged (None -> None), so a None no longer reaches len(None)/list(None) in the downstream length-aware trim (inline #4290791768). Add 3 boundary tests. - omni_connector_runtime._load_custom_func: qualify the docstring (the name-driven kind inference can reject a hook before the structural payload check), escalate an explicitly configured hook that is skipped to a warning naming the register_processor() escape hatch, and keep auto-derived probe candidates at debug (inline #4290791761). Add 2 scan tests. Signed-off-by: RyanYun09 --- .../test_qwen2_5_omni_thinker2talker.py | 66 +++++++++++++++ tests/worker/test_omni_connector_mixin.py | 46 +++++++++++ .../model_runner/omni_connector_runtime.py | 80 ++++++++++++++----- .../stage_input_processors/qwen2_5_omni.py | 10 ++- 4 files changed, 177 insertions(+), 25 deletions(-) diff --git a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py index 3ecae3d5688..ab94b5d08ad 100644 --- a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py +++ b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py @@ -85,5 +85,71 @@ def test_missing_hidden_returns_none(): assert thinker2talker_full_payload(transfer_manager=None, pooling_output={}, request=request) is None +def _make_request_with_raw_token_ids(prompt, output, all_ids, status_name=None): + """Build a request without the ``_make_request`` ``list(...)`` coercion so a + ``None`` token-id field survives to the builder unchanged. + """ + status = SimpleNamespace(name=status_name) if status_name else None + return SimpleNamespace( + request_id="r1", + prompt_token_ids=prompt, + output_token_ids=output, + all_token_ids=all_ids, + status=status, + sampling_params=None, + ) + + +def test_none_prompt_and_output_token_ids_do_not_crash(): + """Regression (linyueqian #4290791768): ``None`` token-id fields must be + normalised to ``[]`` by ``ensure_list`` instead of reaching ``len(None)`` / + ``list(None)`` and raising ``TypeError``. + """ + request = _make_request_with_raw_token_ids(None, None, None) + hidden = torch.arange(2 * 4, dtype=torch.float32).reshape(2, 4) + + # No prompt and no output => target_rows == 0 => graceful None (no crash). + payload = thinker2talker_full_payload( + transfer_manager=None, pooling_output={"hidden": hidden}, request=request + ) + assert payload is None + + +def test_none_output_token_ids_with_prompt_builds_payload(): + """``output_token_ids=None`` (a fresh, pre-decode request) must be treated + as an empty list: prompt-only payload is still shipped. + """ + request = _make_request_with_raw_token_ids((1, 2, 3), None, None) + # Tuple prompt is normalised to list; 3 hidden rows match the prompt span. + hidden = torch.arange(3 * 4, dtype=torch.float32).reshape(3, 4) + + payload = thinker2talker_full_payload( + transfer_manager=None, pooling_output={"hidden": hidden}, request=request + ) + + assert payload is not None + assert payload["ids"]["prompt"] == [1, 2, 3] + assert payload["ids"]["output"] == [] + assert payload["embed"]["prefill"].shape[0] == 3 + + +def test_none_all_token_ids_falls_back_to_prompt_plus_output(): + """``all_token_ids=None`` must fall back to ``prompt + output`` rather than + raising on ``list(None)``. + """ + request = _make_request_with_raw_token_ids([1, 2], [10, 11], None) + hidden = torch.arange(4 * 4, dtype=torch.float32).reshape(4, 4) + + payload = thinker2talker_full_payload( + transfer_manager=None, pooling_output={"hidden": hidden}, request=request + ) + + assert payload is not None + assert payload["ids"]["prompt"] == [1, 2] + assert payload["ids"]["output"] == [10, 11] + assert payload["embed"]["prefill"].shape[0] == 2 + assert payload["hidden_states"]["output"].shape[0] == 2 + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index 9190e1cf3fc..b550c77d1c7 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -551,6 +551,15 @@ def pack(transfer_manager, multimodal_output, request, is_finished=False): return None +def pack_full_payload(transfer_manager, multimodal_output, request, is_finished=False): + """Name suffix says ``producer_full_payload`` while the signature follows the + async-chunk contract (``multimodal_output``) -- the exact shape a + name-driven kind inference mis-classifies before the structural check. + """ + del transfer_manager, multimodal_output, request, is_finished + return None + + class TestLoadCustomFuncSelection(unittest.TestCase): def test_uses_validator_override_from_public_mixin(self): config = SimpleNamespace( @@ -565,6 +574,43 @@ def test_uses_validator_override_from_public_mixin(self): assert selected_path is None assert func is None + def test_warns_when_explicit_hook_is_skipped(self): + """linyueqian #4290791761: an explicitly configured hook that is dropped + (here: rejected before the structural check) is escalated to a + ``warning`` naming the ``register_processor`` escape hatch, instead of + vanishing into a ``debug``-only line. + """ + config = SimpleNamespace( + async_chunk=True, + custom_process_next_stage_input_func=f"{__name__}.pack_full_payload", + ) + + with patch.object(omni_connector_runtime.logger, "warning") as warn: + selected_path, func = MixinHost._load_custom_func(config) + + assert selected_path is None + assert func is None + warn.assert_called_once() + message = warn.call_args.args[0] + assert "explicitly configured connector payload hook" in message + assert "register_processor" in message + + def test_auto_derived_probe_candidates_stay_at_debug(self): + """Auto-derived ``*_full_payload`` / ``*_batch`` probes are the expected + misses of the scan and must not spam ``warning`` when the sibling + builder does not exist. + """ + config = SimpleNamespace( + async_chunk=False, + custom_process_next_stage_input_func=None, + custom_process_input_func="vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav", + ) + + with patch.object(omni_connector_runtime.logger, "warning") as warn: + MixinHost._load_custom_func(config) + + warn.assert_not_called() + def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self): incompatible_paths = [ "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav", diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 64f645245fd..ad7a325a61c 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -617,16 +617,30 @@ def _load_custom_func(cls, model_config: Any) -> tuple[str | None, Any | None]: The loader accepts both supported payload contracts: full-payload builders name their runner payload ``pooling_output`` while async-chunk builders name it ``multimodal_output`` (see - :meth:`_connector_payload_kwarg`). The acceptance gate is the mixin's - structural validator, not the registry's name-driven kind inference, so - an explicitly configured, structurally compatible hook is never - rejected solely because of its name. + :meth:`_connector_payload_kwarg`). For the configured hook the mixin's + structural validator (:meth:`_is_connector_payload_builder`) is the + authoritative acceptance gate. + + Caveat: :func:`resolve_processor` runs the registry's *name-driven* + kind inference before that structural check, so a hook whose name ends + in ``_full_payload`` (for example ``pack_full_payload``) is classified + as a ``producer_full_payload`` builder and validated against *that* + contract. If its signature actually follows a different payload + contract, it is rejected before :meth:`_is_connector_payload_builder` + ever sees it. Such a hook is still reported at ``warning`` level (see + :meth:`_warn_skipped_hook`) rather than silently dropped, and can be + pinned to the intended kind via + ``stage_input_processors.register_processor(path, "producer_async_chunk")``. """ - candidates: list[str] = [] + # Each candidate carries an ``is_explicit`` flag: hooks named in config + # are escalated to ``warning`` when skipped (dropping one disables the + # stage's connector payload transfer), while auto-derived probe + # candidates stay at ``debug`` to avoid noise from the normal scan. + candidates: list[tuple[str, bool]] = [] next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) if isinstance(next_stage_func, str) and next_stage_func: - candidates.append(next_stage_func) + candidates.append((next_stage_func, True)) if not getattr(model_config, "async_chunk", False): input_func = getattr(model_config, "custom_process_input_func", None) @@ -634,16 +648,16 @@ def _load_custom_func(cls, model_config: Any) -> tuple[str | None, Any | None]: try: module_path, func_name = input_func.rsplit(".", 1) if func_name.endswith("_full_payload") or func_name.endswith("_batch"): - candidates.append(f"{module_path}.{func_name}") + candidates.append((f"{module_path}.{func_name}", False)) else: - candidates.append(f"{module_path}.{func_name}_full_payload") - candidates.append(f"{module_path}.{func_name}_batch") - candidates.append(input_func) + candidates.append((f"{module_path}.{func_name}_full_payload", False)) + candidates.append((f"{module_path}.{func_name}_batch", False)) + candidates.append((input_func, False)) except ValueError: - candidates.append(input_func) + candidates.append((input_func, False)) tried: set[str] = set() - for func_path in candidates: + for func_path, is_explicit in candidates: if func_path in tried: continue tried.add(func_path) @@ -654,25 +668,47 @@ def _load_custom_func(cls, model_config: Any) -> tuple[str | None, Any | None]: # authoritative acceptance gate. spec = resolve_processor(func_path) if not cls._is_connector_payload_builder(spec.fn): - logger.debug( - "Skipping incompatible connector payload hook %s; signature=%s", - func_path, - inspect.signature(spec.fn), - ) + try: + detail: Any = inspect.signature(spec.fn) + except (TypeError, ValueError): + detail = "signature unavailable" + cls._warn_skipped_hook(func_path, is_explicit, detail) continue return spec.path, spec.fn except ProcessorValidationError as exc: - logger.debug( - "Skipping incompatible connector payload hook %s: %s", - func_path, - exc, - ) + # A name-driven kind inference can pick the wrong contract + # (e.g. ``pack_full_payload`` -> ``producer_full_payload``) + # before the structural payload check ever runs. + cls._warn_skipped_hook(func_path, is_explicit, exc) continue except Exception: logger.warning("Failed to load custom func: %s", func_path, exc_info=True) return None, None + @classmethod + def _warn_skipped_hook(cls, func_path: str, is_explicit: bool, detail: Any) -> None: + """Report a connector payload hook that was skipped during the scan. + + Auto-derived probe candidates (``*_full_payload`` / ``*_batch``) are + logged at ``debug`` because a missing sibling is the expected case. + An *explicitly* configured hook is escalated to ``warning`` since + silently dropping it disables connector payload transfer for the + stage; the message names the ``register_processor`` escape hatch so a + deliberately named hook can be pinned to the correct kind. + """ + if is_explicit: + logger.warning( + "Skipping explicitly configured connector payload hook %s: %s. " + "If this hook is structurally valid, pin its kind via " + "stage_input_processors.register_processor(%r, 'producer_async_chunk').", + func_path, + detail, + func_path, + ) + else: + logger.debug("Skipping incompatible connector payload hook %s: %s", func_path, detail) + @classmethod def _is_connector_payload_builder(cls, func: Any) -> bool: """Whether *func* matches the mixin payload-builder contract.""" diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 0ab6fad5a22..bb0bf5ff707 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -245,9 +245,13 @@ def thinker2talker_full_payload( ) return None - prompt_token_ids = _common.ensure_list_unchanged(getattr(request, "prompt_token_ids", None)) - output_token_ids = _common.ensure_list_unchanged(getattr(request, "output_token_ids", None)) - all_token_ids = _common.ensure_list_unchanged(getattr(request, "all_token_ids", None) or []) + # Use ``ensure_list`` (``None`` -> ``[]``), not ``ensure_list_unchanged`` + # (``None`` -> ``None``): the downstream arithmetic below indexes and + # takes ``len(...)`` of these lists, so a ``None`` would raise instead of + # degrading gracefully. This matches the pre-refactor behaviour. + prompt_token_ids = _common.ensure_list(getattr(request, "prompt_token_ids", None)) + output_token_ids = _common.ensure_list(getattr(request, "output_token_ids", None)) + all_token_ids = _common.ensure_list(getattr(request, "all_token_ids", None)) if not all_token_ids: all_token_ids = list(prompt_token_ids) + list(output_token_ids) From 946570a095c44bf39390e03cd1258a7d98313fe3 Mon Sep 17 00:00:00 2001 From: RyanYun09 Date: Fri, 16 Oct 2026 15:07:48 +0800 Subject: [PATCH 41/41] [Core] fix pre-commit failures in stage_input_processors tests Fix the hooks that failed on the previous commit so CI passes: - typos: rewrite 'mis-classifies' -> 'misclassifies' in the pack_full_payload docstring (tests/worker/test_omni_connector_mixin.py). - ruff-format: collapse three thinker2talker_full_payload() calls onto single lines (tests/model_executor/stage_input_processors/ test_qwen2_5_omni_thinker2talker.py). No production code changes; all test assertions are preserved. Signed-off-by: RyanYun09 --- .../test_qwen2_5_omni_thinker2talker.py | 14 ++++---------- tests/worker/test_omni_connector_mixin.py | 2 +- 2 files changed, 5 insertions(+), 11 deletions(-) diff --git a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py index ab94b5d08ad..3272a1a05fd 100644 --- a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py +++ b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Light coverage for qwen2_5_omni.thinker2talker_full_payload. Covers the finish-reason-aware stop-row trim contract: when the request @@ -109,9 +109,7 @@ def test_none_prompt_and_output_token_ids_do_not_crash(): hidden = torch.arange(2 * 4, dtype=torch.float32).reshape(2, 4) # No prompt and no output => target_rows == 0 => graceful None (no crash). - payload = thinker2talker_full_payload( - transfer_manager=None, pooling_output={"hidden": hidden}, request=request - ) + payload = thinker2talker_full_payload(transfer_manager=None, pooling_output={"hidden": hidden}, request=request) assert payload is None @@ -123,9 +121,7 @@ def test_none_output_token_ids_with_prompt_builds_payload(): # Tuple prompt is normalised to list; 3 hidden rows match the prompt span. hidden = torch.arange(3 * 4, dtype=torch.float32).reshape(3, 4) - payload = thinker2talker_full_payload( - transfer_manager=None, pooling_output={"hidden": hidden}, request=request - ) + payload = thinker2talker_full_payload(transfer_manager=None, pooling_output={"hidden": hidden}, request=request) assert payload is not None assert payload["ids"]["prompt"] == [1, 2, 3] @@ -140,9 +136,7 @@ def test_none_all_token_ids_falls_back_to_prompt_plus_output(): request = _make_request_with_raw_token_ids([1, 2], [10, 11], None) hidden = torch.arange(4 * 4, dtype=torch.float32).reshape(4, 4) - payload = thinker2talker_full_payload( - transfer_manager=None, pooling_output={"hidden": hidden}, request=request - ) + payload = thinker2talker_full_payload(transfer_manager=None, pooling_output={"hidden": hidden}, request=request) assert payload is not None assert payload["ids"]["prompt"] == [1, 2] diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index b550c77d1c7..3deac22a5ad 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -554,7 +554,7 @@ def pack(transfer_manager, multimodal_output, request, is_finished=False): def pack_full_payload(transfer_manager, multimodal_output, request, is_finished=False): """Name suffix says ``producer_full_payload`` while the signature follows the async-chunk contract (``multimodal_output``) -- the exact shape a - name-driven kind inference mis-classifies before the structural check. + name-driven kind inference misclassifies before the structural check. """ del transfer_manager, multimodal_output, request, is_finished return None