diff --git a/docs/design/index.md b/docs/design/index.md index 99646dc23..748e1d9ca 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -115,3 +115,5 @@ The design contracts separate selection mechanics from backend algorithms: The pre-#5137 pages are preserved in the [legacy module archive](module/archive/README.md) for historical reference and are not active design contracts. + +- [MiniCPM-o 4.5 turn-mode MRv2 performance](minicpm_o45_mrv2_performance.md) diff --git a/docs/design/minicpm_o45_mrv2_performance.md b/docs/design/minicpm_o45_mrv2_performance.md new file mode 100644 index 000000000..64a8c6abf --- /dev/null +++ b/docs/design/minicpm_o45_mrv2_performance.md @@ -0,0 +1,22 @@ +# MiniCPM-o 4.5 MRv2 profiles + +Turn mode uses the stage-level MRv2 contracts from #8184. Thinker emits live +latent metadata outside graph replay; Talker keeps codec history and EOS +control on device. Code2Wav reuses mainline Whole-Euler Flow graphs and shared +prompt state. The previous experimental block compilation, tiled attention, +channels-last and merged-CFM implementations are removed. + +`minicpmo_4_5_turn_mrv2_h200.yaml` selects Talker capacity 16 and 4 GiB KV; +the generic MRv2 profile selects capacity 8 and 2 GiB KV. These configurations +require new end-to-end performance measurements after the mainline integration. +Previously reported numbers do not describe this revised codec backend. + +The generic MRv2 profile and the opt-in native V1 duplex H200 profile enable +TF32 matmuls within Code2Wav forward/capture only; the previous process policy +is restored afterwards. cuDNN's TF32 policy is independent. TF32 changes rounding. +Duplex keeps the mainline V1 session path. Turn results do not establish duplex +performance or interruption correctness. + +The shared asynchronous output snapshot and batched Talker preprocessing also +apply to V1 when async chunking/scheduling is enabled. Default V1 behavior must +therefore be included in end-to-end regression validation. diff --git a/pyproject.toml b/pyproject.toml index 3f7c0b243..bb2f1e0e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -325,6 +325,8 @@ skip_covered = true [tool.typos.default] extend-ignore-identifiers-re = [ + # Canonical PyTorch CUDA device property. + "^shared_memory_per_block_optin$", ".*_thw", ".*thw", "ein", diff --git a/tests/engine/duplex/test_duplex_deploy_config.py b/tests/engine/duplex/test_duplex_deploy_config.py index f7af09d2b..8544c7d90 100644 --- a/tests/engine/duplex/test_duplex_deploy_config.py +++ b/tests/engine/duplex/test_duplex_deploy_config.py @@ -97,6 +97,7 @@ def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str, ("deploy_yaml", "expected"), [ ("minicpmo_4_5.yaml", 16), + ("minicpmo_4_5_duplex_h200.yaml", 16), ("minicpmo_4_5_8x4090.yaml", 1), ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 16), ], diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_stage0_runtime_load.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_stage0_runtime_load.py index 364e56f8f..ce8152293 100644 --- a/tests/model_executor/models/minicpmo_4_5/duplex/test_stage0_runtime_load.py +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_stage0_runtime_load.py @@ -19,6 +19,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _StageModel(torch.nn.Module): + logits_vocab_size = 6562 + def make_empty_intermediate_tensors(self): return None diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 8283bf635..d62abd60f 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -1969,7 +1969,10 @@ def test_whole_euler_query_bucket_keeps_current_first_cache_layout(monkeypatch: (_whole_euler_chunk(1, 8), 6, 2), # narrow padded chunk: capture pads 8 -> 16 (_whole_euler_chunk(1, 16), 16, 0), # consumes the narrow chunk's cache ] - caches = {"exact": (None, None), "bucketed": (None, None)} + caches: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = { + "exact": (None, None), + "bucketed": (None, None), + } for chunk, mel_frames, pad_frames in chunks: chunk["x"][:, :, mel_frames:] = 0.0 results = {} @@ -2319,3 +2322,24 @@ def test_whole_euler_groups_are_acquired_again_after_a_flush(monkeypatch: pytest assert calls == [16, 1, 16, 1] flush_always = True assert wrapper._group_entries(groups, fills) is None + + +@pytest.mark.parametrize("fail", [False, True]) +def test_vocoder_restores_tf32_policy(fail): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import MiniCPMO45Code2Wav + + previous = torch.backends.cuda.matmul.allow_tf32 + + def forward(*args, **kwargs): + assert torch.backends.cuda.matmul.allow_tf32 + if fail: + raise RuntimeError("injected") + return "ok" + + model = SimpleNamespace(_extra_config=lambda: {"token2wav_allow_tf32": True}, _forward_impl=forward) + if fail: + with pytest.raises(RuntimeError, match="injected"): + MiniCPMO45Code2Wav.forward(model) + else: + assert MiniCPMO45Code2Wav.forward(model) == "ok" + assert torch.backends.cuda.matmul.allow_tf32 == previous diff --git a/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py index b3baf871d..5ef8c5561 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py +++ b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py @@ -20,6 +20,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_multimodal_runner_preserves_forwarded_token_identities(monkeypatch): stage_model = torch.nn.Module() + stage_model.logits_vocab_size = 6562 stage_model.make_empty_intermediate_tensors = lambda: None monkeypatch.setattr( "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni.init_vllm_registered_model", diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_host_sync.py b/tests/model_executor/models/minicpmo_4_5/test_talker_host_sync.py new file mode 100644 index 000000000..f1514dddd --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_host_sync.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniCPM-o 4.5 Talker per-step path: batched, host-sync-free, same results.""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + ConditionalChatTTSConfig, +) +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_VOCAB = 8 +_EOS = 7 + + +@dataclass +class _SamplingMetadata: + repetition_penalties: torch.Tensor + prompt_token_ids: torch.Tensor | None = None + no_penalties: bool = False + + +def _make_talker(device: str = "cpu") -> MiniCPMO45OmniTTSForConditionalGeneration: + talker = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + nn.Module.__init__(talker) + talker._num_audio_tokens = _VOCAB + talker._codec_eos_id = _EOS + talker._force_eos_rows = None + talker._mask_eos_rows = None + talker._pending_force_eos_rows = None + talker._penalty_histories = None + talker._request_audio_states = {} + talker._request_condition_states = {} + talker._deferred_cleanup_ids = set() + talker._tts_config = ConditionalChatTTSConfig() + generator = torch.Generator().manual_seed(0) + talker.emb_code = nn.ModuleList([nn.Embedding(_VOCAB, 4)]) + talker.head_code = nn.ModuleList([nn.Linear(4, _VOCAB, bias=False)]) + with torch.no_grad(): + talker.emb_code[0].weight.copy_(torch.randn(_VOCAB, 4, generator=generator)) + talker.head_code[0].weight.copy_(torch.randn(_VOCAB, 4, generator=generator)) + return talker.to(device) + + +def _states() -> dict[str, dict[str, Any]]: + return { + # Mid-stream row: forwards its code, EOS still masked by min_tokens. + "req-live": {"finished": False, "step": 2, "max_tokens": 100, "min_tokens": 5, "recent_codes": [1, 2]}, + # The previous step sampled codec EOS: forward nothing, finish, force EOS. + "req-eos": {"finished": False, "step": 60, "max_tokens": 100, "recent_codes": [4, 4, 5]}, + # Leftover decode of an already-finished request. + "req-done": {"finished": True, "step": 3, "max_tokens": 100}, + # Reaches its codec budget on this step. + "req-cap": {"finished": False, "step": 8, "max_tokens": 10, "recent_codes": list(range(7)) * 3}, + } + + +def _infos(states: dict[str, dict[str, Any]]) -> list[dict[str, Any]]: + return [ + {"request_id": request_id, "audio_state": dict(state), "_omni_is_prefill": False} + for request_id, state in states.items() + ] + + +def _step(talker, infos, hidden, mocker, *, batched: bool, input_ids: torch.Tensor, penalties: torch.Tensor): + if batched: + returned_ids, embeds, updates = talker.preprocess_decode_batch(input_ids=input_ids, req_infos=infos) + assert returned_ids is input_ids + else: + rows = [talker.preprocess(input_ids[row : row + 1], None, **info) for row, info in enumerate(infos)] + embeds = torch.cat([row[1] for row in rows]) + updates = [row[2] for row in rows] + for info, update in zip(infos, updates, strict=True): + info["codes"] = update["codes"] + output = talker.make_omni_output( + hidden, + model_intermediate_buffer=infos, + request_token_spans=[(row, row + 1) for row in range(len(infos))], + ) + logits = talker.compute_logits(output.text_hidden_states) + captured: dict[str, torch.Tensor] = {} + + def _sampler(logits, sampling_metadata): + captured["logits"] = logits.clone() + return SimpleNamespace(sampled_token_ids=logits.argmax(dim=-1, keepdim=True)) + + mocker.patch( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler", + return_value=_sampler, + ) + sampled = talker.sample(logits, _SamplingMetadata(repetition_penalties=penalties)) + return embeds, output, logits, captured["logits"], sampled.sampled_token_ids + + +def test_batched_decode_matches_scalar_preprocess(mocker) -> None: + input_ids = torch.tensor([3, _EOS, _EOS, 6], dtype=torch.int32) + hidden = torch.randn(4, 4, generator=torch.Generator().manual_seed(1)) + penalties = torch.tensor([1.05, 1.2, 1.05, 1.0]) + results = {} + talkers = {} + for batched in (False, True): + talker = _make_talker() + states = _states() + talker._request_audio_states = copy.deepcopy(states) + results[batched] = _step( + talker, _infos(states), hidden, mocker, batched=batched, input_ids=input_ids, penalties=penalties + ) + talkers[batched] = talker + + # Embeddings, masked logits, penalized logits and sampled ids. + for index in (0, 2, 3, 4): + assert torch.equal(results[False][index], results[True][index]) + for result in results.values(): + output = result[1].multimodal_outputs + assert [t.tolist() for t in output["codes"]["audio"]] == [[[3]], [], [], [[6]]] + assert [t.item() for t in output["meta"]["finished"]] == [False, True, True, True] + assert talkers[True]._request_audio_states == talkers[False]._request_audio_states + assert talkers[True]._decode_codec_ids == {} diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_host_sync_cuda.py b/tests/model_executor/models/minicpmo_4_5/test_talker_host_sync_cuda.py new file mode 100644 index 000000000..64e70c343 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_host_sync_cuda.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The MiniCPM-o 4.5 Talker's per-step path must not block the host on the GPU.""" + +from __future__ import annotations + +import copy + +import pytest +import torch +import torch.nn as nn + +from tests.model_executor.models.minicpmo_4_5.test_talker_host_sync import _EOS, _infos, _make_talker, _states, _step +from vllm_omni.utils.device_copy import index_to_device + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +def test_talker_step_does_not_synchronize_cuda(mocker) -> None: + talker = _make_talker("cuda") + states = _states() + talker._request_audio_states = copy.deepcopy(states) + infos = _infos(states) + input_ids = torch.tensor([3, _EOS, _EOS, 6], dtype=torch.int32, device="cuda") + hidden = torch.randn(4, 4, device="cuda") + penalties = torch.tensor([1.05, 1.2, 1.05, 1.0], device="cuda") + # Warm the pinned host allocator and CUDA context outside the check. + index_to_device([1], "cuda") + previous = torch.cuda.get_sync_debug_mode() + try: + torch.cuda.set_sync_debug_mode("error") + with torch.inference_mode(): + *_, sampled = _step(talker, infos, hidden, mocker, batched=True, input_ids=input_ids, penalties=penalties) + finally: + torch.cuda.set_sync_debug_mode(previous) + assert sampled.reshape(-1).tolist()[1:] == [_EOS, _EOS, _EOS] + assert talker._request_audio_states["req-live"]["recent_codes"] == [1, 2, 3] + assert talker._request_audio_states["req-eos"]["finished"] is True + + +def test_talker_condition_upload_does_not_synchronize_cuda() -> None: + talker = _make_talker("cuda") + talker.emb_text = nn.Embedding(16, 4).cuda() + talker.projector_semantic = nn.Linear(6, 4).cuda() + talker._normalize = True + talker._text_eos_id = 14 + talker._tts_bos_id = 15 + token_ids = torch.tensor([2, 3, 4]) + hidden_states = torch.randn(3, 6) + expected = talker._build_condition_embeddings(token_ids.cuda(), hidden_states.cuda()) + previous = torch.cuda.get_sync_debug_mode() + try: + torch.cuda.set_sync_debug_mode("error") + with torch.inference_mode(): + condition = talker._build_condition_embeddings(token_ids, hidden_states) + boundary = talker._build_condition_embeddings(token_ids[:0], hidden_states[:0]) + finally: + torch.cuda.set_sync_debug_mode(previous) + assert torch.equal(condition, expected) + assert boundary.shape == (2, 4) + + +@torch.inference_mode() +def test_v1_batched_decode_and_async_snapshot_match_scalar(mocker): + from vllm_omni.worker.gpu_ar_model_runner import _snapshot_tensor_payload_to_cpu_async + + ids = torch.tensor([3, _EOS, _EOS, 6], dtype=torch.int32, device="cuda") + hidden = torch.arange(16, dtype=torch.float32, device="cuda").reshape(4, 4) + penalties = torch.tensor([1.05, 1.2, 1.05, 1.0], device="cuda") + results, states = [], [] + for batched in (False, True): + talker = _make_talker("cuda") + talker._request_audio_states = copy.deepcopy(_states()) + results.append( + _step(talker, _infos(_states()), hidden, mocker, batched=batched, input_ids=ids, penalties=penalties) + ) + states.append(copy.deepcopy(talker._request_audio_states)) + for index in (0, 2, 3, 4): + torch.testing.assert_close(results[0][index], results[1][index], rtol=0, atol=0) + assert states[0] == states[1] + reference, candidate = [result[1].multimodal_outputs for result in results] + for group in ("codes", "meta"): + for key in reference[group]: + for left, right in zip(reference[group][key], candidate[group][key], strict=True): + torch.testing.assert_close(left, right, rtol=0, atol=0) + snapshot = _snapshot_tensor_payload_to_cpu_async( + {"hidden_states": hidden, "multimodal_outputs": candidate}, + copy_stream=torch.cuda.Stream(), + pin_memory=True, + ) + expected_hidden = hidden.cpu().clone() + snapshot.wait() + torch.testing.assert_close(snapshot.payload["hidden_states"], expected_hidden, rtol=0, atol=0) + for group in ("codes", "meta"): + for key in reference[group]: + for left, right in zip( + reference[group][key], snapshot.payload["multimodal_outputs"][group][key], strict=True + ): + torch.testing.assert_close(left.cpu(), right, rtol=0, atol=0) + hidden.fill_(-1) + torch.testing.assert_close(snapshot.payload["hidden_states"], expected_hidden, rtol=0, atol=0) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_mrv2.py b/tests/model_executor/models/minicpmo_4_5/test_talker_mrv2.py new file mode 100644 index 000000000..2cf8e03d5 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_mrv2.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Model Runner V2 path of MiniCPM-o 4.5's Talker: device-side codec output, +EOS control and the 16-frame codec penalty must match the V1 host path.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import cast + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + _OFFLINE_CODEC_MAX_NEW_TOKENS, + MiniCPMO45OmniTTSForConditionalGeneration, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_EOS = 6561 + + +def _talker(max_reqs: int = 8, max_position_embeddings: int = 4096): + talker = object.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + talker._codec_eos_id = _EOS + talker._tts_config = SimpleNamespace(max_position_embeddings=max_position_embeddings) + talker._mrv2_empty_speech = torch.zeros(max_reqs, dtype=torch.bool) + talker._mrv2_forced_eos = None + talker._mrv2_decode_rows_logged = False + return talker + + +def _batch(rows: list[dict], *, pad_to: int | None = None): + """rows: one dict per request with slot, prompt_len, span (token ids) and prefill flag.""" + ids = [token for row in rows for token in row["span"]] + num_tokens = len(ids) + padded = pad_to or num_tokens + starts = np.cumsum([0] + [len(row["span"]) for row in rows]) + seq_lens = [row["computed"] + len(row["span"]) for row in rows] + return SimpleNamespace( + num_reqs=len(rows), + has_prefill=any(row["prefill"] for row in rows), + is_prefilling_np=np.array([row["prefill"] for row in rows]), + idx_mapping_np=np.array([row["slot"] for row in rows]), + idx_mapping=torch.tensor([row["slot"] for row in rows]), + input_ids=torch.tensor(ids + [0] * (padded - num_tokens), dtype=torch.int32), + logits_indices=torch.tensor(starts[1:] - 1), + seq_lens=torch.tensor(seq_lens + [0] * 2, dtype=torch.int32), + ), padded + + +def _req_states(prompt_lens: dict[int, int], max_reqs: int = 8): + prompt_len = torch.zeros(max_reqs, dtype=torch.int32) + for slot, value in prompt_lens.items(): + prompt_len[slot] = value + return SimpleNamespace(prompt_len=SimpleNamespace(gpu=prompt_len)) + + +def test_mrv2_output_emits_decode_input_codes_with_validity_and_forced_eos() -> None: + talker = _talker() + rows = [ + # Final prefill chunk: no code yet (V1 emits an empty delta). + dict(slot=2, prompt_len=5, computed=0, span=[0, 0, 0, 0, 0], prefill=True), + # Decode with a codec id: emitted this step. + dict(slot=0, prompt_len=4, computed=6, span=[17], prefill=False), + # Decode whose input is codec EOS: request ended, force EOS again. + dict(slot=1, prompt_len=4, computed=8, span=[_EOS], prefill=False), + ] + batch, padded = _batch(rows, pad_to=8) + buffers = [{"audio_state": {"finished": False}}, {}, {}] + hidden = torch.zeros((padded, 4)) + out = talker.make_omni_output_mrv2( + hidden, + input_batch=batch, + req_states=_req_states({cast(int, row["slot"]): cast(int, row["prompt_len"]) for row in rows}), + model_intermediate_buffer=buffers, + ) + assert isinstance(out, OmniOutput) + codes = out.multimodal_outputs["codes"]["audio"] + valid = out.multimodal_outputs["meta"]["codec_frame_valid"] + assert codes.shape == (padded, 1) and valid.shape == (padded,) + assert codes[:, 0].tolist() == batch.input_ids.tolist() + assert valid.tolist() == [False] * 5 + [True, False] + [False] + assert talker.take_mrv2_forced_eos(batch, None, 3).tolist() == [False, False, True] + # Consumed once; a warmup sampler call sees nothing. + assert talker.take_mrv2_forced_eos(batch, None, 3) is None + + +def test_mrv2_output_empty_condition_and_length_cap() -> None: + talker = _talker() + prompt_lens = {3: 6, 4: 4096 - 100} + # Empty Thinker condition: preprocess marks the request finished at prefill. + batch, _ = _batch([dict(slot=3, prompt_len=6, computed=0, span=[0] * 6, prefill=True)]) + talker.make_omni_output_mrv2( + torch.zeros((6, 4)), + input_batch=batch, + req_states=_req_states(prompt_lens), + model_intermediate_buffer=[{"audio_state": {"finished": True}}], + ) + assert talker.take_mrv2_forced_eos(batch, None, 1).tolist() == [True] + # Its later decode rows stay forced and emit nothing, whatever the input id. + batch, _ = _batch([dict(slot=3, prompt_len=6, computed=6, span=[42], prefill=False)]) + out = talker.make_omni_output_mrv2( + torch.zeros((1, 4)), input_batch=batch, req_states=_req_states(prompt_lens), model_intermediate_buffer=[{}] + ) + assert out.multimodal_outputs["meta"]["codec_frame_valid"].tolist() == [False] + assert talker.take_mrv2_forced_eos(batch, None, 1).tolist() == [True] + + # Offline cap: min(2048, context - prompt) samples, the last one EOS. + limit = min(_OFFLINE_CODEC_MAX_NEW_TOKENS, 100) - 1 + for step, forced in ((limit - 1, False), (limit, True)): + batch, _ = _batch( + [dict(slot=4, prompt_len=prompt_lens[4], computed=prompt_lens[4] + step - 1, span=[5], prefill=False)] + ) + out = talker.make_omni_output_mrv2( + torch.zeros((1, 4)), input_batch=batch, req_states=_req_states(prompt_lens), model_intermediate_buffer=[{}] + ) + assert out.multimodal_outputs["meta"]["codec_frame_valid"].tolist() == [True] + assert talker.take_mrv2_forced_eos(batch, None, 1).tolist() == [forced] + + +@pytest.mark.parametrize("forced", [None, [False, True]]) +def test_sampler_adapter_keeps_upstream_counts_and_only_forces_codec_eos(mocker, forced): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import MiniCPMO45TalkerSampler + + output = SimpleNamespace(sampled_token_ids=torch.tensor([[2], [3]]), num_sampled=torch.tensor([1, 0])) + base = mocker.Mock(return_value=output) + base.req_states = object() + mask = None if forced is None else torch.tensor(forced) + talker = SimpleNamespace(_codec_eos_id=7, take_mrv2_forced_eos=mocker.Mock(return_value=mask)) + sampler = MiniCPMO45TalkerSampler(base, talker) + logits, batch = torch.zeros(2, 8), object() + assert sampler(logits, batch) is output + assert output.sampled_token_ids.tolist() == [[2], [3 if forced is None else 7]] + assert output.num_sampled.tolist() == [1, 0] + base.assert_called_once_with(logits, batch) + talker.take_mrv2_forced_eos.assert_called_once_with(batch, base.req_states, 2) diff --git a/tests/model_executor/models/minicpmo_4_5/test_thinker_mrv2.py b/tests/model_executor/models/minicpmo_4_5/test_thinker_mrv2.py new file mode 100644 index 000000000..d34bf614e --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_thinker_mrv2.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MRv2 Thinker preserves multimodal encoding and the llm2tts row ledger.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5 import minicpmo_4_5_omni as omni +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + "profile,capacities,kv_gib", + [ + ("minicpmo_4_5_turn_mrv2.yaml", [16, 8, 8], 2), + ("minicpmo_4_5_turn_mrv2_h200.yaml", [16, 16, 8], 4), + ], +) +def test_mrv2_profile_retains_full_thinker_handoff(profile, capacities, kv_gib, monkeypatch): + from pathlib import Path + + from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config, merge_pipeline_deploy + from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE + from vllm_omni.platforms import current_omni_platform + + # merge_pipeline_deploy resolves the platform again; test the CUDA profile + # consistently even when this CPU test runs on a ROCm host. + monkeypatch.setattr(current_omni_platform, "device_name", "cuda") + deploy = Path(__file__).resolve().parents[4] / "vllm_omni/deploy" / profile + config = _apply_platform_overrides(load_deploy_config(deploy), platform="cuda") + stages = merge_pipeline_deploy(MINICPMO_4_5_PIPELINE, config) + assert [s.yaml_engine_args["use_v2_model_runner"] for s in stages] == [True, True, True] + assert [s.yaml_engine_args["async_chunk"] for s in stages] == [False, True, True] + assert [s.yaml_engine_args["max_num_seqs"] for s in stages] == capacities + assert stages[1].yaml_engine_args["kv_cache_memory_bytes"] == kv_gib * 1024**3 + + +def _model(mocker, *, v2=True, session="turn", async_chunk=False): + thinker = torch.nn.Module() + thinker.make_empty_intermediate_tensors = lambda: None + mocker.patch.object(omni, "init_vllm_registered_model", return_value=thinker) + mocker.patch("vllm_omni.model_executor.models.minicpmo_4_5.duplex.compat.patch_minicpmo_remote_config") + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(), + multimodal_config=None, + model_stage="llm", + use_v2_model_runner=v2, + session_mode=session, + async_chunk=async_chunk, + ) + ) + return omni.MiniCPMO45OmniForConditionalGeneration(vllm_config=config) + + +def test_mrv2_thinker_keeps_native_multimodal_embeddings_and_raw_ids(mocker): + model = _model(mocker) + batch = SimpleNamespace(input_ids=torch.tensor([11, 22, 33]), num_tokens_after_padding=3) + embeddings = torch.randn(3, 4) + runner = SimpleNamespace( + model=model, + supports_mm_inputs=True, + is_first_pp_rank=True, + model_state=SimpleNamespace(dummy_inputs_embeds=lambda _: embeddings), + ) + ids, embeds, _ = OmniGPUModelRunner._prepare_mm_inputs(runner, None, batch, dummy_run=True) + assert ids is batch.input_ids + assert embeds is embeddings + hidden = torch.randn(3, 4) + model.thinker.forward = mocker.Mock(return_value=hidden) + result = model(ids, torch.tensor([5, 6, 7]), inputs_embeds=embeds) + assert result is hidden + torch.testing.assert_close(model.thinker.forward.call_args.kwargs["inputs_embeds"], embeddings) + + +def test_row_ledger_uses_live_batch_after_replay(mocker): + model = _model(mocker) + # Chunked prefill (two tokens), decode (one token), and graph padding. + batch = SimpleNamespace( + input_ids=torch.tensor([11, 12, 41, 0]), + positions=torch.tensor([6, 7, 20, 0]), + ) + hidden = torch.randn(4, 8) + for token in [41, 53]: + batch.input_ids[2] = token + batch.positions[2] += 1 + out = model.make_omni_output_mrv2( + hidden, + input_batch=batch, + req_states=None, + model_intermediate_buffer=[{}, {}], + ) + assert out.text_hidden_states is hidden + assert out.multimodal_outputs["latent"] is hidden + torch.testing.assert_close(out.multimodal_outputs["latent_input_ids"], batch.input_ids[:, None]) + torch.testing.assert_close(out.multimodal_outputs["latent_positions"], batch.positions[:, None]) diff --git a/tests/model_executor/models/minicpmo_4_5/test_vocoder_host_sync.py b/tests/model_executor/models/minicpmo_4_5/test_vocoder_host_sync.py new file mode 100644 index 000000000..6d86b5586 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_vocoder_host_sync.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Shared HiFT must retain its native path unless explicitly opted in.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import hifigan +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import HiFTGenerator + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_shared_hift_keeps_native_istft_without_opt_in(monkeypatch): + n_fft, hop = 16, 4 + window = torch.hann_window(n_fft) + hift = SimpleNamespace( + istft_params={"n_fft": n_fft, "hop_len": hop}, + _get_stft_window=lambda tensor: window, + ) + + def unexpected_cache(*args): + raise AssertionError("shared HiFT must not opt in implicitly") + + monkeypatch.setattr(hifigan, "_istft_without_host_sync", unexpected_cache) + magnitude, phase = torch.rand(2, 9, 20), torch.randn(2, 9, 20) + assert HiFTGenerator._istft(hift, magnitude, phase).shape == (2, 76) + assert not hasattr(hift, "_istft_envelopes") diff --git a/tests/model_executor/models/minicpmo_4_5/test_vocoder_host_sync_cuda.py b/tests/model_executor/models/minicpmo_4_5/test_vocoder_host_sync_cuda.py new file mode 100644 index 000000000..606fbd403 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_vocoder_host_sync_cuda.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The MiniCPM-o 4.5 vocoder stage's per-step path must not block the host.""" + +from __future__ import annotations + +from contextlib import contextmanager + +import pytest +import torch + +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import hifigan +from vllm_omni.worker.gpu_generation_model_runner import _HostCopyBatch + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +@contextmanager +def _no_host_sync(): + previous = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("error") + try: + yield + finally: + torch.cuda.set_sync_debug_mode(previous) + + +def test_istft_is_bitwise_and_sync_free_on_cuda(): + n_fft, hop = 16, 4 + window = torch.hann_window(n_fft, periodic=True, device="cuda") + generator = torch.Generator(device="cuda").manual_seed(0) + magnitude = torch.rand(4, n_fft // 2 + 1, 300, device="cuda", generator=generator) * 3 + phase = torch.randn(4, n_fft // 2 + 1, 300, device="cuda", generator=generator) + spec = torch.complex(magnitude * torch.cos(phase), magnitude * torch.sin(phase)) + envelopes: dict = {} + expected = torch.istft(spec, n_fft, hop, n_fft, window=window) + # The first call per frame count builds and checks the envelope. + hifigan._istft_without_host_sync(spec, n_fft, hop, window, envelopes) + with _no_host_sync(): + actual = hifigan._istft_without_host_sync(spec, n_fft, hop, window, envelopes) + assert torch.equal(actual, expected) + + +def test_host_copy_batch_copies_without_blocking_until_wait(): + to_host = _HostCopyBatch(pin_memory=True) + outputs = [torch.randn(1000, device="cuda") * (row + 1) for row in range(4)] + strided = torch.randn(8, 6, device="cuda").t() + to_host.copy(torch.zeros(1, device="cuda")) + to_host.wait() + with _no_host_sync(): + copies = [to_host.copy(tensor) for tensor in outputs] + copied_strided = to_host.copy(strided) + to_host.wait() + for copy, tensor in zip(copies, outputs, strict=True): + assert copy.device.type == "cpu" and copy.is_pinned() + assert torch.equal(copy, tensor.cpu()) + assert copied_strided.is_contiguous() and torch.equal(copied_strided, strided.cpu()) diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py index 07a709d94..e425cc9c8 100644 --- a/tests/worker_v2/test_omni_ar_model_runner.py +++ b/tests/worker_v2/test_omni_ar_model_runner.py @@ -112,6 +112,7 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history) runner.model = SimpleNamespace(compute_logits=None, logitsprocs_need_output_token_ids=needs_history) runner.model.mrv2_sampling_context = sampling_context + runner.sampler = None runner.sample = MagicMock(return_value=sampler_out) runner.sample.side_effect = lambda *_: sampler_out if sampling_active is needs_history else pytest.fail() logprobs_mock = MagicMock(side_effect=lambda *_: pytest.fail("inside ctx") if sampling_active else {}) @@ -304,3 +305,72 @@ def test_request_reference_codes_preserve_local_axis(prefill_first, padded): assert outputs[index]["codes.ref"].data_ptr() != ref.data_ptr() for i in range(2): assert torch.equal(outputs[i]["codes.audio"], codes[offsets[i] : offsets[i + 1]]) + + +@pytest.mark.parametrize("streaming", [False, True]) +def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch, streaming): + monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) + batch = SimpleNamespace( + num_reqs=2, + query_start_loc_np=np.array([0, 1, 4]), + num_scheduled_tokens=np.array([1, 3]), + num_tokens_after_padding=4, + ) + snapshot = torch.tensor([[10, 11], [20, 21]]) + seen = [] + + def finalize(payload, counts): + seen.append(counts) + return {"codes": {"audio": [payload["snapshot"][0:1].clone(), torch.empty(0, dtype=torch.long)]}} + + result = _async_output( + req_ids=["audio", "partial"], + sampler_output=SamplerOutput(torch.tensor([[99], [100]]), None, None, None), + num_sampled_tokens=torch.tensor([1, 0]), + multimodal_outputs={"snapshot": snapshot}, + input_batch=batch, + async_chunk=streaming, + finalize_multimodal=finalize, + ).get_output() + snapshot.zero_() + assert seen == [[1, 0]] + assert result.sampled_token_ids == [[99], []] + assert result.inter_stage_outputs[0]["codes.audio"].tolist() == [[10, 11]] + assert result.inter_stage_outputs[1]["codes.audio"].numel() == 0 + assert (result.pooler_output is None) == streaming + + +@pytest.mark.parametrize("streaming", [False, True]) +def test_request_owned_snapshot_skips_generic_partition(monkeypatch, streaming): + from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot + + monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) + monkeypatch.setattr( + OmniARModelRunner, + "_build_async_chunk_outputs_from_mm", + lambda *args: pytest.fail("already partitioned payload was repartitioned"), + ) + batch = SimpleNamespace( + num_reqs=2, + query_start_loc_np=np.array([0, 1, 4]), + num_scheduled_tokens=np.array([1, 3]), + num_tokens_after_padding=4, + ) + snapshot = torch.tensor([[10, 11], [20, 21]]) + + def finalize(payload, counts): + return RequestOutputSnapshot([{"codes.audio": payload["snapshot"][0:1].clone()}, None]) + + result = _async_output( + req_ids=["audio", "partial"], + sampler_output=SamplerOutput(torch.tensor([[99], [100]]), None, None, None), + num_sampled_tokens=torch.tensor([1, 0]), + multimodal_outputs={"snapshot": snapshot}, + input_batch=batch, + async_chunk=streaming, + finalize_multimodal=finalize, + ).get_output() + snapshot.zero_() + assert result.inter_stage_outputs[0]["codes.audio"].tolist() == [[10, 11]] + assert result.inter_stage_outputs[1] is None + assert (result.pooler_output is None) == streaming diff --git a/tests/worker_v2/test_omni_sampler.py b/tests/worker_v2/test_omni_sampler.py new file mode 100644 index 000000000..57fd7bedb --- /dev/null +++ b/tests/worker_v2/test_omni_sampler.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from vllm.v1.worker.gpu.sample.output import SamplerOutput + +from vllm_omni.worker_v2.omni_sampler import OmniSampler, OmniSamplingOutput, sample_with_output + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _output(): + counts, rejected = torch.tensor([1, 0]), torch.tensor([0, 0]) + return SamplerOutput(torch.tensor([[7], [0]]), None, None, counts, rejected), counts, rejected + + +@pytest.mark.parametrize("wrapped", [False, True]) +def test_standard_sampling_keeps_counts_grammar_and_lifecycle(wrapped): + base = Mock() + base.req_states = object() + sampler = OmniSampler(base) if wrapped else base + expected = _output() + standard = Mock(return_value=expected) + hidden, batch, grammar = torch.zeros(2, 4), object(), object() + result = sample_with_output(sampler, standard, hidden, batch, base.req_states, grammar) + assert result.sampler_output is expected[0] + assert result.num_sampled is expected[1] and result.num_rejected is expected[2] + assert result.multimodal_outputs is None + standard.assert_called_once_with(hidden, batch, grammar) + sampler.add_request(3, "request") + sampler.apply_staged_writes() + base.add_request.assert_called_once_with(3, "request") + base.apply_staged_writes.assert_called_once_with() + assert sampler.req_states is base.req_states + sampler(hidden, batch) + base.assert_called_once_with(hidden, batch) + + +@pytest.mark.parametrize("payload", [{}, {"codes.audio": torch.tensor([[1, 2]])}]) +def test_audio_sampler_uses_registered_adapter_without_standard_resampling(payload): + expected = OmniSamplingOutput(*_output(), multimodal_outputs=payload) + + class AudioSampler(OmniSampler): + def sample_step(self, hidden, batch, states, grammar, standard): + assert states is self.req_states + if grammar is not None: + raise ValueError("unsupported grammar") + return expected + + sampler = AudioSampler(SimpleNamespace(req_states=object())) + standard = Mock(side_effect=AssertionError("must not resample")) + result = sample_with_output(sampler, standard, torch.zeros(2, 4), object(), sampler.req_states, None) + assert result is expected and result.multimodal_outputs is payload + with pytest.raises(ValueError, match="unsupported grammar"): + sample_with_output(sampler, standard, torch.zeros(2, 4), object(), sampler.req_states, object()) + standard.assert_not_called() + + +def test_adapter_preserves_stock_sampler_staged_write_fast_path(): + from vllm.v1.worker.gpu.sample.sampler import Sampler + + assert OmniSampler(object.__new__(Sampler)).omni_static_staged_writes + assert not OmniSampler(SimpleNamespace()).omni_static_staged_writes + assert OmniSampler(SimpleNamespace(omni_static_staged_writes=True)).omni_static_staged_writes diff --git a/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml b/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml new file mode 100644 index 000000000..c743ee260 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml @@ -0,0 +1,10 @@ +# Opt-in single-H200 native full-duplex codec profile (V1). +# Inherits the native session/stage capacities, sampling and chunk sizes. +# Uses the mainline Whole-Euler Flow and shared prompt-state implementation. +# TF32 changes FP32 GEMM rounding; validate audio quality for your workload. +base_config: minicpmo_4_5.yaml + +connectors: + connector_of_shared_memory: + extra: + token2wav_allow_tf32: true diff --git a/vllm_omni/deploy/minicpmo_4_5_turn_mrv2.yaml b/vllm_omni/deploy/minicpmo_4_5_turn_mrv2.yaml new file mode 100644 index 000000000..a62dcf989 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_turn_mrv2.yaml @@ -0,0 +1,57 @@ +# Opt-in CUDA Model Runner V2 profile for MiniCPM-o 4.5 turn serving: the +# Thinker (stage 0), Talker (stage 1) and Code2Wav (stage 2) run on MRv2. +# The Thinker uses native multimodal encoding and tensor-only FULL graphs; +# downstream stages use the native data plane. These stages support +# turn-based sessions only (duplex needs the V1 chunk adapter). +# +# On MRv2 the Talker keeps its sampled codec id, 16-frame penalty history and +# EOS state on the device (no per-step D2H wait), and Code2Wav schedules +# individual chunks instead of holding a slot per utterance, so more than +# max_num_seqs streams can be live under load. +base_config: minicpmo_4_5_turn.yaml + +connectors: + connector_of_shared_memory: + extra: + hift_graph_capture_batch_sizes: [1, 2, 4, 8] + # TF32 tensor cores for the vocoder's fp32 flow GEMMs: ~25% less GPU + # time per chunk (batch 1-8). A numerics change: the HiFT input mel moves + # by ~3e-4 relative L2; dev-short WER/UTMOS unchanged (64 prompts). + token2wav_allow_tf32: true + +stages: + - stage_id: 0 + model_runner: v2 + # llm2tts consumes the completed Thinker payload via the orchestrator. + # Only Talker -> Code2Wav is a streaming connector edge. Keep the latent + # rows on the full-payload output channel instead of partitioning them + # into an async connector that this stage does not have. + async_chunk: false + - stage_id: 1 + model_runner: v2 + max_num_seqs: 8 + - stage_id: 2 + model_runner: v2 + max_num_seqs: 8 + async_scheduling: false + env: + # Decode waiting first chunks in their own step, but only while every + # stream with a later chunk ready still holds 0.5 s of unplayed audio, + # so first packets do not come at the cost of playback underruns. + VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS: "1" + VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS_SLACK_S: "0.5" + +platforms: + cuda: + stages: + - stage_id: 1 + # Reserve enough KV for eight concurrent 4096-token Talker requests. + kv_cache_memory_bytes: 2147483648 + npu: + model_runner: v1 + xpu: + model_runner: v1 + rocm: + model_runner: v1 + musa: + model_runner: v1 diff --git a/vllm_omni/deploy/minicpmo_4_5_turn_mrv2_h200.yaml b/vllm_omni/deploy/minicpmo_4_5_turn_mrv2_h200.yaml new file mode 100644 index 000000000..349af1108 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_turn_mrv2_h200.yaml @@ -0,0 +1,22 @@ +# Opt-in single-H200 throughput profile for turn-mode MiniCPM-o 4.5. +# Use NVIDIA MPS when sharing the GPU across stages. Warm representative +# streaming shapes before admitting traffic: graph capture can delay +# first audio by seconds. Talker capacity 16 reduces C32 admission-tail +# latency at similar throughput; C16 P99 varied across independent starts. +# This reserves 4 GiB of Talker KV (2 GiB more than the generic MRv2 profile). +# Keep room for lazy graph caches on the 141 GiB H200. +# +# FP32 weights/caches are retained, but the opt-in TF32 policy changes +# rounding. Quality evidence and limits: +# docs/design/minicpm_o45_mrv2_performance.md +base_config: minicpmo_4_5_turn_mrv2.yaml + +stages: + - stage_id: 1 + max_num_seqs: 16 + +platforms: + cuda: + stages: + - stage_id: 1 + kv_cache_memory_bytes: 4294967296 diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index ed81d618d..a353675bf 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -143,6 +143,53 @@ class ResBlock(torch.nn.Module): return folded +def _istft_envelope(window: torch.Tensor, n_fft: int, hop_length: int, n_frames: int) -> torch.Tensor: + """Trimmed overlap-add window envelope that ``torch.istft`` divides by. + + Built with the same ops as ``at::native::istft`` (squared window folded by + ``unfold_backward``, centre padding trimmed). ``torch.istft`` also checks + here that no sample has a vanishing envelope, which reads the result back + to the host; the envelope only depends on the window and the frame count, + so the check runs once per shape instead of once per call. + """ + expected_len = n_fft + hop_length * (n_frames - 1) + squared = window.pow(2).expand(1, n_frames, n_fft) + envelope = torch.ops.aten.unfold_backward(squared, [1, expected_len], 1, n_fft, hop_length) + envelope = envelope[:, n_fft // 2 : expected_len - n_fft // 2] + if bool(envelope.abs().min().lt(1e-11)): + raise RuntimeError("window overlap add min: True") + return envelope + + +def _istft_without_host_sync( + spec: torch.Tensor, + n_fft: int, + hop_length: int, + window: torch.Tensor, + envelopes: dict[tuple, torch.Tensor], +) -> torch.Tensor: + """``torch.istft(spec, n_fft, hop_length, n_fft, window=window)`` without a host sync. + + Mirrors ``at::native::istft`` for this call (``center=True``, one-sided, + unnormalized, ``win_length == n_fft``, no ``length``): the same inverse + real FFT, window product, ``unfold_backward`` overlap-add, trim and + division, so the waveform is bitwise identical. The envelope, and with it + the only device-to-host read, comes from ``envelopes`` after the first + call for a frame count. + """ + batch, _, n_frames = spec.shape + key = (n_frames, n_fft, hop_length, window.device, window.dtype, window.data_ptr()) + envelope = envelopes.get(key) + if envelope is None: + envelope = _istft_envelope(window, n_fft, hop_length, n_frames) + envelopes[key] = envelope + frames = torch.fft.irfft(spec.transpose(1, 2), n=n_fft, dim=-1) + frames = frames * window.view(1, 1, n_fft) + expected_len = n_fft + hop_length * (n_frames - 1) + signal = torch.ops.aten.unfold_backward(frames, [batch, expected_len], 1, n_fft, hop_length) + return signal[:, n_fft // 2 : expected_len - n_fft // 2] / envelope + + def _carry_phase_at_boundary( cum_total: torch.Tensor, phase_acc: torch.Tensor | None, @@ -620,6 +667,10 @@ class HiFTGenerator(nn.Module): ) self.f0_predictor = f0_predictor + def enable_cached_istft(self) -> None: + """Opt into CUDA ISTFT with a cached overlap envelope.""" + self._use_cached_istft = True + def remove_weight_norm(self) -> int: """Fold the generator's frozen weight norms into plain weights. @@ -675,12 +726,28 @@ class HiFTGenerator(nn.Module): magnitude = torch.clip(magnitude, max=1e2) real = magnitude * torch.cos(phase) imag = magnitude * torch.sin(phase) - return torch.istft( - torch.complex(real, imag), + spec = torch.complex(real, imag) + window = self._get_stft_window(magnitude) + if ( + not getattr(self, "_use_cached_istft", False) + or spec.dim() != 3 + or window.dim() != 1 + or window.shape[0] != self.istft_params["n_fft"] + ): + return torch.istft( + spec, + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=window, + ) + envelopes = self.__dict__.setdefault("_istft_envelopes", {}) + return _istft_without_host_sync( + spec, self.istft_params["n_fft"], self.istft_params["hop_len"], - self.istft_params["n_fft"], - window=self._get_stft_window(magnitude), + window, + envelopes, ) def _istft_on_cpu(self, magnitude, phase): 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 09f8d0f65..59998217d 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 @@ -268,6 +268,9 @@ class BatchedToken2Wav(nn.Module): self._trt_stepper = trt_stepper self.flow = token2wav.flow self.hift = token2wav.hift + enable_cached_istft = getattr(self.hift, "enable_cached_istft", None) + if callable(enable_cached_istft): + enable_cached_istft() encoder = getattr(self.flow, "encoder", None) if encoder is not None: _undecorate_dynamo(encoder, "forward_chunk") 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 072d51e50..3f93f0f30 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 @@ -22,6 +22,7 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.utils.device_copy import to_device_nonblocking from .batched_token2wav import ( BatchedToken2Wav, @@ -146,10 +147,12 @@ def _normalize_reference( def _codec_tensor(value: Any, fallback: torch.Tensor) -> torch.Tensor: + # Codec ids arrive from the connector as host data. A pageable host copy + # would block the host once per request per step; stage them pinned. if isinstance(value, torch.Tensor): - return value.reshape(-1).to(device=fallback.device, dtype=torch.long) + return to_device_nonblocking(value.reshape(-1).to(dtype=torch.long), fallback.device) if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): - return torch.as_tensor(value, device=fallback.device, dtype=torch.long).reshape(-1) + return to_device_nonblocking(torch.as_tensor(value, dtype=torch.long).reshape(-1), fallback.device) return fallback.reshape(-1).to(dtype=torch.long) @@ -730,6 +733,33 @@ class MiniCPMO45Code2Wav(nn.Module): inputs_embeds: torch.Tensor | None = None, runtime_additional_information: list[dict[str, Any]] | None = None, **kwargs: Any, + ) -> OmniOutput: + # This stage owns the vocoder process. Restore its previous matmul + # policy after eager execution/capture; cuDNN's policy is independent. + previous_tf32 = torch.backends.cuda.matmul.allow_tf32 + try: + if self._extra_config().get("token2wav_allow_tf32", False): + torch.backends.cuda.matmul.allow_tf32 = True + return self._forward_impl( + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + runtime_additional_information, + **kwargs, + ) + finally: + torch.backends.cuda.matmul.allow_tf32 = previous_tf32 + + @torch.inference_mode() + def _forward_impl( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, ) -> OmniOutput: del positions, intermediate_tensors, inputs_embeds ids = input_ids if isinstance(input_ids, torch.Tensor) else torch.empty(0, dtype=torch.long) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index b34cc5bd0..8056f6c08 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -94,6 +94,19 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp patch_minicpmo_remote_config(config) self.model_stage = vllm_config.model_config.model_stage + self._use_v2_model_runner = bool(getattr(vllm_config.model_config, "use_v2_model_runner", False)) + if ( + self.model_stage == "llm" + and self._use_v2_model_runner + and getattr(vllm_config.model_config, "session_mode", "turn") != "turn" + ): + raise NotImplementedError("MiniCPM-o duplex Thinker requires model_runner: v1") + if ( + self.model_stage == "llm" + and self._use_v2_model_runner + and getattr(vllm_config.model_config, "async_chunk", False) + ): + raise ValueError("MiniCPM-o MRv2 Thinker requires async_chunk: false for its full llm2tts payload") # The Thinker's row ledger needs real token identities even when # embeddings are supplied, including during CUDA graph capture/replay. self.requires_raw_input_tokens = self.model_stage == "llm" @@ -155,6 +168,22 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp if hasattr(self.talker, "init_multi_modal"): self.talker.init_multi_modal(config) self.model = self.talker + # The runner looks this hook up on the wrapper. Without it every + # decode row runs scalar preprocess, which reads its codec id with + # a blocking ``.item()``. + batch_decode = getattr(self.talker, "preprocess_decode_batch", None) + if callable(batch_decode): + self.preprocess_decode_batch = batch_decode + # Model Runner V2 hooks: device-side codec output, EOS control and + # codec penalty (see MiniCPMO45OmniTTSForConditionalGeneration). + for hook_name in ("preprocess_decode_batch_mrv2", "make_omni_output_mrv2", "mrv2_custom_sampler"): + hook = getattr(self.talker, hook_name, None) + if callable(hook): + setattr(self, hook_name, hook) + self.mrv2_decode_preprocess_is_identity = bool( + getattr(self.talker, "mrv2_decode_preprocess_is_identity", False) + ) + self.logits_vocab_size = int(self.talker.logits_vocab_size) else: raise ValueError(f"Invalid model stage: {self.model_stage}. Must be one of: 'llm', 'tts'") @@ -173,7 +202,13 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp # preprocess for duplex audio, while the Talker converts the # tts_token_ids/tts_hidden_states handoff into its conditioning # embeddings and initializes request-local codec generation state. - self.has_preprocess = self.model_stage in {"llm", "tts"} + # Turn-mode Thinker inputs use the native multimodal encoder/cache on + # MRv2. Marking it as a custom-preprocess model disables that encoder + # path in the runner. Duplex still uses the V1 preprocess hook. + self.has_preprocess = self.model_stage == "tts" or not self._use_v2_model_runner + # Neither AR stage has a postprocess, so step outputs can use the + # runner's async snapshot instead of a blocking per-step D2H. + self.use_async_omni_output = self.model_stage in {"llm", "tts"} if self.model_stage == "llm" and getattr(vllm_config.model_config, "session_mode", "turn") == "duplex": # Build the Stage-0 duplex runtime (remote-code processor and @@ -617,6 +652,12 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp if text_hidden_states.ndim == 3 and text_hidden_states.shape[0] == 1: text_hidden_states = text_hidden_states.squeeze(0) + if getattr(self, "_use_v2_model_runner", False): + # FULL graphs return tensors. The row ledger is reconstructed + # from this step's live input batch after replay, never from + # Python values or token buffers captured during warmup. + return text_hidden_states + # Return hidden states with latent in multimodal_outputs for stage_input_processors multimodal_outputs = {"latent": text_hidden_states} # Keep per-forward row identities alongside the latent payload. @@ -691,6 +732,27 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp raise ValueError(f"Unsupported model stage: {self.model_stage}") + def make_omni_output_mrv2(self, model_outputs, *, input_batch, req_states, model_intermediate_buffer): + """Attach the Thinker row identities used by the llm2tts bridge.""" + if self.model_stage != "llm": + return self.talker.make_omni_output_mrv2( + model_outputs, + input_batch=input_batch, + req_states=req_states, + model_intermediate_buffer=model_intermediate_buffer, + ) + if any(isinstance(info, dict) and info.get("duplex") for info in model_intermediate_buffer): + raise NotImplementedError("MiniCPM-o duplex Thinker requires model_runner: v1") + num_tokens = model_outputs.shape[0] + return OmniOutput( + text_hidden_states=model_outputs, + multimodal_outputs={ + "latent": model_outputs, + "latent_input_ids": input_batch.input_ids[:num_tokens].reshape(-1, 1), + "latent_positions": input_batch.positions[:num_tokens].reshape(-1, 1), + }, + ) + def make_omni_output(self, model_outputs, **kwargs): if self.model_stage != "tts": return model_outputs diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 320332fa6..206783903 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -14,8 +14,10 @@ Pipeline: from collections.abc import Iterable, Mapping, Sequence from dataclasses import replace +from types import SimpleNamespace from typing import Any +import numpy as np import torch import torch.nn as nn import torch.nn.functional as F @@ -34,6 +36,8 @@ from vllm_omni.model_executor.models.minicpmo_4_5 import ( ) from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform +from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking +from vllm_omni.worker_v2.omni_sampler import OmniSampler logger = init_logger(__name__) @@ -105,6 +109,38 @@ def _restore_weight_norm_weight(weight_g: torch.Tensor, weight_v: torch.Tensor) return torch._weight_norm(weight_v, weight_g, dim=0) +class _HostCodecIds: + """A batch of device codec ids, copied to pinned host memory without a host sync. + + The copy is queued before the Talker forward, so ``values()`` waits only + for the previous step's sampler (the same point vLLM's penalty path already + waits for), not for this step's forward. + """ + + __slots__ = ("_event", "_host", "_values") + + def __init__(self, ids: torch.Tensor) -> None: + self._event: torch.cuda.Event | None = None + self._host: torch.Tensor | None = None + self._values: list[int] | None = None + ids = ids.detach().reshape(-1) + if ids.device.type == "cuda": + self._host = torch.empty(ids.shape, dtype=ids.dtype, device="cpu", pin_memory=True) + self._host.copy_(ids, non_blocking=True) + self._event = torch.cuda.Event() + self._event.record() + else: + self._values = ids.tolist() + + def values(self) -> list[int]: + if self._values is None: + assert self._event is not None and self._host is not None + self._event.synchronize() + self._values = self._host.tolist() + self._event = self._host = None + return self._values + + def _apply_batched_repetition_penalty( logits: torch.Tensor, histories: Sequence[torch.Tensor], @@ -149,7 +185,8 @@ def _apply_batched_repetition_penalty( # The vocabulary fixes the output size. CUDA bincount still reads the # maximum id back to the host even when minlength is provided. encoded = encoded_rows[0] if len(encoded_rows) == 1 else torch.cat(encoded_rows) - encoded = encoded.to(device=logits.device) + # Pinned, non-blocking: a pageable H2D would wait for queued GPU work. + encoded = to_device_nonblocking(encoded, logits.device) frequencies = torch.zeros((end - start) * vocab_size, dtype=torch.long, device=logits.device) frequencies.scatter_add_(0, encoded, torch.ones_like(encoded)) frequencies = frequencies.reshape(end - start, vocab_size) @@ -159,6 +196,150 @@ def _apply_batched_repetition_penalty( return penalized +def _apply_codec_window_penalty_gpu( + logits: torch.Tensor, + slots: torch.Tensor, + all_token_ids: torch.Tensor, + total_len: torch.Tensor, + prompt_len: torch.Tensor, + penalty: torch.Tensor, + *, + window_size: int, +) -> None: + """In place: ``_apply_batched_repetition_penalty`` on the runner's device token history. + + Model Runner V2 keeps every request's sampled ids (``all_token_ids``) and + lengths on the device, so the 16-frame window is gathered there instead of + being rebuilt on the host each step. Row ``r`` is scored over its last + ``window_size`` output ids, ``[max(total_len - window, prompt_len), + total_len)``: the codes sampled so far, exactly the history the V1 path + builds in ``make_omni_output``. Same arithmetic as the V1 helper, so the + penalized logits are identical. + """ + num_rows, vocab_size = logits.shape + if num_rows == 0: + return + slots = slots.long() + end = total_len.index_select(0, slots).long() + start = torch.maximum(end - window_size, prompt_len.index_select(0, slots).long()) + offsets = torch.arange(window_size, device=logits.device, dtype=torch.long) + positions = end.unsqueeze(1) - window_size + offsets.unsqueeze(0) + valid = positions >= start.unsqueeze(1) + rows = all_token_ids.index_select(0, slots) + history = rows.gather(1, positions.clamp_min(0)).long() + # Invalid positions count into a spare column that is dropped below. + history = torch.where(valid & (history >= 0) & (history < vocab_size), history, vocab_size) + frequencies = torch.zeros((num_rows, vocab_size + 1), dtype=torch.long, device=logits.device) + frequencies.scatter_add_(1, history, torch.ones_like(history)) + alpha = torch.pow( + penalty.index_select(0, slots).to(dtype=logits.dtype).unsqueeze(1), + frequencies[:, :vocab_size].to(dtype=logits.dtype), + ) + logits.copy_(torch.where(logits < 0, logits * alpha, logits / alpha)) + + +class _CodecWindowPenaltiesState: + """MRv2 sampler penalties for the MiniCPM-o Talker. + + Replaces the sampler's presence-based repetition penalty (whole prompt and + output) with MiniCPMTTS's frequency penalty over the last 16 codes, scored + on the device. Frequency and presence penalties, if a request sets them, + still go through the upstream state, which keeps the output bin counts. + """ + + def __init__(self, base: Any, *, window_size: int) -> None: + from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor + + self.base = base + self.req_states = base.req_states + self.window_size = int(window_size) + max_num_reqs = int(self.req_states.max_num_reqs) + self.repetition_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32) + self.repetition_penalty.np.fill(1.0) + self.repetition_penalty.copy_to_uva() + self.use_window = np.zeros(max_num_reqs, dtype=bool) + self.use_penalty = np.zeros(max_num_reqs, dtype=bool) + + @property + def output_bin_counts(self) -> torch.Tensor: + return self.base.output_bin_counts + + def add_request(self, req_idx: int, sampling_params: Any) -> None: + repetition = float(getattr(sampling_params, "repetition_penalty", 1.0)) + self.repetition_penalty.np[req_idx] = repetition + self.use_window[req_idx] = repetition != 1.0 + self.base.add_request( + req_idx, + SimpleNamespace( + repetition_penalty=1.0, + frequency_penalty=float(getattr(sampling_params, "frequency_penalty", 0.0)), + presence_penalty=float(getattr(sampling_params, "presence_penalty", 0.0)), + ), + ) + self.use_penalty[req_idx] = bool(self.use_window[req_idx] or self.base.use_penalty[req_idx]) + + def apply_staged_writes(self) -> None: + self.repetition_penalty.copy_to_uva() + self.base.apply_staged_writes() + + def apply_penalties( + self, + logits: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + idx_mapping_np: np.ndarray, + input_ids: torch.Tensor, + expanded_local_pos: torch.Tensor, + ) -> None: + if np.any(self.use_window[idx_mapping_np]): + _apply_codec_window_penalty_gpu( + logits, + expanded_idx_mapping, + self.req_states.all_token_ids.gpu, + self.req_states.total_len.gpu, + self.req_states.prompt_len.gpu, + self.repetition_penalty.gpu, + window_size=self.window_size, + ) + self.base.apply_penalties(logits, expanded_idx_mapping, idx_mapping_np, input_ids, expanded_local_pos) + + +def _install_mrv2_talker_sampler(sampler: Any, talker: "MiniCPMO45OmniTTSForConditionalGeneration") -> Any: + """Wrap the runner's MRv2 sampler with codec penalties and EOS control. + + The upstream sampler already applies ``min_tokens`` (codec EOS is a stage + stop token), temperature, top-k/top-p and seeded sampling on the device. + The Talker adds its 16-frame codec penalty and overrides rows the model + terminates (``_mrv2_forced_eos``) with codec EOS after sampling, exactly + where V1's ``_force_eos_on_sampled_ids`` does. + """ + + sampler.penalties_state = _CodecWindowPenaltiesState(sampler.penalties_state, window_size=_CODEC_PENALTY_WINDOW) + talker._mrv2_empty_speech = torch.zeros( + int(sampler.req_states.max_num_reqs), dtype=torch.bool, device=sampler.req_states.device + ) + logger.info( + "MiniCPM-o Talker: MRv2 sampler with device-side %d-frame codec penalty and EOS control", + _CODEC_PENALTY_WINDOW, + ) + return MiniCPMO45TalkerSampler(sampler, talker) + + +class MiniCPMO45TalkerSampler(OmniSampler): + omni_static_staged_writes = True + + def __init__(self, base_sampler, talker): + super().__init__(base_sampler) + self.talker = talker + + def __call__(self, logits: torch.Tensor, input_batch: Any) -> Any: + forced = self.talker.take_mrv2_forced_eos(input_batch, self.req_states, logits.shape[0]) + output = self.base_sampler(logits, input_batch) + if forced is not None: + sampled = output.sampled_token_ids + sampled.masked_fill_(forced.view(-1, *([1] * (sampled.ndim - 1))), int(self.talker._codec_eos_id)) + return output + + class _MiniCPMTTSProjector(nn.Module): """Checkpoint-compatible hidden-state projector used by MiniCPMTTS.""" @@ -188,6 +369,9 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): self._mask_eos_rows: list[bool] | None = None self._pending_force_eos_rows: list[bool] | None = None self._penalty_histories: list[torch.Tensor] | None = None + # Host copies of this step's decode input ids, keyed by request id: + # make_omni_output reads them for EOS detection and penalty history. + self._decode_codec_ids: dict[str, tuple[_HostCodecIds, int]] = {} self._request_audio_states: dict[str, dict[str, Any]] = {} # Mirrors upstream TTSStreamingGenerator._chunk_info: one committed # condition plus, during a rollover, one immutable recompute recipe. @@ -212,9 +396,17 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): self.has_preprocess = True self.has_postprocess = False # Same-step codes travel through make_omni_output from the previous - # sampled id (decode preprocess embeds that id via emb_code). - self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("codes", "audio")} + # sampled id (decode preprocess embeds that id via emb_code). They are + # CPU transport deltas, so no buffer key is GPU-resident. self._init_native_talker(prefix) + # Model Runner V2 keeps the sampled id, codec history and EOS state on + # the GPU (see make_omni_output_mrv2 and mrv2_custom_sampler). + self._use_v2_model_runner = bool(getattr(vllm_config.model_config, "use_v2_model_runner", False)) + # Per request slot: the Thinker handed over an empty condition. + self._mrv2_empty_speech: torch.Tensor | None = None + # Rows the sampler must force to codec EOS, computed with this step's output. + self._mrv2_forced_eos: torch.Tensor | None = None + self._mrv2_decode_rows_logged = False def _init_native_talker(self, prefix: str) -> None: if self._tts_config is None: @@ -255,11 +447,7 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): def _boundary_embeddings(self) -> torch.Tensor: """Embed the ```` tail every condition ends with.""" - ids = torch.tensor( - [self._text_eos_id, self._tts_bos_id], - device=self.emb_text.weight.device, - dtype=torch.long, - ) + ids = index_to_device([self._text_eos_id, self._tts_bos_id], self.emb_text.weight.device) return self.emb_text(ids) def _build_condition_embeddings( @@ -278,8 +466,9 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): return self._boundary_embeddings() device = self.emb_text.weight.device dtype = self.emb_text.weight.dtype - token_ids = tts_token_ids.to(device=device, dtype=torch.long).reshape(-1) - hidden = tts_hidden_states.to(device=device, dtype=dtype) + # Pinned, non-blocking H2D: a pageable copy would stall every row in the step. + token_ids = to_device_nonblocking(tts_token_ids, device).to(dtype=torch.long).reshape(-1) + hidden = to_device_nonblocking(tts_hidden_states, device).to(dtype=dtype) if hidden.shape[0] != token_ids.shape[0] and token_ids.shape[0] != 1: raise ValueError( "MiniCPM-o Talker condition length mismatch: " @@ -289,7 +478,7 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): hidden_embeds = self.projector_semantic(hidden) if self._normalize: hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1) - audio_bos = self.emb_text(torch.tensor([self._tts_bos_id], device=device, dtype=torch.long)) + audio_bos = self.emb_text(index_to_device([self._tts_bos_id], device)) condition = text_embeds + hidden_embeds if native_duplex: # Match MiniCPMTTS.generate_chunk's streaming condition. @@ -565,6 +754,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): # Decode: vLLM's previous sampled codec id is this step's input. # Embed it with the codec table (not the 32k Llama embed_tokens) and # hand the same id to make_omni_output so Code2Wav sees it this step. + # The runner sends one-token decode rows through preprocess_decode_batch, + # which defers this ``.item()`` EOS check to make_omni_output. code = input_ids.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1)[-1:] embeds = self.emb_code[0](code) code_id = int(code.item()) @@ -580,6 +771,185 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): delta = torch.tensor([[code_id]], dtype=torch.long, device="cpu") return input_ids, embeds, {"codes": {"audio": delta}} + def _decode_codec_id_map(self) -> dict[str, tuple[_HostCodecIds, int]]: + pending = getattr(self, "_decode_codec_ids", None) + if pending is None: + pending = {} + self._decode_codec_ids = pending + return pending + + def preprocess_decode_batch( + self, + *, + input_ids: torch.Tensor, + req_infos: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, list[dict[str, Any]]]: + """Embed every one-token decode row at once, without reading ids on the host. + + Same result per row as scalar ``preprocess``, except that its per-row + ``.item()`` moves to ``make_omni_output``: one non-blocking copy of all + rows' ids is queued here, before the forward, and read there to build the + CPU delta and apply the codec-EOS check. Nothing reads the request's + ``finished`` state between the two points. + """ + num_rows = len(req_infos) + if input_ids.numel() != num_rows: + raise ValueError(f"MiniCPM-o Talker batched decode needs one id per row: {input_ids.numel()} != {num_rows}") + weight = self.emb_code[0].weight + codes = input_ids.reshape(-1).to(device=weight.device, dtype=torch.long) + embeds = self.emb_code[0](codes) + # The runner and connector consume CPU codec ids. + empty_codes = torch.empty(0, dtype=torch.long, device="cpu") + request_states = self._request_audio_states + updates: list[dict[str, Any]] = [] + live_rows: list[tuple[int, str]] = [] + zero_rows: list[int] = [] + for row, info in enumerate(req_infos): + request_id = str(info.get("request_id", "0")) + state = info.get("audio_state") + if info.get("_omni_is_prefill", False) or not isinstance(state, dict): + # A row without Talker state builds its condition like a prefill. + _, row_embeds, update = self.preprocess(input_ids=input_ids[row : row + 1], input_embeds=None, **info) + embeds[row : row + 1] = row_embeds + updates.append(update) + continue + stored = request_states.get(request_id) + if isinstance(stored, dict): + state = stored + if state.get("finished"): + # Leftover decode of a finished request: shape-correct zeros. + zero_rows.append(row) + updates.append({"codes": {"audio": empty_codes}}) + continue + live_rows.append((row, request_id)) + # The id is not on the host yet: make_omni_output builds this row's + # delta from the host copy. The empty entry keeps no stale delta. + updates.append({"codes": {"audio": empty_codes}}) + if zero_rows: + embeds.index_fill_(0, index_to_device(zero_rows, embeds.device), 0.0) + if live_rows: + host_ids = _HostCodecIds(codes) + pending = self._decode_codec_id_map() + for row, request_id in live_rows: + pending[request_id] = (host_ids, row) + return input_ids, embeds, updates + + # ------------------------------------------------------------------ + # Model Runner V2 + # ------------------------------------------------------------------ + # + # Under MRv2 a decode row's input id is the previous step's sample, which + # the runner gathers on the device (``last_sampled_tokens``) and embeds + # into its static input buffer with ``embed_input_ids`` (``emb_code``). + # That is the whole decode preprocess, so decode rows skip the per-row + # hook. Everything V1 derives on the host from that id each step -- the + # codec delta, codec-EOS detection, the step count behind the length cap + # and the 16-frame penalty history -- is derived on the device from the + # runner's request state instead, and nothing waits for a D2H copy. + + #: Decode rows need no per-row preprocess beyond the runner's embedding. + mrv2_decode_preprocess_is_identity = True + + @property + def logits_vocab_size(self) -> int: + return int(self._num_audio_tokens) + + def preprocess_decode_batch_mrv2( + self, + *, + input_ids: torch.Tensor, + input_embeds: torch.Tensor, + req_infos: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[dict[str, Any]]]: + """Decode rows are ``emb_code`` of their input id, already in ``input_embeds``.""" + num_rows = len(req_infos) + empty = input_embeds.new_empty((num_rows, 0)) + return input_ids, input_embeds, empty, empty, [{} for _ in range(num_rows)] + + def mrv2_custom_sampler(self, sampler: Any) -> tuple[Any, None]: + return _install_mrv2_talker_sampler(sampler, self), None + + def make_omni_output_mrv2( + self, + model_outputs: torch.Tensor | OmniOutput, + *, + input_batch: Any, + req_states: Any, + model_intermediate_buffer: list[dict[str, Any]], + ) -> OmniOutput: + """Device-side ``make_omni_output`` for Model Runner V2. + + Emits every token row's input id as ``codes.audio`` plus + ``meta.codec_frame_valid``: true only for a decode row whose input is a + codec id of a request that has not ended. That id is the code V1 hands + to Code2Wav on this step; the stage payload builder drops the other + rows on the host, after the asynchronous output copy. Rows the model + terminates this step (empty condition, codec EOS already sampled, or + the offline length cap) are recorded for the sampler. + """ + if isinstance(model_outputs, OmniOutput): + return model_outputs + hidden = model_outputs + empty_speech = self._mrv2_empty_speech + if empty_speech is None: + raise RuntimeError("MiniCPM-o Talker MRv2 output requires its MRv2 sampler (mrv2_custom_sampler)") + num_reqs = int(input_batch.num_reqs) + device = hidden.device + if input_batch.has_prefill: + # Prefill rows only: record whether the Thinker handed over an + # empty condition (preprocess marks such a request finished). + slots: list[int] = [] + flags: list[bool] = [] + for row in np.flatnonzero(input_batch.is_prefilling_np[:num_reqs]).tolist(): + info = model_intermediate_buffer[row] + if info.get("native_duplex") is True: + raise NotImplementedError("MiniCPM-o native duplex Talker requires model_runner: v1") + state = info.get("audio_state") + flags.append(bool(state.get("finished")) if isinstance(state, dict) else False) + slots.append(int(input_batch.idx_mapping_np[row])) + if slots: + empty_speech.index_copy_( + 0, + index_to_device(slots, device), + index_to_device(flags, device, dtype=torch.bool), + ) + num_tokens = int(hidden.shape[0]) + token_ids = input_batch.input_ids[:num_tokens] + last_rows = input_batch.logits_indices[:num_reqs].long() + slot_ids = input_batch.idx_mapping[:num_reqs].long() + prompt_len = req_states.prompt_len.gpu.index_select(0, slot_ids).long() + # Codes emitted before this step's input (V1's ``state["step"]``). + step = input_batch.seq_lens[:num_reqs].long() - prompt_len + last_ids = token_ids.index_select(0, last_rows) + empty = empty_speech.index_select(0, slot_ids) + decode = step > 0 + eos_input = decode & (last_ids == int(self._codec_eos_id)) + valid = decode & ~eos_input & ~empty + # MiniCPMTTS.generate's max_new_token, clamped to the Talker context + # (see ``preprocess``); the sample after the last allowed code is EOS. + remaining = int(self._tts_config.max_position_embeddings) - prompt_len + limit = torch.clamp(remaining, min=1, max=_OFFLINE_CODEC_MAX_NEW_TOKENS) - 1 + self._mrv2_forced_eos = empty | eos_input | (step >= limit) + frame_valid = torch.zeros(num_tokens, dtype=torch.bool, device=device) + frame_valid.index_copy_(0, last_rows, valid) + if not self._mrv2_decode_rows_logged and not input_batch.has_prefill: + self._mrv2_decode_rows_logged = True + logger.info("MiniCPM-o Talker: MRv2 device-side codec output active (no host read of sampled ids)") + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={ + "codes": {"audio": token_ids.to(dtype=torch.long).reshape(num_tokens, 1)}, + "meta": {"codec_frame_valid": frame_valid}, + }, + ) + + def take_mrv2_forced_eos(self, input_batch: Any, req_states: Any, num_rows: int) -> torch.Tensor | None: + """This step's forced-EOS rows, once; ``None`` outside a model step (warmup).""" + forced, self._mrv2_forced_eos = self._mrv2_forced_eos, None + if forced is None or int(forced.shape[0]) != int(num_rows): + return None + return forced + def make_omni_output( self, model_outputs: torch.Tensor | OmniOutput, @@ -606,6 +976,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): mask_eos_rows = [False] * len(infos) empty_history = torch.empty(0, dtype=torch.long, device="cpu") penalty_histories = [empty_history for _ in infos] + pending_codec_ids = self._decode_codec_id_map() + codec_eos_id = int(self._codec_eos_id) for index, info in enumerate(infos): info_dict = info if isinstance(info, dict) else {} native_duplex = info_dict.get("native_duplex") is True @@ -657,9 +1029,20 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): if not isinstance(state, dict): state = dict(info.get("audio_state", {}) or {}) self._request_audio_states[request_id] = state - empty_speech = bool(state.get("finished")) codes = info.get("codes", {}) audio = codes.get("audio") if isinstance(codes, Mapping) else None + pending = pending_codec_ids.pop(request_id, None) + if pending is not None: + # A preprocess_decode_batch row: build what scalar preprocess + # would have stored, from this step's host copy of the ids. + host_ids, row = pending + code_id = host_ids.values()[row] + if code_id == codec_eos_id: + state["finished"] = True + audio = None + else: + audio = torch.tensor([[code_id]], dtype=torch.long, device="cpu") + empty_speech = bool(state.get("finished")) if isinstance(audio, torch.Tensor) and audio.numel() > 0: audio = audio.to(device="cpu", dtype=torch.long) codec_deltas[index] = audio.reshape(-1, 1) @@ -725,9 +1108,11 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): def _flush_deferred_cleanup(self) -> None: request_audio_states = getattr(self, "_request_audio_states", {}) request_condition_states = getattr(self, "_request_condition_states", {}) + decode_codec_ids = getattr(self, "_decode_codec_ids", {}) for request_id in self._deferred_cleanup_ids: request_audio_states.pop(request_id, None) request_condition_states.pop(request_id, None) + decode_codec_ids.pop(request_id, None) self._deferred_cleanup_ids.clear() def _dummy_hidden_states( @@ -789,18 +1174,21 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): # we just forced (it is in the stage's ``stop_token_ids``), leaving an # all -inf row and a request that never releases. self._pending_force_eos_rows = force_eos if need_force else None + # Row masks go up through pinned memory and are applied with where / + # masked_fill: a pageable H2D or boolean-mask indexing blocks the host. if need_force or need_mask: logits = logits.clone() eos_id = int(self._codec_eos_id) if need_force: assert force_eos is not None - forced = torch.tensor(force_eos, dtype=torch.bool, device=logits.device) - logits[forced] = float("-inf") - logits[forced, eos_id] = 0.0 + forced = index_to_device(force_eos, logits.device, dtype=torch.bool) + eos_only = torch.full_like(logits[:1], float("-inf")) + eos_only[:, eos_id] = 0.0 + logits = torch.where(forced.unsqueeze(1), eos_only, logits) if need_mask: assert mask_eos is not None - masked = torch.tensor(mask_eos, dtype=torch.bool, device=logits.device) - logits[masked, eos_id] = float("-inf") + masked = index_to_device(mask_eos, logits.device, dtype=torch.bool) + logits[:, eos_id].masked_fill_(masked, float("-inf")) return logits def sample(self, logits, sampling_metadata): @@ -867,8 +1255,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): sampled = getattr(output, "sampled_token_ids", None) if not isinstance(sampled, torch.Tensor) or sampled.shape[0] != len(force_eos): return output - rows = torch.tensor(force_eos, dtype=torch.bool, device=sampled.device) - sampled[rows] = int(self._codec_eos_id) + rows = index_to_device(force_eos, sampled.device, dtype=torch.bool) + sampled.masked_fill_(rows.view(-1, *([1] * (sampled.ndim - 1))), int(self._codec_eos_id)) return output def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index e75393845..f6d71d44d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -47,6 +47,7 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( owns_tokenizer=True, requires_multimodal_data=True, engine_output_type="latent", + supports_native_mrv2_data_plane=True, sampling_constraints={ "detokenize": True, # The llm2tts bridge discards this boundary and every row @@ -65,6 +66,9 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( custom_process_input_func=f"{_PROC}.llm2tts", custom_process_next_stage_input_func=f"{_PROC}.tts2code2wav_full_payload", async_chunk_process_next_stage_input_func=f"{_PROC}.tts2code2wav_async_chunk", + # Takes effect only when the deploy selects model_runner v2 for + # this stage (turn sessions only; duplex stays on V1). + supports_native_mrv2_data_plane=True, sampling_constraints={ "detokenize": False, # MiniCPM-o 4.5 codec EOS is tts_config.num_audio_tokens - 1. @@ -81,6 +85,7 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="audio", model_arch="MiniCPMO45Code2Wav", + supports_native_mrv2_data_plane=True, sync_process_input_func=f"{_PROC}.tts2code2wav_token_only", sampling_constraints={"detokenize": True}, requires_full_payload_input=True, diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index 5855cf079..fe8ba2666 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -3,7 +3,7 @@ """MiniCPM-o 4.5 Thinker-to-Talker and Talker-to-Code2Wav bridges.""" import logging -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence from typing import Any import torch @@ -203,6 +203,14 @@ def _extract_codec_delta(pooling_output: Any, request_id: str) -> list[int]: return [] codes = pooling_output.get("codes") audio = codes.get("audio") if isinstance(codes, Mapping) else pooling_output.get("codes.audio") + valid = meta.get("codec_frame_valid") if isinstance(meta, Mapping) else None + if valid is None: + valid = pooling_output.get("meta.codec_frame_valid") + if isinstance(valid, torch.Tensor) and isinstance(audio, torch.Tensor): + # Model Runner V2 emits one id per token row plus its validity + # (decode rows of live requests); keep only the codec frames. + rows = valid.detach().to(device="cpu").reshape(-1).bool() + audio = audio.detach().to(device="cpu").reshape(rows.numel(), -1)[rows] return _codec_scalars(audio) if isinstance(pooling_output, Sequence) and not isinstance( pooling_output, @@ -227,7 +235,7 @@ def _drop_codec_state(transfer_manager: Any, request_id: str) -> None: else: request_payload.pop(request_id, None) code_accumulators = getattr(transfer_manager, "code_prompt_token_ids", None) - if hasattr(code_accumulators, "pop"): + if isinstance(code_accumulators, dict): code_accumulators.pop(request_id, None) @@ -316,7 +324,8 @@ def tts2code2wav_async_chunk( container[_MINICPMO45_ASYNC_STATE] = state pending = state["pending"] - pending.extend(_extract_codec_delta(multimodal_output, request_id)) + delta = _extract_codec_delta(multimodal_output, request_id) + pending.extend(delta) pending_text_utf8 = state.setdefault("pending_text_utf8", []) current_text_utf8 = ( segment_text_utf8.detach().to(device="cpu", dtype=torch.uint8).reshape(-1).tolist() @@ -602,7 +611,7 @@ def _decode_native_duplex_token_ids( request_id: str, ) -> str | None: decode_token_ids = getattr(streaming_context, "source_token_decoder", None) - if not isinstance(decode_token_ids, Callable): + if not callable(decode_token_ids): return None decode_ids = [int(token_id) for token_id in token_ids] try: @@ -1006,7 +1015,9 @@ def llm2tts( if data_plane_metadata is not None: model_intermediate_buffer["duplex"] = data_plane_metadata meta["native_duplex_segment_text"] = thinker_text - meta.setdefault("override_keys", []).extend( + override_keys = meta.setdefault("override_keys", []) + assert isinstance(override_keys, list) + override_keys.extend( [ "llm_output_text", ["meta", "native_duplex_segment_text"], diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index 28f363a1b..32622e8ef 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -35,7 +35,8 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import partition_flat_payload from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner -from vllm_omni.worker_v2.output_snapshot import PackedOutputSnapshot, pack_output_snapshot +from vllm_omni.worker_v2.omni_sampler import sample_with_output +from vllm_omni.worker_v2.output_snapshot import PackedOutputSnapshot, RequestOutputSnapshot, pack_output_snapshot logger = init_logger(__name__) _ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT = 64 @@ -187,11 +188,13 @@ class OmniARModelRunner(OmniGPUModelRunner): ), ) with sampling_context: - sampler_output, num_sampled, num_rejected = self.sample( - text_hidden, - input_batch, - grammar_output, + sampling_output = sample_with_output( + self.sampler, self.sample, text_hidden, input_batch, self.req_states, grammar_output ) + sampler_output = sampling_output.sampler_output + num_sampled, num_rejected = sampling_output.num_sampled, sampling_output.num_rejected + if sampling_output.multimodal_outputs is not None: + multimodal_outputs = sampling_output.multimodal_outputs run_eager_mtp = getattr(self.model_state, "run_eager_mtp", None) if multimodal_outputs and run_eager_mtp is not None: run_eager_mtp( @@ -255,11 +258,12 @@ class OmniARModelRunner(OmniGPUModelRunner): num_sampled_tokens=num_sampled, main_stream=self.main_stream, copy_stream=self.output_copy_stream, - text_hidden=text_hidden if need_pooler else None, + text_hidden=text_hidden if need_pooler and sampling_output.include_hidden_states else None, multimodal_outputs=multimodal_outputs if need_pooler else None, input_batch=input_batch if need_pooler else None, async_chunk=bool(getattr(self.model_config, "async_chunk", False)), finalize_output=(None if materialize_native else self._finalize_native_data_plane_output), + finalize_multimodal=sampling_output.finalize_multimodal, check_ep_fault=self.check_ep_fault, routed_experts=routed_experts, extra_multimodal_outputs=extra_outputs, @@ -268,7 +272,8 @@ class OmniARModelRunner(OmniGPUModelRunner): _guard_graph_replay_for_pooler_copy( self.main_stream, async_output.copy_event, - need_pooler=need_pooler, + need_pooler=need_pooler + and (sampling_output.include_hidden_states or not sampling_output.owns_multimodal_outputs), async_chunk=bool(getattr(self.model_config, "async_chunk", False)), ) @@ -671,6 +676,7 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): input_batch: Any | None = None, async_chunk: bool = False, finalize_output: Any | None = None, + finalize_multimodal: Any | None = None, check_ep_fault: bool = False, routed_experts: RoutedExpertsTensors | None = None, extra_multimodal_outputs: tuple[dict[str, Any], torch.cuda.Event] | None = None, @@ -682,10 +688,13 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): self.copy_event = copy_event if copy_event is not None else torch.cuda.Event(blocking=True) self._async_chunk = bool(async_chunk) self._finalize_output = finalize_output + self._finalize_multimodal = finalize_multimodal self._has_fault: torch.Tensor | None = None # Snapshot input_batch metadata needed for pooler_output slicing - self._need_pooler = text_hidden is not None or (self._async_chunk and bool(multimodal_outputs)) + self._need_pooler = text_hidden is not None or ( + (self._async_chunk or finalize_multimodal is not None) and bool(multimodal_outputs) + ) self._query_start_loc_np: np.ndarray | None = None self._num_scheduled_tokens: np.ndarray | None = None self._num_reqs: int = 0 @@ -759,7 +768,7 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): self._hidden_cpu: torch.Tensor | None = None self._mm_cpu: dict[str, Any] = {} self._mm_snapshot: dict[str, Any] = {} - if self._need_pooler and self._async_chunk: + if self._need_pooler and (self._async_chunk or self._finalize_multimodal is not None): # CUDA graph replay reuses the model's output buffers. Take # ownership directly in pinned host memory on the output copy # stream so deferred finalization never performs a blocking @@ -836,16 +845,27 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): # * pooler_output -> sync/full-payload path (inline pooling_output bridge) # * multimodal_outputs -> wire multimodal_output, which the async_chunk # stage-input processor (talker2code2wav_async_chunk) reads for codes. - if self._need_pooler and self._async_chunk: - pooler_inter, pooler_client = OmniARModelRunner._build_async_chunk_outputs_from_mm( - self._mm_snapshot, - self._query_start_loc_np, - self._num_scheduled_tokens, - self._num_reqs, - self._total_tokens, - self._padded_total_tokens, - ) - self.model_runner_output.pooler_output = None + if self._need_pooler and (self._async_chunk or self._finalize_multimodal is not None): + if self._finalize_multimodal is not None: + self._mm_snapshot = self._finalize_multimodal(self._mm_snapshot, num_sampled_tokens) + pooler_inter: list[dict[str, Any] | None] | None + pooler_client: list[dict[str, Any] | None] | None + if isinstance(self._mm_snapshot, RequestOutputSnapshot): + pooler_inter, pooler_client = self._mm_snapshot.inter_stage, self._mm_snapshot.client + if len(pooler_inter) != self._num_reqs or ( + pooler_client is not None and len(pooler_client) != self._num_reqs + ): + raise ValueError("Model-owned output snapshot does not match the request batch") + else: + pooler_inter, pooler_client = OmniARModelRunner._build_async_chunk_outputs_from_mm( + self._mm_snapshot, + self._query_start_loc_np, + self._num_scheduled_tokens, + self._num_reqs, + self._total_tokens, + self._padded_total_tokens, + ) + self.model_runner_output.pooler_output = None if self._async_chunk else pooler_inter self.model_runner_output.inter_stage_outputs = pooler_inter self.model_runner_output.multimodal_outputs = ( [_ensure_tensor_values(_async_copy_mm_value(p)) if p else {} for p in pooler_client] diff --git a/vllm_omni/worker_v2/omni_sampler.py b/vllm_omni/worker_v2/omni_sampler.py new file mode 100644 index 000000000..1f1235aa3 --- /dev/null +++ b/vllm_omni/worker_v2/omni_sampler.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""One sampler registration contract for token and post-sampling audio output.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch +from vllm.v1.core.sched.output import GrammarOutput +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.sample.output import SamplerOutput +from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.states import RequestState + +from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot + +StandardSample = Callable[ + [torch.Tensor, InputBatch, GrammarOutput | None], tuple[SamplerOutput, torch.Tensor, torch.Tensor] +] + + +@dataclass +class OmniSamplingOutput: + sampler_output: SamplerOutput + num_sampled: torch.Tensor + num_rejected: torch.Tensor + # None preserves forward's payload. An empty dict explicitly replaces it. + # Tensors must remain valid until the runner finishes its output copy. + multimodal_outputs: dict[str, Any] | None = None + include_hidden_states: bool = True + owns_multimodal_outputs: bool = False + # Invoked on CPU-owned snapshots after D2H, with effective per-request counts. + finalize_multimodal: Callable[[dict[str, Any], list[int]], dict[str, Any] | RequestOutputSnapshot] | None = None + + +class OmniSampler: + """Register through ModelState.custom_sampler, just like an upstream sampler. + + Delegate request lifecycle/state to the original sampler without copying + its attributes. Models can override __call__ for logits-only adjustments, + or sample_step when sampling also produces audio. Both use one registered + sampler and one result contract; there is no separate ModelState step hook. + """ + + def __init__(self, base_sampler: Sampler): + self.base_sampler = base_sampler + + @property + def omni_static_staged_writes(self) -> bool: + # Preserve the runner's no-new-requests fast path for the stock sampler. + return type(self.base_sampler) is Sampler or bool( + getattr(self.base_sampler, "omni_static_staged_writes", False) + ) + + def __getattr__(self, name: str) -> Any: + return getattr(self.base_sampler, name) + + def __call__(self, logits: torch.Tensor, input_batch: InputBatch) -> SamplerOutput: + return self.base_sampler(logits, input_batch) + + def sample_step( + self, + hidden_states: torch.Tensor, + input_batch: InputBatch, + req_states: RequestState, + grammar_output: GrammarOutput | None, + standard_sample: StandardSample, + ) -> OmniSamplingOutput: + # standard_sample calls this adapter's __call__ through the runner. + return OmniSamplingOutput(*standard_sample(hidden_states, input_batch, grammar_output)) + + +def sample_with_output( + sampler: Sampler | OmniSampler | None, + standard_sample: StandardSample, + hidden_states: torch.Tensor, + input_batch: InputBatch, + req_states: RequestState, + grammar_output: GrammarOutput | None, +) -> OmniSamplingOutput: + if isinstance(sampler, OmniSampler): + return sampler.sample_step(hidden_states, input_batch, req_states, grammar_output, standard_sample) + return OmniSamplingOutput(*standard_sample(hidden_states, input_batch, grammar_output)) diff --git a/vllm_omni/worker_v2/output_snapshot.py b/vllm_omni/worker_v2/output_snapshot.py index 31187ddf7..a8ae384db 100644 --- a/vllm_omni/worker_v2/output_snapshot.py +++ b/vllm_omni/worker_v2/output_snapshot.py @@ -6,12 +6,25 @@ from __future__ import annotations from collections import defaultdict from collections.abc import Callable +from dataclasses import dataclass from typing import Any import torch from torch.utils._pytree import tree_flatten, tree_unflatten +@dataclass(frozen=True) +class RequestOutputSnapshot: + """Already partitioned, CPU-owned payloads in the current request order. + + Model finalizers may return this after D2H completes to avoid recursively + cloning and partitioning data whose request ownership is already known. + """ + + inter_stage: list[dict[str, Any] | None] + client: list[dict[str, Any] | None] | None = None + + class PackedOutputSnapshot(dict): """A normal payload mapping with an internal batched-copy plan.