diff --git a/pyproject.toml b/pyproject.toml index 84a950719..5135661ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,8 @@ dev = [ "FlagEmbedding==1.4.0", # for hunyuan_image3 tests "kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0) "s3tokenizer==0.3.0", # for tts test + # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo, + # whose top-level stepaudio2 package conflicts with this distribution. "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml) ] diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 50fcbcf4e..ba0b75dca 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.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 """Unit tests for Omni AR streaming-session async placeholder handling.""" @@ -64,6 +64,9 @@ def _run_resumable_segment_stop( session: Request, *, session_finished: bool = False, + handle_stopped=None, + chunk_transfer_adapter=None, + inter_stage_output=None, ): sched = MagicMock() sched.requests = {session.request_id: session} @@ -75,11 +78,14 @@ def _run_resumable_segment_stop( return [42], True sched._update_request_with_output.side_effect = stop_request - sched._handle_stopped_request.return_value = session_finished # vLLM 0.26 returns (kv_xfer_params, ec_xfer_params); an unconfigured # MagicMock iterates empty and fails to unpack at the call site. sched._free_request.return_value = (None, None) - sched.chunk_transfer_adapter = None + if handle_stopped is None: + sched._handle_stopped_request.return_value = session_finished + else: + sched._handle_stopped_request.side_effect = handle_stopped + sched.chunk_transfer_adapter = chunk_transfer_adapter sched.running = [session] sched.waiting_for_transfer_free = set() sched.transfer_triggered_requests = set() @@ -105,6 +111,7 @@ def _run_resumable_segment_stop( model_runner_output.cudagraph_stats = None model_runner_output.req_id_to_index = {session.request_id: 0} model_runner_output.routed_experts = None + model_runner_output.inter_stage_outputs = [inter_stage_output] if inter_stage_output is not None else None return OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) @@ -160,6 +167,41 @@ def test_update_from_output_settles_in_flight_tokens() -> None: assert session.num_in_flight_tokens == 0 +def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> None: + """The old segment's transfer must not observe the next segment's reset. + + A queued streaming update can replace the same mutable Request while the + stop output is still being handled. The connector needs the confirmed + token count from before that replacement to avoid dropping the boundary. + """ + session = _make_request() + session.status = RequestStatus.RUNNING + session.resumable = True + session.num_computed_tokens = 26 + adapter = MagicMock() + adapter._confirmed_num_computed_tokens.return_value = 26 + inter_stage_output = {"codes": {"audio": [7]}} + + def replace_with_next_segment(request: Request) -> bool: + request.num_computed_tokens = 0 + request.status = RequestStatus.WAITING + return False + + _run_resumable_segment_stop( + session, + handle_stopped=replace_with_next_segment, + chunk_transfer_adapter=adapter, + inter_stage_output=inter_stage_output, + ) + + adapter.save_async.assert_called_once_with( + inter_stage_output, + session, + True, + confirmed_num_computed_tokens=26, + ) + + def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> None: """A decode step that neither stops nor carries an inter-stage payload. diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index d0fd00376..2dbe967e7 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.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 import threading from collections import deque @@ -9,9 +9,10 @@ from unittest.mock import Mock, patch import pytest import torch from pytest_mock import MockerFixture +from vllm.sampling_params import SamplingParams from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats -from vllm.v1.request import RequestStatus +from vllm.v1.request import Request, RequestStatus from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler @@ -337,6 +338,121 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad assert len(adapter._pending_save_reqs) == 1 +def test_segment_boundary_starts_new_send_watermark_before_background_flush(build_adapter): + """A queued boundary owns the end of its deduplication generation. + + The next segment can start before the save thread sends the old boundary. + Sending that old task later must not erase the new segment's watermark. + """ + adapter, _ = build_adapter(stage_id=1) + request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream") + request.resumable = True + request.num_computed_tokens = 0 + adapter.requests_num_chunks_sent["ext-stream"] = 26 + + adapter.save_async( + multimodal_output=None, + request=request, + is_segment_finished=True, + confirmed_num_computed_tokens=26, + ) + + assert len(adapter._pending_save_reqs) == 1 + boundary_task = adapter._pending_save_reqs.popleft() + assert "ext-stream" not in adapter.requests_num_chunks_sent + + request.num_computed_tokens = 3 + adapter.save_async( + multimodal_output=None, + request=request, + is_segment_finished=False, + ) + + assert len(adapter._pending_save_reqs) == 1 + assert adapter.requests_num_chunks_sent["ext-stream"] == 3 + + adapter._send_single_request(boundary_task) + + assert adapter.requests_num_chunks_sent["ext-stream"] == 3 + + +def test_background_send_uses_enqueued_request_snapshot(build_adapter): + """A queued segment must not observe later in-place request mutations.""" + adapter, _ = build_adapter(stage_id=1) + request = Request( + request_id="req-stream", + prompt_token_ids=[1, 2], + sampling_params=SamplingParams(max_tokens=8), + pooling_params=None, + resumable=True, + ) + request.external_req_id = "ext-stream" + ref_audio = torch.tensor([0.1, -0.1]) + request.additional_information = {"codes": {"ref": ref_audio}, "meta": {"segment": "old"}} + request.append_output_token_ids([7]) + seen_requests = [] + + def recording_processor(**kwargs): + queued = kwargs["request"] + seen_requests.append( + ( + queued.additional_information["meta"]["segment"], + list(queued.prompt_token_ids), + list(queued.output_token_ids), + list(queued.all_token_ids), + queued.additional_information["codes"]["ref"] is ref_audio, + ) + ) + return OmniPayloadStruct() + + adapter.custom_process_next_stage_input_func = recording_processor + adapter.save_async( + multimodal_output=None, + request=request, + is_segment_finished=True, + ) + request.additional_information["meta"]["segment"] = "next" + request.prompt_token_ids.append(3) + request.append_output_token_ids([8]) + adapter.save_async( + multimodal_output=None, + request=request, + is_segment_finished=True, + ) + request.additional_information["meta"]["segment"] = "later" + request.prompt_token_ids.append(4) + request.append_output_token_ids([9]) + + assert len(adapter._pending_save_reqs) == 2 + first_task = adapter._pending_save_reqs.popleft() + second_task = adapter._pending_save_reqs.popleft() + assert first_task["request"].additional_information is not request.additional_information + assert first_task["request"].additional_information["meta"] is not request.additional_information["meta"] + + adapter._send_single_request(first_task) + adapter._send_single_request(second_task) + + assert seen_requests == [ + ("old", [1, 2], [7], [1, 2, 7], True), + ("next", [1, 2, 3], [7, 8], [1, 2, 7, 8], True), + ] + + +def test_save_without_custom_processor_does_not_snapshot_request(build_adapter, mocker): + adapter, _ = build_adapter(stage_id=1) + request = _req("req-direct", RequestStatus.RUNNING, external_req_id="ext-direct") + snapshot = mocker.patch.object( + adapter, + "_snapshot_processor_request", + side_effect=AssertionError("unexpected snapshot"), + ) + + adapter.save_async(multimodal_output=None, request=request) + + snapshot.assert_not_called() + assert adapter._pending_save_reqs.popleft()["request"] is request + + def test_send_single_request_terminal_chunk_still_flushes_processor(build_adapter, monkeypatch): """A terminal stop is not a segment boundary (#5383), but the producer-side processor must still receive the flush signal on the terminal chunk. diff --git a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py index 50474a644..e688b9360 100644 --- a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py +++ b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniCPM-o 4.5 Daily-Omni + Seed-TTS accuracy regression coverage. Daily-Omni settings follow the MiniCPM interleaved AV recipe that reaches @@ -21,13 +22,25 @@ from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import ( build_acc_benchmark_cli_argv, find_vllm_cli, ) -from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import SERVER_PARAMS as DUPLEX_TEST_PARAMS +from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import ( + DEPLOY_CONFIG as _DUPLEX_DEPLOY_CONFIG, +) +from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import ( + SERVER_PARAMS as DUPLEX_TEST_PARAMS, +) from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config _MODEL = os.environ.get("VLLM_TEST_MINICPMO_4_5_MODEL", "openbmb/MiniCPM-o-4_5") _DEPLOY_CONFIG = get_deploy_config_path("minicpmo_4_5.yaml") +_DUPLEX_BF16_DEPLOY_CONFIG = modify_stage_config( + _DUPLEX_DEPLOY_CONFIG, + updates={ + "base_config": _DEPLOY_CONFIG, + "connectors.connector_of_shared_memory.extra.code2wav_bfloat16_attention_cache": True, + }, +) _RESULT_DIR = Path( os.environ.get( "ACC_BENCH_RESULT_DIR", @@ -91,6 +104,20 @@ seed_test_params = [ ] +duplex_accuracy_test_params = [ + *DUPLEX_TEST_PARAMS, + pytest.param( + OmniServerParams( + model=_MODEL, + stage_config_path=_DUPLEX_BF16_DEPLOY_CONFIG, + use_stage_cli=False, + server_args=list(_SEED_TTS_SERVER_ARGS), + ), + id="three-stage-single-gpu-bf16-attention-cache", + ), +] + + def _require_vllm_cli() -> None: try: find_vllm_cli() @@ -174,7 +201,7 @@ def test_minicpmo_4_5_seed_tts_wer_bench(omni_server) -> None: @hardware_test(res={"cuda": "H100"}, num_cards=1) -@pytest.mark.parametrize("omni_server", DUPLEX_TEST_PARAMS, indirect=True) +@pytest.mark.parametrize("omni_server", duplex_accuracy_test_params, indirect=True) def test_minicpmo_4_5_duplex_seed_tts_wer_bench(omni_server) -> None: """Gate Seed-TTS WER through the explicit Realtime TTS contract.""" _require_vllm_cli() diff --git a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py index 17e0cb2eb..d5355887d 100644 --- a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py +++ b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Strict MiniCPM-o 4.5 Realtime duplex E2E scenario harness. This script is intentionally scenario-based instead of a generic chat client. @@ -1351,6 +1354,9 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: try: await ws.send(json.dumps(_session_update_event(args))) await _wait_for(state, lambda: state.count("session.created") > 0, timeout_s=20, label="session.created") + start_barrier = getattr(args, "start_barrier", None) + if start_barrier is not None: + await start_barrier.wait() turn_transcripts = _turn_transcripts(args.first_turn_transcript, turns=args.turns) turn_specs = list(zip(turn_transcripts, turn_durations, strict=True)) diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py index 331f0d20f..910a5a180 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Concurrent MiniCPM-o Realtime duplex and resumable-session E2E driver.""" from __future__ import annotations @@ -14,6 +17,7 @@ from types import SimpleNamespace from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit import websockets +from websockets.asyncio.client import ClientConnection from websockets.exceptions import ConnectionClosed SCRIPT_DIR = Path(__file__).resolve().parent @@ -42,6 +46,67 @@ async def run_demo(args): return await _scenario_module().run_demo(args) +class _SynchronizedStartGate: + """One-shot, failure-propagating start gate for Python 3.10+.""" + + def __init__(self, parties: int, *, timeout_s: float): + if parties < 1: + raise ValueError("synchronized start parties must be positive") + if timeout_s <= 0: + raise ValueError("synchronized start timeout must be positive") + self._parties = parties + self._timeout_s = timeout_s + self._arrived = 0 + self._released = False + self._failure: BaseException | None = None + self._event = asyncio.Event() + self._lock = asyncio.Lock() + + async def wait(self) -> None: + async with self._lock: + if self._failure is not None: + raise RuntimeError(f"synchronized start aborted: {self._failure!r}") from self._failure + if self._released: + return + self._arrived += 1 + if self._arrived == self._parties: + self._released = True + self._event.set() + return + + try: + await asyncio.wait_for(self._event.wait(), timeout=self._timeout_s) + except asyncio.TimeoutError as exc: + failure = TimeoutError( + f"synchronized start timed out after {self._timeout_s:.1f}s " + f"({self._arrived}/{self._parties} sessions ready)" + ) + await self.abort(failure) + if self._released and self._failure is None: + return + raise failure from exc + + if self._failure is not None: + raise RuntimeError(f"synchronized start aborted: {self._failure!r}") from self._failure + + async def abort(self, failure: BaseException) -> None: + async with self._lock: + if self._released or self._failure is not None: + return + self._failure = failure + self._event.set() + + +async def _run_demo_with_start_gate(args: SimpleNamespace): + start_gate = getattr(args, "start_barrier", None) + try: + return await run_demo(args) + except BaseException as exc: + if start_gate is not None: + await start_gate.abort(exc) + raise + + def _with_resume_mode(url: str) -> str: parts = urlsplit(url) query = dict(parse_qsl(parts.query, keep_blank_values=True)) @@ -101,14 +166,24 @@ def _validate_semantic_isolation( return True if len(expected_tokens) != len(results): return False - for result, expected_token in zip(results, expected_tokens, strict=True): + normalized_expected_tokens = [token.strip().casefold() for token in expected_tokens] + if any(not token for token in normalized_expected_tokens): + return False + for result_index, (result, expected_token) in enumerate(zip(results, normalized_expected_tokens, strict=True)): details = result.get("transcript_integrity") transcripts = ( [str(item.get("transcript", "")) for item in details if isinstance(item, dict)] if isinstance(details, list) else [] ) - if expected_token not in "".join(transcripts): + joined_transcripts = "".join(transcripts).casefold() + if expected_token not in joined_transcripts: + return False + if any( + other_token in joined_transcripts + for other_index, other_token in enumerate(normalized_expected_tokens) + if other_index != result_index + ): return False return True @@ -130,10 +205,14 @@ async def _receive_until(ws, event_type: str, *, timeout_s: float) -> tuple[dict return await asyncio.wait_for(receive(), timeout=timeout_s) +def _server_event_sequences(events: list[dict[str, object]]) -> list[int]: + return [sequence for event in events if isinstance(sequence := event.get("server_event_seq"), int)] + + async def _open_admission_session( args: argparse.Namespace, session_id: str, -) -> tuple[object, dict[str, object]]: +) -> tuple[ClientConnection, dict[str, object]]: url = _url_with_model( args.url, args.model, @@ -159,7 +238,7 @@ async def _open_admission_session( return ws, created -async def _close_admission_session(ws, *, timeout_s: float) -> None: +async def _close_admission_session(ws: ClientConnection, *, timeout_s: float) -> None: await ws.send(json.dumps({"type": "session.close"})) await _receive_until(ws, "session.closed", timeout_s=timeout_s) await ws.close() @@ -169,8 +248,8 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str, if limit < 1: raise ValueError("admission limit must be positive") prefix = f"admission-{uuid.uuid4().hex}" - accepted: list[tuple[object, dict[str, object]]] = [] - replacement: tuple[object, dict[str, object]] | None = None + accepted: list[tuple[ClientConnection, dict[str, object]]] = [] + replacement: tuple[ClientConnection, dict[str, object]] | None = None overflow_code = None try: for index in range(limit): @@ -274,14 +353,7 @@ async def _resume_probe( generation = created.get("attachment_generation") if not isinstance(token, str) or not isinstance(incarnation, int): raise RuntimeError("session.created omitted resumable credentials") - last_seq = max( - ( - event.get("server_event_seq", 0) - for event in first_events - if isinstance(event.get("server_event_seq"), int) - ), - default=0, - ) + last_seq = max(_server_event_sequences(first_events), default=0) delay_s = args.expire_after_s if expect_expired else args.resume_after_ms / 1000 if delay_s > 0: @@ -324,7 +396,7 @@ async def _resume_probe( await second.send(json.dumps({"type": "session.close"})) closed, close_events = await _receive_until(second, "session.closed", timeout_s=args.timeout_s) - replay_sequences = [event["server_event_seq"] for event in replay if isinstance(event.get("server_event_seq"), int)] + replay_sequences = _server_event_sequences(replay) return { "ok": ( resumed.get("session_id") == session_id @@ -386,14 +458,7 @@ async def _takeover_probe( generation = created.get("attachment_generation") if not isinstance(token, str) or not isinstance(incarnation, int) or not isinstance(generation, int): raise RuntimeError("session.created omitted takeover credentials") - last_seq = max( - ( - event.get("server_event_seq", 0) - for event in first_events - if isinstance(event.get("server_event_seq"), int) - ), - default=0, - ) + last_seq = max(_server_event_sequences(first_events), default=0) second = await websockets.connect(_with_resume_mode(url), max_size=64 * 1024 * 1024) await second.send( @@ -457,7 +522,9 @@ async def _takeover_probe( await first.close() -def _demo_args(args: argparse.Namespace, index: int) -> SimpleNamespace: +def _demo_args( + args: argparse.Namespace, index: int, start_barrier: _SynchronizedStartGate | None = None +) -> SimpleNamespace: validation_mode = "response-required" if args.response_required else "model-policy" input_wav = args.session_input_wav[index] if args.session_input_wav else args.input_wav return SimpleNamespace( @@ -477,14 +544,16 @@ def _demo_args(args: argparse.Namespace, index: int) -> SimpleNamespace: omit_transcript_hints=True, validation_mode=validation_mode, temperature=args.temperature, - scenario="sequential", + scenario=getattr(args, "scenario", "sequential"), + silence_ms=getattr(args, "silence_ms", 500), require_audio=args.response_required, require_distinct_inputs=False, - expect_empty_turn=[], + expect_empty_turn=list(getattr(args, "expect_empty_turn", []) or []), short_ack_ms=350, turns=args.turns, timeout_s=args.timeout_s, model_policy_settle_ms=args.model_policy_settle_ms, + start_barrier=start_barrier, ) @@ -511,8 +580,16 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]: ) lifecycle_result = await run_lifecycle_probes(args) + start_barrier = ( + _SynchronizedStartGate( + args.sessions, + timeout_s=min(float(args.timeout_s), 30.0), + ) + if getattr(args, "synchronized_start", False) + else None + ) session_results = await asyncio.gather( - *(run_demo(_demo_args(args, index)) for index in range(args.sessions)), + *(_run_demo_with_start_gate(_demo_args(args, index, start_barrier)) for index in range(args.sessions)), return_exceptions=True, ) failures = [repr(result) for result in session_results if isinstance(result, BaseException)] @@ -596,7 +673,14 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--turns", type=int, default=1) parser.add_argument("--first-turn-ms", type=int, default=1400) parser.add_argument("--turn-duration-ms", type=int, action="append", default=[]) + parser.add_argument("--expect-empty-turn", type=int, action="append", default=[]) parser.add_argument("--response-required", action="store_true") + parser.add_argument( + "--scenario", + choices=["sequential", "listen-only-overlap"], + default="sequential", + ) + parser.add_argument("--silence-ms", type=int, default=500) parser.add_argument("--temperature", type=float, default=None) parser.add_argument("--disconnect-session-index", type=int) parser.add_argument("--takeover-session-index", type=int) @@ -604,6 +688,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--expire-session-index", type=int) parser.add_argument("--expire-after-s", type=float, default=6.0) parser.add_argument("--verify-admission-limit", type=int) + parser.add_argument( + "--synchronized-start", + action="store_true", + help="Wait until every session is created before starting input.", + ) parser.add_argument("--model-policy-settle-ms", type=int, default=600) parser.add_argument("--timeout-s", type=float, default=120.0) args = parser.parse_args() @@ -615,6 +704,13 @@ def parse_args() -> argparse.Namespace: parser.error("provide exactly one --session-expected-token per session") if args.session_expected_token and not args.session_input_wav: parser.error("--session-expected-token requires --session-input-wav") + normalized_expected_tokens = [token.strip().casefold() for token in args.session_expected_token] + for token_index, token in enumerate(normalized_expected_tokens): + if any( + token in other_token or other_token in token + for other_token in normalized_expected_tokens[token_index + 1 :] + ): + parser.error("--session-expected-token values must not overlap after normalization") return args diff --git a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py index b2e40c58f..5fd65bbcb 100644 --- a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py +++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import base64 import importlib.util @@ -758,6 +761,34 @@ def test_realtime_duplex_multi_session_rejects_cross_session_response_identity() ) +def test_realtime_duplex_multi_session_rejects_overlapping_expected_tokens(monkeypatch, capsys): + demo = _load_multi_demo_module() + monkeypatch.setattr( + demo.sys, + "argv", + [ + "multi.py", + "--sessions", + "2", + "--input-wav", + "fallback.wav", + "--session-input-wav", + "a.wav", + "--session-input-wav", + "b.wav", + "--session-expected-token", + "cat", + "--session-expected-token", + "cater", + ], + ) + + with pytest.raises(SystemExit): + demo.parse_args() + + assert "must not overlap after normalization" in capsys.readouterr().err + + def test_realtime_duplex_multi_session_reads_nested_terminal_identity(): demo = _load_multi_demo_module() @@ -768,6 +799,67 @@ def test_realtime_duplex_multi_session_reads_nested_terminal_identity(): ) +def test_realtime_duplex_synchronized_start_gate_releases_all_participants(): + demo = _load_multi_demo_module() + + async def exercise(): + gate = demo._SynchronizedStartGate(3, timeout_s=0.5) + released = [] + + async def participant(index): + await gate.wait() + released.append(index) + + await asyncio.wait_for( + asyncio.gather(*(participant(index) for index in range(3))), + timeout=1.0, + ) + return released + + assert set(asyncio.run(exercise())) == {0, 1, 2} + + +def test_realtime_duplex_synchronized_start_gate_aborts_waiters_on_failure(monkeypatch): + demo = _load_multi_demo_module() + + async def fake_run_demo(args): + if args.fail: + raise RuntimeError("session creation failed") + await args.start_barrier.wait() + return {"ok": True} + + monkeypatch.setattr(demo, "run_demo", fake_run_demo) + + async def exercise(): + gate = demo._SynchronizedStartGate(2, timeout_s=0.5) + return await asyncio.wait_for( + asyncio.gather( + demo._run_demo_with_start_gate(SimpleNamespace(start_barrier=gate, fail=True)), + demo._run_demo_with_start_gate(SimpleNamespace(start_barrier=gate, fail=False)), + return_exceptions=True, + ), + timeout=1.0, + ) + + failed, aborted = asyncio.run(exercise()) + assert isinstance(failed, RuntimeError) + assert str(failed) == "session creation failed" + assert isinstance(aborted, RuntimeError) + assert "synchronized start aborted" in str(aborted) + assert "session creation failed" in str(aborted) + + +def test_realtime_duplex_synchronized_start_gate_has_bounded_wait(): + demo = _load_multi_demo_module() + + async def exercise(): + gate = demo._SynchronizedStartGate(2, timeout_s=0.01) + await gate.wait() + + with pytest.raises(TimeoutError, match=r"1/2 sessions ready"): + asyncio.run(exercise()) + + def test_realtime_duplex_demo_resolves_distinct_turn_inputs(): demo = _load_demo_module() primary = Path("first.wav") diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py index 9579a9aac..a0e7ed4db 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py @@ -1,3 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import subprocess +import sys +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace @@ -5,8 +11,10 @@ import pytest import torch import torch.nn as nn +import vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav as batched_token2wav_module from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( BatchedToken2Wav, + _token2wav_sdpa_context, plan_token2wav_encode_slices, relpos_encode_token_budget, ) @@ -52,6 +60,9 @@ class _FakeEstimator(nn.Module): self.blocks = [_FakeBlock()] self.cfg_batches: list[int] = [] self.speaker_order: list[list[float]] = [] + self.attention_cache_dtypes: list[torch.dtype | None] = [] + self.register_buffer("att_cache_buffer", torch.ones(1), persistent=False) + self.register_buffer("cnn_cache_buffer", torch.ones(1), persistent=False) def t_embedder(self, time): return time[:, None] @@ -66,7 +77,8 @@ class _FakeEstimator(nn.Module): cnn_out, att_out, ): - del time, mask, cnn_cache, att_cache + del time, mask, cnn_cache + self.attention_cache_dtypes.append(att_cache.dtype if isinstance(att_cache, torch.Tensor) else None) self.cfg_batches.append(inputs.shape[0]) self.speaker_order.append(inputs[:, 2, 0].tolist()) marker = inputs[:, 1, 0] @@ -100,16 +112,17 @@ class _FakeHiFT(nn.Module): super().__init__() self.calls: list[int] = [] - def forward(self, mel, source): + def forward(self, batch, device): + mel = batch["speech_feat"].transpose(1, 2).to(device) + return self.inference(mel, mel.new_zeros((mel.shape[0], 1, 0))) + + def inference(self, mel, source): del source self.calls.append(mel.shape[0]) speech = mel[:, 0].repeat_interleave(3, dim=1) generated_source = speech[:, None] return speech, generated_source - def inference(self, mel, source): - return self.forward(mel, source) - class _FakeToken2Wav: def __init__(self): @@ -140,13 +153,14 @@ class _FakeToken2Wav: raise AssertionError("sequential __call__ fallback must never be called") -def _config(minimum: int = 1): +def _config(minimum: int = 1, initial: int = 0): return SimpleNamespace( model_config=SimpleNamespace( model="/fake/model", stage_connector_config={ "extra": { "code2wav_min_batch_size": minimum, + "code2wav_initial_batch_size": initial, "prompt_cache_id": "shared", "prompt_wav": "/fake/prompt.wav", } @@ -155,14 +169,45 @@ def _config(minimum: int = 1): ) -def _model(): +def _model(initial: int = 0, minimum: int = 1): token2wav = _FakeToken2Wav() backend = BatchedToken2Wav(token2wav) - model = MiniCPMO45Code2Wav(vllm_config=_config()) + _enable_fake_ragged_kernel(backend) + model = MiniCPMO45Code2Wav(vllm_config=_config(minimum=minimum, initial=initial)) model.backend = backend return model, token2wav +def _enable_fake_ragged_kernel(adapter: BatchedToken2Wav) -> None: + """Keep fake-model tests focused on ragged grouping and row mapping.""" + estimator = adapter.flow.decoder.estimator + estimator.in_proj = nn.Identity() + + def fake_ragged_kernel( + estimator, + estimator_input, + time_embedding, + attn_mask, + cnn_cache, + att_cache, + cnn_cache_buffer, + att_cache_buffer, + valid_lengths, + ): + del valid_lengths + return estimator.blocks_forward_chunk( + estimator_input, + time_embedding, + attn_mask, + cnn_cache, + att_cache, + cnn_cache_buffer, + att_cache_buffer, + ) + + adapter._blocks_forward_chunk_ragged = fake_ragged_kernel # type: ignore[method-assign] + + def _info( request_id: str, chunk_seq: int, @@ -256,6 +301,400 @@ def test_decode_batch_splits_overlong_prefill_inside_relpos_budget(): assert audios[0].numel() > 0 +def test_ragged_overlong_prefill_uses_relpos_safe_exact_slices(): + token2wav = _FakeToken2Wav() + token2wav.flow.encoder.pre_lookahead_layer = SimpleNamespace(pre_lookahead_len=3) + adapter = BatchedToken2Wav(token2wav) + _enable_fake_ragged_kernel(adapter) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 2) + setup_encodes = len(token2wav.flow.encoder.calls) + adapter._max_encode_token_frames = lambda _states: 4 + + tokens = [ + torch.arange(10, dtype=torch.long), + torch.arange(4, dtype=torch.long), + ] + last_chunks = [True, False] + audios, next_states = adapter.decode_ragged_batch( + tokens, + prompt, + states, + last_chunks=last_chunks, + ) + + assert token2wav.flow.encoder.calls[setup_encodes:] == [1, 1, 1, 1] + assert token2wav.flow.encoder.last_chunk_calls[setup_encodes:] == [False, False, True, False] + assert len(audios) == len(next_states) == 2 + assert all(audio.numel() > 0 for audio in audios) + + for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): + reference_token2wav = _FakeToken2Wav() + reference_token2wav.flow.encoder.pre_lookahead_layer = SimpleNamespace(pre_lookahead_len=3) + reference_adapter = BatchedToken2Wav(reference_token2wav) + reference_prompt = reference_adapter.prepare_prompt("shared", "/fake/prompt.wav") + reference_states = reference_adapter.setup_batch(reference_prompt, 1) + reference_adapter._max_encode_token_frames = lambda _states: 4 + reference_audio, reference_state = reference_adapter.decode_batch( + row_tokens.unsqueeze(0), + reference_prompt, + reference_states, + last_chunk=last_chunk, + ) + torch.testing.assert_close(audios[row], reference_audio[0]) + for cache_name, expected in reference_state[0].flow_cache.items(): + torch.testing.assert_close( + next_states[row].flow_cache[cache_name], + expected, + ) + for cache_name, expected in reference_state[0].hift_cache.items(): + torch.testing.assert_close( + next_states[row].hift_cache[cache_name], + expected, + ) + + +def test_ragged_diffusion_batches_eight_rows_and_preserves_exact_results(): + token2wav = _FakeToken2Wav() + adapter = BatchedToken2Wav(token2wav) + _enable_fake_ragged_kernel(adapter) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 8) + token2wav.flow.encoder.calls.clear() + token2wav.flow.decoder.estimator.cfg_batches.clear() + token2wav.hift.calls.clear() + tokens = [ + *(torch.tensor([10 + row, 11 + row, 12 + row]) for row in range(6)), + torch.tensor([20, 21]), + torch.tensor([30, 31]), + ] + last_chunks = [False] * 6 + [True, True] + + audios, next_states = adapter.decode_ragged_batch( + tokens, + prompt, + states, + last_chunks=last_chunks, + ) + + assert token2wav.flow.encoder.calls == [6, 2] + assert token2wav.flow.decoder.estimator.cfg_batches == [16, 16] + assert token2wav.hift.calls == [6, 2] + assert len(audios) == len(next_states) == 8 + + for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): + reference_adapter = BatchedToken2Wav(_FakeToken2Wav()) + reference_prompt = reference_adapter.prepare_prompt("shared", "/fake/prompt.wav") + reference_states = reference_adapter.setup_batch(reference_prompt, 1) + reference_audio, reference_state = reference_adapter.decode_batch( + row_tokens.unsqueeze(0), + reference_prompt, + reference_states, + last_chunk=last_chunk, + ) + torch.testing.assert_close(audios[row], reference_audio[0]) + for cache_name, expected in reference_state[0].flow_cache.items(): + torch.testing.assert_close( + next_states[row].flow_cache[cache_name], + expected, + ) + for cache_name, expected in reference_state[0].hift_cache.items(): + torch.testing.assert_close( + next_states[row].hift_cache[cache_name], + expected, + ) + + +def test_npu_patch_preserves_ragged_estimator_contract(): + script = r""" +import torch + +from tests.model_executor.models.minicpmo_4_5.test_code2wav_batching import ( + _FakeToken2Wav, + _enable_fake_ragged_kernel, +) +from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import BatchedToken2Wav +import vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav as npu_patch + + +class _FailGraphRunner: + def run(self, *args, **kwargs): + raise AssertionError("ragged decode must bypass the exact-shape NPU Graph") + + +npu_patch.apply_minicpmo_4_5_code2wav_patch() +adapter = BatchedToken2Wav(_FakeToken2Wav()) +_enable_fake_ragged_kernel(adapter) +prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") +states = adapter.setup_batch(prompt, 2) +npu_patch._backend_graph_runners[adapter] = _FailGraphRunner() +audios, next_states = adapter.decode_ragged_batch( + [ + torch.tensor([10, 11, 12]), + torch.tensor([20, 21]), + ], + prompt, + states, + last_chunks=[False, True], +) +assert len(audios) == len(next_states) == 2 +""" + completed = subprocess.run( + [sys.executable, "-c", script], + cwd=Path(__file__).resolve().parents[4], + capture_output=True, + text=True, + ) + + assert completed.returncode == 0, completed.stderr + + +def test_decode_cfm_enters_platform_sdpa_context(monkeypatch): + entered: list[str] = [] + + @contextmanager + def recording_context(): + entered.append("enter") + yield + entered.append("exit") + + monkeypatch.setattr( + batched_token2wav_module, + "_token2wav_sdpa_context", + lambda _device: recording_context(), + ) + adapter = BatchedToken2Wav(_FakeToken2Wav()) + _enable_fake_ragged_kernel(adapter) + adapter._decode_cfm( + torch.ones((2, 1, 2)), + torch.ones((2, 1)), + torch.zeros((2, 1, 2)), + cnn_cache=None, + att_cache=None, + valid_lengths=[2, 1], + ) + + assert entered == ["enter", "exit"] + + +def test_ragged_decode_bypasses_exact_shape_accelerators(): + adapter = BatchedToken2Wav(_FakeToken2Wav()) + _enable_fake_ragged_kernel(adapter) + + class ForbiddenAccelerator: + def replay(self, *args): + pytest.fail("ragged decode must not use the exact-shape CFM graph") + + def step(self, **kwargs): + pytest.fail("ragged decode must not use the exact-shape TensorRT stepper") + + accelerator = ForbiddenAccelerator() + adapter._cfm_graph_wrapper = accelerator + adapter._trt_stepper = accelerator + adapter._decode_cfm( + torch.ones((2, 1, 2)), + torch.ones((2, 1)), + torch.zeros((2, 1, 2)), + cnn_cache=None, + att_cache=None, + valid_lengths=[2, 1], + ) + + +def test_ragged_decode_fails_when_estimator_has_no_ragged_kernel(): + adapter = BatchedToken2Wav(_FakeToken2Wav()) + + with pytest.raises( + RuntimeError, + match=r'"reason":"ragged_kernel_unavailable"', + ): + adapter._decode_cfm( + torch.ones((2, 1, 2)), + torch.ones((2, 1)), + torch.zeros((2, 1, 2)), + cnn_cache=None, + att_cache=None, + valid_lengths=[2, 1], + ) + + +def test_cfm_graph_receives_bfloat16_cache_in_compute_dtype(): + adapter = BatchedToken2Wav(_FakeToken2Wav(), bfloat16_attention_cache=True) + seen_dtypes: list[torch.dtype] = [] + + class RecordingGraph: + def replay(self, estimator_input, time_embedding, cnn_cache, att_cache, cnn_out, att_out): + del time_embedding, cnn_cache + seen_dtypes.append(att_cache.dtype) + return estimator_input[:, :1], cnn_out, att_out + + adapter._cfm_graph_wrapper = RecordingGraph() + adapter._estimator_step( + adapter.flow.decoder.estimator, + x=torch.ones((2, 1, 2)), + mu=torch.ones((2, 1, 2)), + time=torch.zeros(2), + speakers=torch.ones((2, 1)), + cond=torch.zeros((2, 1, 2)), + cnn_cache=torch.zeros((1, 2, 2, 1)), + att_cache=torch.zeros((1, 2, 1, 1, 2), dtype=torch.bfloat16), + ) + + assert seen_dtypes == [torch.float32] + + +def test_npu_sdpa_context_uses_platform_patch(monkeypatch): + entered: list[str] = [] + + @contextmanager + def recording_context(): + entered.append("enter") + yield + entered.append("exit") + + monkeypatch.setitem( + sys.modules, + "vllm_omni.platforms.npu.models.step_audio2_token2wav", + SimpleNamespace(npu_token2wav_sdpa_context=recording_context), + ) + with _token2wav_sdpa_context(SimpleNamespace(type="npu")): + assert entered == ["enter"] + + assert entered == ["enter", "exit"] + + +def test_ragged_outputs_fail_closed_on_middle_hole(): + adapter = BatchedToken2Wav(_FakeToken2Wav()) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 3) + audios: list[torch.Tensor | None] = [ + torch.zeros(1), + None, + torch.ones(1), + ] + next_states = [states[0], None, states[2]] + + with pytest.raises( + RuntimeError, + match=r'"reason":"incomplete_ragged_output","audio_rows":\[1\],"state_rows":\[1\]', + ): + adapter._require_complete_ragged_outputs(audios, next_states) + + +@pytest.mark.parametrize( + "bfloat16_attention_cache", + [False, True], + ids=["fp32-cache", "bf16-cache"], +) +def test_ragged_dit_two_step_decode_matches_per_row_exact_decode( + bfloat16_attention_cache: bool, +): + from cosyvoice2.flow.decoder_dit import DiT + + torch.manual_seed(17) + template = DiT( + in_channels=4, + out_channels=1, + depth=2, + num_heads=2, + head_dim=2, + hidden_size=4, + ).eval() + template_state = {name: value.detach().clone() for name, value in template.state_dict().items()} + rand_noise = torch.randn(1, 1, 32) + + def make_adapter() -> BatchedToken2Wav: + token2wav = _FakeToken2Wav() + estimator = DiT( + in_channels=4, + out_channels=1, + depth=2, + num_heads=2, + head_dim=2, + hidden_size=4, + ).eval() + estimator.load_state_dict(template_state) + token2wav.flow.decoder.estimator = estimator + token2wav.flow.decoder.rand_noise = rand_noise.clone() + return BatchedToken2Wav( + token2wav, + bfloat16_attention_cache=bfloat16_attention_cache, + ) + + ragged_adapter = make_adapter() + ragged_prompt = ragged_adapter.prepare_prompt("shared", "/fake/prompt.wav") + ragged_states = ragged_adapter.setup_batch(ragged_prompt, 2) + + exact_adapters = [make_adapter(), make_adapter()] + exact_prompts = [adapter.prepare_prompt("shared", "/fake/prompt.wav") for adapter in exact_adapters] + exact_states = [ + adapter.setup_batch(prompt, 1) for adapter, prompt in zip(exact_adapters, exact_prompts, strict=True) + ] + tolerance = (2e-2, 2e-3) if bfloat16_attention_cache else (1e-5, 1e-6) + steps = [ + ( + [torch.tensor([10, 11, 12]), torch.tensor([20, 21, 22])], + [False, True], + ), + ( + [torch.tensor([13, 14]), torch.tensor([22, 23, 24])], + [True, True], + ), + ] + first_step_history_width: int | None = None + + for step, (tokens, last_chunks) in enumerate(steps): + if step == 1: + assert first_step_history_width is not None + assert all( + int(state.flow_cache["estimator_att_cache"].shape[4]) == first_step_history_width + for state in ragged_states + ) + ragged_audios, ragged_states = ragged_adapter.decode_ragged_batch( + tokens, + ragged_prompt, + ragged_states, + last_chunks=last_chunks, + ) + for row, (adapter, prompt) in enumerate(zip(exact_adapters, exact_prompts, strict=True)): + exact_audios, exact_states[row] = adapter.decode_batch( + tokens[row].unsqueeze(0), + prompt, + exact_states[row], + last_chunk=last_chunks[row], + ) + torch.testing.assert_close( + ragged_audios[row], + exact_audios[0], + rtol=tolerance[0], + atol=tolerance[1], + ) + for cache_name, expected in exact_states[row][0].flow_cache.items(): + actual = ragged_states[row].flow_cache[cache_name] + assert actual.dtype == expected.dtype + torch.testing.assert_close( + actual, + expected, + rtol=tolerance[0], + atol=tolerance[1], + ) + for cache_name, expected in exact_states[row][0].hift_cache.items(): + torch.testing.assert_close( + ragged_states[row].hift_cache[cache_name], + expected, + rtol=tolerance[0], + atol=tolerance[1], + ) + + estimator_att = ragged_states[0].flow_cache["estimator_att_cache"] + assert estimator_att.shape[4] > 0 + if step == 0: + first_step_history_width = int(estimator_att.shape[4]) + else: + assert first_step_history_width is not None + assert int(estimator_att.shape[4]) > first_step_history_width + + def test_fade_in_out_limits_overlap_to_available_previous_audio(): speech = torch.arange(6, dtype=torch.float32).reshape(1, -1) previous = torch.full((1, 3), 2.0) @@ -295,6 +734,69 @@ def test_estimator_cache_stack_split_round_trip_preserves_cfg_rows(): ) +def test_bfloat16_estimator_attention_cache_materializes_only_current_timestep(): + token2wav = _FakeToken2Wav() + adapter = BatchedToken2Wav(token2wav, bfloat16_attention_cache=True) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 2) + + assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in states) + assert ( + states[0].flow_cache["estimator_att_cache"].data_ptr() != states[1].flow_cache["estimator_att_cache"].data_ptr() + ) + + stacked = adapter._stack_flow_cache(states) + assert stacked["estimator_att_cache"].dtype == torch.bfloat16 + restored = adapter._split_flow_cache(stacked, 2) + for original, round_tripped in zip(states, restored, strict=True): + assert round_tripped["estimator_att_cache"].dtype == torch.bfloat16 + torch.testing.assert_close( + round_tripped["estimator_att_cache"], + original.flow_cache["estimator_att_cache"], + ) + + _, next_states = adapter.decode_batch( + torch.tensor([[10, 11], [20, 21]]), + prompt, + states, + last_chunk=False, + ) + assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in next_states) + assert token2wav.flow.decoder.estimator.attention_cache_dtypes[-2:] == [torch.float32, torch.float32] + + +def test_bfloat16_estimator_attention_cache_materializes_each_timestep(): + adapter = BatchedToken2Wav(_FakeToken2Wav(), bfloat16_attention_cache=True) + + _, estimator_cnn, estimator_att = adapter._decode_cfm( + torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]), + torch.ones((2, 1)), + torch.zeros((2, 1, 2)), + cnn_cache=None, + att_cache=None, + ) + + assert estimator_cnn.dtype == torch.float32 + assert estimator_att.dtype == torch.bfloat16 + assert estimator_att.shape[:3] == (2, 1, 4) + + +def test_float32_estimator_attention_cache_materializes_each_timestep(): + adapter = BatchedToken2Wav(_FakeToken2Wav()) + + _, estimator_cnn, estimator_att = adapter._decode_cfm( + torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]), + torch.ones((2, 1)), + torch.zeros((2, 1, 2)), + cnn_cache=None, + att_cache=None, + ) + + assert estimator_cnn.dtype == torch.float32 + assert estimator_att.dtype == torch.float32 + assert estimator_att.shape[:3] == (2, 1, 4) + + def test_model_preserves_output_slots_and_prefers_runtime_codes(): model, token2wav = _model() output = _forward( @@ -386,6 +888,28 @@ def test_initial_empty_segment_marker_initializes_stream_without_audio(): assert "duplex" in model._states +def test_initial_empty_segment_markers_respect_initial_batch_limit(): + model, token2wav = _model(initial=2) + names = ["a", "b", "c", "d", "e"] + boundaries = [] + for name in names: + boundary = _info(name, 0, []) + boundary["meta"].update( + { + "code_flat_numel": 0, + "tts_is_last_chunk": True, + "turn_end": False, + } + ) + boundaries.append(boundary) + + output = _forward(model, boundaries) + + assert token2wav.flow.encoder.calls == [2, 2, 1] + assert [audio.numel() for audio in output.multimodal_outputs["model_outputs"]] == [0] * len(names) + assert set(model._states) == set(names) + + def test_shared_runtime_prompt_recreates_missing_file_before_second_owner(tmp_path, monkeypatch): monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path)) model, _ = _model() @@ -647,11 +1171,56 @@ def test_singleton_and_mixed_shape_buckets_use_same_batched_backend_without_fall assert token2wav.hift.calls[-2:] == [1, 1] +def test_initial_batch_limit_allows_later_full_batch(): + model, token2wav = _model(initial=4) + names = [f"request-{index}" for index in range(8)] + + for chunk_seq in (0, 1): + infos = [_info(name, chunk_seq, [1, 2]) for name in names] + _forward(model, infos) + assert token2wav.hift.calls[-2:] == [4, 4] + + token2wav.hift.calls.clear() + _forward(model, [_info(name, 2, [1, 2]) for name in names]) + assert token2wav.hift.calls == [8] + + +def test_initial_batch_partition_preserves_minimum(): + model, token2wav = _model(initial=4, minimum=3) + names = [f"request-{index}" for index in range(6)] + + _forward(model, [_info(name, 0, [1, 2]) for name in names]) + + assert token2wav.hift.calls == [3, 3] + + +def test_initial_batch_partition_rejects_impossible_remainder(): + model, _ = _model(initial=4, minimum=3) + names = [f"request-{index}" for index in range(5)] + + with pytest.raises(RuntimeError, match="initial_batch_partition_below_minimum"): + _forward(model, [_info(name, 0, [1, 2]) for name in names]) + + +@pytest.mark.parametrize("chunk_seqs", [[0, 2, 2, 2], [0, 0, 0, 2]]) +def test_initial_steady_partition_rejects_undersized_wave(chunk_seqs): + model, _ = _model(initial=4, minimum=3) + bucket = [SimpleNamespace(chunk_seq=chunk_seq) for chunk_seq in chunk_seqs] + + with pytest.raises(RuntimeError, match="decode_wave_below_minimum"): + list(model._iter_decode_batches([bucket])) + + def test_backend_failure_does_not_commit_any_request_state(monkeypatch): model, _ = _model() _forward( model, - [_info(name, 0, [index + 1, index + 2]) for index, name in enumerate(("a", "b", "c", "d"))], + [ + _info("a", 0, [1, 2], cache_epoch=0), + _info("b", 0, [2, 3], cache_epoch=0), + _info("c", 0, [3, 4], cache_epoch=1), + _info("d", 0, [4, 5], cache_epoch=1), + ], ) before = dict(model._states) original = model.backend.decode_batch @@ -671,8 +1240,8 @@ def test_backend_failure_does_not_commit_any_request_state(monkeypatch): [ _info("a", 1, [5, 6]), _info("b", 1, [7, 8]), - _info("c", 1, [9, 10, 11]), - _info("d", 1, [12, 13, 14]), + _info("c", 1, [9, 10], cache_epoch=1), + _info("d", 1, [12, 13], cache_epoch=1), ], ) assert call_count == 2 diff --git a/tests/platforms/npu/test_graph_tools.py b/tests/platforms/npu/test_graph_tools.py index 7ec36dfe9..703b8f630 100644 --- a/tests/platforms/npu/test_graph_tools.py +++ b/tests/platforms/npu/test_graph_tools.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 contextlib import contextmanager from types import SimpleNamespace @@ -196,6 +196,7 @@ def test_code2wav_patch_reads_stage_additional_config( model = SimpleNamespace( backend=None, + _extra_config=lambda: {}, vllm_config=SimpleNamespace( additional_config={ "code2wav_enable_npu_graph": enabled, @@ -224,6 +225,47 @@ def test_code2wav_patch_reads_stage_additional_config( assert model.backend not in code2wav_patch._backend_graph_runners +@pytest.mark.parametrize("value", [True, "true", "false", "0", "off"]) +def test_code2wav_patch_rejects_bfloat16_attention_cache(monkeypatch, value): + built = False + + model = SimpleNamespace( + backend=None, + _extra_config=lambda: {"code2wav_bfloat16_attention_cache": value}, + ) + + def build_backend(_instance): + nonlocal built + built = True + + monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend) + + with pytest.raises(ValueError, match="currently supported only on CUDA"): + code2wav_patch._patched_build_backend(model) + + assert not built + + +@pytest.mark.parametrize("value", [False, None, 0, ""]) +def test_code2wav_patch_allows_disabled_bfloat16_attention_cache(monkeypatch, value): + built = False + + model = SimpleNamespace( + backend=None, + _extra_config=lambda: {"code2wav_bfloat16_attention_cache": value}, + ) + + def build_backend(instance): + nonlocal built + built = True + instance.backend = object() + + monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend) + code2wav_patch._patched_build_backend(model) + + assert built + + @pytest.mark.parametrize("with_cache", [False, True]) def test_code2wav_estimator_graph_dispatch_is_platform_owned(monkeypatch, with_cache): calls = [] diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index f24edae38..dd7a75821 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from collections import defaultdict @@ -463,7 +466,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): request.resumable = False stopped = True + confirmed_num_computed_tokens = None if stopped: + if self.chunk_transfer_adapter is not None: + confirmed_num_computed_tokens = self.chunk_transfer_adapter._confirmed_num_computed_tokens(request) if model_runner_output.routed_experts is not None: routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request) @@ -555,6 +561,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): inter_stage_output, request, is_segment_finished, + confirmed_num_computed_tokens=confirmed_num_computed_tokens, ) self._remove_stopped_requests_from_queues( diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index ae19a41a3..d572453fd 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import time 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 9bcc74257..f79192faa 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 @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import copy import importlib import time from collections import defaultdict, deque @@ -10,6 +11,7 @@ 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 @@ -118,6 +120,38 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): num_placeholders = int(getattr(request, "num_output_placeholders", 0) or 0) return max(0, num_computed - num_placeholders) + @staticmethod + def _snapshot_processor_request(request: Request) -> Request: + """Snapshot mutable processor inputs at save-queue admission time.""" + + def snapshot_container(value: Any) -> Any: + if isinstance(value, dict): + return {key: snapshot_container(item) for key, item in value.items()} + if isinstance(value, list): + return [snapshot_container(item) for item in value] + if isinstance(value, tuple): + return tuple(snapshot_container(item) for item in value) + return value + + snapshot = copy.copy(request) + for name in ( + "additional_information", + "prompt_token_ids", + ): + if hasattr(request, name): + setattr(snapshot, name, snapshot_container(getattr(request, name))) + + for private_name, public_name in ( + ("_all_token_ids", "all_token_ids"), + ("_output_token_ids", "output_token_ids"), + ): + token_ids = getattr(request, private_name, None) + if isinstance(token_ids, list): + frozen_ids = token_ids.copy() + setattr(snapshot, private_name, frozen_ids) + setattr(snapshot, public_name, ConstantList(frozen_ids)) + return snapshot + @staticmethod def _refresh_generation_chunk_prefill_state(request: Request) -> None: request.num_prompt_tokens = len(request.prompt_token_ids) @@ -169,6 +203,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): multimodal_output: dict[str, Any] | None = None, request: Request | None = None, is_segment_finished: bool = False, + confirmed_num_computed_tokens: int | None = None, ): """Build and enqueue one chunk for asynchronous sending. @@ -186,10 +221,13 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): multimodal_output: Per-request multimodal output dictionary request: Request object is_segment_finished: whether the segment of request is finished + confirmed_num_computed_tokens: committed token count captured + before a streaming transition can mutate ``request`` """ is_finished = request.is_finished() and not request.resumable - confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request) + if confirmed_num_computed_tokens is None: + confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request) # If the request is preempted, skip the already saved chunks. if confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(request.external_req_id, 0): @@ -204,11 +242,17 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.requests_num_chunks_sent[request.external_req_id] = confirmed_num_computed_tokens task = { "multimodal_output": multimodal_output, - "request": request, + "request": ( + self._snapshot_processor_request(request) if self.custom_process_next_stage_input_func else request + ), "is_finished": is_finished, "is_segment_finished": is_segment_finished, } self._pending_save_reqs.append(task) + if is_segment_finished: + # The queued FIFO item now owns the old segment. Start the next + # segment's deduplication watermark before the worker sends it. + self.requests_num_chunks_sent.pop(request.external_req_id, None) with self._save_cond: self._save_cond.notify() @@ -425,7 +469,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if is_segment_finished: self.code_prompt_token_ids.pop(external_req_id, None) - self.requests_num_chunks_sent.pop(external_req_id, None) self.ramp_chunk_count.pop(external_req_id, None) cached_ic = getattr(self, "_cached_ic", None) if cached_ic is not None: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index f8692f349..2bffeb49e 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.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 """Strict, state-explicit batching for MiniCPM-o 4.5 Token2wav.""" from __future__ import annotations @@ -7,7 +7,7 @@ from __future__ import annotations from collections.abc import Mapping from contextlib import nullcontext from dataclasses import dataclass -from typing import Any +from typing import Any, cast import torch import torch.nn as nn @@ -94,6 +94,17 @@ def _autocast_disabled(device: torch.device): return nullcontext() +def _token2wav_sdpa_context(device: torch.device): + if device.type != "npu": + return nullcontext() + + from vllm_omni.platforms.npu.models.step_audio2_token2wav import ( + npu_token2wav_sdpa_context, + ) + + return npu_token2wav_sdpa_context() + + def tensor_signature(value: torch.Tensor) -> tuple[tuple[int, ...], str, str]: return tuple(value.shape), str(value.dtype), value.device.type @@ -152,6 +163,7 @@ class BatchedToken2Wav(nn.Module): connector_config: Mapping[str, int] | None = None, hift_graph_config: Mapping[str, Any] | None = None, cfm_graph_config: Mapping[str, Any] | None = None, + bfloat16_attention_cache: bool = False, ): super().__init__() self._token2wav = token2wav @@ -185,6 +197,10 @@ class BatchedToken2Wav(nn.Module): del speech, source torch.accelerator.empty_cache() self.float16 = bool(token2wav.float16) + self._estimator_att_compute_dtype = torch.float16 if self.float16 else torch.float32 + self._estimator_att_cache_dtype = ( + torch.bfloat16 if bfloat16_attention_cache else self._estimator_att_compute_dtype + ) self.n_timesteps = int(token2wav.n_timesteps) self.mel_cache_len = int(token2wav.mel_cache_len) self.source_cache_len = int(token2wav.source_cache_len) @@ -310,9 +326,7 @@ class BatchedToken2Wav(nn.Module): return max(1, int(stride)) if stride is not None else 2 def _max_encode_token_frames(self, states: list[BatchedToken2WavState]) -> int: - att = None - if states: - att = self._stack_flow_cache(states).get("conformer_att_cache") + att = states[0].flow_cache.get("conformer_att_cache") if states else None offset1 = int(att.shape[3] // 2) if att is not None else 0 lookahead = self._pre_lookahead_len() or 0 return relpos_encode_token_budget( @@ -382,8 +396,10 @@ class BatchedToken2Wav(nn.Module): cond: torch.Tensor, cnn_cache: torch.Tensor | None, att_cache: torch.Tensor | None, + attn_mask: torch.Tensor | None = None, + valid_lengths: list[int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if self._trt_stepper is not None: + if self._trt_stepper is not None and valid_lengths is None: out, new_cnn, new_att = self._trt_stepper.step( x=x, mu=mu, @@ -399,32 +415,139 @@ class BatchedToken2Wav(nn.Module): speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache) - if self._cfm_graph_wrapper is not None: - if cnn_cache is None: - cnn_cache = torch.zeros_like(cnn_out) - if att_cache is None: - att_cache = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) + old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks) + if isinstance(old_att, torch.Tensor) and old_att.dtype != estimator_input.dtype: + old_att = old_att.to(dtype=estimator_input.dtype) + if self._cfm_graph_wrapper is not None and valid_lengths is None: + graph_cnn = torch.zeros_like(cnn_out) if cnn_cache is None else cnn_cache + graph_att = ( + estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + if att_cache is None + else old_att + ) return self._cfm_graph_wrapper.replay( estimator_input, time_embedding, - cnn_cache, - att_cache, + graph_cnn, + graph_att, + cnn_out, + att_out, + ) + if valid_lengths is not None: + if not hasattr(estimator, "in_proj"): + raise RuntimeError('MiniCPMO45Code2WavBatchError {"reason":"ragged_kernel_unavailable"}') + result = self._blocks_forward_chunk_ragged( + estimator, + estimator_input, + time_embedding, + attn_mask, + old_cnn, + old_att, + cnn_out, + att_out, + valid_lengths, + ) + else: + result = estimator.blocks_forward_chunk( + estimator_input, + time_embedding, + attn_mask, + old_cnn, + old_att, cnn_out, att_out, ) - old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) - old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks) - result = estimator.blocks_forward_chunk( - estimator_input, - time_embedding, - None, - old_cnn, - old_att, - cnn_out, - att_out, - ) return result, cnn_out, att_out + @staticmethod + def _gather_causal_cache( + history: torch.Tensor, + valid_lengths: torch.Tensor, + width: int, + ) -> torch.Tensor: + indices = valid_lengths[:, None] + torch.arange(width, device=history.device)[None, :] + return history.gather( + 2, + indices[:, None, :].expand(-1, int(history.shape[1]), -1), + ) + + def _blocks_forward_chunk_ragged( + self, + estimator: nn.Module, + estimator_input: torch.Tensor, + time_embedding: torch.Tensor, + attn_mask: torch.Tensor | None, + cnn_cache: Any, + att_cache: Any, + cnn_cache_buffer: torch.Tensor, + att_cache_buffer: torch.Tensor, + valid_lengths: list[int], + ) -> torch.Tensor: + """Run one padded DiT batch while capturing exact per-row CNN state.""" + lengths = torch.tensor( + (*valid_lengths, *valid_lengths), + device=estimator_input.device, + dtype=torch.long, + ) + x = estimator.in_proj(estimator_input.transpose(1, 2)) + for block_index, block in enumerate(estimator.blocks): + ( + shift_msa, + scale_msa, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + shift_conv, + scale_conv, + gate_conv, + ) = block.adaLN_modulation(time_embedding).chunk(9, dim=-1) + + normalized = block.norm1(x) * (1 + scale_msa) + shift_msa + x_att, new_att = block.attn.forward_chunk( + normalized, + att_cache[block_index], + attn_mask, + ) + x = x + gate_msa * x_att + + conv_input = block.norm3(x) * (1 + scale_conv) + shift_conv + old_cnn = cnn_cache[block_index] + if old_cnn is None: + width = int(block.conv.block[1].causal_padding[0]) + old_cnn = conv_input.new_zeros( + ( + int(conv_input.shape[0]), + int(block.conv.in_channels + block.conv.out_channels), + width, + ) + ) + old_cnn1, old_cnn2 = old_cnn.split((block.conv.in_channels, block.conv.out_channels), dim=1) + + conv1_input = block.conv.block[0](conv_input) + conv1_output, _ = block.conv.block[1].forward_chunk(conv1_input, old_cnn1) + new_cnn1 = self._gather_causal_cache( + torch.cat((old_cnn1, conv1_input), dim=2), + lengths, + int(old_cnn1.shape[2]), + ) + + conv2_input = block.conv.block[2:6](conv1_output) + conv2_output, _ = block.conv.block[6].forward_chunk(conv2_input, old_cnn2) + new_cnn2 = self._gather_causal_cache( + torch.cat((old_cnn2, conv2_input), dim=2), + lengths, + int(old_cnn2.shape[2]), + ) + x = x + gate_conv * block.conv.block[7](conv2_output) + x = x + gate_mlp * block.mlp(block.norm2(x) * (1 + scale_mlp) + shift_mlp) + + cnn_cache_buffer[block_index].copy_(torch.cat((new_cnn1, new_cnn2), dim=1)) + att_cache_buffer[block_index][:, :, : int(new_att.shape[2]), :].copy_(new_att) + + return estimator.final_layer(x, time_embedding).transpose(1, 2) + def _decode_cfm( self, mu: torch.Tensor, @@ -433,7 +556,12 @@ class BatchedToken2Wav(nn.Module): *, cnn_cache: torch.Tensor | None, att_cache: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + valid_lengths: list[int] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | list[torch.Tensor], + ]: decoder = self.flow.decoder estimator = decoder.estimator batch_size = int(mu.shape[0]) @@ -458,36 +586,114 @@ class BatchedToken2Wav(nn.Module): mu_cfg = torch.cat((mu, torch.zeros_like(mu)), dim=0) speakers_cfg = torch.cat((speakers, torch.zeros_like(speakers)), dim=0) cond_cfg = torch.cat((cond, torch.zeros_like(cond)), dim=0) - next_cnn: list[torch.Tensor] = [] - next_att: list[torch.Tensor] = [] - dt = timeline[1] - timeline[0] - for step in range(self.n_timesteps): - old_cnn = cnn_cache[step] if cnn_cache is not None else None - old_att = att_cache[step] if att_cache is not None else None - estimate, step_cnn, step_att = self._estimator_step( - estimator, - x=torch.cat((x, x), dim=0), - mu=mu_cfg, - time=torch.cat((time, time), dim=0), - speakers=speakers_cfg, - cond=cond_cfg, - cnn_cache=old_cnn, - att_cache=old_att, + attn_mask = None + if valid_lengths is not None: + if len(valid_lengths) != batch_size: + raise ValueError(f"valid length count {len(valid_lengths)} != batch {batch_size}") + cfg_lengths = torch.tensor( + (*valid_lengths, *valid_lengths), + device=mu.device, ) - conditional, unconditional = estimate.split(batch_size, dim=0) - velocity = (1.0 + decoder.inference_cfg_rate) * conditional - decoder.inference_cfg_rate * unconditional - x = x + dt * velocity - time = time + dt - if step + 1 < self.n_timesteps: - dt = timeline[step + 2] - time[0] - next_cnn.append(step_cnn) - next_att.append(step_att) - return x, torch.stack(next_cnn), torch.stack(next_att) + positions = torch.arange(int(mu.shape[2]), device=mu.device) + valid_queries = positions.unsqueeze(0) < cfg_lengths.unsqueeze(1) + current_keys = valid_queries.unsqueeze(1).expand(-1, int(mu.shape[2]), -1) + old_keys = torch.ones( + (2 * batch_size, int(mu.shape[2]), offset), + dtype=torch.bool, + device=mu.device, + ) + attn_mask = valid_queries.unsqueeze(2) & torch.cat((current_keys, old_keys), dim=2) + next_cnn: list[torch.Tensor] = [] + next_att_cache: torch.Tensor | None = None + ragged_att_cache: list[torch.Tensor] | None = None + dt = timeline[1] - timeline[0] + with _token2wav_sdpa_context(mu.device): + for step in range(self.n_timesteps): + old_cnn = cnn_cache[step] if cnn_cache is not None else None + old_att = att_cache[step] if att_cache is not None else None + estimate, step_cnn, step_att = self._estimator_step( + estimator, + x=torch.cat((x, x), dim=0), + mu=mu_cfg, + time=torch.cat((time, time), dim=0), + speakers=speakers_cfg, + cond=cond_cfg, + cnn_cache=old_cnn, + att_cache=old_att, + attn_mask=attn_mask, + valid_lengths=valid_lengths, + ) + conditional, unconditional = estimate.split(batch_size, dim=0) + velocity = (1.0 + decoder.inference_cfg_rate) * conditional - decoder.inference_cfg_rate * unconditional + x = x + dt * velocity + time = time + dt + if step + 1 < self.n_timesteps: + dt = timeline[step + 2] - time[0] + next_cnn.append(step_cnn) + if valid_lengths is not None: + if ragged_att_cache is None: + ragged_att_cache = [ + torch.empty( + ( + self.n_timesteps, + int(step_att.shape[0]), + 2, + int(step_att.shape[2]), + valid_length + offset, + int(step_att.shape[4]), + ), + device=step_att.device, + dtype=self._estimator_att_cache_dtype, + ) + for valid_length in valid_lengths + ] + current_width = int(mu.shape[2]) + for row, (valid_length, row_cache) in enumerate(zip(valid_lengths, ragged_att_cache, strict=True)): + for cfg_row, source_row in enumerate((row, batch_size + row)): + row_cache[step, :, cfg_row, :, :valid_length].copy_( + step_att[:, source_row, :, :valid_length] + ) + if offset: + row_cache[step, :, cfg_row, :, valid_length:].copy_( + step_att[:, source_row, :, current_width:] + ) + del step_att + continue + if next_att_cache is None: + next_att_cache = torch.empty( + (self.n_timesteps, *step_att.shape), + device=step_att.device, + dtype=self._estimator_att_cache_dtype, + ) + next_att_cache[step].copy_(step_att) + del step_att + if ragged_att_cache is not None: + return x, torch.stack(next_cnn), ragged_att_cache + assert next_att_cache is not None + return x, torch.stack(next_cnn), next_att_cache - @staticmethod - def _split_flow_cache(cache: dict[str, torch.Tensor], batch_size: int) -> list[dict[str, torch.Tensor]]: + def _split_flow_cache(self, cache: dict[str, torch.Tensor], batch_size: int) -> list[dict[str, torch.Tensor]]: result: list[dict[str, torch.Tensor]] = [] for row in range(batch_size): + estimator_att = cache["estimator_att_cache"] + if estimator_att.dtype == self._estimator_att_cache_dtype: + request_att = torch.cat( + ( + estimator_att[:, :, row : row + 1], + estimator_att[:, :, batch_size + row : batch_size + row + 1], + ), + dim=2, + ).detach() + else: + request_shape = list(estimator_att.shape) + request_shape[2] = 2 + request_att = torch.empty( + request_shape, + device=estimator_att.device, + dtype=self._estimator_att_cache_dtype, + ) + request_att[:, :, 0:1].copy_(estimator_att[:, :, row : row + 1]) + request_att[:, :, 1:2].copy_(estimator_att[:, :, batch_size + row : batch_size + row + 1]) result.append( { "conformer_cnn_cache": cache["conformer_cnn_cache"][row : row + 1].detach().clone(), @@ -499,29 +705,23 @@ class BatchedToken2Wav(nn.Module): ), dim=2, ).detach(), - "estimator_att_cache": torch.cat( - ( - cache["estimator_att_cache"][:, :, row : row + 1], - cache["estimator_att_cache"][:, :, batch_size + row : batch_size + row + 1], - ), - dim=2, - ).detach(), + "estimator_att_cache": request_att, } ) return result - @staticmethod - def _stack_flow_cache(states: list[BatchedToken2WavState]) -> dict[str, torch.Tensor]: + def _stack_flow_cache(self, states: list[BatchedToken2WavState]) -> dict[str, torch.Tensor]: flows = [state.flow_cache for state in states] conditional_cnn = [flow["estimator_cnn_cache"][:, :, 0:1] for flow in flows] unconditional_cnn = [flow["estimator_cnn_cache"][:, :, 1:2] for flow in flows] conditional_att = [flow["estimator_att_cache"][:, :, 0:1] for flow in flows] unconditional_att = [flow["estimator_att_cache"][:, :, 1:2] for flow in flows] + estimator_att = torch.cat((*conditional_att, *unconditional_att), dim=2) return { "conformer_cnn_cache": torch.cat([flow["conformer_cnn_cache"] for flow in flows], dim=0), "conformer_att_cache": torch.cat([flow["conformer_att_cache"] for flow in flows], dim=1), "estimator_cnn_cache": torch.cat((*conditional_cnn, *unconditional_cnn), dim=2), - "estimator_att_cache": torch.cat((*conditional_att, *unconditional_att), dim=2), + "estimator_att_cache": estimator_att, } def setup_batch( @@ -729,3 +929,188 @@ class BatchedToken2Wav(nn.Module): ] audios = [emitted[row].reshape(-1).to(dtype=torch.float32) for row in range(batch_size)] return audios, next_states + + @staticmethod + def _require_complete_ragged_outputs( + audios: list[torch.Tensor | None], + next_states: list[BatchedToken2WavState | None], + ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]: + missing_audio_rows = [row for row, audio in enumerate(audios) if audio is None] + missing_state_rows = [row for row, state in enumerate(next_states) if state is None] + if missing_audio_rows or missing_state_rows: + raise RuntimeError( + "MiniCPMO45Code2WavBatchError " + f'{{"reason":"incomplete_ragged_output","audio_rows":{missing_audio_rows},' + f'"state_rows":{missing_state_rows}}}' + ) + return ( + cast(list[torch.Tensor], audios), + cast(list[BatchedToken2WavState], next_states), + ) + + def decode_ragged_batch( + self, + tokens: list[torch.Tensor], + features: PromptFeatures, + states: list[BatchedToken2WavState], + *, + last_chunks: list[bool], + ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]: + batch_size = len(tokens) + if batch_size != len(states) or batch_size != len(last_chunks): + raise ValueError("ragged token, state, and final-flag batches must have the same size") + if batch_size == 0: + return [], [] + + lookahead = self._pre_lookahead_len() + for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): + num_frames = int(row_tokens.numel()) + if lookahead is not None and not last_chunk and num_frames <= lookahead: + raise RuntimeError( + "MiniCPMO45Code2WavBatchError " + f'{{"reason":"chunk_below_lookahead_window","row":{row},' + f'"frames":{num_frames},"minimum":{lookahead + 1}}}' + ) + + max_frames = self._max_encode_token_frames(states) + if any(int(row_tokens.numel()) > max_frames for row_tokens in tokens): + # decode_batch owns the RelPos-safe encoder slicing used by + # non-async Talker dumps. Preserve that path for overlong rows; + # ordinary short duplex chunks continue through one padded DiT. + exact_groups: dict[tuple[int, bool], list[int]] = {} + for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): + exact_groups.setdefault((int(row_tokens.numel()), last_chunk), []).append(row) + audios: list[torch.Tensor | None] = [None] * batch_size + next_states: list[BatchedToken2WavState | None] = [None] * batch_size + for (_, last_chunk), rows in exact_groups.items(): + group_audio, group_states = self.decode_batch( + torch.stack([tokens[row] for row in rows], dim=0), + features, + [states[row] for row in rows], + last_chunk=last_chunk, + ) + for group_row, row in enumerate(rows): + audios[row] = group_audio[group_row] + next_states[row] = group_states[group_row] + return self._require_complete_ragged_outputs(audios, next_states) + encoder_groups: dict[tuple[int, bool], list[int]] = {} + for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): + encoder_groups.setdefault((int(row_tokens.numel()), last_chunk), []).append(row) + + hidden_rows: list[torch.Tensor | None] = [None] * batch_size + conformer_cnn_rows: list[torch.Tensor | None] = [None] * batch_size + conformer_att_rows: list[torch.Tensor | None] = [None] * batch_size + for (_, last_chunk), rows in encoder_groups.items(): + group_states = [states[row] for row in rows] + group_cache = self._stack_flow_cache(group_states) + group_tokens = torch.stack([tokens[row] for row in rows], dim=0) + with self._autocast(group_tokens.device): + hidden, conformer_cnn, conformer_att = self._encode_chunk( + group_tokens, + last_chunk=last_chunk, + cnn_cache=group_cache["conformer_cnn_cache"], + att_cache=group_cache["conformer_att_cache"], + ) + for group_row, row in enumerate(rows): + hidden_rows[row] = hidden[group_row : group_row + 1] + conformer_cnn_rows[row] = conformer_cnn[group_row : group_row + 1] + conformer_att_rows[row] = conformer_att[:, group_row : group_row + 1] + + missing_hidden_rows = [row for row, value in enumerate(hidden_rows) if value is None] + if missing_hidden_rows: + raise RuntimeError( + f'MiniCPMO45Code2WavBatchError {{"reason":"incomplete_encoder_output","rows":{missing_hidden_rows}}}' + ) + resolved_hidden = cast(list[torch.Tensor], hidden_rows) + hidden_lengths = [int(value.shape[1]) for value in resolved_hidden] + max_hidden_length = max(hidden_lengths) + padded_hidden = resolved_hidden[0].new_zeros((batch_size, max_hidden_length, int(resolved_hidden[0].shape[2]))) + for row, hidden in enumerate(resolved_hidden): + padded_hidden[row, : int(hidden.shape[1])].copy_(hidden[0]) + + flow_cache = self._stack_flow_cache(states) + speakers = features.speaker_embedding.expand(batch_size, -1) + with self._autocast(padded_hidden.device): + projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + cond = torch.zeros_like(padded_hidden).transpose(1, 2).contiguous() + chunk_mel, estimator_cnn, estimator_att = self._decode_cfm( + padded_hidden.transpose(1, 2).contiguous(), + projected_speakers, + cond, + cnn_cache=flow_cache["estimator_cnn_cache"], + att_cache=flow_cache["estimator_att_cache"], + valid_lengths=hidden_lengths, + ) + + prompt_len = int(features.mels.shape[1]) + assert isinstance(estimator_att, list) + new_flow: list[dict[str, torch.Tensor]] = [] + for row, row_estimator_att in enumerate(estimator_att): + conformer_cnn = conformer_cnn_rows[row] + conformer_att = conformer_att_rows[row] + assert conformer_cnn is not None and conformer_att is not None + if row_estimator_att.shape[4] > prompt_len + 100: + row_estimator_att = torch.cat( + ( + row_estimator_att[..., :prompt_len, :], + row_estimator_att[..., -100:, :], + ), + dim=4, + ) + if conformer_att.shape[3] > prompt_len + 100: + conformer_att = torch.cat( + ( + conformer_att[..., :prompt_len, :], + conformer_att[..., -100:, :], + ), + dim=3, + ) + new_flow.append( + { + "conformer_cnn_cache": conformer_cnn.detach().clone(), + "conformer_att_cache": conformer_att.detach().clone(), + "estimator_cnn_cache": torch.cat( + ( + estimator_cnn[:, :, row : row + 1], + estimator_cnn[:, :, batch_size + row : batch_size + row + 1], + ), + dim=2, + ).detach(), + "estimator_att_cache": row_estimator_att, + } + ) + + audios: list[torch.Tensor | None] = [None] * batch_size + next_states: list[BatchedToken2WavState | None] = [None] * batch_size + vocoder_groups: dict[tuple[int, bool], list[int]] = {} + for row, last_chunk in enumerate(last_chunks): + vocoder_groups.setdefault((hidden_lengths[row], last_chunk), []).append(row) + for (hidden_length, last_chunk), rows in vocoder_groups.items(): + old_mel = torch.cat([states[row].hift_cache["mel"] for row in rows], dim=0) + old_source = torch.cat([states[row].hift_cache["source"] for row in rows], dim=0) + old_speech = torch.cat([states[row].hift_cache["speech"] for row in rows], dim=0) + group_mel = torch.cat( + [chunk_mel[row : row + 1, :, :hidden_length] for row in rows], + dim=0, + ) + mel = torch.cat((old_mel, group_mel), dim=2) + speech, source = self._hift_inference(mel, old_source) + if old_speech.shape[-1] > 0: + window = self.speech_window.to(device=speech.device, dtype=speech.dtype) + speech = self._fade_in_out(speech, old_speech, window) + next_hift = { + "mel": mel[..., -self.mel_cache_len :].detach(), + "source": source[..., -self.source_cache_len :].detach(), + "speech": speech[..., -self.source_cache_len :].detach(), + } + emitted = speech if last_chunk else speech[..., : -self.source_cache_len] + for group_row, row in enumerate(rows): + audios[row] = emitted[group_row].reshape(-1).to(dtype=torch.float32) + next_states[row] = BatchedToken2WavState( + flow_cache=new_flow[row], + hift_cache={ + name: value[group_row : group_row + 1].detach().clone() for name, value in next_hift.items() + }, + ) + + return self._require_complete_ragged_outputs(audios, next_states) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 4ebe0b28c..dfd207e7c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.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 """Strict batched codec-to-waveform stage for MiniCPM-o 4.5.""" from __future__ import annotations @@ -135,7 +135,7 @@ class _WorkItem: class MiniCPMO45Code2Wav(nn.Module): - """LLM_GENERATION model that admits only true exact-shape GPU batches.""" + """LLM_GENERATION model with request-owned state and compatible batching.""" input_modalities = "audio" have_multimodal_outputs = True @@ -184,6 +184,11 @@ class MiniCPMO45Code2Wav(nn.Module): self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1)) if self._min_batch_size < 1: raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1") + self._initial_batch_size = int(extra.get("code2wav_initial_batch_size", 0)) + if self._initial_batch_size < 0: + raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be >= 0") + if self._initial_batch_size and self._initial_batch_size < self._min_batch_size: + raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be 0 or >= code2wav_min_batch_size") self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) self._prompt_wav_override = extra.get("prompt_wav") @@ -482,13 +487,63 @@ class MiniCPMO45Code2Wav(nn.Module): return ( item.prompt_cache_id, item.prompt_wav, - int(item.tokens.numel()), cache_signature, - item.last_chunk, - item.tts_is_last_chunk, item.cache_epoch, ) + def _iter_limited_batches( + self, + buckets: Iterable[list[_WorkItem]], + *, + maximum: int, + ) -> Iterable[list[_WorkItem]]: + for bucket in buckets: + if not bucket: + continue + if not maximum: + yield bucket + continue + batch_count = (len(bucket) + maximum - 1) // maximum + if len(bucket) < batch_count * self._min_batch_size: + raise _batch_error( + "initial_batch_partition_below_minimum", + size=len(bucket), + minimum=self._min_batch_size, + maximum=maximum, + ) + batch_size, larger_batches = divmod(len(bucket), batch_count) + start = 0 + for batch_index in range(batch_count): + stop = start + batch_size + (batch_index < larger_batches) + yield bucket[start:stop] + start = stop + + def _iter_decode_batches( + self, + buckets: Iterable[list[_WorkItem]], + ) -> Iterable[list[_WorkItem]]: + for bucket in buckets: + if not self._initial_batch_size: + yield bucket + continue + initial = [item for item in bucket if item.chunk_seq <= 1] + steady = [item for item in bucket if item.chunk_seq > 1] + + undersized = [ + {"wave": wave, "size": len(items)} + for wave, items in (("initial", initial), ("steady", steady)) + if items and len(items) < self._min_batch_size + ] + if undersized: + raise _batch_error( + "decode_wave_below_minimum", + minimum=self._min_batch_size, + waves=undersized, + ) + yield from self._iter_limited_batches([initial], maximum=self._initial_batch_size) + if steady: + yield steady + @torch.inference_mode() def forward( self, @@ -624,7 +679,10 @@ class MiniCPMO45Code2Wav(nn.Module): (item.prompt_cache_id, item.prompt_wav), [], ).append(item) - for bucket in initial_marker_buckets.values(): + for bucket in self._iter_limited_batches( + initial_marker_buckets.values(), + maximum=self._initial_batch_size, + ): try: features = self.backend.prepare_prompt( bucket[0].prompt_cache_id, @@ -656,7 +714,7 @@ class MiniCPMO45Code2Wav(nn.Module): prompt_wav=item.prompt_wav, token2wav=state, ) - for bucket in buckets.values(): + for bucket in self._iter_decode_batches(buckets.values()): batch_size = len(bucket) try: features = self.backend.prepare_prompt( @@ -667,13 +725,23 @@ class MiniCPMO45Code2Wav(nn.Module): states = self.backend.setup_batch(features, batch_size) else: states = [item.previous.token2wav for item in bucket if item.previous is not None] - tokens = torch.stack([item.tokens for item in bucket], dim=0) - audios, next_states = self.backend.decode_batch( - tokens, - features, - states, - last_chunk=bucket[0].last_chunk, - ) + token_lengths = {int(item.tokens.numel()) for item in bucket} + last_chunk_values = {item.last_chunk for item in bucket} + if len(token_lengths) > 1 or len(last_chunk_values) > 1: + audios, next_states = self.backend.decode_ragged_batch( + [item.tokens for item in bucket], + features, + states, + last_chunks=[item.last_chunk for item in bucket], + ) + else: + tokens = torch.stack([item.tokens for item in bucket], dim=0) + audios, next_states = self.backend.decode_batch( + tokens, + features, + states, + last_chunk=bucket[0].last_chunk, + ) except Exception as exc: self._prune_unowned_runtime_prompts() if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "): @@ -822,4 +890,5 @@ class MiniCPMO45Code2Wav(nn.Module): connector_config=self._connector_config, hift_graph_config=self._hift_graph_config, cfm_graph_config=self._cfm_graph_config, + bfloat16_attention_cache=bool(extra.get("code2wav_bfloat16_attention_cache", False)), ) diff --git a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py index c6effe6f8..0d34c93a5 100644 --- a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py +++ b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.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 """Inject MiniCPM-o Code2Wav NPUGraph acceleration on Ascend.""" from __future__ import annotations @@ -7,6 +7,7 @@ from __future__ import annotations import os from collections.abc import Mapping from contextlib import nullcontext +from typing import cast from weakref import WeakKeyDictionary import torch @@ -24,6 +25,7 @@ _original_decode_batch = None _backend_graph_runners: WeakKeyDictionary[object, NPUExactGraphRunner] = WeakKeyDictionary() _ENABLE_KEY = "code2wav_enable_npu_graph" _MAX_GRAPHS_KEY = "code2wav_max_npu_graphs" +_BF16_ATTENTION_CACHE_KEY = "code2wav_bfloat16_attention_cache" def _config_bool(value: object, default: bool) -> bool: @@ -104,10 +106,18 @@ def _patched_estimator_step( cond, cnn_cache, att_cache, + attn_mask=None, + valid_lengths=None, ): assert _original_estimator_step is not None graph_runner = _backend_graph_runners.get(self) - if graph_runner is None or self._trt_stepper is not None or self._cfm_graph_wrapper is not None: + if ( + graph_runner is None + or self._trt_stepper is not None + or self._cfm_graph_wrapper is not None + or attn_mask is not None + or valid_lengths is not None + ): return _original_estimator_step( self, estimator, @@ -118,6 +128,8 @@ def _patched_estimator_step( cond=cond, cnn_cache=cnn_cache, att_cache=att_cache, + attn_mask=attn_mask, + valid_lengths=valid_lengths, ) if (cnn_cache is None) != (att_cache is None): raise ValueError("estimator CNN and attention caches must both be present or absent") @@ -198,8 +210,15 @@ def _patched_build_backend(self) -> None: if self.backend is not None: return + extra = self._extra_config() + if bool(extra.get(_BF16_ATTENTION_CACHE_KEY, False)): + raise ValueError( + "MiniCPM-o Code2Wav code2wav_bfloat16_attention_cache is " + "currently supported only on CUDA; leave it unset or false on NPU." + ) + config = _graph_config(self) - max_graphs = max(0, int(config.get(_MAX_GRAPHS_KEY, 32))) + max_graphs = max(0, int(cast(int | str, config.get(_MAX_GRAPHS_KEY, 32)))) graph_enabled = max_graphs > 0 and _config_bool(config.get(_ENABLE_KEY), False) if graph_enabled: # NPUOmniPlatform enables internal format for quantized LLM kernels.