diff --git a/tests/model_executor/models/minicpmo_4_5/test_fused_residual_layer_norm.py b/tests/model_executor/models/minicpmo_4_5/test_fused_residual_layer_norm.py new file mode 100644 index 000000000..aa967f970 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_fused_residual_layer_norm.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""``residual_layer_norm`` matches the eager gated residual + LayerNorm chain.""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.model_executor.models.minicpmo_4_5.ops import residual_layer_norm + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.mark.parametrize("device", ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])) +def test_gated_residual_modulated_layer_norm(device: str) -> None: + torch.manual_seed(0) + channels = 64 + x = torch.randn(3, 7, channels, device=device) + y = torch.randn(3, 7, channels, device=device) + gate, scale, shift = (torch.randn(channels, device=device) for _ in range(3)) + expected_residual = x + gate * y + expected = F.layer_norm(expected_residual, (channels,), eps=1e-6) * (1 + scale) + shift + residual = x.clone() + out = residual_layer_norm(residual, y, gate=gate, weight=1 + scale, bias=shift, eps=1e-6, residual_out=residual) + torch.testing.assert_close(residual, expected_residual, rtol=1e-6, atol=1e-6) + torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5) diff --git a/tests/model_executor/models/minicpmo_4_5/test_incremental_fbank.py b/tests/model_executor/models/minicpmo_4_5/test_incremental_fbank.py new file mode 100644 index 000000000..0a4079e22 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_incremental_fbank.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""``duplex_incremental_fbank`` is bitwise equal to the remote ``StreamingMelProcessorExact`` (skipped without it).""" + +import importlib +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch +from transformers.dynamic_module_utils import init_hf_modules +from transformers.utils import HF_MODULES_CACHE + +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import MiniCPMO45Stage0DuplexRuntime + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_SR = 16000 + + +@pytest.fixture(scope="module") +def remote() -> ModuleType: # openbmb/MiniCPM-o-4_5's remote code, from the transformers module cache + candidates = Path(HF_MODULES_CACHE, "transformers_modules").glob("*/*/processing_minicpmo.py") + paths = [path for path in candidates if "class StreamingMelProcessorExact" in path.read_text()] + if not paths: + pytest.skip("MiniCPM-o 4.5 remote processing code is not in the transformers module cache") + init_hf_modules() + path = max(paths, key=lambda path: path.stat().st_mtime) + return importlib.import_module(".".join(path.relative_to(HF_MODULES_CACHE).with_suffix("").parts)) + + +def _mel(remote: ModuleType, *, incremental: bool) -> Any: + """Stage 0's streaming configuration, switched over by the runtime as a session's processor is.""" + mel = remote.StreamingMelProcessorExact( + feature_extractor=remote.MiniCPMAAudioProcessor(), + chunk_ms=1000, + first_chunk_ms=1035, + cnn_redundancy_ms=20, + enable_sliding_window=True, + slide_trigger_seconds=30.0, + slide_stride_seconds=10.0, + ) + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime._incremental_fbank = incremental + runtime._maybe_enable_incremental_fbank(SimpleNamespace(_streaming_mel_processor=mel)) + assert (type(mel) is not remote.StreamingMelProcessorExact) is incremental + return mel + + +@pytest.mark.parametrize("chunks", ["stage0", "random"]) +def test_streams_are_bitwise_equal(remote, chunks: str) -> None: + # Stage 0's chunk sizes, or random ones (unaligned buffer ends); a 66 s stream slides the window twice. + rng = np.random.default_rng(7) + audio = (rng.standard_normal(66 * _SR) * 0.05).astype(np.float32) + audio[_SR * 3 : _SR * 7] *= 20 # loud stretch for the dynamic range + audio[_SR * 12 : _SR * 14] = 0.0 # digital silence + audio[_SR * 45 :] *= 1e-3 # quiet tail + ref, inc = _mel(remote, incremental=False), _mel(remote, incremental=True) + position, slides = 0, 0 + while position < len(audio): + size = ref.get_chunk_size() if chunks == "stage0" else int(rng.integers(1, 2 * _SR)) + dropped = ref.left_samples_dropped + expected, expected_info = ref.process(audio[position : position + size]) + got, got_info = inc.process(audio[position : position + size]) + assert torch.equal(got, expected) and got_info == expected_info, position + assert torch.equal(inc._extract_full(), remote.StreamingMelProcessorExact._extract_full(inc)), position + position += size + slides += ref.left_samples_dropped != dropped + assert slides >= 2 diff --git a/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling.py b/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling.py new file mode 100644 index 000000000..43de261c2 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stage-0 duplex sampling: batched rows against a row-by-row reference, the deferred path against both.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from vllm_omni.model_executor.models.minicpmo_4_5 import minicpmo_4_5_omni +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +Model = minicpmo_4_5_omni.MiniCPMO45OmniForConditionalGeneration +CHUNK_EOS, LISTEN, TTS_PAD, VOCAB = 5, 6, 10, 40 +_NAMES = ("unit", "chunk_eos", "listen", "tts_bos", "turn_eos", "chunk_tts_eos", "tts_pad") +TOKEN_IDS = {f"{name}_token_id": token for name, token in zip(_NAMES, (1, CHUNK_EOS, LISTEN, 7, 8, 9, TTS_PAD))} + + +class _CharTokenizer: + """Token t decodes to (t % 3) + 1 copies of one CJK character; special ids decode to "".""" + + clean_up_tokenization_spaces, bad_token_ids, all_special_ids = False, list[int](), sorted(TOKEN_IDS.values()) + + def __len__(self): + return VOCAB + + def decode(self, ids, skip_special_tokens=False): + return "".join("" if t in self.all_special_ids else chr(0x4E00 + t) * (t % 3 + 1) for t in map(int, ids)) + + def batch_decode(self, batch, skip_special_tokens=False): + return [self.decode(ids) for ids in batch] + + +def _setup(seed: int, *, all_greedy: bool = False, max_chars: int = 6, turns_open: bool = True): + """Six rows: sampled, listen in an open turn, speak-length cap, boundary chunk_eos, greedy, penalized.""" + torch.manual_seed(seed) + steps = [torch.randn(6, VOCAB) * 2 for _ in range(2)] + for logits in steps: + logits[3, CHUNK_EOS] += 50.0 + logits[1, LISTEN] += 40.0 + states = {row: SimpleNamespace(generated_tokens=[], current_turn_ended=True) for row in range(6)} + states[1].current_turn_ended, states[5].generated_tokens = not turns_open, [11, 11, 12] + model = Model.__new__(Model) + model.model_stage, model.max_new_speak_tokens_per_chunk, model.max_speak_chars_per_chunk = "llm", 5, max_chars + model._minicpmo45_native_duplex_token_ids_cache = dict(TOKEN_IDS) + model._minicpmo45_tokenizer_cache, model._minicpmo45_duplex_state_for_row = _CharTokenizer(), states.get + model._minicpmo45_duplex_payload_for_row = model._minicpmo45_duplex_row_request_max_tokens = lambda row: None + md = SimpleNamespace(all_greedy=all_greedy, top_p=torch.tensor([0.9, 1.0, 0.8, 0.8, 0.8, 0.85])) + md.generators = {row: torch.Generator().manual_seed(100 + row) for row in range(6)} + md.output_token_ids = [[11, 12], [2], [1, 2, 3, 4], [2, 3], [7, 2], [20, 21, 22]] + md.temperature, md.top_k = torch.tensor([0.8, 1.1, 0.7, 0.7, 0.0, 0.9]), torch.tensor([10, 0, 100, 100, 100, 20]) + return model, states, md, steps + + +def _reference(logits, md, states, row_params) -> list[int]: + """Row by row: speak-length cap, boundary draw, then the stage-2 draw (no cut, no listen rewrite).""" + out = [] + for row, (temperature, top_k, top_p) in enumerate(row_params): + generator, recent = md.generators[row], md.output_token_ids[row] + if len(recent) >= 4 or torch.rand((), generator=generator) < torch.softmax(logits[row], -1)[CHUNK_EOS]: + out.append(CHUNK_EOS) + continue + row_logits = logits[row : row + 1].clone() + row_logits[0, [TTS_PAD, CHUNK_EOS]] = float("-inf") + for token in set((states[row].generated_tokens or recent)[-MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE :]): + row_logits[0, token] /= 1.05 + if temperature <= 0: + out.append(int(row_logits.argmax())) + continue + scaled = row_logits / temperature + candidates = Model._duplex_top_k_top_p_candidates(scaled, top_k=top_k, top_p=top_p) + # The candidate-space distribution is the dense top-k / top-p filter's. + dense = torch.softmax(Model._top_k_top_p_filter(scaled, top_k=int(top_k), top_p=top_p), -1) + torch.testing.assert_close(_expand(scaled, candidates), dense) + uniform = torch.rand((), generator=generator) + out.append(int(Model._duplex_draw_candidates(scaled, candidates, uniform.view(1)))) + return out + + +def _expand(logits, candidates) -> torch.Tensor: + """The candidates as a dense distribution: the tie slot spread over its lowest-index tied tokens.""" + probs, ids, kth, ties = candidates + dense = torch.zeros(logits.shape, dtype=probs.dtype).scatter_(1, ids[:, :-1], probs[:, :-1]) + tied = logits == kth + kept = tied & (tied.cumsum(dim=-1) <= ties) + return dense + kept * (probs[:, -1:] / ties.clamp(min=1)) + + +def _run(model, md, steps) -> tuple[list[list[int]], list[bool]]: + """Two steps through the Stage-0 sampler; also whether each left a deferred commit.""" + tokens, deferred = [], [] + for logits in steps: + out = model._sample_minicpmo45_native_duplex_stage0(logits.clone(), md, duplex_rows=list(range(6))) + tokens.append(out.sampled_token_ids.squeeze(-1).tolist()) + deferred.append(getattr(model, "_minicpmo45_duplex_pending_samples", None) is not None) + model._commit_minicpmo45_duplex_pending_samples() + for row, token in enumerate(tokens[-1]): + md.output_token_ids[row].append(token) + return tokens, deferred + + +@pytest.mark.parametrize("seed", range(4)) +def test_batched_rows_match_the_row_by_row_reference(seed): + model, states, md, steps = _setup(seed, max_chars=10**6, turns_open=False) + _, ref_states, ref_md, _ = _setup(seed) + params = model._minicpmo45_duplex_row_params(md, 6) + kwargs = dict(row_idxs=list(range(6)), token_ids=TOKEN_IDS, row_params=params) + rows = model._sample_minicpmo45_native_duplex_rows(steps[0], md, **kwargs) + assert rows == _reference(steps[0], ref_md, ref_states, params) + assert all(torch.equal(md.generators[r].get_state(), ref_md.generators[r].get_state()) for r in range(6)) + + +@pytest.mark.parametrize("all_greedy", [False, True]) +@pytest.mark.parametrize("seed", range(3)) +def test_deferred_rows_match_the_synchronous_rows(seed, all_greedy, monkeypatch): + sync_model, sync_states, sync_md, steps = _setup(seed, all_greedy=all_greedy) + monkeypatch.setattr(sync_model, "_sample_minicpmo45_native_duplex_rows_deferred", lambda *a, **k: None) + expected, _ = _run(sync_model, sync_md, steps) + model, states, md, _ = _setup(seed, all_greedy=all_greedy) + assert _run(model, md, steps) == (expected, [True, True]) + assert {row: vars(s) for row, s in states.items()} == {row: vars(s) for row, s in sync_states.items()} + assert all(torch.equal(md.generators[r].get_state(), sync_md.generators[r].get_state()) for r in range(6)) + + +def _tied_logits(device: str) -> torch.Tensor: + """Rows whose k-th logit is tied well past k, with the nucleus cutting inside and before the tie.""" + logits = torch.full((4, 64), -3.0, device=device) + logits[0, [3, 9, 17, 30, 41, 50]] = 1.0 # six tied for k=2 + logits[1, 7], logits[1, [2, 12, 22, 32, 42]] = 4.0, 1.0 # a leader, then five tied for k=3 + logits[2] = torch.randn(64, generator=torch.Generator().manual_seed(0)).to(device) + logits[3, ::4] = 0.5 # sixteen tied for k=4 + return logits + + +@pytest.mark.parametrize("top_p", [1.0, 0.9, 0.5, 0.2]) +@pytest.mark.parametrize("top_k", [2, 3, 4, 0]) +def test_candidate_draws_follow_the_dense_distribution(top_k, top_p): + # A midpoint grid of uniforms integrates the inverse CDF exactly: each token + # is drawn in proportion to its dense probability, ties included. + logits, n = _tied_logits("cpu"), 4096 + expected = torch.softmax(Model._top_k_top_p_filter(logits, top_k=top_k, top_p=top_p), dim=-1) + rows = logits.repeat(n, 1) + grid = ((torch.arange(n, dtype=torch.float32) + 0.5) / n).repeat_interleave(logits.shape[0]) + draws = Model._duplex_draw_candidates( + rows, Model._duplex_top_k_top_p_candidates(rows, top_k=top_k, top_p=top_p), grid + ) + counts = torch.zeros_like(expected).index_put_( + (torch.arange(rows.shape[0]) % logits.shape[0], draws), torch.ones(rows.shape[0]), accumulate=True + ) + torch.testing.assert_close(counts / n, expected, atol=2.0 / n, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("top_p", [0.2, 0.25, 0.5, 0.8, 1.0]) +@pytest.mark.parametrize("n", [4, 10, 20, 64, 100]) +def test_tied_nucleus_boundaries_match_the_dense_filter(n, top_p, dtype): + # Uniform ties put every top_p on (or next to) a candidate boundary. + logits = torch.full((3, n + 8), -30.0) + logits[0, :n], logits[1, 3 : 3 + n], logits[2, :2], logits[2, 4 : 4 + n] = 20.0, 0.5, 2.0, 1.0 + logits = logits.to(dtype) + # The candidates compute in float32; bf16 widens exactly, so ties survive. + expected = torch.softmax(Model._top_k_top_p_filter(logits.float(), top_k=3, top_p=top_p), dim=-1) + actual = _expand(logits, Model._duplex_top_k_top_p_candidates(logits, top_k=3, top_p=top_p)) + assert torch.equal(actual > 0, expected > 0) + torch.testing.assert_close(actual, expected) + + +class _NoHostReads(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if func in (torch.ops.aten._local_scalar_dense.default, torch.ops.aten.is_nonzero.default): + raise AssertionError(f"host read: {func}") + return func(*args, **(kwargs or {})) + + +@pytest.mark.parametrize( + "device", + ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"))], +) +def test_candidate_sampling_reads_nothing_back_to_the_host(device): + logits = _tied_logits(device) + with _NoHostReads(): + for top_k, top_p in [(2, 0.5), (4, 1.0), (0, 0.8)]: + candidates = Model._duplex_top_k_top_p_candidates(logits, top_k=top_k, top_p=top_p) + Model._duplex_draw_candidates(logits, candidates, torch.rand(logits.shape[0], device=device)) diff --git a/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_batch.py b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_batch.py new file mode 100644 index 000000000..4414a1e47 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_batch.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The cross-session batched streaming audio encoder against the per-session path.""" + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from torch import nn +from transformers.models.whisper.modeling_whisper import WhisperConfig + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import MiniCPMO45OmniForConditionalGeneration +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, + MiniCPMWhisperEncoder, + MultiModalProjector, +) +from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import StreamingAudioChunk + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +N_MELS, POOL = 16, 5 +# Per round, each session's unit frames (None: sits the round out): 22 / 24 frames mirror 1020 / 1040 ms +# at 1/10 scale, odd lengths pad conv2, 25 frames leave the pooler a remainder, and 40 positions force resets. +_SCHEDULE: list[list[int | None]] = [ + [22, 22, 23, 21], + [24, None, 24, 25], + [24, 24, None, 24], + [24, 23, 24, 24], + [24, 24, 24, None], + [25, 24, 23, 24], +] + + +def _thinker() -> Any: + """The audio half of the thinker: real modules and methods, toy sizes.""" + torch.manual_seed(0) + config = WhisperConfig(num_mel_bins=N_MELS, d_model=32, encoder_layers=2, encoder_attention_heads=2) + config.update(dict(encoder_ffn_dim=128, max_source_positions=40, dropout=0.0, attention_dropout=0.0)) + config._attn_implementation = "sdpa" + thinker = MiniCPMO45OmniLLMForConditionalGeneration.__new__(MiniCPMO45OmniLLMForConditionalGeneration) + nn.Module.__init__(thinker) + thinker.apm = MiniCPMWhisperEncoder(config).eval() + thinker.audio_projection_layer = MultiModalProjector(in_dim=32, out_dim=24) # takes ffn // 4, as the checkpoint + thinker.audio_avg_pooler = nn.AvgPool1d(POOL, stride=POOL) + thinker.audio_encoder_layer, thinker.audio_past_key_values = -1, None + thinker.config = SimpleNamespace(audio_pool_step=POOL, duplex_audio_kv_page_positions=16) + return thinker + + +def test_batched_rounds_match_sequential_sessions() -> None: + thinker = _thinker() + legacy: list[Any] = [None] * 4 + batched: list[Any] = [None] * 4 + resets = 0 + for round_index, frames_per_session in enumerate(_SCHEDULE): + chunks, active, expected = [], [], [] + for session, frames in enumerate(frames_per_session): + if frames is None: + continue + mel = torch.randn((1, N_MELS, frames), generator=torch.Generator().manual_seed(100 * round_index + session)) + extra = dict(prefix_extra_frames=0 if legacy[session] is None else 2, suffix_extra_frames=2) + features = {"audio_features": mel, "audio_feature_lens": [torch.tensor([frames])]} + thinker.audio_past_key_values = legacy[session] + with torch.no_grad(): + nested = thinker.get_audio_embedding_streaming(features, use_extra_context=True, **extra) + legacy[session] = thinker.audio_past_key_values + expected.append(torch.cat([t for row in nested for t in row])) + chunks.append(StreamingAudioChunk(mel, batched[session], **extra)) + active.append(session) + idle = {s: (cache, cache.length if cache else 0) for s, cache in enumerate(batched) if s not in active} + outputs, caches = thinker.get_audio_embedding_streaming_batch(chunks) + for session, output, cache, reference in zip(active, outputs, caches, expected, strict=True): + torch.testing.assert_close(output, reference, rtol=1e-5, atol=1e-5) + resets += batched[session] is not None and cache is not batched[session] + batched[session] = cache + want = legacy[session].self_attention_cache + assert cache.length == want.get_seq_length() + for got, ref in zip(cache.to_legacy_cache().self_attention_cache.layers, want.layers, strict=True): + torch.testing.assert_close((got.keys, got.values), (ref.keys, ref.values)) + # Not in the batch: the same cache object, the same committed length. + assert all(batched[s] is c and (c.length if c else 0) == n for s, (c, n) in idle.items()) + assert resets >= 2, "the schedule must cross the max_source_positions reset" + + +def test_graph_build_honors_eager_without_engine_config() -> None: + # The thinker keeps no vllm_config: the gate reads the flag __init__ stored. + thinker = _thinker() + thinker._enforce_eager = True + assert thinker.build_streaming_audio_graph_encoder(unit_frames=24) is False + assert thinker._duplex_audio_cuda_graph_encoder is None + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_graph_build_captures_on_cuda() -> None: + thinker = _thinker().cuda() + thinker._enforce_eager = False + thinker.config.duplex_audio_encoder_cuda_graph_batch_sizes = [1, 2] + thinker.config.duplex_audio_encoder_cuda_graph_cache_buckets = [8] + assert thinker.build_streaming_audio_graph_encoder(unit_frames=24) is True + assert thinker._duplex_audio_cuda_graph_encoder is not None + + +@pytest.mark.parametrize( + ("session_mode", "device", "captures"), + [("duplex", "cuda", 1), ("turn", "cuda", 0), ("duplex", "cpu", 0)], +) +def test_duplex_audio_graphs_capture_while_the_model_loads(monkeypatch, session_mode, device, captures) -> None: + # Captured in load_weights, the graphs count against the KV budget instead of past it. + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + nn.Module.__init__(model) + model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode=session_mode)) + model.thinker = nn.Linear(1, 1) + calls: list[None] = [] + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(build_audio_cuda_graph=lambda: calls.append(None)) + monkeypatch.setattr( + MiniCPMO45OmniForConditionalGeneration, "_module_device", staticmethod(lambda module: torch.device(device)) + ) + model._capture_duplex_audio_graphs() + assert len(calls) == captures diff --git a/tests/model_executor/models/minicpmo_4_5/test_vision_fused.py b/tests/model_executor/models/minicpmo_4_5/test_vision_fused.py new file mode 100644 index 000000000..ae2755a5c --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_vision_fused.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused packed SigLIP layers match the eager packed encode; prefetched frames serve only their own append.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5 import vision_fused +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import MiniCPMO45Stage0DuplexRuntime +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + SiglipVisionConfig, + SiglipVisionTransformer, + _vision_encode_paths, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_fused_layers_match_eager_packed_layers() -> None: + torch.manual_seed(0) + sizes = dict(hidden_size=32, intermediate_size=72, num_hidden_layers=3, num_attention_heads=4) + config = SiglipVisionConfig(**sizes, image_size=28, patch_size=2, attention_dropout=0.0) + config._attn_implementation = "sdpa" + vpm = SiglipVisionTransformer(config).eval() + for norm in [module for module in vpm.modules() if isinstance(module, torch.nn.LayerNorm)]: + torch.nn.init.normal_(norm.weight, mean=1.0, std=0.2) + torch.nn.init.normal_(norm.bias, std=0.2) + seq_groups, hidden = [(0, 3, 16), (48, 2, 15)], torch.randn((78, 32), generator=torch.Generator().manual_seed(1)) + with torch.inference_mode(): + expected = hidden.clone() + for layer in vpm.encoder.layers: + expected = layer.forward_packed(expected, seq_groups) + fused = vision_fused.encode_packed_fused(vpm, hidden.clone(), seq_groups) + torch.testing.assert_close(fused, vpm.post_layernorm(expected), rtol=1e-5, atol=1e-5) + + +def test_prefetched_frames_are_chunked_and_keyed_by_append() -> None: + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.sessions, runtime._prefetched_vision = {}, {} + runtime.processor = SimpleNamespace(process_image=lambda *args, **kwargs: None) + runtime.stage_model = runtime.thinker = SimpleNamespace(config=SimpleNamespace(vision_batch_size=2)) + runtime._decode_video_frames_payload = lambda payload: list(payload["video_frames"]) + runtime._preprocess_frames = lambda frames: [f"p{frame}" for frame in frames] + # Each encoded frame is tagged with its tower call's size: vision_batch_size appends per call. + runtime._encode_processed_vision_batch = lambda items: [[f"e{item}/{len(items)}"] for item in items] + frames = ["a", "b", "c"] + appends = [ + {"session_id": f"s{i}", "epoch": 0, "seq": 1, "payload": {"video_frames": frames[: i + 1]}} for i in range(3) + ] + runtime.prefetch_vision(appends) + assert runtime.take_prefetched_vision({**appends[1], "payload": {"video_frames": ["x", "y"]}}) is None # changed + assert runtime.take_prefetched_vision(appends[1]) is None # and the entry is gone + assert runtime.take_prefetched_vision({**appends[2], "seq": "2"}) is None # another append + last, payload = appends[2], appends[2]["payload"] + assert runtime.frame_kwargs(last, payload) == {"encoded_frames": [["epa/3"], ["epb/3"], ["epc/3"]]} + assert runtime.frame_kwargs(last, payload) == {"video_frames": frames} # taken once + + +def test_packed_vision_adapter_protocol() -> None: + shape, out_shape = (1, 3, 16), torch.Size([1, 4, 8]) + adapter = vision_fused._PackedVisionAdapter(lambda px, runs: px[:, :4, :8], [(2, 2, 1)], shape, out_shape) + assert adapter.supports_encoder_cudagraph is True + cfg = adapter.get_encoder_cudagraph_config() + assert cfg.modalities == ["image"] + assert cfg.buffer_keys == ["pixels"] + assert cfg.out_hidden_size == 8 + inputs = adapter.prepare_encoder_cudagraph_capture_inputs(4, 1, 0, torch.device("cpu"), torch.float32) + assert "pixels" in inputs.values + dest: list[torch.Tensor | None] = [None] + adapter.postprocess_encoder_output({"default": torch.ones(1, 4, 8)}, [0], [4], dest, clone=True) + assert dest[0] is not None and dest[0].shape == (1, 4, 8) + + +@pytest.mark.parametrize( + ("overrides", "encoder_graphs", "expected"), + [ + ({}, True, (False, False)), # default: the padded batch keeps the vpm encoder graph + ({"vision_cuda_graph": True}, True, (True, True)), + ({"vision_cuda_graph": True}, False, (False, False)), # encoder_cuda_graph=False / --enforce-eager + ({"vision_packed_encode": True}, True, (True, False)), # packed eager, for A/B + ({"vision_packed_encode": False, "vision_cuda_graph": True}, True, (False, False)), + ], +) +def test_vision_encode_paths(overrides, encoder_graphs, expected) -> None: + assert _vision_encode_paths(SimpleNamespace(**overrides), encoder_graphs=encoder_graphs) == expected diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index 65fb52ad0..d8b6d29ca 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/tests/worker/test_native_duplex_hooks.py @@ -415,7 +415,8 @@ def test_minicpmo_non_duplex_sample_skips_duplex_row_scan(): ) assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler" - assert calls == [] + # The empty row set is still published: the model then skips its per-row prompt scan. + assert calls == ["prepare"] def test_minicpmo_duplex_sample_clears_stale_rows_once_without_scanning(): @@ -442,7 +443,7 @@ def test_minicpmo_duplex_sample_clears_stale_rows_once_without_scanning(): assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler" assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler" - assert calls == [()] + assert calls == [(), ()] def test_generic_ar_runner_has_no_minicpmo_sampler_state_or_typeerror_probe(): @@ -2890,6 +2891,7 @@ def test_minicpmo_stage0_native_sampler_forced_listen_yields_floor(): ) sampled = model.sample(logits, sampling_metadata) + model._commit_minicpmo45_duplex_pending_samples() # the runner commits before the next step assert sampled is not None assert sampled.sampled_token_ids.tolist() == [[151705]] @@ -3155,6 +3157,7 @@ def test_minicpmo_stage0_turn_eos_is_forwarded_not_pending(): logits[0, 151717] = 30.0 sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201])) + model._commit_minicpmo45_duplex_pending_samples() # the runner commits before the next step assert sampled.sampled_token_ids.tolist() == [[151717]] assert state.pending_terminator_token is None @@ -3166,6 +3169,7 @@ def test_minicpmo_stage0_turn_eos_is_forwarded_not_pending(): logits = torch.full((1, 151723), -100.0) logits[0, 151718] = 30.0 sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201, 151717])) + model._commit_minicpmo45_duplex_pending_samples() assert sampled.sampled_token_ids.tolist() == [[151718]] assert state.pending_terminator_token == 151718 diff --git a/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml b/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml index dbabc9172..4c71b93f5 100644 --- a/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml @@ -1,6 +1,6 @@ # 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. +# Enables all Stage 0 (multimodal duplex prefill) and Stage 2 (code2wav DiT) optimizations. # Ordinary TF32 for CFM DiT GEMMs (cuBLAS allow_tf32 + Triton tf32 dots). # HiFT stays IEEE FP32. Validate audio quality for your workload. base_config: minicpmo_4_5.yaml @@ -9,3 +9,28 @@ connectors: connector_of_shared_memory: extra: token2wav_allow_tf32: true + # Stage 2: Fused DiT transformer body (LayerNorm + QKV + Residual) + cfm_fused_body: true + # Stage 2: Flow condition encoder CUDA graph + cfm_encoder_cuda_graph: true + # Stage 2: Pre-allocated persistent slot pool and row offset merging + cfm_slot_pool: true + cfm_row_offset_merge: true + # Stage 2: Compress DiT attention cache to bf16 + code2wav_bfloat16_attention_cache: true + +stages: + - stage_id: 0 + enforce_eager: false + hf_overrides: + # Stage 0: Incremental sliding window Fbank feature extraction + duplex_incremental_fbank: true + # Stage 0: Fused packed SigLIP encoder layers + vision_fused_layers: true + # Stage 0: Packed SigLIP tower & resampler CUDA graph + vision_cuda_graph: true + # Stage 0: Steady Whisper streaming audio encoder CUDA graph + duplex_audio_encoder_cuda_graph: true + duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions: true + # Stage 0: Pinned memory async H2D copy for streaming audio chunks + duplex_audio_encoder_pinned_h2d: true diff --git a/vllm_omni/model_executor/duplex_sampling.py b/vllm_omni/model_executor/duplex_sampling.py index dd48431f6..5546ae27c 100644 --- a/vllm_omni/model_executor/duplex_sampling.py +++ b/vllm_omni/model_executor/duplex_sampling.py @@ -22,6 +22,20 @@ class DuplexSamplingRow: seq: int | None payload: dict[str, object] | None max_tokens: int | None + # Host copies of the row's sampling parameters (None when the runner has none). + temperature: float | None = None + top_k: int | None = None + top_p: float | None = None + + +def _host_value(values: Any, row_idx: int, cast: type) -> Any: + """``values[row_idx]`` of a host array as ``cast``, or None when there is none.""" + if values is None: + return None + try: + return cast(values[row_idx]) + except (IndexError, TypeError, ValueError): + return None class DuplexSamplingHelper: @@ -53,15 +67,19 @@ class DuplexSamplingHelper: else: self.active_request_ids.discard(req_id) - def update_states(self, runner: object, scheduler_output: object) -> None: + def update_states(self, runner: Any, scheduler_output: Any) -> None: self.active_request_ids.difference_update(str(req_id) for req_id in scheduler_output.finished_req_ids) for request in scheduler_output.scheduled_new_reqs: self.refresh_active_request(runner, str(request.req_id)) - def rows(self, runner: object) -> tuple[DuplexSamplingRow, ...]: + def rows(self, runner: Any) -> tuple[DuplexSamplingRow, ...]: rows: list[DuplexSamplingRow] = [] - req_ids = [str(req_id) for req_id in getattr(runner.input_batch, "req_ids", [])] + input_batch = runner.input_batch + req_ids = [str(req_id) for req_id in getattr(input_batch, "req_ids", [])] requests = getattr(runner, "requests", {}) + temperature_cpu = getattr(input_batch, "temperature_cpu", None) + top_k_cpu = getattr(input_batch, "top_k_cpu", None) + top_p_cpu = getattr(input_batch, "top_p_cpu", None) for row_idx, req_id in enumerate(req_ids): if req_id not in self.active_request_ids: continue @@ -73,7 +91,8 @@ class DuplexSamplingHelper: if not isinstance(session_id, str) or not session_id: session_id = None try: - seq = int(duplex.get("seq")) + seq_raw = duplex.get("seq") + seq = int(seq_raw) if seq_raw is not None else None except (TypeError, ValueError): seq = None payload = duplex.get("payload") @@ -93,6 +112,9 @@ class DuplexSamplingHelper: seq=seq, payload=payload, max_tokens=max_tokens if max_tokens > 0 else None, + temperature=_host_value(temperature_cpu, row_idx, float), + top_k=_host_value(top_k_cpu, row_idx, int), + top_p=_host_value(top_p_cpu, row_idx, float), ) ) return tuple(rows) @@ -142,8 +164,9 @@ class DuplexSamplingRunnerMixin: return helper = getattr(self, "_duplex_sampling_helper", None) rows = helper.rows(self) if helper is not None and helper.active_request_ids else () - if rows or (helper is not None and helper.hook_active): - prepare_duplex_sampling(logits, prepared_sampling_metadata, rows) + # Publish the row set even when empty, so the model skips its per-row + # prompt-token scan (a host sync per row) in turn-based serving. + prepare_duplex_sampling(logits, prepared_sampling_metadata, rows) if helper is not None: helper.hook_active = bool(rows) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/incremental_fbank.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/incremental_fbank.py new file mode 100644 index 000000000..6538b2d7e --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/incremental_fbank.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Incremental Stage-0 streaming fbank (HF override ``duplex_incremental_fbank``, default off). + +The remote ``StreamingMelProcessorExact._extract_full`` recomputes the CPU Whisper log-mel over the whole +streaming buffer (up to its 30 s slide) on every unit. This keeps the ``|STFT|^2`` rows already computed and +recomputes only the right edge (frames whose window reaches new samples or the reflect pad) and, after a slide, +the new left edge. The mel projection, log, dynamic range and scaling still run over the whole buffer with the +original's shapes and memory layout, so the features stay bitwise equal: the mel GEMM is not row independent +(MKL's remainder kernel), and an STFT batch's last frames may take a remainder path, so recomputes start on +``_ALIGN_FRAMES`` boundaries and their last ``_ALIGN_FRAMES`` frames are not kept. +""" + +from __future__ import annotations + +import copy +from typing import Any + +import numpy as np +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# Recompute start granularity, trailing recompute frames not kept, and the +# left-edge recompute after a slide (>= the 2 frames reading the pad). +_ALIGN_FRAMES = 16 +_FEATURE_EXTRACTOR_ATTRS = ("n_fft", "hop_length", "n_samples", "sampling_rate", "mel_filters", "set_spac_log_norm") +_FEATURE_EXTRACTOR_ATTRS += ("dynamic_log_norm", "dynamic_range_db", "log_floor_db") +_INCREMENTAL_CLASSES: dict[type, type] = {} +_SELF_CHECKED: dict[tuple[type, type], bool] = {} + + +def _stft_power(buffer: np.ndarray, start: int, stop: int, n_fft: int, hop: int, window: torch.Tensor) -> np.ndarray: + """``(stop - start, n_fft // 2 + 1)`` rows of ``|STFT|^2``, framed as the original's centered STFT. + + The batch is ``[start, stop]`` minus its last frame, like the original's + ``stft(...)[..., :-1].abs() ** 2``; edge reflect padding is built explicitly. + """ + half = n_fft // 2 + lo, hi = start * hop - half, stop * hop + half + length = len(buffer) + segment = buffer[max(lo, 0) : min(hi, length)] + if lo < 0 or hi > length: + segment = np.pad(segment, (max(0, -lo), max(0, hi - length)), mode="reflect") + spec = torch.stft( + torch.from_numpy(segment)[None], n_fft=n_fft, hop_length=hop, window=window, center=False, return_complex=True + ) + magnitudes = spec[..., :-1].abs() ** 2 # [1, n_freq, stop - start], frame-major in memory + return magnitudes[0].T.numpy() + + +def _log_mel(power: torch.Tensor, mel_filters: torch.Tensor, feature_extractor: Any) -> torch.Tensor: + """``_torch_extract_fbank_features`` from the ``|STFT|^2`` rows, op for op and layout for layout.""" + magnitudes = power.unsqueeze(0).transpose(1, 2) # [1, n_freq, T], strides of stft(...)[..., :-1] + log_spec = torch.clamp(mel_filters.T @ magnitudes, min=1e-10).log10() + if feature_extractor.dynamic_log_norm: + max_val_bt = log_spec.max(dim=2, keepdim=True)[0].max(dim=1, keepdim=True)[0] + log_spec = torch.maximum(log_spec, max_val_bt - feature_extractor.dynamic_range_db) + else: + log_spec = torch.maximum(log_spec, torch.tensor(feature_extractor.log_floor_db, dtype=log_spec.dtype)) + return (log_spec + 4.0) / 4.0 + + +class _IncrementalFbankMixin: + """Overrides of the remote ``StreamingMelProcessorExact``, mixed in by ``enable_incremental_fbank``.""" + + #: ``|STFT|^2`` rows (frame t at row t); rows ``[0, _ifb_valid)`` are exact for the current buffer. + _ifb_power: np.ndarray | None = None + _ifb_valid = 0 + #: ``left_samples_dropped`` and the buffer length when the rows were kept. + _ifb_dropped = 0 + _ifb_length = 0 + #: The ``n_fft`` samples ending where the last kept frame's window ends. + _ifb_guard: np.ndarray | None = None + _ifb_window: torch.Tensor | None = None + _ifb_mel_filters: torch.Tensor | None = None + + def reset(self) -> None: + super().reset() + self._ifb_valid = 0 + + def restore_snapshot(self, snapshot: dict) -> None: + super().restore_snapshot(snapshot) + self._ifb_valid = 0 + + def _extract_full(self) -> torch.Tensor: + fe, buffer = self.feature_extractor, self.buffer + if not ( + isinstance(buffer, np.ndarray) + and buffer.ndim == 1 + and buffer.dtype == np.float32 + and fe.n_fft <= len(buffer) <= fe.n_samples + and getattr(fe, "dither", 0.0) == 0.0 + and fe.sampling_rate == self.sample_rate + ): + # Too short (the original raises), truncated, dithered or misconfigured: the original. + self._ifb_valid = 0 + return super()._extract_full() + # The original's log-floor switch, side effect on the feature extractor included. + if len(buffer) < 5 * self.sample_rate: + fe.set_spac_log_norm(log_floor_db=-10) + else: + fe.set_spac_log_norm(dynamic_range_db=8) + return _log_mel(self._ifb_update_power(buffer, fe), self._ifb_mel_filters, fe) + + def _ifb_update_power(self, buffer: np.ndarray, fe: Any) -> torch.Tensor: + n_fft, hop = int(fe.n_fft), int(fe.hop_length) + if self._ifb_window is None: + self._ifb_window = torch.hann_window(n_fft) + self._ifb_mel_filters = torch.from_numpy(fe.mel_filters).to(torch.float32) + power = self._ifb_power + if power is None or power.ctypes.data % 64: + # Rows from a torch allocation: 64-byte aligned like the original's fresh + # |X|^2 tensor, which the GEMM reads (a deepcopy's numpy array may not be). + aligned = torch.empty(int(fe.n_samples) // hop + 1, n_fft // 2 + 1, dtype=torch.float32).numpy() + if power is not None: + aligned[: self._ifb_valid] = power[: self._ifb_valid] + self._ifb_power = power = aligned + frames = len(buffer) // hop # the original drops the STFT's last frame + start = self._ifb_rebase(buffer, n_fft, hop) // _ALIGN_FRAMES * _ALIGN_FRAMES + power[start:frames] = _stft_power(buffer, start, frames, n_fft, hop, self._ifb_window) + # Keep frames whose window lies inside the buffer, minus the recompute's tail. + valid = min((len(buffer) - n_fft // 2) // hop + 1, frames - _ALIGN_FRAMES) + guard_end = (valid - 1) * hop + n_fft // 2 + if guard_end < n_fft: + valid = 0 + else: + self._ifb_guard = buffer[guard_end - n_fft : guard_end].copy() + self._ifb_valid = valid + self._ifb_dropped = int(self.left_samples_dropped) + self._ifb_length = len(buffer) + return torch.from_numpy(power[:frames]) + + def _ifb_rebase(self, buffer: np.ndarray, n_fft: int, hop: int) -> int: + """Rows still exact for ``buffer``: the kept rows (moved down after a slide), or 0.""" + valid = self._ifb_valid + if valid == 0: + return 0 + shift = int(self.left_samples_dropped) - self._ifb_dropped + guard_end = (valid - 1) * hop + n_fft // 2 - shift + if ( + shift < 0 + or shift % hop + or len(buffer) + shift < self._ifb_length + or guard_end < n_fft + or not np.array_equal(buffer[guard_end - n_fft : guard_end], self._ifb_guard) + ): + return 0 # not the same stream (or not a hop-aligned slide of it): start over + if shift == 0: + return valid + drop = shift // hop + valid -= drop + if valid <= _ALIGN_FRAMES: + return 0 + power = self._ifb_power + power[:valid] = power[drop : drop + valid] + # The new left edge's frames read the reflect pad. + power[:_ALIGN_FRAMES] = _stft_power(buffer, 0, _ALIGN_FRAMES, n_fft, hop, self._ifb_window) + return valid + + +def _self_check(mel_processor: Any, incremental_cls: type) -> bool: # Any: remote-code processor + """Compare the incremental and the original features on a probe stream (with a slide) once per class.""" + fe = mel_processor.feature_extractor + saved = {name: getattr(fe, name) for name in ("dynamic_log_norm", "dynamic_range_db", "log_floor_db")} + probe = copy.copy(mel_processor) + probe.__class__ = incremental_cls + hop = int(fe.hop_length) + audio = (0.1 * np.random.default_rng(0).standard_normal(8 * int(mel_processor.sample_rate))).astype(np.float32) + try: + # Below and above the 5 s dynamic-range switch, then a slide of 100 hops. + for dropped, end in ((0, len(audio) // 3 + 37), (0, len(audio) - 11), (100 * hop, len(audio))): + probe.buffer, probe.left_samples_dropped = audio[dropped:end], dropped + expected = type(mel_processor)._extract_full(probe) + if not torch.equal(incremental_cls._extract_full(probe), expected): + return False + return True + except Exception: + logger.warning("MiniCPM-o incremental fbank self-check raised", exc_info=True) + return False + finally: + for name, value in saved.items(): + setattr(fe, name, value) + + +def enable_incremental_fbank(mel_processor: Any) -> bool: # Any: remote-code StreamingMelProcessorExact + """Switch one session's streaming mel processor to the incremental ``_extract_full`` (idempotent). + + The instance's class becomes a subclass of its remote class with the mixin + in front, so copies keep it. Returns False, leaving the full recompute, for + a processor of another shape or one that differs on the probe stream. + """ + if isinstance(mel_processor, _IncrementalFbankMixin): + return True + cls = type(mel_processor) + fe = getattr(mel_processor, "feature_extractor", None) + if not callable(getattr(cls, "_extract_full", None)) or not all( + hasattr(fe, name) for name in _FEATURE_EXTRACTOR_ATTRS + ): + return False + if not all(hasattr(mel_processor, name) for name in ("buffer", "sample_rate", "left_samples_dropped")): + return False + incremental_cls = _INCREMENTAL_CLASSES.get(cls) + if incremental_cls is None: + incremental_cls = type(f"Incremental{cls.__name__}", (_IncrementalFbankMixin, cls), {"__module__": __name__}) + _INCREMENTAL_CLASSES[cls] = incremental_cls + key = (cls, type(fe)) + if key not in _SELF_CHECKED: + _SELF_CHECKED[key] = _self_check(mel_processor, incremental_cls) + if not _SELF_CHECKED[key]: + logger.warning( + "MiniCPM-o incremental fbank differs from %s._extract_full on the probe stream; " + "the streaming fbank stays on the full recompute", + cls.__name__, + ) + if not _SELF_CHECKED[key]: + return False + mel_processor.__class__ = incremental_cls + return True diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py index c65c8dadb..3eacd802a 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py @@ -5,15 +5,19 @@ from __future__ import annotations import base64 import copy +import os import time -from collections.abc import Callable +from collections.abc import Callable, Generator +from concurrent.futures import ThreadPoolExecutor from contextlib import suppress from dataclasses import dataclass, field from threading import Lock +from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast import numpy as np from numpy.typing import NDArray +from vllm.logger import init_logger from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy @@ -22,9 +26,32 @@ if TYPE_CHECKING: from PIL import Image from torch import nn +logger = init_logger(__name__) + _MINICPMO45_SPECIAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS _MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS _MINICPMO45_PROCESSOR_LOAD_LOCK = Lock() +_MINICPMO45_FRAME_POOL_LOCK = Lock() +_MINICPMO45_FRAME_POOL: ThreadPoolExecutor | None = None + + +def _optional_int(value: Any) -> int | None: + try: + return int(value) if value is not None else None + except (TypeError, ValueError): + return None + + +def _frame_preprocess_pool() -> ThreadPoolExecutor | None: + """Shared pool for camera-frame JPEG decode and slicing (both release the GIL); ``None`` on one CPU.""" + global _MINICPMO45_FRAME_POOL + workers = min(4, os.cpu_count() or 1) + if workers <= 1: + return None + with _MINICPMO45_FRAME_POOL_LOCK: + if _MINICPMO45_FRAME_POOL is None: + _MINICPMO45_FRAME_POOL = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="minicpmo45-frames") + return _MINICPMO45_FRAME_POOL @dataclass @@ -55,8 +82,13 @@ class _MiniCPMO45Stage0SessionState: prepared_inputs_embeds: torch.Tensor | None = None prepared_input_token_ids: list[int] = field(default_factory=list) prepared_result: dict[str, object] = field(default_factory=dict) - #: Opaque transformers cache, handed back to the model unchanged. + #: Streaming-encoder state: a ``StreamingAudioKVCache`` on the batched + #: path, else the opaque transformers cache handed back to the model. audio_past_key_values: object | None = None + #: Append identity of the last prefill that did work (successful or not). + attempted_append_identity: tuple[int | None, int] | None = None + #: This step's batched prefill result, taken by the append's own preprocess. + staged_prefill: tuple[tuple[int | None, int], dict[str, object]] | None = None pending_terminator_token: int | None = None last_terminator_token: int | None = None pending_speech_context: bool = False @@ -65,6 +97,18 @@ class _MiniCPMO45Stage0SessionState: generated_tokens: list[int] = field(default_factory=list) +@dataclass(frozen=True, slots=True) +class _MiniCPMO45AudioEncodeRequest: + """One streaming unit a prefill needs encoded; ``_stage_prefill_steps`` yields it.""" + + #: Any: transformers BatchFeature (or the dict fallback of _process_streaming_audio). + batch_feature: Any + state: _MiniCPMO45Stage0SessionState | None + + +_PrefillSteps = Generator[_MiniCPMO45AudioEncodeRequest, "torch.Tensor | None", dict[str, object]] + + class MiniCPMO45Stage0DuplexRuntime: """Build scheduler-owned MiniCPM-o 4.5 Stage0 duplex inputs.""" @@ -87,6 +131,8 @@ class MiniCPMO45Stage0DuplexRuntime: self.model_path = model_path self.device = device self.sessions: dict[tuple[str, int], _MiniCPMO45Stage0SessionState] = {} + # (session_id, epoch, seq) -> (frame payload, per-frame embeddings) of this runner step. + self._prefetched_vision: dict[tuple[str, int | None, int], tuple[tuple[object, ...], list[Any]]] = {} self.thinker = getattr(stage_model, "thinker", None) or getattr(stage_model, "model", None) or stage_model self.processor = ( getattr(stage_model, "processor", None) @@ -101,6 +147,49 @@ class MiniCPMO45Stage0DuplexRuntime: self._init_token_ids() if self.tokenizer is not None: self._require_special_token_ids() + # ``duplex_incremental_fbank`` (HF override, default off): see incremental_fbank.py. + self._incremental_fbank = ( + getattr(getattr(stage_model, "config", None), "duplex_incremental_fbank", False) is True + ) + # Built with the model, before its weights load and it switches to eval mode: + # the streaming encoder's CUDA graphs are captured later, by build_audio_cuda_graph(). + self._audio_graph_attempted = False + + def build_audio_cuda_graph(self) -> bool: + """Capture the streaming-encoder CUDA graphs once weights are loaded (idempotent; failures stay eager). + + The steady unit's mel-frame count comes from a probe processor configured like a session's: the + shared processor's default chunk is no model unit, and a first unit has its own size. + """ + thinker = getattr(self, "thinker", None) + config = getattr(thinker, "config", None) + attempted = getattr(self, "_audio_graph_attempted", False) + self._audio_graph_attempted = True + if attempted or not bool(getattr(config, "duplex_audio_encoder_cuda_graph", True)): + return False + apm = getattr(self.thinker, "apm", None) + if getattr(apm, "training", False) is True and callable(getattr(apm, "eval", None)): + apm.eval() # inference only; the loader's model.eval() covers it as well + build = getattr(self.thinker, "build_streaming_audio_graph_encoder", None) + if not callable(build) or self.processor is None: + return False + try: + probe = self._configure_streaming_processor(SimpleNamespace(streaming_processor=None)) + if probe is None: + return False + for _ in range(2): # the first unit, then the steady one + silence = np.zeros(self._streaming_chunk_size(probe), dtype=np.float32) + feature = probe.process_audio_streaming(silence, reset=False) + # A BatchFeature-like mapping or a (features, metadata) tuple, by processor version. + frames = int((feature["audio_features"] if isinstance(feature, dict) else feature[0]).shape[-1]) + except Exception: + logger.info( + "MiniCPM-o could not derive the steady streaming-audio unit shape; " + "the streaming audio encoder stays eager", + exc_info=True, + ) + return False + return bool(build(unit_frames=frames)) def _stage_runtime_ready(self) -> bool: return self.processor is not None and self.tokenizer is not None and self.thinker is not None @@ -138,7 +227,7 @@ class MiniCPMO45Stage0DuplexRuntime: reset_streaming = getattr(processor, "reset_streaming", None) if callable(reset_streaming): reset_streaming() - return processor + return self._maybe_enable_incremental_fbank(processor) configure_streaming = getattr(processor, "configure_streaming", None) if callable(configure_streaming): configure_streaming( @@ -147,6 +236,19 @@ class MiniCPMO45Stage0DuplexRuntime: slide_trigger_seconds=30.0, slide_stride_seconds=10.0, ) + return self._maybe_enable_incremental_fbank(processor) + + def _maybe_enable_incremental_fbank(self, processor: Any) -> Any: # Any: remote-code processor + """Put the processor's streaming mel front end on the incremental path (``duplex_incremental_fbank``).""" + mel_processor = getattr(processor, "_streaming_mel_processor", None) + if getattr(self, "_incremental_fbank", False) and mel_processor is not None: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.incremental_fbank import enable_incremental_fbank + + if not enable_incremental_fbank(mel_processor): + logger.warning_once( + "MiniCPM-o duplex_incremental_fbank: %s is not supported; the streaming fbank recomputes in full", + type(mel_processor).__name__, + ) return processor def _prepare_session_context( @@ -208,24 +310,90 @@ class MiniCPMO45Stage0DuplexRuntime: state.current_turn_ended = False state.pending_speech_response_open = True - def _stage_prefill_embeddings_only( + def _stage_prefill_embeddings_only(self, state, audio_waveform, **kwargs: Any) -> dict[str, object]: + """Run :meth:`_stage_prefill_steps` for one append, encoding its units one at a time.""" + return self._run_prefill_steps([self._stage_prefill_steps(state, audio_waveform, **kwargs)])[0] + + @staticmethod + def take_staged_prefill( + state: _MiniCPMO45Stage0SessionState, epoch: int | None, seq: int | None + ) -> dict[str, object] | None: + """Pop the append :meth:`stage_prefill_batch` built for ``(epoch, seq)`` this step, if any.""" + staged = state.staged_prefill + if staged is None: + return None + state.staged_prefill = None + if seq is None or staged[0] != (epoch, seq): + return None + return staged[1] + + def batches_audio_encoder(self) -> bool: + """Whether the thinker can encode several sessions' audio units in one call.""" + return self._batched_audio_encoder() is not None + + def needs_prefill(self, state: _MiniCPMO45Stage0SessionState | None, epoch: int | None, seq: int | None) -> bool: + """Whether an append has not been built (or attempted) for this session yet.""" + if seq is None or state is None: + return seq is not None + built = state.prepared_append_identity == (epoch, seq) and state.prepared_inputs_embeds is not None + return not built and state.attempted_append_identity != (epoch, seq) + + def stage_prefill_batch( + self, + appends: list[tuple[_MiniCPMO45Stage0SessionState, NDArray[np.float32] | None, dict[str, Any]]], + ) -> None: + """Build distinct sessions' appends (each with a ``seq``), batching their ``k``-th unit encodes. + + Results are staged on the session for the append's own ``preprocess`` this step. + """ + if len({id(state) for state, _, _ in appends}) != len(appends): + raise ValueError("stage_prefill_batch needs one append per session") + steps = [self._stage_prefill_steps(state, audio, **kwargs) for state, audio, kwargs in appends] + for (state, _, kwargs), result in zip(appends, self._run_prefill_steps(steps), strict=True): + state.staged_prefill = ((kwargs.get("epoch"), kwargs["seq"]), result) + + def _run_prefill_steps(self, steps: list[_PrefillSteps]) -> list[dict[str, object]]: + """Drive prefills to completion, encoding each round's pending units together.""" + results: list[dict[str, object] | None] = [None] * len(steps) + pending: list[tuple[int, _MiniCPMO45AudioEncodeRequest]] = [] + + def advance(index: int, value: torch.Tensor | None, *, first: bool = False) -> None: + try: + request = next(steps[index]) if first else steps[index].send(value) + except StopIteration as stop: + results[index] = stop.value + return + pending.append((index, request)) + + for index in range(len(steps)): + advance(index, None, first=True) + while pending: + round_, pending = pending, [] + embeds = self._stage_audio_embeddings_batch([request for _, request in round_]) + for (index, _), embed in zip(round_, embeds, strict=True): + advance(index, embed) + return cast("list[dict[str, object]]", results) + + def _stage_prefill_steps( self, state: _MiniCPMO45Stage0SessionState, audio_waveform: NDArray[np.float32] | None, *, video_frames: list[Image.Image] | None = None, + encoded_frames: list[list[torch.Tensor]] | None = None, epoch: int | None = None, seq: int | None = None, is_speech: bool = False, final: bool = False, stage0_window: dict[str, object] | None = None, stage0_reanchor: dict[str, object] | None = None, - ) -> dict[str, object]: + ) -> _PrefillSteps: """Build scheduler-owned Stage0 input embeddings for one audio append. Unlike the legacy worker-control path, this method never calls an eager model forward. The normal vLLM runner consumes the returned embeddings and owns attention metadata, block tables, KV cache, and sampling. + Yields each unit's encoder request; the caller sends back its embeddings. """ start_time = time.time() processor = self._configure_streaming_processor(state) @@ -239,6 +407,7 @@ class MiniCPMO45Stage0DuplexRuntime: result["inputs_embeds"] = state.prepared_inputs_embeds result["input_token_ids"] = list(state.prepared_input_token_ids) return result + state.attempted_append_identity = append_identity self._require_special_token_ids() if isinstance(stage0_window, dict): completed_ids = stage0_window.get("completed_token_ids") @@ -252,11 +421,12 @@ class MiniCPMO45Stage0DuplexRuntime: state.pending_terminator_token = completed_terminator if audio_waveform is None or len(audio_waveform) == 0: return self._stage_prefill_result(False, start_time, "empty audio") + num_frames = len(encoded_frames) if encoded_frames is not None else len(video_frames or ()) state.audio_buffer = np.concatenate([state.audio_buffer, np.asarray(audio_waveform, dtype=np.float32)]) chunk_size = self._streaming_chunk_size(processor) self._pad_first_audio_chunk_if_needed(state, processor) if len(state.audio_buffer) < chunk_size: - if video_frames: + if num_frames: # Frames belong to the unit this append closes, and the # scheduler already reserved their vision tokens against this # append. A frame arriving before its unit can close has @@ -265,7 +435,7 @@ class MiniCPMO45Stage0DuplexRuntime: return self._stage_prefill_result( False, start_time, - f"{len(video_frames)} video frame(s) on an append that closes no model unit " + f"{num_frames} video frame(s) on an append that closes no model unit " f"(need {chunk_size} samples, only {len(state.audio_buffer)}); " "send frames on unit boundaries", ) @@ -278,10 +448,11 @@ class MiniCPMO45Stage0DuplexRuntime: # carry them, mirroring official streaming_prefill (feed , then # every frame_list entry as its own block, then the audio). frame_blocks: list[list[torch.Tensor]] = [] - if video_frames: + if num_frames: self._require_vision_token_ids() - encoded_frames = self._stage_vision_embeddings(video_frames) - if encoded_frames is None or len(encoded_frames) != len(video_frames): + if encoded_frames is None: + encoded_frames = self._stage_vision_embeddings(video_frames or []) + if encoded_frames is None or len(encoded_frames) != num_frames: return self._stage_prefill_result(False, start_time, "streaming vision embedding failed") if any(not slices for slices in encoded_frames): return self._stage_prefill_result(False, start_time, "streaming vision embedding failed") @@ -312,7 +483,7 @@ class MiniCPMO45Stage0DuplexRuntime: ): with suppress(Exception): setattr(batch_feature, name, value) - audio_embeds = self._stage_audio_embeddings(batch_feature, state=state) + audio_embeds = yield _MiniCPMO45AudioEncodeRequest(batch_feature, state) if audio_embeds is None: if units_built == 0: return self._stage_prefill_result(False, start_time, "streaming audio embedding returned empty") @@ -623,7 +794,8 @@ class MiniCPMO45Stage0DuplexRuntime: def _embed_token(self, token_id: int) -> torch.Tensor: import torch - token = torch.tensor([int(token_id)], dtype=torch.long, device=self._model_device()) + # A device-side fill: ``torch.tensor(..., device=...)`` is a blocking copy, 2-3 per unit. + token = torch.full((1,), int(token_id), dtype=torch.long, device=self._model_device()) embedder = self._token_embedder() embeds = embedder(token) return self._as_2d_tensor(embeds) @@ -759,7 +931,11 @@ class MiniCPMO45Stage0DuplexRuntime: state: _MiniCPMO45Stage0SessionState | None = None, ) -> torch.Tensor | None: if hasattr(batch_feature, "to"): - batch_feature = batch_feature.to(self.device) + batch_feature = self._audio_feature_to_device(batch_feature) + if state is not None and hasattr(state.audio_past_key_values, "to_legacy_cache"): + # A batched-path session whose unit this path has to take: hand the + # history over as the transformers cache the per-session path grows. + state.audio_past_key_values = state.audio_past_key_values.to_legacy_cache() self._ensure_dynamic_cache_compat() has_audio_cache = state is not None and hasattr(self.thinker, "audio_past_key_values") previous_audio_past_key_values = ( @@ -799,6 +975,94 @@ class MiniCPMO45Stage0DuplexRuntime: if has_audio_cache: self.thinker.audio_past_key_values = previous_audio_past_key_values + def _stage_audio_embeddings_batch( + self, + requests: list[_MiniCPMO45AudioEncodeRequest], + ) -> list[torch.Tensor | None]: + """Encode one unit per request (distinct sessions); those the batched encoder takes share one pass.""" + results: list[torch.Tensor | None] = [None] * len(requests) + encoder = self._batched_audio_encoder() + batched: list[tuple[int, _MiniCPMO45AudioEncodeRequest, Any]] = [] + for index, request in enumerate(requests): + chunk = self._batched_audio_chunk(request) if encoder is not None else None + if chunk is None: + results[index] = self._stage_audio_embeddings(request.batch_feature, state=request.state) + else: + batched.append((index, request, chunk)) + if batched: + assert encoder is not None + embeds, caches = encoder([chunk for _, _, chunk in batched]) + for (index, request, _), embed, cache in zip(batched, embeds, caches, strict=True): + assert request.state is not None + request.state.audio_past_key_values = cache + results[index] = embed + return results + + def _batched_audio_encoder(self) -> Callable[[list[Any]], tuple[list[Any], list[Any]]] | None: + encoder = getattr(self, "_batched_audio_encoder_fn", False) + if encoder is False: + if not getattr(self, "_audio_graph_attempted", False): + self.build_audio_cuda_graph() + thinker = getattr(self, "thinker", None) + batch = getattr(thinker, "get_audio_embedding_streaming_batch", None) + supported = getattr(thinker, "supports_streaming_audio_batch", None) + encoder = batch if callable(batch) and callable(supported) and supported() else None + self._batched_audio_encoder_fn = encoder + return encoder + + @staticmethod + def _batched_audio_chunk(request: _MiniCPMO45AudioEncodeRequest) -> Any | None: # Any: StreamingAudioChunk + import torch + + from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import ( + StreamingAudioChunk, + StreamingAudioKVCache, + ) + + cache = getattr(request.state, "audio_past_key_values", None) + if request.state is None or (cache is not None and not isinstance(cache, StreamingAudioKVCache)): + return None + feature = request.batch_feature + try: + features, lengths = feature["audio_features"], feature["audio_feature_lens"] + except (KeyError, TypeError): + return None + if not isinstance(features, torch.Tensor) or features.ndim != 3 or features.shape[0] != 1: + return None + if isinstance(lengths, (list, tuple)) and len(lengths) == 1: + lengths = lengths[0] + if not isinstance(lengths, torch.Tensor) or lengths.numel() != 1: + return None + # Same extra-context frames the per-session call passes. + return StreamingAudioChunk( + features=features, + cache=cache, + prefix_extra_frames=0 if int(getattr(feature, "chunk_idx", 0)) == 0 else 2, + suffix_extra_frames=2, + feature_length=int(lengths.reshape(-1)[0]), + ) + + def _audio_feature_to_device(self, batch_feature: Any) -> Any: # Any: transformers BatchFeature + """Move one unit's log-mel features to the encoder without a host sync; integer lengths stay on the host.""" + import torch + + from vllm_omni.utils.device_copy import to_device_nonblocking + + data = getattr(batch_feature, "data", None) + streaming = any( + callable(getattr(target, "get_audio_embedding_streaming", None)) + for target in (self.stage_model, self.thinker) + ) + if torch.device(self.device).type != "cuda" or not isinstance(data, dict) or not streaming: + return batch_feature.to(self.device) + batch_feature.data = { + key: to_device_nonblocking(value, self.device) + if isinstance(value, torch.Tensor) and torch.is_floating_point(value) + else value + for key, value in data.items() + } + return batch_feature + @staticmethod def _decode_ref_audio_from_session_config(session_config: dict[str, object]) -> NDArray[np.float32] | None: from vllm_omni.model_executor.models.minicpmo_4_5.duplex.input import decode_native_ref_audio_from_config @@ -1092,7 +1356,8 @@ class MiniCPMO45Stage0DuplexRuntime: return [1] return [2] + [1] * (frame_count - 1) - def _encode_processed_vision(self, processed: Any) -> list[torch.Tensor] | None: # Any: BatchFeature + def _encode_processed_vision_batch(self, processed: list[Any]) -> list[list[torch.Tensor]] | None: # BatchFeature + """Encode every slice of the ``processed`` frames in one vision-tower call.""" targets = (self.stage_model, self.thinker, getattr(self.stage_model, "model", None)) for target in targets: if target is None: @@ -1106,18 +1371,23 @@ class MiniCPMO45Stage0DuplexRuntime: vpm_param = next(vpm.parameters()) device, dtype = vpm_param.device, vpm_param.dtype - pixel_nested = processed["pixel_values"] - tgt_nested = processed["tgt_sizes"] flat_pixels: list[torch.Tensor] = [] flat_tgt: list[torch.Tensor] = [] - for image_slices, image_tgt in zip(pixel_nested, tgt_nested): - for slice_pixels in image_slices: - flat_pixels.append(slice_pixels.to(device=device, dtype=dtype)) - tgt_tensor = image_tgt if hasattr(image_tgt, "reshape") else torch.tensor(image_tgt) - flat_tgt.append(tgt_tensor.reshape(-1, 2)) + slices_per_frame: list[int] = [] + for frame in processed: + num_slices = 0 + for image_slices, image_tgt in zip(frame["pixel_values"], frame["tgt_sizes"]): + for slice_pixels in image_slices: + flat_pixels.append(slice_pixels.to(device=device, dtype=dtype)) + num_slices += 1 + tgt_tensor = image_tgt if hasattr(image_tgt, "reshape") else torch.tensor(image_tgt) + flat_tgt.append(torch.as_tensor(tgt_tensor).reshape(-1, 2)) + if num_slices == 0: + return None + slices_per_frame.append(num_slices) if not flat_pixels: return None - tgt_sizes = torch.cat(flat_tgt, dim=0).to(device=device, dtype=torch.int64) + tgt_sizes = torch.cat(flat_tgt, dim=0).to(device="cpu", dtype=torch.int64) with torch.no_grad(): hidden = get_hidden({"pixel_values": flat_pixels, "tgt_sizes": tgt_sizes}) except Exception: # noqa: BLE001 - prefill fails with a reason @@ -1132,7 +1402,10 @@ class MiniCPMO45Stage0DuplexRuntime: if embed_dtype is not None and block_2d.dtype != embed_dtype: block_2d = block_2d.to(dtype=embed_dtype) out.append(block_2d) - return out + if len(out) != sum(slices_per_frame): + return None + blocks = iter(out) + return [[next(blocks) for _ in range(num_slices)] for num_slices in slices_per_frame] return None def _stage_vision_embeddings(self, frames: list[Image.Image]) -> list[list[torch.Tensor]] | None: @@ -1146,17 +1419,96 @@ class MiniCPMO45Stage0DuplexRuntime: Frames are processed one at a time because ``process_image([a, b])`` packs both PILs into one batch item. """ + processed = self._preprocess_frames(frames) + return None if processed is None else self._encode_processed_vision_batch(processed) + + def _preprocess_frames(self, frames: list[Image.Image]) -> list[Any] | None: # Any: BatchFeature + """The CPU half of ``_stage_vision_embeddings``: each frame's processed slices; ``None`` on failure.""" process_image = getattr(self.processor, "process_image", None) if not callable(process_image): return None - out: list[list[torch.Tensor]] = [] + out: list[Any] = [] for frame, max_slices in zip(frames, self._official_max_slice_nums(len(frames))): try: - processed = process_image([frame], max_slice_nums=max_slices) + out.append(process_image([frame], max_slice_nums=max_slices)) except Exception: # noqa: BLE001 - prefill fails with a reason return None - encoded = self._encode_processed_vision(processed) - if encoded is None: - return None - out.append(encoded) return out + + def frame_kwargs(self, duplex: dict[str, Any], payload: dict[str, object]) -> dict[str, Any]: + """``encoded_frames`` prefetched for this append, else its decoded ``video_frames`` (``ValueError``).""" + encoded = self.take_prefetched_vision(duplex) + if encoded is not None: + return {"encoded_frames": encoded} + return {"video_frames": self._decode_video_frames_payload(payload)} + + @staticmethod + def _vision_prefetch_key(duplex: dict[str, Any]) -> tuple[str, int | None, int] | None: + """``(session_id, epoch, seq)`` of an append, parsed like ``preprocess`` does; ``None`` if unkeyed.""" + session_id, seq = str(duplex.get("session_id") or ""), _optional_int(duplex.get("seq")) + return (session_id, _optional_int(duplex.get("epoch")), seq) if session_id and seq is not None else None + + def prefetch_vision(self, appends: list[dict[str, Any]]) -> None: + """Encode this step's appends' camera frames in ``vision_batch_size``-append tower calls. + + Decode and slicing run on the frame pool, the next chunk's overlapping this chunk's encode. + Results are keyed by append and checked against its frame payload on use. + """ + self._prefetched_vision.clear() + jobs: list[tuple[tuple[str, int | None, int], tuple[object, ...]]] = [] + for duplex in appends: + key = self._vision_prefetch_key(duplex) + payload = duplex.get("payload") + frames = payload.get("video_frames") if isinstance(payload, dict) else None + if key is None or not isinstance(frames, list) or not frames: + continue + state = self.sessions.get(key[0]) + if state is None or state.prepared_append_identity != key[1:] or state.prepared_inputs_embeds is None: + jobs.append((key, tuple(frames))) # else already built: preprocess replays the cached append + if not jobs or not callable(getattr(self.processor, "process_image", None)): + return + + def preprocess(frames: tuple[object, ...]) -> list[Any] | None: + try: + decoded = self._decode_video_frames_payload({"video_frames": list(frames)}) + except ValueError: + return None + return self._preprocess_frames(decoded) if decoded else None + + pool = _frame_preprocess_pool() if len(jobs) > 1 else None + + def start(chunk: list[tuple[Any, tuple[object, ...]]]) -> list[Any]: + return [pool.submit(preprocess, frames) if pool is not None else frames for _, frames in chunk] + + def finish(started: list[Any]) -> list[list[Any] | None]: + return [item.result() if pool is not None else preprocess(item) for item in started] + + size = self._vision_prefetch_chunk_size() + chunks = [jobs[begin : begin + size] for begin in range(0, len(jobs), size)] + started = start(chunks[0]) + for index, chunk in enumerate(chunks): + # Submit the next chunk before this chunk's encode so they overlap. + following = start(chunks[index + 1]) if index + 1 < len(chunks) else [] + ready = [(job, processed) for job, processed in zip(chunk, finish(started)) if processed is not None] + items = [item for _, processed in ready for item in processed] + encoded = self._encode_processed_vision_batch(items) if ready else None + blocks = iter(encoded or ()) + for (key, frames), processed in ready if encoded is not None else (): + self._prefetched_vision[key] = (frames, [next(blocks) for _ in processed]) + started = following + + def _vision_prefetch_chunk_size(self) -> int: + """Jobs per ``prefetch_vision`` chunk: the checkpoint's ``vision_batch_size``, else 16.""" + for target in (self.stage_model, self.thinker, getattr(self.stage_model, "model", None)): + size = getattr(getattr(target, "config", None), "vision_batch_size", None) + if isinstance(size, int) and size > 0: + return size + return 16 + + def take_prefetched_vision(self, duplex: dict[str, Any]) -> list[list[torch.Tensor]] | None: + """Pop this append's prefetched per-frame embeddings if its frame payload still matches.""" + key = self._vision_prefetch_key(duplex) + frames_payload, encoded = self._prefetched_vision.pop(key, (None, None)) if key is not None else (None, None) + payload = duplex.get("payload") + frames = payload.get("video_frames") if isinstance(payload, dict) else None + return encoded if isinstance(frames, list) and tuple(frames) == frames_payload else None 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 8056f6c08..8a5c337f9 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 @@ -17,8 +17,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Iterable +from collections.abc import Callable, Iterable from contextlib import suppress +from dataclasses import dataclass from functools import cached_property from typing import Any @@ -45,10 +46,35 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights from vllm_omni.platforms import current_omni_platform +from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking logger = init_logger(__name__) +@dataclass(slots=True) +class _MiniCPMO45PendingSamples: + """A deferred Stage-0 duplex step awaiting its host commit.""" + + host: torch.Tensor + event: Any + row_idxs: list[int] + stage2: list[bool] + rewinds: dict[int, Callable[[], None]] + token_ids: dict[str, int] + row_sessions: dict[int, str] | None + row_payloads: dict[int, dict[str, Any]] | None + + +def _generator_rewind(generator: torch.Generator) -> Callable[[], None]: + """A call that puts ``generator`` back where it stands now: its Philox offset, else its full state.""" + try: + offset = generator.get_offset() + return lambda: generator.set_offset(offset) + except RuntimeError: # not a Philox generator (CPU) + state = generator.get_state() + return lambda: generator.set_state(state) + + @MULTIMODAL_REGISTRY.register_processor( MiniCPMO45OmniLLMMultiModalProcessor, info=MiniCPMO45OmniLLMProcessingInfo, @@ -244,7 +270,14 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp ) -> None: """Apply MiniCPM duplex policy before the standard model sampler.""" del sampling_metadata + # The previous step's deferred decisions update the session state read below. + self._commit_minicpmo45_duplex_pending_samples() self._minicpmo45_active_duplex_rows = [row.row_idx for row in rows] + self._minicpmo45_duplex_row_sampling_host = { + row.row_idx: (row.temperature, row.top_k, row.top_p) + for row in rows + if row.temperature is not None and row.top_k is not None and row.top_p is not None + } self._minicpmo45_duplex_row_sessions = { row.row_idx: row.session_id for row in rows if row.session_id is not None } @@ -306,8 +339,9 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp if force_listen and segment_key in force_listen_segments: continue if force_listen: - logits[row_idx, :] = float("-inf") - logits[row_idx, listen_id] = 0.0 + # fill_ keeps the scalars on the host; an element assignment copies a scalar tensor and syncs. + logits[row_idx].fill_(float("-inf")) + logits[row_idx].narrow(0, listen_id, 1).fill_(0.0) force_listen_segments.add(segment_key) # -------------------- Device utilities -------------------- @@ -418,53 +452,20 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": "bad_duplex_payload"}} - session_key = session_id - state = helper.sessions.get(session_key) - if state is None: - from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( - _MiniCPMO45Stage0SessionState, - ) - - state = _MiniCPMO45Stage0SessionState(session_id=session_id) - helper.sessions[session_key] = state - session_config = duplex.get("session_config") - session_config = dict(session_config) if isinstance(session_config, dict) else {} - runtime_config = duplex.get("runtime_config") - runtime_config = dict(runtime_config) if isinstance(runtime_config, dict) else {} - if hasattr(helper.thinker, "audio_past_key_values"): - helper.thinker.audio_past_key_values = None - helper._configure_streaming_processor(state) - helper._prepare_session_context(state, session_config, runtime_config=runtime_config) - - audio_waveform = helper._decode_audio_payload(payload) - try: - video_frames = helper._decode_video_frames_payload(payload) - except ValueError as exc: - embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) - return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": str(exc)}} - seq = duplex.get("seq") - try: - seq = int(seq) if seq is not None else None - except (TypeError, ValueError): - seq = None - epoch = duplex.get("epoch") - try: - epoch = int(epoch) if epoch is not None else None - except (TypeError, ValueError): - epoch = None - result = helper._stage_prefill_embeddings_only( - state, - audio_waveform, - video_frames=video_frames, - epoch=epoch, - seq=seq, - is_speech=bool(payload.get("is_speech", False)), - final=bool(duplex.get("final")), - stage0_window=(duplex.get("stage0_window") if isinstance(duplex.get("stage0_window"), dict) else None), - stage0_reanchor=( - duplex.get("stage0_reanchor") if isinstance(duplex.get("stage0_reanchor"), dict) else None - ), - ) + # A deferred sample of this session updates the latches its append reads. + self._commit_minicpmo45_duplex_pending_samples(session_ids={session_id}) + state = self._minicpmo45_duplex_session_state(helper, session_id, duplex) + prefill_kwargs = self._minicpmo45_duplex_prefill_kwargs(duplex, payload) + seq = prefill_kwargs["seq"] + result = helper.take_staged_prefill(state, prefill_kwargs["epoch"], seq) + if result is None: # not built by this step's preprocess_batch + audio_waveform = helper._decode_audio_payload(payload) + try: + prefill_kwargs.update(helper.frame_kwargs(duplex, payload)) + except ValueError as exc: + embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) + return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": str(exc)}} + result = helper._stage_prefill_embeddings_only(state, audio_waveform, **prefill_kwargs) update_result = dict(result) if result.get("stage0_window_replaced") is True: window = duplex.get("stage0_window", {}) @@ -542,12 +543,105 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp if len(input_token_ids) < span_len: input_token_ids.extend([pad_token_id] * (span_len - len(input_token_ids))) if input_token_ids: - req_input_ids = torch.tensor(input_token_ids, dtype=input_ids.dtype, device=input_ids.device) + # Pinned + non-blocking: a pageable copy would wait for all queued GPU work. + req_input_ids = index_to_device(input_token_ids, input_ids.device, input_ids.dtype) update_result["duplex_prompt_token_ids"] = full_input_token_ids else: req_input_ids = torch.full_like(input_ids, helper._required_token_id("unit_token_id")) return req_input_ids, req_embeds, {"duplex": update_result} + def preprocess_batch( + self, + *, + req_ids: list[str], + model_intermediate_buffer: dict[str, dict[str, Any]], + device: torch.device, + ) -> None: + """Runner hook before ``preprocess``: batch this step's duplex appends across sessions. + + With two or more appends, camera frames share vision-tower calls (``prefetch_vision``) and + unbuilt appends share streaming-encoder passes (``stage_prefill_batch``); ``preprocess`` + takes the staged results and builds or reports anything skipped here. + """ + del device + if self.model_stage != "llm": + return + infos = [model_intermediate_buffer.get(req_id) for req_id in req_ids] + duplexes = [info.get("duplex") for info in infos if isinstance(info, dict)] + appends = [duplex for duplex in duplexes if isinstance(duplex, dict) and duplex.get("data_plane") is True] + helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None) + frame_appends = [d for d in appends if isinstance(d.get("payload"), dict) and d["payload"].get("video_frames")] + if len(frame_appends) >= 2: + try: + helper = self._duplex_data_plane_helper() + helper.prefetch_vision(frame_appends) + except Exception: # noqa: BLE001 - preprocess encodes each append itself instead + logger.warning("MiniCPM-o duplex frame prefetch failed; encoding per request", exc_info=True) + elif helper is not None: + helper.prefetch_vision([]) # drop entries the previous step did not consume + if helper is None or len(req_ids) < 2 or not helper.batches_audio_encoder(): + return + candidates: list[tuple[str, dict[str, Any], dict[str, Any], dict[str, Any]]] = [] + seen: set[str] = set() + for duplex in appends: + session_id = str(duplex.get("session_id") or "") + payload = duplex.get("payload") + repeated = session_id in seen + seen.add(session_id) + if not session_id or not isinstance(payload, dict) or repeated: + continue # a second append of one session keeps its order: preprocess + kwargs = self._minicpmo45_duplex_prefill_kwargs(duplex, payload) + if helper.needs_prefill(helper.sessions.get(session_id), kwargs["epoch"], kwargs["seq"]): + candidates.append((session_id, duplex, payload, kwargs)) + if len(candidates) < 2: + return + self._commit_minicpmo45_duplex_pending_samples(session_ids={candidate[0] for candidate in candidates}) + staged = [] + for session_id, duplex, payload, kwargs in candidates: + try: + audio_waveform = helper._decode_audio_payload(payload) + kwargs.update(helper.frame_kwargs(duplex, payload)) + except ValueError: + continue + staged.append((self._minicpmo45_duplex_session_state(helper, session_id, duplex), audio_waveform, kwargs)) + if len(staged) >= 2: + helper.stage_prefill_batch(staged) + + def _minicpmo45_duplex_session_state(self, helper, session_id: str, duplex: dict[str, Any]): + """The Stage-0 state of ``session_id``, created with its session context on first use.""" + state = helper.sessions.get(session_id) + if state is None: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + + state = _MiniCPMO45Stage0SessionState(session_id=session_id) + helper.sessions[session_id] = state + session_config = duplex.get("session_config") + session_config = dict(session_config) if isinstance(session_config, dict) else {} + runtime_config = duplex.get("runtime_config") + runtime_config = dict(runtime_config) if isinstance(runtime_config, dict) else {} + if hasattr(helper.thinker, "audio_past_key_values"): + helper.thinker.audio_past_key_values = None + helper._configure_streaming_processor(state) + helper._prepare_session_context(state, session_config, runtime_config=runtime_config) + return state + + @staticmethod + def _minicpmo45_duplex_prefill_kwargs(duplex: dict[str, Any], payload: dict[str, Any]) -> dict[str, Any]: + """Keyword arguments of ``_stage_prefill_embeddings_only`` for one append, frames aside.""" + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _optional_int + + window, reanchor = duplex.get("stage0_window"), duplex.get("stage0_reanchor") + return { + "epoch": _optional_int(duplex.get("epoch")), + "seq": _optional_int(duplex.get("seq")), + "is_speech": bool(payload.get("is_speech", False)), + "final": bool(duplex.get("final")), + "stage0_window": window if isinstance(window, dict) else None, + "stage0_reanchor": reanchor if isinstance(reanchor, dict) else None, + } + def _duplex_data_plane_helper(self): helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None) if helper is not None: @@ -694,13 +788,11 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp if isinstance(key, str) and isinstance(value, int) and value >= 0 } if special_keys: + # Host tensors: every consumer reads them on the host, and a pageable + # host->device copy per row and key would wait for the whole forward. multimodal_outputs["meta"] = { key: [ - torch.tensor( - [int(value)], - dtype=torch.long, - device=text_hidden_states.device, - ) + torch.tensor([int(value)], dtype=torch.long) if isinstance(value, int) and value >= 0 else None for duplex_info in duplex_rows @@ -808,6 +900,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp ) -> SamplerOutput | None: if self.model_stage != "llm" or logits.ndim != 2 or logits.shape[0] == 0: return None + self._commit_minicpmo45_duplex_pending_samples() token_ids = self._minicpmo45_native_duplex_token_ids() unit_id = token_ids.get("unit_token_id", -1) if unit_id < 0: @@ -823,120 +916,399 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp chunk_terminators = self._minicpmo45_chunk_terminator_token_ids(token_ids) output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or [] - sampled_ids: list[int] = [] - for row_idx in range(logits.shape[0]): + num_rows = logits.shape[0] + row_params = self._minicpmo45_duplex_row_params(sampling_metadata, num_rows) + + # Lookahead after a chunk terminator: re-emit it and skip sampling. + sampled_ids: list[int] = [-1] * num_rows + pending_rows: list[int] = [] + for row_idx in range(num_rows): accepted = output_token_ids[row_idx] if row_idx < len(output_token_ids) else [] last_accepted = next((int(t) for t in reversed(accepted) if isinstance(t, int) and t >= 0), None) if last_accepted in chunk_terminators: - # Async scheduling runs one lookahead frame after the chunk - # terminator was sampled but before the scheduler observes - # the segment stop. The scheduler discards this frame's - # token, so decide nothing here: re-emit the terminator and - # leave the model-owned policy state exactly as the accepted - # history left it. The next append re-injects that terminator. - sampled_ids.append(last_accepted) - continue - row_logits = logits[row_idx : row_idx + 1].clone() - sampled = self._sample_minicpmo45_native_duplex_row( - row_logits, - sampling_metadata, - row_idx=row_idx, - token_ids=token_ids, - ) - self._record_minicpmo45_duplex_terminator(row_idx, sampled, token_ids) - sampled_ids.append(sampled) + sampled_ids[row_idx] = last_accepted + else: + pending_rows.append(row_idx) + + if pending_rows: + kwargs = {"row_idxs": pending_rows, "token_ids": token_ids, "row_params": row_params} + deferred = self._sample_minicpmo45_native_duplex_rows_deferred(logits, sampling_metadata, **kwargs) + if deferred is not None: + sampled = index_to_device(sampled_ids, logits.device, torch.int32) + sampled[index_to_device(pending_rows, logits.device)] = deferred.to(torch.int32) + return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None) + synchronous = self._sample_minicpmo45_native_duplex_rows(logits, sampling_metadata, **kwargs) + for row_idx, sampled in zip(pending_rows, synchronous, strict=True): + sampled_ids[row_idx] = sampled + self._record_minicpmo45_duplex_terminator(row_idx, sampled, token_ids) + return SamplerOutput( - sampled_token_ids=torch.tensor(sampled_ids, device=logits.device, dtype=torch.int32).unsqueeze(-1), + sampled_token_ids=index_to_device(sampled_ids, logits.device, torch.int32).unsqueeze(-1), logprobs_tensors=None, ) - def _sample_minicpmo45_native_duplex_row( + def _minicpmo45_duplex_row_params( + self, + sampling_metadata: SamplingMetadata, + num_rows: int, + ) -> list[tuple[float, float, float]]: + """Every row's (temperature, top_k, top_p): host copies when present, else one device read per parameter.""" + host = getattr(self, "_minicpmo45_duplex_row_sampling_host", None) or {} + columns = [] + for column, (name, default) in enumerate((("temperature", 0.7), ("top_k", 100), ("top_p", 0.8))): + value = getattr(sampling_metadata, name, None) + if not isinstance(value, torch.Tensor): + columns.append([self._sampling_metadata_value(sampling_metadata, name, 0, default)] * num_rows) + elif value.numel() == num_rows and host.keys() >= set(range(num_rows)): + columns.append([float(host[row][column]) for row in range(num_rows)]) + elif value.numel() == 0: + columns.append([default] * num_rows) + else: + flat = [float(v) for v in value.reshape(-1).tolist()] + columns.append([flat[min(row, len(flat) - 1)] for row in range(num_rows)]) + return list(zip(*columns, strict=True)) + + def _minicpmo45_policy_limit(self, name: str, default: int) -> int: + """A per-chunk duplex limit set on the model (``0``/``None`` mean ``default``).""" + return int(getattr(self, name, default) or default) + + def _minicpmo45_duplex_forbidden_index(self, token_ids: dict[str, int], vocab: int, device) -> torch.Tensor | None: + """``_minicpmo45_native_forbidden_token_ids`` inside the vocab, as a cached device index.""" + key = (vocab, str(device)) + cache = getattr(self, "_minicpmo45_duplex_forbidden_index_cache", None) + if cache is None or cache[0] != key: + forbidden = self._minicpmo45_native_forbidden_token_ids(token_ids) + valid = sorted({token_id for token_id in forbidden if 0 <= token_id < vocab}) + cache = (key, torch.tensor(valid, dtype=torch.long, device=device) if valid else None) + self._minicpmo45_duplex_forbidden_index_cache = cache + return cache[1] + + def _minicpmo45_duplex_cut_tables(self, token_ids: dict[str, int], vocab: int, device): + """Per-token decoded length and "never cut" mask for ``_maybe_cut_minicpmo45_native_duplex_text_chunk``. + + Byte-level decoding is additive after a prefix ending on a character boundary: + ``len(decode(chunk + [t])) == len(decode(chunk)) + length[t]``. None (no device cut) without batch + decode, or with the whitespace clean-up that rewrites across token boundaries. + """ + key = (vocab, str(device)) + cache = getattr(self, "_minicpmo45_duplex_cut_tables_cache", None) + if cache is not None and cache[0] == key: + return cache[1] + tables = None + tokenizer = self._minicpmo45_tokenizer() + batch_decode = getattr(tokenizer, "batch_decode", None) + if callable(batch_decode) and not getattr(tokenizer, "clean_up_tokenization_spaces", False): + try: + num_tokens = min(vocab, len(tokenizer)) + texts = batch_decode([[token_id] for token_id in range(num_tokens)], skip_special_tokens=True) + except Exception: + texts = None + if texts is not None and len(texts) == num_tokens: + lengths = torch.zeros(vocab, dtype=torch.int32) + lengths[:num_tokens] = torch.tensor([len(text) for text in texts], dtype=torch.int32) + never_cut = torch.ones(vocab, dtype=torch.bool) + never_cut[:num_tokens] = False + special = [t for t in self._minicpmo45_native_special_token_ids(token_ids) if 0 <= t < vocab] + if special: + never_cut[torch.tensor(special, dtype=torch.long)] = True + tables = (lengths.to(device), never_cut.to(device)) + self._minicpmo45_duplex_cut_tables_cache = (key, tables) + return tables + + def _minicpmo45_duplex_cut_bases( + self, + row_idxs: list[int], + recent_tokens_by_row: dict[int, list[int]], + token_ids: dict[str, int], + max_chars: int, + ) -> list[int] | None: + """Each row's decoded current-chunk length (-1: never cut), or None when a row needs the host decode. + + A chunk ending in U+FFFD may end inside a character the next token's bytes complete: not additive. + """ + decode = getattr(self._minicpmo45_tokenizer(), "decode", None) + if max_chars <= 0 or not callable(decode): + return [-1] * len(row_idxs) + bases = [] + for row_idx in row_idxs: + chunk = self._minicpmo45_current_chunk_tokens(recent_tokens_by_row[row_idx], token_ids) + try: + text = decode(chunk, skip_special_tokens=True) + except Exception: + return None + if not isinstance(text, str) or text.endswith("\ufffd"): + return None + bases.append(len(text)) + return bases + + @staticmethod + def _minicpmo45_duplex_recent_tokens( + sampling_metadata: SamplingMetadata, row_idxs: list[int] + ) -> dict[int, list[int]]: + """Each row's accepted output tokens (the current chunk's history), placeholders dropped.""" + output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or [] + rows = {row_idx: output_token_ids[row_idx] if row_idx < len(output_token_ids) else () for row_idx in row_idxs} + return {row: [int(t) for t in tokens if isinstance(t, int) and t >= 0] for row, tokens in rows.items()} + + def _minicpmo45_duplex_capped_rows(self, row_idxs: list[int], recent: dict[int, list[int]]) -> set[int]: + """Rows at their per-chunk speak-length cap: chunk_eos without sampling anything.""" + limit = self._minicpmo45_policy_limit( + "max_new_speak_tokens_per_chunk", MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK + ) + request_max = {row_idx: self._minicpmo45_duplex_row_request_max_tokens(row_idx) for row_idx in row_idxs} + row_limit = {row: limit if request is None else min(limit, request) for row, request in request_max.items()} + return {row for row, row_cap in row_limit.items() if len(recent[row]) >= max(1, row_cap - 1)} + + def _minicpmo45_duplex_boundary_hits( + self, logits: torch.Tensor, rows: list[int], chunk_eos_id: int, sampling_metadata: SamplingMetadata + ) -> torch.Tensor: + """Whether each row's boundary draw (the released StreamDecoder's first pass) picks chunk_eos. + + The drawn token only matters when it is chunk_eos, so a Bernoulli at + ``P(sample == chunk_eos)`` drawn with the row's generator is the same + decision, without a full-vocabulary softmax and multinomial per row. + """ + rows_logits = logits[index_to_device(rows, logits.device)] + if getattr(sampling_metadata, "all_greedy", False): + return torch.argmax(rows_logits, dim=-1) == chunk_eos_id + generators = getattr(sampling_metadata, "generators", {}) or {} + draws = [torch.rand((), generator=generators.get(row_idx), device=logits.device) for row_idx in rows] + return torch.stack(draws) < self._duplex_boundary_chunk_eos_probs(rows_logits, chunk_eos_id) + + def _minicpmo45_duplex_stage2_draws( + self, + logits: torch.Tensor, + rows: list[int], + recent: dict[int, list[int]], + token_ids: dict[str, int], + row_params: list[tuple[float, float, float]], + sampling_metadata: SamplingMetadata, + rewinds: dict[int, Callable[[], None]] | None = None, + ) -> torch.Tensor: + """Each row's stage-2 token: forbidden ids masked, repetition penalty, then greedy or a candidate draw. + + ``rewinds`` receives, per row position, a call undoing that row's draw. + """ + device, vocab = logits.device, logits.shape[-1] + all_greedy = bool(getattr(sampling_metadata, "all_greedy", False)) + generators = getattr(sampling_metadata, "generators", {}) or {} + stage2_logits = logits[index_to_device(rows, device)].clone() + forbidden_index = self._minicpmo45_duplex_forbidden_index(token_ids, vocab, device) + if forbidden_index is not None: + stage2_logits.index_fill_(1, forbidden_index, float("-inf")) + history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE + penalty_rows: list[int] = [] + penalty_cols: list[int] = [] + for local_idx, row_idx in enumerate(rows): + state = self._minicpmo45_duplex_state_for_row(row_idx) + repetition_tokens = getattr(state, "generated_tokens", None) or recent[row_idx] + penalized = [token_id for token_id in set(repetition_tokens[-history_size:]) if 0 <= token_id < vocab] + penalty_rows.extend([local_idx] * len(penalized)) + penalty_cols.extend(penalized) + if penalty_rows: + # Distinct (row, col) pairs, so one indexed division equals dividing each once. + rows_t, cols_t = index_to_device(penalty_rows, device), index_to_device(penalty_cols, device) + stage2_logits[rows_t, cols_t] = stage2_logits[rows_t, cols_t] / 1.05 + + samples: list[torch.Tensor | None] = [None] * len(rows) + greedy = [local for local, row_idx in enumerate(rows) if all_greedy or float(row_params[row_idx][0]) <= 0] + if greedy: + greedy_sample = torch.argmax(stage2_logits[index_to_device(greedy, device)], dim=-1) + for pos, local_idx in enumerate(greedy): + samples[local_idx] = greedy_sample[pos : pos + 1] + drawn = [local_idx for local_idx in range(len(rows)) if samples[local_idx] is None] + # One uniform per row from its generator, in row order, then a batched draw per (top_k, top_p) group. + uniforms: dict[int, torch.Tensor] = {} + groups: dict[tuple[int, float], list[int]] = {} + for local_idx in drawn: + generator = generators.get(rows[local_idx]) + if rewinds is not None: + rewinds[local_idx] = _generator_rewind(generator) + uniforms[local_idx] = torch.rand((), generator=generator, device=device) + _temperature, top_k, top_p = row_params[rows[local_idx]] + groups.setdefault((int(top_k), float(top_p)), []).append(local_idx) + if drawn: + temps = torch.tensor([float(row_params[row_idx][0]) for row_idx in rows], dtype=stage2_logits.dtype) + temps = to_device_nonblocking(temps, device).view(-1, 1) + for (top_k, top_p), members in groups.items(): + # Device-side row gathers: a python-list index would stage a + # pageable host copy and block the host on queued GPU work. + members_t = index_to_device(members, device) + group_logits = stage2_logits[members_t] / temps[members_t] + candidates = self._duplex_top_k_top_p_candidates(group_logits, top_k=top_k, top_p=top_p) + group_uniforms = torch.stack([uniforms[local_idx] for local_idx in members]) + draws = self._duplex_draw_candidates(group_logits, candidates, group_uniforms) + for i, local_idx in enumerate(members): + samples[local_idx] = draws[i : i + 1] + return torch.cat(samples) # type: ignore[arg-type] + + def _sample_minicpmo45_native_duplex_rows( self, logits: torch.Tensor, sampling_metadata: SamplingMetadata, *, - row_idx: int, + row_idxs: list[int], token_ids: dict[str, int], - ) -> int: + row_params: list[tuple[float, float, float]], + ) -> list[int]: + """Sample duplex rows with at most two host reads: the boundary decisions, then the stage-2 tokens. + + Per row: chunk_eos at the speak-length cap, else the boundary draw, else + the stage-2 draw, the chunk cut and the listen->tts_bos rewrite. Each + row's own generator draws as it would one row at a time; rows sharing + the process-wide generator draw every boundary before any stage-2 token. + """ chunk_eos_id = token_ids.get("chunk_eos_token_id", -1) - generator = getattr(sampling_metadata, "generators", {}).get(row_idx) - output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or [] - raw_recent_tokens = output_token_ids[row_idx] if row_idx < len(output_token_ids) else [] - recent_tokens = [int(token_id) for token_id in raw_recent_tokens if isinstance(token_id, int) and token_id >= 0] - temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7)) - top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100)) - top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8)) - state = self._minicpmo45_duplex_state_for_row(row_idx) - if chunk_eos_id >= 0 and chunk_eos_id < logits.shape[-1]: - max_speak_tokens = int( - getattr( - self, - "max_new_speak_tokens_per_chunk", - MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK, - ) - or MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK + recent = self._minicpmo45_duplex_recent_tokens(sampling_metadata, row_idxs) + resolved: dict[int, int] = {} + if 0 <= chunk_eos_id < logits.shape[-1]: + resolved = dict.fromkeys(self._minicpmo45_duplex_capped_rows(row_idxs, recent), chunk_eos_id) + boundary = [row_idx for row_idx in row_idxs if row_idx not in resolved] + if boundary: + hits = self._minicpmo45_duplex_boundary_hits(logits, boundary, chunk_eos_id, sampling_metadata) + resolved.update((row_idx, chunk_eos_id) for row_idx, hit in zip(boundary, hits.tolist()) if hit) + stage2_rows = [row_idx for row_idx in row_idxs if row_idx not in resolved] + if stage2_rows: + draws = self._minicpmo45_duplex_stage2_draws( + logits, stage2_rows, recent, token_ids, row_params, sampling_metadata ) - request_max_tokens = self._minicpmo45_duplex_row_request_max_tokens(row_idx) - effective_max_speak_tokens = max_speak_tokens - if request_max_tokens is not None: - effective_max_speak_tokens = min(effective_max_speak_tokens, request_max_tokens) - if len(recent_tokens) >= max(1, effective_max_speak_tokens - 1): - return int(chunk_eos_id) + for row_idx, sampled in zip(stage2_rows, draws.tolist(), strict=True): + self._record_minicpmo45_duplex_generation_token(row_idx, sampled) + sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk(sampled, recent[row_idx], token_ids) + resolved[row_idx] = self._finalize_minicpmo45_native_duplex_sample(row_idx, sampled, token_ids) + return [int(resolved[row_idx]) for row_idx in row_idxs] - # Match the released StreamDecoder: first sample the original - # distribution only to preserve the model's own chunk boundary. - # If it does not choose chunk_eos, mask that token before the - # normal text/listen/turn sampling pass below. - if getattr(sampling_metadata, "all_greedy", False): - boundary_sample = int(torch.argmax(logits, dim=-1).item()) - else: - boundary_probs = F.softmax(logits, dim=-1) - boundary_sample = int(torch.multinomial(boundary_probs, num_samples=1, generator=generator).item()) - if boundary_sample == chunk_eos_id: - return int(chunk_eos_id) + def _sample_minicpmo45_native_duplex_rows_deferred( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + *, + row_idxs: list[int], + token_ids: dict[str, int], + row_params: list[tuple[float, float, float]], + ) -> torch.Tensor | None: + """``_sample_minicpmo45_native_duplex_rows`` decided on the device without a host sync, or ``None``. - logits = logits.clone() - forbidden = self._minicpmo45_native_forbidden_token_ids(token_ids) - if forbidden: - valid_forbidden = [token_id for token_id in forbidden if 0 <= token_id < logits.shape[-1]] - if valid_forbidden: - logits[:, valid_forbidden] = float("-inf") - - generated_tokens = getattr(state, "generated_tokens", None) - repetition_tokens = generated_tokens if generated_tokens else recent_tokens - repetition_penalty = 1.05 - if repetition_penalty != 1.0 and repetition_tokens: - history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE - for token_id in set(repetition_tokens[-history_size:]): - if token_id < 0 or token_id >= logits.shape[-1]: - continue - logits[0, token_id] /= repetition_penalty - - if getattr(sampling_metadata, "all_greedy", False) or temperature <= 0: - sampled = int(torch.argmax(logits, dim=-1).item()) - self._record_minicpmo45_duplex_generation_token(row_idx, sampled) - sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk( - sampled, - recent_tokens, - token_ids, + Host state is committed a step later (``_commit_minicpmo45_duplex_pending_samples``). Stage 2 + is drawn for every uncapped row; a row the boundary resolves has its generator rewound at + commit, so generators advance as on the synchronous path. Process-wide generator rows can not + be rewound alone and keep the step synchronous. + """ + device, vocab = logits.device, logits.shape[-1] + chunk_eos_id = token_ids.get("chunk_eos_token_id", -1) + has_chunk_eos = 0 <= chunk_eos_id < vocab + all_greedy = bool(getattr(sampling_metadata, "all_greedy", False)) + generators = getattr(sampling_metadata, "generators", {}) or {} + recent = self._minicpmo45_duplex_recent_tokens(sampling_metadata, row_idxs) + capped = self._minicpmo45_duplex_capped_rows(row_idxs, recent) if has_chunk_eos else set() + stage2_rows = [row_idx for row_idx in row_idxs if row_idx not in capped] + drawn_rows = [row_idx for row_idx in stage2_rows if not (all_greedy or float(row_params[row_idx][0]) <= 0)] + if any(generators.get(row_idx) is None for row_idx in drawn_rows): + return None + cut = None + if has_chunk_eos and stage2_rows: + max_chars = self._minicpmo45_policy_limit( + "max_speak_chars_per_chunk", MiniCPMO45DuplexPolicy.DEFAULT_MAX_SPEAK_CHARS_PER_CHUNK ) - return self._finalize_minicpmo45_native_duplex_sample(row_idx, sampled, token_ids) + bases = self._minicpmo45_duplex_cut_bases(stage2_rows, recent, token_ids, max_chars) + if bases is None: + return None + if any(base >= 0 for base in bases): + tables = self._minicpmo45_duplex_cut_tables(token_ids, vocab, device) + if tables is None: + return None + cut = (tables, bases, max_chars) - logits = logits / temperature - logits = self._top_k_top_p_filter(logits, top_k=top_k, top_p=top_p) - probs = F.softmax(logits, dim=-1) - sampled = int(torch.multinomial(probs, num_samples=1, generator=generator).item()) - self._record_minicpmo45_duplex_generation_token(row_idx, sampled) - sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk( - sampled, - recent_tokens, - token_ids, + final = torch.full((len(row_idxs),), max(chunk_eos_id, 0), dtype=torch.long, device=device) + position = {row_idx: pos for pos, row_idx in enumerate(row_idxs)} + rewinds: dict[int, Callable[[], None]] = {} + if stage2_rows: + positions_t = index_to_device([position[row_idx] for row_idx in stage2_rows], device) + hit = ( + self._minicpmo45_duplex_boundary_hits(logits, stage2_rows, chunk_eos_id, sampling_metadata) + if has_chunk_eos + else torch.zeros(len(stage2_rows), dtype=torch.bool, device=device) + ) + rewinds_into = rewinds if has_chunk_eos else None + raw = self._minicpmo45_duplex_stage2_draws( + logits, stage2_rows, recent, token_ids, row_params, sampling_metadata, rewinds_into + ) + token = raw + if cut is not None: + (lengths, never_cut), bases, max_chars = cut + base_t = index_to_device(bases, device) + cut_t = (base_t >= 0) & ~never_cut[raw] & (base_t + lengths[raw].long() >= max_chars) + token = torch.where(cut_t, torch.full_like(raw, chunk_eos_id), raw) + listen_id = token_ids.get("listen_token_id", -1) + if 0 <= listen_id: + # The host rewrite of a sampled listen (to tts_bos inside an open turn), decided per row. + rewritten = [ + self._finalize_minicpmo45_native_duplex_sample(r, listen_id, token_ids) for r in stage2_rows + ] + if any(value != listen_id for value in rewritten): + token = torch.where(token == listen_id, index_to_device(rewritten, device), token) + if has_chunk_eos: + token = torch.where(hit, torch.full_like(token, chunk_eos_id), token) + final[positions_t] = token + + # One pinned copy of (final, raw stage-2 sample, boundary hit) per row + # for the commit; the host reads it once the runner has waited for it. + packed = torch.stack([final, torch.full_like(final, -1), torch.zeros_like(final)], dim=1) + if stage2_rows: + packed[positions_t, 1] = raw + packed[positions_t, 2] = hit.long() + host = torch.empty(packed.shape, dtype=packed.dtype, pin_memory=device.type == "cuda") + host.copy_(packed, non_blocking=True) + event = torch.cuda.Event() if device.type == "cuda" else None + if event is not None: + event.record() + self._minicpmo45_duplex_pending_samples = _MiniCPMO45PendingSamples( + host=host, + event=event, + row_idxs=list(row_idxs), + stage2=[row_idx not in capped for row_idx in row_idxs], + rewinds={position[stage2_rows[local_idx]]: rewind for local_idx, rewind in rewinds.items()}, + token_ids=token_ids, + row_sessions=getattr(self, "_minicpmo45_duplex_row_sessions", None), + row_payloads=getattr(self, "_minicpmo45_duplex_row_payloads", None), ) - return self._finalize_minicpmo45_native_duplex_sample( - row_idx, - sampled, - token_ids, + return final + + def _commit_minicpmo45_duplex_pending_samples(self, *, session_ids: set[str] | None = None) -> None: + """Apply the host state updates of the last deferred step, in the synchronous path's order. + + With ``session_ids``, only when that step sampled one of them: a + caller building those sessions' next appends needs their latches, and + must not wait for an unrelated step otherwise. + """ + pending = getattr(self, "_minicpmo45_duplex_pending_samples", None) + if pending is None: + return + if session_ids is not None: + sessions = pending.row_sessions or {} + if not any(sessions.get(row_idx) in session_ids for row_idx in pending.row_idxs): + return + self._minicpmo45_duplex_pending_samples = None + if pending.event is not None: + pending.event.synchronize() + values = pending.host.tolist() + saved = ( + getattr(self, "_minicpmo45_duplex_row_sessions", None), + getattr(self, "_minicpmo45_duplex_row_payloads", None), ) + self._minicpmo45_duplex_row_sessions = pending.row_sessions + self._minicpmo45_duplex_row_payloads = pending.row_payloads + try: + for pos, row_idx in enumerate(pending.row_idxs): + _final, raw, hit = values[pos] + if pending.stage2[pos] and hit: + pending.rewinds.get(pos, lambda: None)() + elif pending.stage2[pos]: + self._record_minicpmo45_duplex_generation_token(row_idx, int(raw)) + for pos, row_idx in enumerate(pending.row_idxs): + self._record_minicpmo45_duplex_terminator(row_idx, int(values[pos][0]), pending.token_ids) + finally: + self._minicpmo45_duplex_row_sessions, self._minicpmo45_duplex_row_payloads = saved def _maybe_cut_minicpmo45_native_duplex_text_chunk( self, @@ -1227,7 +1599,8 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp kth = torch.topk(logits, top_k, dim=-1).values[..., -1, None] logits = logits.masked_fill(logits < kth, float("-inf")) if 0.0 < top_p < 1.0: - sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) + # Stable: tied logits enter the nucleus lowest index first. + sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1, stable=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_remove = cumulative_probs > top_p sorted_remove[..., 1:] = sorted_remove[..., :-1].clone() @@ -1237,6 +1610,80 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp logits = logits.masked_fill(remove, float("-inf")) return logits + @staticmethod + def _duplex_top_k_top_p_candidates( + logits: torch.Tensor, *, top_k: int, top_p: float + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """``softmax(_top_k_top_p_filter(logits))`` on its support as ``(probs, indices, kth, ties)``, no host read. + + Nucleus order (value descending, index ascending), O(V + k log k). All logits tied at the k-th value + stay; their count is data-dependent, so they share the last slot (index -1): the ``ties`` lowest-index + ones the nucleus keeps, at one probability each. + """ + vocab = logits.shape[-1] + k = int(top_k) if top_k else 0 + if 0 < k < vocab: + values, indices = torch.topk(logits, k, dim=-1) + kth = values[..., -1:] + values = values.masked_fill(values <= kth, float("-inf")) + perm = torch.argsort(indices, dim=-1) + values, order = torch.sort(values.gather(-1, perm), descending=True, dim=-1, stable=True) + indices = indices.gather(-1, perm).gather(-1, order) + tied = (logits == kth).sum(dim=-1, keepdim=True) + else: + values, indices = torch.sort(logits, descending=True, dim=-1, stable=True) + kth = torch.full_like(values[..., :1], float("-inf")) + tied = torch.zeros_like(indices[..., :1]) + # In float32, as softmax does: exponentials past the largest logit, then one division by their sum. + values, kth_value = values.float(), kth.float() + top = torch.maximum(values[..., :1], kth_value) + weights, tie_weight = (values - top).exp(), torch.where(tied > 0, (kth_value - top).exp(), 0.0) + total = weights.sum(dim=-1, keepdim=True) + tied * tie_weight + probs, tie_prob = weights / total, tie_weight / total + ties = tied + if 0.0 < float(top_p) < 1.0: + cumulative = probs.cumsum(dim=-1) + remove = cumulative > float(top_p) + remove[..., 1:] = remove[..., :-1].clone() + remove[..., 0] = False + probs = probs.masked_fill(remove, 0.0) + # Tied token j stays while the mass before it, ``before + j * tie_prob``, is at most top_p (the + # first candidate always does); the floor's estimate of the last j is corrected on that mass. + before = cumulative[..., -1:] + last = ((float(top_p) - before) / tie_prob.clamp(min=torch.finfo(torch.float32).tiny)).floor() + last = torch.where(before + (last + 1) * tie_prob <= float(top_p), last + 1, last) + last = torch.where(before + last * tie_prob > float(top_p), last - 1, last) + kept = torch.where(tie_prob > 0, last + 1, torch.zeros_like(last)) + ties = torch.minimum(kept.clamp(min=0), tied.float()).long() + probs = torch.cat([probs, ties * tie_prob], dim=-1) + probs = probs / probs.sum(dim=-1, keepdim=True) + return probs, torch.cat([indices, torch.full_like(indices[..., :1], -1)], dim=-1), kth, ties + + @staticmethod + def _duplex_draw_candidates( + logits: torch.Tensor, + candidates: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + uniforms: torch.Tensor, + ) -> torch.Tensor: + """Per row, the inverse-CDF draw of ``_duplex_top_k_top_p_candidates(logits)`` at a uniform in [0, 1).""" + probs, indices, kth, ties = candidates + cdf = probs.cumsum(dim=-1) + target = uniforms.view(-1, 1).to(cdf.dtype) * cdf[..., -1:] + slot = torch.searchsorted(cdf, target, right=True).clamp(max=probs.shape[-1] - 1) + token = indices.gather(-1, slot) + # In the tie slot, the j-th lowest-index tied token, j uniform over the kept ties. + block = probs[..., -1:] + share = (target - (cdf[..., -1:] - block)) / block.clamp(min=torch.finfo(cdf.dtype).tiny) * ties + j = torch.minimum(share.floor().clamp(min=0), (ties - 1).clamp(min=0).to(share.dtype)).long() + tied = logits == kth + tie_token = (tied & (tied.cumsum(dim=-1) == j + 1)).int().argmax(dim=-1, keepdim=True) + return torch.where(slot == probs.shape[-1] - 1, tie_token, token).squeeze(-1) + + @staticmethod + def _duplex_boundary_chunk_eos_probs(logits: torch.Tensor, chunk_eos_id: int) -> torch.Tensor: + """Per-row ``softmax(logits)[:, chunk_eos_id]`` as ``exp(l_eos - logsumexp(l))``.""" + return torch.exp(logits[:, chunk_eos_id] - torch.logsumexp(logits, dim=-1)) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights for all components of the omni model.""" loaded_weights = set() @@ -1266,4 +1713,19 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp talker_loaded = add_prefix_to_loaded_weights(talker_loaded, "talker") loaded_weights.update(talker_loaded) + self._capture_duplex_audio_graphs() return loaded_weights + + def _capture_duplex_audio_graphs(self) -> None: + """Capture the duplex streaming audio encoder graphs while the model loads. + + Inside load_model their memory counts as model memory, so the KV budget leaves room for it. + A lazy capture on the first duplex request allocates past gpu_memory_utilization instead, + which OOMs the stages sharing the GPU. + """ + model_config = getattr(self.vllm_config, "model_config", None) + if self.thinker is None or getattr(model_config, "session_mode", "turn") != "duplex": + return + if self._module_device(self.thinker).type != "cuda": + return + self._duplex_data_plane_helper().build_audio_cuda_graph() diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 086bf74e8..74aafcef4 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -18,6 +18,7 @@ # limitations under the License. import inspect +import itertools import math import os import warnings @@ -138,6 +139,13 @@ def _encode_tokens(tokenizer: Any, prompt: str) -> list[int]: from vllm.utils.tensor_schema import TensorSchema, TensorShape +from vllm_omni.model_executor.models.minicpmo_4_5 import streaming_audio_encoder as streaming_audio +from vllm_omni.model_executor.models.minicpmo_4_5 import streaming_audio_encoder_graph as streaming_audio_graph +from vllm_omni.model_executor.models.minicpmo_4_5.vision_fused import ( + VisionGraphEncoder, + encode_packed_fused, + supports_fused_layers, +) from vllm_omni.model_executor.models.model_local_kv import ( ModelLocalKVScope, ModelLocalKVSpec, @@ -952,6 +960,8 @@ class SiglipVisionEmbeddings(nn.Module): self.num_patches = self.num_patches_per_side**2 self.num_positions = self.num_patches self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + # (h, w, device) -> position ids of that patch grid, see layout_position_ids. + self._layout_position_ids: dict[tuple[int, int, str], torch.Tensor] = {} def _create_grid_position_ids( self, @@ -1018,6 +1028,19 @@ class SiglipVisionEmbeddings(nn.Module): return position_ids.to(device=device) + def layout_position_ids(self, patch_grid_height: int, patch_grid_width: int, device: torch.device) -> torch.Tensor: + """Position ids of one ``(h, w)`` patch grid on ``device``: ``_create_position_ids``'s, cached.""" + cache = self._layout_position_ids + key = (int(patch_grid_height), int(patch_grid_width), str(device)) + ids = cache.get(key) + if ids is None: + if len(cache) >= 1024: + cache.clear() + boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side) + ids = self._create_grid_position_ids(key[0], key[1], boundaries).to(device=device) + cache[key] = ids + return ids + def forward( self, pixel_values: torch.FloatTensor, @@ -1112,6 +1135,51 @@ class SiglipAttention(nn.Module): return attn_output, attn_weights + def _attend_unpadded(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + """Attention over ``(n, L, heads, head_dim)`` items with no padding; same layout out.""" + query, key, value = (x.transpose(1, 2) for x in (query, key, value)) + attn_weights = torch.matmul(query, key.transpose(2, 3)) * self.scale + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + return torch.matmul(attn_weights, value).transpose(1, 2) + + def forward_packed(self, hidden_states: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]]) -> torch.Tensor: + """Self-attention over packed ``(tokens, embed)`` items, one mask-free call per ``seq_groups`` run. + + ``seq_groups`` holds contiguous ``(start_token, num_items, seq_len)`` runs of unpadded items. + """ + states = [projection(hidden_states) for projection in (self.q_proj, self.k_proj, self.v_proj)] + attn_output = torch.empty_like(states[0]) + for start, num_items, seq_len in seq_groups: + end, shape = start + num_items * seq_len, (num_items, seq_len, self.num_heads, self.head_dim) + attn_output[start:end].view(shape).copy_(self._attend_unpadded(*(x[start:end].view(shape) for x in states))) + return self.out_proj(attn_output) + + +class SiglipSdpaAttention(SiglipAttention): + """SigLIP attention through ``F.scaled_dot_product_attention`` (adapted from vllm-omni#5385).""" + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if output_attentions: + return super().forward(hidden_states, attention_mask=attention_mask, output_attentions=output_attentions) + batch_size, q_len, _ = hidden_states.size() + shape = (batch_size, q_len, self.num_heads, self.head_dim) + states = [projection(hidden_states).view(shape) for projection in (self.q_proj, self.k_proj, self.v_proj)] + attn_output = self._attend_unpadded(*states, attn_mask=attention_mask) + return self.out_proj(attn_output.reshape(batch_size, q_len, self.embed_dim)), None + + def _attend_unpadded( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: torch.Tensor | None = None + ) -> torch.Tensor: + query, key, value = (x.transpose(1, 2) for x in (query, key, value)) + return F.scaled_dot_product_attention( + query, key, value, attn_mask=attn_mask, dropout_p=self.dropout if self.training else 0.0, scale=self.scale + ).transpose(1, 2) + class SiglipFlashAttention2(SiglipAttention): """ @@ -1329,6 +1397,10 @@ class SiglipFlashAttention2(SiglipAttention): (max_seqlen_in_batch_q, max_seqlen_in_batch_k), ) + def _attend_unpadded(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + # Flash attention consumes the (n, L, heads, head_dim) layout directly. + return self._flash_attention_forward(query, key, value, None, query.shape[1]) + # Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Siglip class SiglipMLP(nn.Module): @@ -1352,7 +1424,8 @@ class SiglipEncoderLayer(nn.Module): super().__init__() self.embed_dim = config.hidden_size self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" - self.self_attn = SiglipAttention(config) if not self._use_flash_attention_2 else SiglipFlashAttention2(config) + attention_classes = {"flash_attention_2": SiglipFlashAttention2, "sdpa": SiglipSdpaAttention} + self.self_attn = attention_classes.get(config._attn_implementation, SiglipAttention)(config) self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.mlp = SiglipMLP(config) self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) @@ -1404,6 +1477,15 @@ class SiglipEncoderLayer(nn.Module): return outputs + def forward_packed(self, hidden_states: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]]) -> torch.Tensor: + """``forward`` over packed ``(tokens, embed)`` items; see ``SiglipAttention.forward_packed``.""" + residual = hidden_states + hidden_states = self.self_attn.forward_packed(self.layer_norm1(hidden_states), seq_groups) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp(self.layer_norm2(hidden_states)) + return residual + hidden_states + class SiglipPreTrainedModel(PreTrainedModel): """ @@ -1570,8 +1652,14 @@ class SiglipEncoder(nn.Module): class SiglipVisionTransformer(SiglipPreTrainedModel): config_class = SiglipVisionConfig main_input_name = "pixel_values" + # Transformers 5 spells the flash flag ``_supports_flash_attn``; SDPA is the + # default vision kernel (see MiniCPMO45OmniLLMForConditionalGeneration). + _supports_flash_attn = True _supports_flash_attn_2 = True + _supports_sdpa = True _no_split_modules = [] + # ``forward_packed`` through ``vision_fused.encode_packed_fused``; set by the model. + fused_layers = False def __init__(self, config: SiglipVisionConfig): super().__init__(config) @@ -1687,6 +1775,35 @@ class SiglipVisionTransformer(SiglipPreTrainedModel): attentions=encoder_outputs.attentions, ) + def forward_packed(self, pixel_values: torch.Tensor, layouts: Sequence[tuple[int, int, int]]) -> torch.Tensor: + """The padded ``forward`` per item, without padding, over items packed along the patch axis. + + ``pixel_values`` is ``(1, C, patch, tokens * patch)``, every item's strip in order (the stride-``patch`` + embedding cuts at item boundaries); ``layouts`` lists ``(h, w, count)`` runs. Returns ``(tokens, embed)``. + """ + embeddings = self.embeddings + device = embeddings.position_embedding.weight.device + position_ids = [ + embeddings.layout_position_ids(height, width, device).repeat(count) for height, width, count in layouts + ] + position_ids = position_ids[0] if len(position_ids) == 1 else torch.cat(position_ids) + patch_embeds = embeddings.patch_embedding(pixel_values).flatten(2).transpose(1, 2)[0] + hidden_states = patch_embeds + embeddings.position_embedding(position_ids) + + seq_groups: list[tuple[int, int, int]] = [] + start = 0 + # Adjacent equal-length grids share one attention call. + for seq_len, runs in itertools.groupby(layouts, key=lambda layout: layout[0] * layout[1]): + count = sum(run[2] for run in runs) + seq_groups.append((start, count, seq_len)) + start += count * seq_len + + if self.fused_layers and hidden_states.is_cuda and supports_fused_layers(self): + return encode_packed_fused(self, hidden_states, seq_groups) + for encoder_layer in self.encoder.layers: + hidden_states = encoder_layer.forward_packed(hidden_states, seq_groups) + return self.post_layernorm(hidden_states) + # ============== Resampler Classes ============== # coding=utf-8 @@ -1873,6 +1990,154 @@ class Resampler(nn.Module): def _repeat(self, query, N: int): return query.unsqueeze(1).repeat(1, N, 1) + def supports_packed(self) -> bool: + attn = self.attn + plain = not attn.batch_first and attn._qkv_same_embed_dim and attn.in_proj_weight is not None + return plain and attn.bias_k is None and attn.bias_v is None and not attn.add_zero_attn + + def forward_packed(self, x: torch.Tensor, layouts: Sequence[tuple[int, int, int]]) -> torch.Tensor: + """``forward`` without key padding for packed ``(tokens, kv_dim)`` features; returns ``(B, Q, D)``. + + Each ``(h, w, count)`` run of equal grids attends in one mask-free call (host ints: no device sync). + """ + max_h = max(height for height, _, _ in layouts) + max_w = max(width for _, width, _ in layouts) + if max_h > self.max_size[0] or max_w > self.max_size[1]: + self.max_size = [max(max_h, self.max_size[0]), max(max_w, self.max_size[1])] + self._set_2d_pos_cache(self.max_size, x.device) + if self.pos_embed.device != x.device: + self.pos_embed = self.pos_embed.to(x.device) + + x = self.ln_kv(self.kv_proj(x)) + key_input = torch.empty_like(x) + start = 0 + for height, width, count in layouts: + end = start + count * height * width + pos_embed = self.pos_embed[:height, :width, :].reshape(height * width, -1).to(x.dtype) + shape = (count, height * width, -1) + torch.add(x[start:end].view(shape), pos_embed, out=key_input[start:end].view(shape)) + start = end + + attn = self.attn + num_heads, head_dim = attn.num_heads, self.embed_dim // attn.num_heads + w_q, w_k, w_v = attn.in_proj_weight.chunk(3) + b_q, b_k, b_v = attn.in_proj_bias.chunk(3) if attn.in_proj_bias is not None else (None, None, None) + query = F.linear(self.ln_q(self.query), w_q, b_q) + keys = F.linear(key_input, w_k, b_k) + values = F.linear(x, w_v, b_v) + num_queries = query.shape[0] + query_heads = query.view(num_queries, num_heads, head_dim).transpose(0, 1) + eager = getattr(self, "attn_implementation", "sdpa") == "eager" + + outputs: list[torch.Tensor] = [] + start = 0 + for height, width, count in layouts: + seq_len = height * width + end = start + count * seq_len + shape = (count, seq_len, num_heads, head_dim) + key_heads, value_heads = (t[start:end].view(shape).transpose(1, 2) for t in (keys, values)) + q = query_heads.unsqueeze(0).expand(count, -1, -1, -1).contiguous() + if eager: + # Mirrors nn.MultiheadAttention's need_weights path used by ``forward``. + weights = torch.matmul(q / math.sqrt(head_dim), key_heads.transpose(-2, -1)) + out = torch.matmul(F.softmax(weights, dim=-1), value_heads) + else: + out = F.scaled_dot_product_attention(q, key_heads, value_heads) + outputs.append(out.transpose(1, 2).reshape(count, num_queries, self.embed_dim)) + start = end + + out = outputs[0] if len(outputs) == 1 else torch.cat(outputs) + out = self.ln_post(attn.out_proj(out)) + return out @ self.proj + + +def _vision_encode_paths(config: Any, *, encoder_graphs: bool) -> tuple[bool, bool]: + """``(vision_packed_encode, vision_cuda_graph)`` from the HF config. + + The packed path (_encode_vision_packed, vision_fused.py) serves the opt-in + ``vision_fused_layers`` and ``vision_cuda_graph``; otherwise the padded + batch runs and replays the SigLIP graph of ``vpm._encoder_graph``. + ``encoder_graphs`` (``encoder_cuda_graph`` and not ``--enforce-eager``) + gates the packed graphs too. ``vision_packed_encode`` overrides the path, + e.g. to measure packed eager encoding. + """ + graph = bool(getattr(config, "vision_cuda_graph", False)) and encoder_graphs + packed = getattr(config, "vision_packed_encode", None) + if packed is None: + return bool(getattr(config, "vision_fused_layers", False)) or graph, graph + return bool(packed), graph and bool(packed) + + +def _packed_vision_layouts( + pixel_values: Sequence[torch.Tensor], + tgt_sizes: torch.Tensor | Sequence[Sequence[int]], + patch_size: int, +) -> list[tuple[int, int]] | None: + """Host ``(h, w)`` patch grid per slice, or ``None`` on a mismatch; the packed path's only sync.""" + sizes = tgt_sizes.tolist() if isinstance(tgt_sizes, torch.Tensor) else [list(size) for size in tgt_sizes] + if len(sizes) != len(pixel_values) or not sizes: + return None + layouts: list[tuple[int, int]] = [] + for pixels, size in zip(pixel_values, sizes): + if len(size) != 2: + return None + height, width = int(size[0]), int(size[1]) + if ( + min(height, width) <= 0 + or pixels.ndim != 3 + or pixels.shape[-2:] != (patch_size, height * width * patch_size) + ): + return None + layouts.append((height, width)) + return layouts + + +def _plan_packed_vision_chunks( + layouts: Sequence[tuple[int, int]], max_items: int +) -> list[tuple[list[int], list[tuple[int, int, int]]]]: + """Per ``max_items``-slice chunk, its slice indices and ``(h, w, count)`` runs, equal grids adjacent. + + Sorting by ``h * w`` first also puts equal sequence lengths together, sharing an attention call. + """ + order = sorted(range(len(layouts)), key=lambda i: (layouts[i][0] * layouts[i][1], layouts[i])) + chunks: list[tuple[list[int], list[tuple[int, int, int]]]] = [] + for begin in range(0, len(order), max(1, max_items)): + indices = order[begin : begin + max(1, max_items)] + grids = itertools.groupby(layouts[index] for index in indices) + chunks.append((indices, [(height, width, len(list(run))) for (height, width), run in grids])) + return chunks + + +def _encode_vision_packed( + vpm: "SiglipVisionTransformer", + resampler: Resampler, + pixel_values: Sequence[torch.Tensor], + layouts: Sequence[tuple[int, int]], + max_items: int, + graph_encoder: VisionGraphEncoder | None = None, +) -> torch.Tensor: + """SigLIP + resampler over unpadded slices packed by grid; ``(num_slices, query_num, embed)`` in input order. + + No slice is padded, so its output does not depend on the others; single-grid chunks replay graphs. + """ + weight = vpm.embeddings.patch_embedding.weight + outputs: list[torch.Tensor] = [] + packed_order: list[int] = [] + for indices, runs in _plan_packed_vision_chunks(layouts, max_items): + parts = [pixel_values[index] for index in indices] + pixels = parts[0] if len(parts) == 1 else torch.cat(parts, dim=-1) + pixels = pixels.to(device=weight.device, dtype=weight.dtype).unsqueeze(0) + out = graph_encoder.encode(pixels, *runs[0]) if graph_encoder is not None and len(runs) == 1 else None + if out is None: + out = resampler.forward_packed(vpm.forward_packed(pixels, runs), runs) + outputs.append(out) + packed_order.extend(indices) + out = outputs[0] if len(outputs) == 1 else torch.cat(outputs) + if packed_order != list(range(len(packed_order))): + inverse = torch.tensor(packed_order, dtype=torch.long).argsort() + out = out.index_select(0, inverse.to(device=out.device, non_blocking=True)) + return out + class MultiheadAttention(nn.MultiheadAttention): def __init__( @@ -3944,7 +4209,10 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S multimodal_config = vllm_config.model_config.multimodal_config self.config = config - self.multimodal_config = multimodal_config + # Duplex session cap of this deployment; sizes the audio-encoder graph + # grid when ``duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions``. + self._duplex_max_sessions = int(getattr(vllm_config.model_config, "duplex_max_sessions", 1) or 1) + self._enforce_eager = bool(vllm_config.model_config.enforce_eager) # Model-local opt-out for A/B measurements; --enforce-eager always wins. encoder_graphs = ( bool(getattr(config, "encoder_cuda_graph", True)) and not vllm_config.model_config.enforce_eager @@ -3970,10 +4238,15 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S # Initialize vision encoder (SigLIP) if multimodal_config.get_limit_per_prompt("image"): # Set attention implementation - if hasattr(vllm_config, "model_config") and hasattr(vllm_config.model_config, "_attn_implementation"): + # (an engine override, else the vision config, else SDPA: eager math without the (B, H, L, L) scores) + if getattr(getattr(vllm_config, "model_config", None), "_attn_implementation", None) is not None: config.vision_config._attn_implementation = vllm_config.model_config._attn_implementation else: - config.vision_config._attn_implementation = "eager" + vision_implementation = getattr(config.vision_config, "_attn_implementation", None) + config.vision_config._attn_implementation = vision_implementation or "sdpa" + if config.vision_config._attn_implementation == "flash_attention_2" and not is_flash_attn_2_available(): + logger.warning("flash_attention_2 requested for the MiniCPM-o vision tower but unavailable; using sdpa") + config.vision_config._attn_implementation = "sdpa" self.vpm = SiglipVisionTransformer(config.vision_config) if encoder_graphs: @@ -4029,8 +4302,13 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S kv_dim=vision_dim, # Vision encoder dimension adaptive=True, # Enable adaptive resampling ) + self.resampler.attn_implementation = config.vision_config._attn_implementation else: self.resampler = None + if self.vpm is not None: + self.vpm.fused_layers = bool(getattr(config, "vision_fused_layers", False)) + self.vision_packed_encode, self.vision_cuda_graph = _vision_encode_paths(config, encoder_graphs=encoder_graphs) + self._vision_graph_encoder: VisionGraphEncoder | None = None # Initialize audio encoder (APM) and audio projection if getattr(config, "init_audio", True) and hasattr(config, "audio_config") and config.audio_config is not None: @@ -4309,6 +4587,19 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S pixel_values = data["pixel_values"] tgt_sizes = data["tgt_sizes"] + vpm = self.vpm + packed = getattr(self, "vision_packed_encode", False) and hasattr(vpm, "forward_packed") + if packed and hasattr(self.resampler, "forward_packed") and self.resampler.supports_packed(): + layouts = _packed_vision_layouts(pixel_values, tgt_sizes, vpm.embeddings.patch_size) + if layouts is not None: + graph_encoder = None + if getattr(self, "vision_cuda_graph", False) and pixel_values[0].is_cuda: + if self._vision_graph_encoder is None: + self._vision_graph_encoder = VisionGraphEncoder(vpm, self.resampler) + graph_encoder = self._vision_graph_encoder + max_items = max(1, int(self.config.vision_batch_size)) + return _encode_vision_packed(vpm, self.resampler, pixel_values, layouts, max_items, graph_encoder) + B = len(pixel_values) P = pixel_values[0].shape[-2] L = max(item.shape[-1] for item in pixel_values) @@ -4488,8 +4779,8 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S audio_embeds = self._encode_audio_features(wavforms, audio_attention_mask) _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens) - - num_audio_tokens = feature_lens_after_pooling + # One host read of every length instead of an implicit .item() per slice bound. + num_audio_tokens = feature_lens_after_pooling.tolist() final_audio_embeds = list[torch.Tensor]() idx = 0 @@ -4590,18 +4881,23 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S device=wavforms.device, ) + # As in get_audio_embedding: every layer's output only to pick the last one keeps 25 activations alive. + selects_final_layer = self.audio_encoder_layer == -1 audio_outputs = self.apm( wavforms, past_key_values=self.audio_past_key_values, use_cache=True, - output_hidden_states=True, + output_hidden_states=not selects_final_layer, attention_mask=audio_attention_mask, use_extra_context=use_extra_context, prefix_extra_frames=prefix_extra_frames, suffix_extra_frames=suffix_extra_frames, cnn_min_length=cnn_min_length, ) - audio_states = audio_outputs.hidden_states[self.audio_encoder_layer] + if selects_final_layer: + audio_states = audio_outputs.last_hidden_state + else: + audio_states = audio_outputs.hidden_states[self.audio_encoder_layer] self.audio_past_key_values = audio_outputs.past_key_values audio_embeds = self.audio_projection_layer(audio_states) @@ -4610,16 +4906,97 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S audio_embeds = audio_embeds.transpose(1, 2) _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens) + # One host read of the lengths instead of a device sync per slice bound. + pooled_lens = feature_lens_after_pooling.tolist() final_audio_embeds: list[list[torch.Tensor]] = [] idx = 0 for i in range(len(audio_feature_lens_raw)): target_audio_embeds = [] for _ in range(len(audio_feature_lens_raw[i])): - target_audio_embeds.append(audio_embeds[idx, : feature_lens_after_pooling[idx], :]) + target_audio_embeds.append(audio_embeds[idx, : pooled_lens[idx], :]) idx += 1 final_audio_embeds.append(target_audio_embeds) return final_audio_embeds + def supports_streaming_audio_batch(self) -> bool: + """Whether ``get_audio_embedding_streaming_batch`` can serve this encoder.""" + if self.apm is None or self.audio_projection_layer is None or self.audio_avg_pooler is None: + return False + requested = getattr(self.config, "duplex_audio_encoder_batching", None) + if requested is not None and not bool(requested): + return False + reason = streaming_audio.streaming_batch_unsupported_reason( + self.apm, audio_encoder_layer=self.audio_encoder_layer + ) + if reason is not None: + logger.info("MiniCPM-o streaming audio encoder stays per-session: %s", reason) + return False + return requested is not None or self.apm.conv1.weight.device.type in ("cuda", "cpu") + + def build_streaming_audio_graph_encoder(self, *, unit_frames: int) -> bool: + """Capture CUDA graphs of the streaming encoder's steady unit. Failures stay eager.""" + self._duplex_audio_cuda_graph_encoder = None + config = self.config + enabled = bool(getattr(config, "duplex_audio_encoder_cuda_graph", True)) and not self._enforce_eager + if not self.supports_streaming_audio_batch() or not enabled: + return False + if self.apm.conv1.weight.device.type != "cuda": + logger.info("MiniCPM-o streaming audio encoder CUDA graph needs a CUDA device; staying eager") + return False + batch_sizes = getattr(config, "duplex_audio_encoder_cuda_graph_batch_sizes", None) + if getattr(config, "duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions", False) is True: + # Overrides the explicit list: the grid follows duplex max_sessions. + batch_sizes = streaming_audio_graph.batch_sizes_for_sessions(getattr(self, "_duplex_max_sessions", 1)) + allocated = torch.accelerator.memory_allocated(self.apm.conv1.weight.device) + try: + wrapper = streaming_audio_graph.StreamingAudioGraphEncoder( + self.apm, + self.audio_projection_layer, + self.audio_avg_pooler, + unit_frames=int(unit_frames), + pool_step=int(config.audio_pool_step), + batch_sizes=batch_sizes, + cache_buckets=getattr(config, "duplex_audio_encoder_cuda_graph_cache_buckets", None), + page_positions=self._duplex_audio_kv_page_positions(), + pinned_h2d=getattr(config, "duplex_audio_encoder_pinned_h2d", False) is True, + ) + wrapper.capture() + except Exception: + logger.exception("MiniCPM-o streaming audio encoder CUDA graph capture failed; staying eager") + return False + self._duplex_audio_cuda_graph_encoder = wrapper + logger.info( + "Captured MiniCPM-o streaming audio encoder CUDA graphs: unit_frames=%d batch_sizes=%s " + "cache_buckets=%s pinned_h2d=%s, took %.2f GiB", + unit_frames, + wrapper.batch_sizes, + wrapper.cache_buckets, + wrapper.pinned_h2d, + (torch.accelerator.memory_allocated(self.apm.conv1.weight.device) - allocated) / (1 << 30), + ) + return True + + def get_audio_embedding_streaming_batch( + self, + chunks: Sequence[Any], + ) -> tuple[list[torch.Tensor | None], list[Any]]: + """Batched ``get_audio_embedding_streaming`` (embeddings, updated caches); graphs take steady rows.""" + graph = getattr(self, "_duplex_audio_cuda_graph_encoder", None) + if graph is not None: + return graph.encode(chunks) + return streaming_audio.encode_streaming_audio_batch( + self.apm, + self.audio_projection_layer, + self.audio_avg_pooler, + chunks, + pool_step=int(self.config.audio_pool_step), + page_positions=self._duplex_audio_kv_page_positions(), + ) + + def _duplex_audio_kv_page_positions(self) -> int: + page_positions = getattr(self.config, "duplex_audio_kv_page_positions", None) + return int(page_positions or streaming_audio.DEFAULT_KV_PAGE_POSITIONS) + def _process_audio_input( self, audio_input: MiniCPMOAudioInputs, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder.py b/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder.py new file mode 100644 index 000000000..f0d0fa824 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Batched streaming Whisper encoder for MiniCPM-o duplex sessions.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass, field + +import torch +from torch import nn + +#: Growth step of a session buffer, in encoder positions (a 1 s unit is 50): +#: at most five reallocations per 1500-position cycle, ~12 MiB average slack. +DEFAULT_KV_PAGE_POSITIONS = 250 + +_SUPPORTED_ATTENTION = ("sdpa", "eager") +_SUPPORTED_DTYPES = (torch.bfloat16, torch.float32) + + +@dataclass(eq=False, slots=True) +class StreamingAudioKVCache: + """One session's streaming-encoder self-attention cache, written in place. + + Layout ``[layers, 2 (key/value), capacity, embed_dim]``: token-major, so appending a unit is one contiguous + copy per layer and attention reads a strided ``[1, heads, length, head_dim]`` view without materializing it. + Capacity grows in ``page_positions`` steps up to ``max_positions``. + """ + + num_layers: int + embed_dim: int + num_heads: int + max_positions: int + page_positions: int = DEFAULT_KV_PAGE_POSITIONS + length: int = field(default=0, init=False) + _buffer: torch.Tensor | None = field(default=None, init=False, repr=False) + + def __post_init__(self) -> None: + if min(self.num_layers, self.embed_dim, self.num_heads, self.max_positions, self.page_positions) <= 0: + raise ValueError("streaming audio KV cache dimensions must be positive") + + @classmethod + def for_encoder(cls, encoder: nn.Module, page_positions: int) -> StreamingAudioKVCache: + config, positions = encoder.config, int(encoder.embed_positions.weight.shape[0]) + return cls( + len(encoder.layers), int(config.d_model), int(config.encoder_attention_heads), positions, page_positions + ) + + @property + def capacity(self) -> int: + return 0 if self._buffer is None else int(self._buffer.shape[2]) + + def reserve(self, positions: int, *, dtype: torch.dtype, device: torch.device) -> None: + """Make room for ``positions`` entries, keeping the committed prefix.""" + if positions <= self.capacity: + return + pages = math.ceil(positions / self.page_positions) * self.page_positions + capacity = max(positions, min(self.max_positions, pages)) + buffer = torch.empty((self.num_layers, 2, capacity, self.embed_dim), dtype=dtype, device=device) + if self._buffer is not None and self.length: + buffer[:, :, : self.length].copy_(self._buffer[:, :, : self.length]) + self._buffer = buffer + + def history(self, length: int) -> torch.Tensor: + """``[layers, 2, length, embed_dim]`` rows ``[0, length)``, every layer at once.""" + assert self._buffer is not None + return self._buffer[:, :, :length] + + def commit(self, offset: int, new_kv: torch.Tensor) -> None: + """Write ``[layers, 2, positions, embed_dim]`` rows at ``offset``; the caller advances ``length``.""" + assert self._buffer is not None + positions = int(new_kv.shape[2]) + self._buffer[:, :, offset : offset + positions].copy_(new_kv) + + def layer_views(self, start: int, end: int) -> tuple[tuple[torch.Tensor, ...], ...]: + """Per layer, ``(key_rows, value_rows, key_heads, value_heads)``: the ``[end - start, embed_dim]`` + write targets and the ``[1, heads, end, head_dim]`` attention inputs over ``[0, end)``.""" + assert self._buffer is not None + head_dim = self.embed_dim // self.num_heads + heads = ( + self._buffer[:, :, :end] + .view(self.num_layers, 2, end, self.num_heads, head_dim) + .transpose(2, 3) + .unsqueeze(2) + ) + rows = self._buffer[:, :, start:end] + return rows[:, 0].unbind(0), rows[:, 1].unbind(0), heads[:, 0].unbind(0), heads[:, 1].unbind(0) + + def to_legacy_cache(self): + """The committed history as the ``EncoderDecoderCache`` the per-session path grows.""" + from transformers.cache_utils import DynamicCache, EncoderDecoderCache + + self_attention = DynamicCache() + if self.length: + _, _, keys, values = self.layer_views(0, self.length) + for layer in range(self.num_layers): + self_attention.update(keys[layer].contiguous(), values[layer].contiguous(), layer) + return EncoderDecoderCache(self_attention, DynamicCache()) + + +@dataclass(frozen=True, slots=True) +class StreamingAudioChunk: + """One session's ready unit: log-mel frames plus the state it continues.""" + + #: ``[n_mels, frames]`` or ``[1, n_mels, frames]`` log-mel features, host or device. + features: torch.Tensor + #: Session cache; ``None`` starts a new history. + cache: StreamingAudioKVCache | None + prefix_extra_frames: int = 0 + suffix_extra_frames: int = 0 + use_extra_context: bool = True + #: ``audio_feature_lens`` of the unit when it differs from ``features.shape[-1]``. + feature_length: int | None = None + + +@dataclass(slots=True) +class _Row: + index: int + frames: int + start: int + length: int + pooled_length: int + cache: StreamingAudioKVCache + past: int = 0 + offset: int = 0 + + @property + def total(self) -> int: + return self.past + self.length + + +def _extra_context_trim(frames: int) -> int: + return (int(frames) + 1) // 2 if frames > 0 else 0 + + +def streaming_batch_unsupported_reason(encoder: nn.Module, *, audio_encoder_layer: int | None) -> str | None: + """Why ``encoder`` cannot take the batched path, or ``None`` when it can.""" + if audio_encoder_layer != -1: + return "only the final encoder layer output is batched" + if encoder.training: + return "training mode (dropout / layerdrop)" + weight = getattr(getattr(encoder, "conv1", None), "weight", None) + if weight is None or weight.dtype not in _SUPPORTED_DTYPES: + # float16 needs the per-session inf/nan clamp, which is per tensor. + return f"unsupported encoder dtype {getattr(weight, 'dtype', None)}" + layers = getattr(encoder, "layers", None) + if not layers: + return "encoder has no layers" + for layer in layers: + attention = getattr(layer, "self_attn", None) + if attention is None or not all( + hasattr(attention, name) for name in ("q_proj", "k_proj", "v_proj", "out_proj", "scaling") + ): + return "encoder layers are not Whisper self-attention layers" + implementation = getattr(encoder.config, "_attn_implementation", None) or "eager" + if implementation not in _SUPPORTED_ATTENTION: + return f"attention implementation {implementation!r} is not batched" + return None + + +def _attend(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor, implementation: str): + """Whisper self-attention as the per-session path calls it; ``query`` already carries ``scaling``.""" + if implementation == "sdpa": + return torch.nn.functional.scaled_dot_product_attention( + query, key, value, attn_mask=mask, dropout_p=0.0, scale=1.0 + ) + weights = torch.matmul(query, key.transpose(2, 3)) + mask + return torch.matmul(torch.nn.functional.softmax(weights, dim=-1), value) + + +def _stack_features(rows: list[_Row], features: list[torch.Tensor], device: torch.device) -> torch.Tensor: + """Zero-padded ``[B, mels, max_frames]`` on ``device`` with one host-to-device copy.""" + max_frames = max(row.frames for row in rows) + on_host = all(feature.device.type == "cpu" for feature in features) + pin = on_host and device.type == "cuda" + batch = torch.zeros( + (len(rows), int(features[0].shape[0]), max_frames), + dtype=features[0].dtype, + device="cpu" if on_host else device, + pin_memory=pin, + ) + for slot, (row, feature) in enumerate(zip(rows, features, strict=True)): + batch[slot, :, : row.frames].copy_(feature) + # Pinned staging keeps the host from waiting on queued GPU work; the + # caching host allocator holds the buffer until the copy completes. + return batch.to(device, non_blocking=pin) + + +def _position_rows(encoder: nn.Module, rows: list[_Row]) -> torch.Tensor: + """Learned positions ``past .. past + n`` per row; past the table, repeat its last row.""" + table = encoder.embed_positions.weight + capacity = int(table.shape[0]) + pieces: list[torch.Tensor] = [] + for row in rows: + end = row.past + row.length + if end <= capacity: + pieces.append(table[row.past : end]) + else: + pieces.append(table[row.past :]) + pieces.append(table[-1:].expand(end - max(row.past, capacity), -1)) + return pieces[0] if len(pieces) == 1 else torch.cat(pieces) + + +@torch.inference_mode() +def encode_streaming_audio_batch( + encoder: nn.Module, + projection: nn.Module, + pooler: nn.Module, + chunks: Sequence[StreamingAudioChunk], + *, + pool_step: int, + page_positions: int = DEFAULT_KV_PAGE_POSITIONS, +) -> tuple[list[torch.Tensor | None], list[StreamingAudioKVCache | None]]: + """Encode one streaming unit per session in one encoder pass.""" + implementation = getattr(encoder.config, "_attn_implementation", None) or "eager" + weight = encoder.conv1.weight + dtype, device = weight.dtype, weight.device + max_positions = int(encoder.embed_positions.weight.shape[0]) + embed_dim, num_heads = int(encoder.config.d_model), int(encoder.config.encoder_attention_heads) + + outputs: list[torch.Tensor | None] = [None] * len(chunks) + caches: list[StreamingAudioKVCache | None] = [chunk.cache for chunk in chunks] + rows: list[_Row] = [] + features: list[torch.Tensor] = [] + for index, chunk in enumerate(chunks): + feature = chunk.features + if feature.ndim == 3: + if feature.shape[0] != 1: + raise ValueError("a streaming chunk carries exactly one audio") + feature = feature[0] + frames = int(feature.shape[-1]) + conv_length = (frames - 1) // 2 + 1 + prefix = _extra_context_trim(chunk.prefix_extra_frames) if chunk.use_extra_context else 0 + suffix = _extra_context_trim(chunk.suffix_extra_frames) if chunk.use_extra_context else 0 + length = conv_length - prefix - suffix + if frames <= 0 or length <= 0: + continue + feature_length = frames if chunk.feature_length is None else int(chunk.feature_length) + pooled_length = ((feature_length - 1) // 2 + 1 - pool_step) // pool_step + 1 + cache = chunk.cache + if cache is None or cache.length + length >= max_positions: + cache = StreamingAudioKVCache.for_encoder(encoder, page_positions) + rows.append(_Row(index, frames, prefix, length, pooled_length, cache, past=cache.length)) + features.append(feature) + if not rows: + return outputs, caches + + offset = 0 + for row in rows: + row.offset = offset + offset += row.length + row.cache.reserve(row.total, dtype=dtype, device=device) + + # CNN front end on the zero-padded batch, transformer layers on the packed rows. + mel = _stack_features(rows, features, device).to(dtype=dtype) + hidden = nn.functional.gelu(encoder.conv1(mel)) + max_frames = int(mel.shape[-1]) + for slot, row in enumerate(rows): + if row.frames < max_frames: + # conv2 must see zeros past each row's end, as its own padding would in the per-session call. + hidden[slot, :, row.frames :].zero_() + hidden = nn.functional.gelu(encoder.conv2(hidden)).permute(0, 2, 1) + pieces = [hidden[slot, row.start : row.start + row.length] for slot, row in enumerate(rows)] + hidden = pieces[0] if len(pieces) == 1 else torch.cat(pieces) + hidden = hidden + _position_rows(encoder, rows) + + # Every per-row view is built once here, not per layer: at batch sizes + # that matter the host cost of view ops, not the GPU, bounds this loop. + views = [row.cache.layer_views(row.past, row.total) for row in rows] + lengths = [row.length for row in rows] + total_rows = sum(lengths) + head_dim = embed_dim // num_heads + # The per-session path passes an all-zero additive mask; keep it for the same attention kernel. + masks = [torch.zeros((1, 1, row.length, row.total), dtype=dtype, device=device) for row in rows] + for layer_index, layer in enumerate(encoder.layers): + attention = layer.self_attn + residual = hidden + normed = layer.self_attn_layer_norm(hidden) + query = attention.q_proj(normed) * attention.scaling + key, value = attention.k_proj(normed), attention.v_proj(normed) + destinations = [view[0][layer_index] for view in views] + [view[1][layer_index] for view in views] + torch._foreach_copy_(destinations, [*key.split(lengths), *value.split(lengths)]) + queries = query.view(1, total_rows, num_heads, head_dim).transpose(1, 2).split(lengths, dim=2) + attended = [ + _attend(row_query, view[2][layer_index], view[3][layer_index], mask, implementation) + for row_query, view, mask in zip(queries, views, masks, strict=True) + ] + # [1, heads, rows, head_dim] per row -> packed [rows, embed_dim]. + attended_rows = (attended[0] if len(attended) == 1 else torch.cat(attended, dim=2)).transpose(1, 2) + hidden = residual + attention.out_proj(attended_rows.reshape(total_rows, embed_dim)) + residual = hidden + hidden = layer.activation_fn(layer.fc1(layer.final_layer_norm(hidden))) + hidden = residual + layer.fc2(hidden) + hidden = encoder.layer_norm(hidden) + embeds = projection(hidden) + if all(row.length % pool_step == 0 for row in rows): + pooled = pooler(embeds.transpose(0, 1).unsqueeze(0)).squeeze(0).transpose(0, 1) + for row in rows: + start = row.offset // pool_step + count = min(row.length // pool_step, row.pooled_length) + outputs[row.index] = pooled[start : start + count] + else: + for row in rows: + segment = embeds[row.offset : row.offset + row.length].transpose(0, 1).unsqueeze(0) + pooled = pooler(segment).squeeze(0).transpose(0, 1) + outputs[row.index] = pooled[: row.pooled_length] + + for row in rows: + row.cache.length = row.total + caches[row.index] = row.cache + return outputs, caches diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder_graph.py b/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder_graph.py new file mode 100644 index 000000000..3c8833d60 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder_graph.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA graphs of MiniCPM-o's streaming Whisper encoder for the steady unit shape.""" + +from __future__ import annotations + +import functools +import itertools +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from torch import nn + +from .streaming_audio_encoder import ( + DEFAULT_KV_PAGE_POSITIONS, + StreamingAudioChunk, + StreamingAudioKVCache, + _attend, + _extra_context_trim, + encode_streaming_audio_batch, +) + +if TYPE_CHECKING: + from torch.cuda import CUDAGraph + +_NEG_INF = float("-inf") + +# stage0.py builds every unit but a session's first with these trims. +STEADY_PREFIX_EXTRA_FRAMES = 2 +STEADY_SUFFIX_EXTRA_FRAMES = 2 + +DEFAULT_GRAPH_BATCH_SIZES: tuple[int, ...] = (1, 2, 4, 8, 16, 32) +# Shared storage scales with the largest bucket only. 900 stays under ~3 GiB +# at batch 32 on MiniCPM-o 4.5 (bf16); history past it falls back to eager. +DEFAULT_GRAPH_CACHE_BUCKETS: tuple[int, ...] = (250, 500, 750, 900) + + +def batch_sizes_for_sessions(max_sessions: int) -> tuple[int, ...]: + """Powers of two below ``max_sessions``, plus ``max_sessions`` itself.""" + top = max(1, int(max_sessions)) + sizes = [1 << i for i in range(top.bit_length()) if (1 << i) < top] + return normalize_buckets([*sizes, top]) + + +def select_bucket(value: int, buckets: Sequence[int]) -> int | None: + return min((bucket for bucket in buckets if value <= bucket), default=None) + + +def normalize_buckets(values: Sequence[int]) -> tuple[int, ...]: + return tuple(sorted({int(v) for v in values if int(v) > 0})) + + +def steady_unit_length(unit_frames: int) -> int: + conv_length = (int(unit_frames) - 1) // 2 + 1 + trim = _extra_context_trim(STEADY_PREFIX_EXTRA_FRAMES) + _extra_context_trim(STEADY_SUFFIX_EXTRA_FRAMES) + return conv_length - trim + + +def row_is_steady(chunk: StreamingAudioChunk, *, expected_frames: int) -> bool: + frames = int(chunk.features.shape[-1]) + return ( + chunk.cache is not None + and chunk.use_extra_context + and (int(chunk.prefix_extra_frames), int(chunk.suffix_extra_frames)) + == (STEADY_PREFIX_EXTRA_FRAMES, STEADY_SUFFIX_EXTRA_FRAMES) + and frames == expected_frames + and (chunk.feature_length is None or int(chunk.feature_length) == frames) + ) + + +@dataclass(slots=True) +class _Group: + index: int + cache_in: StreamingAudioKVCache + cache_out: StreamingAudioKVCache + past: int + + +class StreamingAudioGraphEncoder: + """Pre-captured CUDA graphs of the streaming encoder's steady unit shape.""" + + def __init__( + self, + encoder: nn.Module, + projection: nn.Module, + pooler: nn.Module, + *, + unit_frames: int, + pool_step: int, + batch_sizes: Sequence[int] | None = None, + cache_buckets: Sequence[int] | None = None, + page_positions: int = DEFAULT_KV_PAGE_POSITIONS, + pinned_h2d: bool = False, + ) -> None: + weight, config = encoder.conv1.weight, encoder.config + if weight.device.type != "cuda": + raise ValueError("StreamingAudioGraphEncoder requires a CUDA encoder") + self.encoder, self.projection, self.pooler = encoder, projection, pooler + self.device, self.dtype = weight.device, weight.dtype + self.pool_step, self.page_positions, self.pinned_h2d = int(pool_step), int(page_positions), bool(pinned_h2d) + self.embed_dim, self.num_heads = int(config.d_model), int(config.encoder_attention_heads) + self.max_positions = int(encoder.embed_positions.weight.shape[0]) + self.implementation = getattr(config, "_attn_implementation", None) or "eager" + self.unit_frames = int(unit_frames) + self.unit_length = steady_unit_length(self.unit_frames) + nominal_pooled = ((self.unit_frames - 1) // 2 + 1 - self.pool_step) // self.pool_step + 1 + self.pooled_length = min(nominal_pooled, self.unit_length // self.pool_step) + if self.unit_length <= 0 or self.pooled_length <= 0: + raise ValueError(f"unit_frames={unit_frames} pool_step={pool_step} leaves no steady unit output") + self.batch_sizes = normalize_buckets(batch_sizes or DEFAULT_GRAPH_BATCH_SIZES) + self.cache_buckets = normalize_buckets(cache_buckets or DEFAULT_GRAPH_CACHE_BUCKETS) + if not self.batch_sizes or not self.cache_buckets: + raise ValueError("StreamingAudioGraphEncoder needs at least one batch size and one cache bucket") + self.max_batch, self.max_cache_bucket = max(self.batch_sizes), max(self.cache_buckets) + # (batch, cache length) -> (graph, pooled output); every graph reads prefix views of one storage. + self._graphs: dict[tuple[int, int], tuple[CUDAGraph, torch.Tensor]] = {} + + def _views(self, batch: int, cache_len: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + total = cache_len + self.unit_length + return ( + self._mel[:batch], + self._cache[:, :, :batch, :total, :], + self._mask[:batch, :, :, :total], + self._position_ids[:batch], + ) + + def capture(self) -> None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("cannot capture the streaming audio encoder graph mid-capture") + from vllm.platforms import current_platform + + batch, total = self.max_batch, self.max_cache_bucket + self.unit_length + zeros = functools.partial(torch.zeros, dtype=self.dtype, device=self.device) + self._mel = zeros((batch, int(self.encoder.config.num_mel_bins), self.unit_frames)) + self._cache = zeros((len(self.encoder.layers), 2, batch, total, self.embed_dim)) + self._mask = zeros((batch, 1, self.unit_length, total)) + self._position_ids = torch.zeros((batch, self.unit_length), dtype=torch.long, device=self.device) + pool = current_platform.get_global_graph_pool() + current_stream = torch.cuda.current_stream(self.device) + for key in itertools.product(self.batch_sizes, self.cache_buckets): + views = self._views(*key) + warmup_stream = torch.cuda.Stream(device=self.device) + warmup_stream.wait_stream(current_stream) + with torch.cuda.stream(warmup_stream), torch.inference_mode(): + for _ in range(3): + self._forward(*views, key[1]) + current_stream.wait_stream(warmup_stream) + graph = torch.cuda.CUDAGraph() + with torch.inference_mode(), torch.cuda.graph(graph, pool=pool): + self._graphs[key] = (graph, self._forward(*views, key[1])) + + def _forward( + self, + mel: torch.Tensor, + cache: torch.Tensor, + mask: torch.Tensor, + position_ids: torch.Tensor, + cache_len: int, + ) -> torch.Tensor: + encoder = self.encoder + batch = mel.shape[0] + hidden = nn.functional.gelu(encoder.conv1(mel)) + hidden = nn.functional.gelu(encoder.conv2(hidden)).permute(0, 2, 1) + prefix = _extra_context_trim(STEADY_PREFIX_EXTRA_FRAMES) + hidden = hidden[:, prefix : prefix + self.unit_length, :] + hidden = hidden + encoder.embed_positions(position_ids) + head_dim = self.embed_dim // self.num_heads + for layer_index, layer in enumerate(encoder.layers): + attention = layer.self_attn + residual = hidden + normed = layer.self_attn_layer_norm(hidden) + query = attention.q_proj(normed) * attention.scaling + key, value = attention.k_proj(normed), attention.v_proj(normed) + cache[layer_index, 0, :, cache_len : cache_len + self.unit_length, :].copy_(key) + cache[layer_index, 1, :, cache_len : cache_len + self.unit_length, :].copy_(value) + q = query.view(batch, self.unit_length, self.num_heads, head_dim).transpose(1, 2) + k_all = cache[layer_index, 0].view(batch, -1, self.num_heads, head_dim).transpose(1, 2) + v_all = cache[layer_index, 1].view(batch, -1, self.num_heads, head_dim).transpose(1, 2) + attended = _attend(q, k_all, v_all, mask, self.implementation) + attended = attended.transpose(1, 2).reshape(batch, self.unit_length, self.embed_dim) + hidden = residual + attention.out_proj(attended) + residual = hidden + hidden = layer.activation_fn(layer.fc1(layer.final_layer_norm(hidden))) + hidden = residual + layer.fc2(hidden) + hidden = encoder.layer_norm(hidden) + embeds = self.projection(hidden) + return self.pooler(embeds.transpose(1, 2)).transpose(1, 2) + + def encode( + self, + chunks: Sequence[StreamingAudioChunk], + ) -> tuple[list[torch.Tensor | None], list[StreamingAudioKVCache | None]]: + outputs: list[torch.Tensor | None] = [None] * len(chunks) + caches: list[StreamingAudioKVCache | None] = [chunk.cache for chunk in chunks] + steady: list[_Group] = [] + eager: list[int] = [] + for index, chunk in enumerate(chunks): + cache_in = chunk.cache + reset = cache_in is not None and cache_in.length + self.unit_length >= self.max_positions + past = 0 if reset or cache_in is None else cache_in.length + if not row_is_steady(chunk, expected_frames=self.unit_frames) or past > self.max_cache_bucket: + eager.append(index) + continue + cache_out = StreamingAudioKVCache.for_encoder(self.encoder, self.page_positions) if reset else cache_in + steady.append(_Group(index, cache_in, cache_out, past)) + for start in range(0, len(steady), self.max_batch): + self._replay(steady[start : start + self.max_batch], chunks, outputs, caches) + if eager: + eager_outputs, eager_caches = encode_streaming_audio_batch( + self.encoder, + self.projection, + self.pooler, + [chunks[i] for i in eager], + pool_step=self.pool_step, + page_positions=self.page_positions, + ) + for index, output, cache in zip(eager, eager_outputs, eager_caches, strict=True): + outputs[index], caches[index] = output, cache + return outputs, caches + + @torch.inference_mode() + def _replay( + self, + group: list[_Group], + chunks: Sequence[StreamingAudioChunk], + outputs: list[torch.Tensor | None], + caches: list[StreamingAudioKVCache | None], + ) -> None: + batch = select_bucket(len(group), self.batch_sizes) + cache_len = select_bucket(max((row.past for row in group), default=0), self.cache_buckets) + assert batch is not None and cache_len is not None + mel, cache, mask, position_ids = self._views(batch, cache_len) + mask.fill_(0.0) + if len(group) < batch: + mask[len(group) :, :, :, :cache_len].fill_(_NEG_INF) + position_ids[len(group) :].copy_(torch.arange(self.unit_length, device=self.device)) + for slot, row in enumerate(group): + feature = chunks[row.index].features + feature = (feature[0] if feature.ndim == 3 else feature).to(dtype=self.dtype) + if self.pinned_h2d and feature.device.type == "cpu": + mel[slot].copy_(feature.pin_memory(), non_blocking=True) + else: + mel[slot].copy_(feature.to(device=self.device)) + pad = cache_len - row.past + if pad: + mask[slot, :, :, :pad].fill_(_NEG_INF) + if row.past: + cache[:, :, slot, pad:cache_len, :].copy_(row.cache_in.history(row.past)) + position_ids[slot].copy_(torch.arange(row.past, row.past + self.unit_length, device=self.device)) + graph, pooled = self._graphs[(batch, cache_len)] + graph.replay() + for slot, row in enumerate(group): + outputs[row.index] = pooled[slot, : self.pooled_length].clone() + row.cache_out.reserve(row.past + self.unit_length, dtype=self.dtype, device=self.device) + row.cache_out.commit(row.past, cache[:, :, slot, cache_len : cache_len + self.unit_length, :]) + row.cache_out.length = row.past + self.unit_length + caches[row.index] = row.cache_out diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/vision_fused.py b/vllm_omni/model_executor/models/minicpmo_4_5/vision_fused.py new file mode 100644 index 000000000..4bba6e113 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/vision_fused.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused packed SigLIP layers and CUDA graphs for MiniCPM-o 4.5 vision encode.""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Sequence +from copy import copy +from math import prod +from typing import Any, ClassVar, Literal + +import torch +import torch.nn.functional as F +from torch import nn +from vllm.config import ModelConfig, VllmConfig +from vllm.config.multimodal import MultiModalConfig +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph +from vllm.platforms import current_platform +from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager +from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + EncoderCudaGraphConfig, + EncoderCudaGraphReplayBuffers, + EncoderItemSpec, +) + +from .ops import residual_layer_norm + +logger = init_logger(__name__) + +# Batch buckets of the vision graphs; the checkpoint's vision_batch_size (16) +# bounds a chunk, so it is the largest bucket. +DEFAULT_GRAPH_BATCH_SIZES = (1, 2, 4, 8, 16) +# Distinct (grid, bucket) graphs kept at once; duplex uses one grid. +_MAX_GRAPHS = 16 +# A (grid, bucket) is captured on its third encode, so one-off shapes (the +# startup profile run, a single turn-mode image) stay eager. +_CAPTURE_AFTER = 2 + + +def supports_fused_layers(vpm: nn.Module) -> bool: + """Whether ``encode_packed_fused`` covers ``vpm`` (SigLIP with a GELU MLP).""" + layers = getattr(getattr(vpm, "encoder", None), "layers", None) + if not layers or getattr(vpm, "post_layernorm", None) is None: + return False + if getattr(getattr(vpm, "config", None), "_attn_implementation", None) == "eager": + return False + attn, mlp = getattr(layers[0], "self_attn", None), getattr(layers[0], "mlp", None) + return all( + getattr(attn, name, None) is not None for name in ("q_proj", "k_proj", "v_proj", "out_proj") + ) and getattr(getattr(mlp, "config", None), "hidden_act", None) in ("gelu_pytorch_tanh", "gelu") + + +def packed_qkv(attn: nn.Module) -> tuple[torch.Tensor, torch.Tensor | None]: + """``(3E, E)`` q/k/v weight (and bias) of ``attn``, stacked once.""" + cached = attn.__dict__.get("_packed_qkv") + projections = (attn.q_proj, attn.k_proj, attn.v_proj) + if cached is not None and all( + p.weight.data_ptr() == cached[0][i * p.weight.shape[0]].data_ptr() for i, p in enumerate(projections) + ): + return cached + with torch.no_grad(): + weight = torch.cat([p.weight for p in projections]) + biases = [p.bias for p in projections] + bias = torch.cat(biases) if all(b is not None for b in biases) else None + rows = int(projections[0].weight.shape[0]) + for index, projection in enumerate(projections): + projection.weight.data = weight[index * rows : (index + 1) * rows] + if bias is not None: + projection.bias.data = bias[index * rows : (index + 1) * rows] + attn.__dict__["_packed_qkv"] = (weight, bias) + return weight, bias + + +def _attend(attn: nn.Module, qkv: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]]) -> torch.Tensor: + """``(tokens, E)`` attention output of packed ``(tokens, 3E)`` q/k/v, each item attending to itself.""" + heads, head_dim = attn.num_heads, attn.head_dim + embed = heads * head_dim + outputs = [] + for start, count, seq_len in seq_groups: + run = qkv[start : start + count * seq_len].view(count, seq_len, 3, heads, head_dim) + query, key, value = (run[:, :, index].transpose(1, 2) for index in range(3)) + out = F.scaled_dot_product_attention(query, key, value, scale=attn.scale) + outputs.append(out.transpose(1, 2).reshape(count * seq_len, embed)) + return outputs[0] if len(outputs) == 1 else torch.cat(outputs) + + +def _mlp_up(mlp: nn.Module, hidden: torch.Tensor) -> torch.Tensor: + fc1 = mlp.fc1 + if mlp.config.hidden_act == "gelu_pytorch_tanh" and fc1.bias is not None and hidden.is_cuda: + # cuBLASLt's GELU epilogue is the tanh approximation. + return torch._addmm_activation(fc1.bias, hidden, fc1.weight.t(), use_gelu=True) + return mlp.activation_fn(fc1(hidden)) + + +def _add_norm(residual: torch.Tensor, norm: nn.Module, y=None, y_bias=None, residual_out=None) -> torch.Tensor: + """``norm(residual + y + y_bias)`` for a ``(tokens, E)`` ``y``, the sum also written to ``residual_out``.""" + y = None if y is None else y.unsqueeze(0) + return residual_layer_norm( + residual, y, y_bias=y_bias, weight=norm.weight, bias=norm.bias, eps=norm.eps, residual_out=residual_out + ) + + +def encode_packed_fused( + vpm: nn.Module, hidden_states: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]] +) -> torch.Tensor: + """The encoder layers and ``post_layernorm`` of ``vpm.forward_packed`` over ``(tokens, E)`` embeddings. + + The residual stream is updated in place in a contiguous copy of + ``hidden_states`` (the patch embedding hands over a transposed view). + """ + layers = vpm.encoder.layers + residual = hidden_states.contiguous().unsqueeze(0) + normed = _add_norm(residual, layers[0].layer_norm1) + for index, layer in enumerate(layers): + attn, mlp = layer.self_attn, layer.mlp + qkv_weight, qkv_bias = packed_qkv(attn) + attended = _attend(attn, F.linear(normed[0], qkv_weight, qkv_bias), seq_groups) + y = F.linear(attended, attn.out_proj.weight) + normed = _add_norm(residual, layer.layer_norm2, y, attn.out_proj.bias, residual_out=residual) + last = index + 1 == len(layers) + norm = vpm.post_layernorm if last else layers[index + 1].layer_norm1 + y = F.linear(_mlp_up(mlp, normed[0]), mlp.fc2.weight) + normed = _add_norm(residual, norm, y, mlp.fc2.bias, residual_out=None if last else residual) + return normed[0] + + +def _copy_pixels_padded(dst: torch.Tensor, src: torch.Tensor) -> None: + width_px = int(src.shape[-1]) + dst[..., :width_px].copy_(src) + if width_px < int(dst.shape[-1]): + dst[..., width_px:].zero_() + + +class _PackedVisionAdapter(SupportsEncoderCudaGraph): + """Adapter exposing the packed SigLIP tower + resampler to vLLM's EncoderCudaGraphManager.""" + + supports_encoder_cudagraph: ClassVar[Literal[True]] = True + + def __init__(self, forward: Any, runs: Any, pixel_shape: Any, output_shape: torch.Size): + self.forward, self.runs, self.pixel_shape, self.output_shape = forward, runs, pixel_shape, output_shape + self.tokens = prod(output_shape[:-1]) or 1 + + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=["pixels"], + out_hidden_size=self.output_shape[-1], + padding_logics={"pixels": _copy_pixels_padded}, + ) + + def get_encoder_cudagraph_budget_range(self, vllm_config: VllmConfig) -> tuple[int, int]: + return self.tokens, self.tokens + + def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]: + return [EncoderItemSpec(input_size=self.tokens, output_tokens=self.tokens)] + + def select_encoder_cudagraph_items(self, mm_kwargs: dict[str, Any], indices: list[int]) -> dict[str, Any]: + return dict(mm_kwargs) + + def prepare_encoder_cudagraph_capture_inputs(self, *args: Any, **kwargs: Any) -> EncoderCudaGraphCaptureInputs: + device = kwargs.get("device") or args[3] + dtype = kwargs.get("dtype") or args[4] + return EncoderCudaGraphCaptureInputs({"pixels": torch.zeros(self.pixel_shape, dtype=dtype, device=device)}) + + def prepare_encoder_cudagraph_replay_buffers( + self, mm_kwargs: dict[str, Any], *args: Any, **kwargs: Any + ) -> EncoderCudaGraphReplayBuffers: + return EncoderCudaGraphReplayBuffers(mm_kwargs) + + def encoder_cudagraph_forward(self, inputs: dict[str, torch.Tensor], path: str = "default") -> torch.Tensor: + return self.forward(inputs["pixels"], self.runs) + + def encoder_eager_forward(self, mm_kwargs: dict[str, Any], path: str = "default") -> torch.Tensor: + return self.encoder_cudagraph_forward(mm_kwargs, path) + + def postprocess_encoder_output( + self, + outputs: dict[str, torch.Tensor], + indices: list[int], + per_item_out_tokens: list[int], + dest: dict[int, torch.Tensor] | list[torch.Tensor | None], + clone: bool = False, + batch_mm_kwargs: dict[str, Any] | None = None, + ) -> None: + dest[0] = outputs["default"].clone() + + +class VisionGraphEncoder: + """CUDA graphs of the packed SigLIP tower + resampler for single-grid chunks.""" + + def __init__( + self, + vpm: nn.Module, + resampler: nn.Module, + *, + batch_sizes: Sequence[int] = DEFAULT_GRAPH_BATCH_SIZES, + vllm_config: VllmConfig | None = None, + ): + self.vpm, self.resampler, self.vllm_config = vpm, resampler, vllm_config + self.batch_sizes = tuple(sorted({int(b) for b in batch_sizes if int(b) > 0})) + # (h, w, bucket) -> EncoderCudaGraphManager, least recently used first. + self._managers: OrderedDict[tuple[int, int, int], EncoderCudaGraphManager] = OrderedDict() + self._failed: set[tuple[int, int, int]] = set() + self._seen: dict[tuple[int, int, int], int] = {} + + def encode(self, pixels: torch.Tensor, height: int, width: int, count: int) -> torch.Tensor | None: + if torch.cuda.is_current_stream_capturing(): + return None + bucket = min((size for size in self.batch_sizes if count <= size), default=None) + if bucket is None: + return None + key = (int(height), int(width), bucket) + manager = self._managers.get(key) + if manager is None: + if key in self._failed: + return None + seen = self._seen.get(key, 0) + if seen < _CAPTURE_AFTER: + if len(self._seen) >= 4 * _MAX_GRAPHS: + self._seen.clear() + self._seen[key] = seen + 1 + return None + manager = self._capture(key, pixels) + if manager is None: + return None + self._managers.move_to_end(key) + return manager.execute({"pixels": pixels})[0][:count] + + def _run(self, pixels: torch.Tensor, runs: list[tuple[int, int, int]]) -> torch.Tensor: + return self.resampler.forward_packed(self.vpm.forward_packed(pixels, runs), runs) + + def _capture(self, key: tuple[int, int, int], like: torch.Tensor) -> EncoderCudaGraphManager | None: + height, width, bucket = key + runs, patch = [(height, width, bucket)], int(like.shape[-2]) + pixel_shape = (*like.shape[:-1], bucket * height * width * patch) + dummy_pixels = torch.zeros(pixel_shape, device=like.device, dtype=like.dtype) + try: + stream = torch.cuda.Stream(device=like.device) + stream.wait_stream(torch.cuda.current_stream(like.device)) + with torch.cuda.stream(stream), torch.inference_mode(): + out = self._run(dummy_pixels, runs) + for _ in range(1): + self._run(dummy_pixels, runs) + torch.cuda.current_stream(like.device).wait_stream(stream) + adapter = _PackedVisionAdapter(self._run, runs, pixel_shape, out.shape) + config = copy(self.vllm_config) if self.vllm_config is not None else VllmConfig() + if getattr(config, "model_config", None) is None: + config.model_config = ModelConfig.__new__(ModelConfig) + if getattr(config.model_config, "multimodal_config", None) is None: + config.model_config.multimodal_config = MultiModalConfig() + cc = config.compilation_config = copy(config.compilation_config) + cc.encoder_cudagraph_token_budgets = [adapter.tokens] + cc.encoder_cudagraph_max_vision_items_per_batch, cc.encoder_cudagraph_max_frames_per_batch = 1, 0 + config.parallel_config = copy(config.parallel_config) + config.parallel_config.tensor_parallel_size = 1 + + manager = EncoderCudaGraphManager(config, like.device, like.dtype, adapter) + with torch.cuda.stream(stream): + manager.capture(graph_pool=current_platform.get_global_graph_pool()) + torch.cuda.current_stream(like.device).wait_stream(stream) + except Exception: + logger.warning("MiniCPM-o vision CUDA graph capture failed for %s; staying eager", key, exc_info=True) + self._failed.add(key) + return None + if len(self._managers) >= _MAX_GRAPHS: + _, evicted = self._managers.popitem(last=False) + evicted.clear() + self._managers[key] = manager + logger.info("Captured MiniCPM-o vision CUDA graph grid=%dx%d batch=%d via manager", height, width, bucket) + return manager