diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_scheduler_token_budget.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_scheduler_token_budget.py new file mode 100644 index 000000000..03072cb33 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_scheduler_token_budget.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""What a duplex append reserves in the scheduler, and why that number. + +The reservation is not advisory: ``build_duplex_data_plane_prompt`` turns it +into ``[scheduler_token_id] * budget``, and a unit whose embeddings outnumber +its slots has its tail dropped with a warning rather than failing +(``MiniCPMO45OmniModel.get_input_embeddings``). Over-reserving only wastes +slots, so every uncertainty here has to resolve upwards. + +The camera side depends on HD slicing, which depends on the frame size +*relative to the checkpoint's normalization tile*. The tile is configuration, +not a constant, so the cases below drive the arithmetic at more than one tile +size and check it against ``MiniCPMVImageProcessor.get_sliced_grid`` -- the +same grid search the checkpoint's own processor runs. +""" + +from __future__ import annotations + +import base64 +import io +from dataclasses import dataclass + +import pytest +from PIL import Image + +from vllm_omni.engine.duplex.config import DuplexSessionConfig +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import ( + PRIVATE_RUNTIME_CONFIG_KEYS, + _apply_default_scheduler_policy, + _duplex_vision_tile_pixels, + _duplex_vision_tokens, + _model_vision_tile_pixels, + duplex_scheduler_token_budget, +) +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import MiniCPMVImageProcessor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +TOKENS_PER_BLOCK = 66 +#: ``_official_max_slice_nums`` gives the first frame of a stacked pair 2 and +#: every other frame 1. +BASE_FRAME_MAX_SLICES = 2 +#: What the released MiniCPM-o 4.5 checkpoint configures. +TILE_448 = 448 * 448 + +SIZES = [(224, 224), (320, 180), (448, 448), (448, 449), (640, 480), (960, 540), (1280, 720), (540, 960)] + + +def _frame(size: tuple[int, int], fmt: str = "JPEG") -> str: + buffer = io.BytesIO() + Image.new("RGB", size).save(buffer, format=fmt) + return base64.b64encode(buffer.getvalue()).decode() + + +def _blocks_the_model_makes(size: tuple[int, int], max_slice_nums: int, scale_resolution: int) -> int: + """Source tile plus the HD patches, straight from the model's grid search.""" + processor = MiniCPMVImageProcessor(scale_resolution=scale_resolution) + grid = processor.get_sliced_grid(image_size=size, max_slice_nums=max_slice_nums) + return 1 + grid[0] * grid[1] if grid else 1 + + +@pytest.mark.parametrize("scale_resolution", [336, 448, 560]) +@pytest.mark.parametrize("size", SIZES, ids=[f"{w}x{h}" for w, h in SIZES]) +@pytest.mark.parametrize("fmt", ["JPEG", "PNG"]) +def test_a_stacked_pair_reserves_what_the_model_slices_it_into( + size: tuple[int, int], fmt: str, scale_resolution: int +) -> None: + """The tile comes from the checkpoint, so the arithmetic has to follow it.""" + frame = _frame(size, fmt) + expected_blocks = _blocks_the_model_makes(size, BASE_FRAME_MAX_SLICES, scale_resolution) + _blocks_the_model_makes( + size, 1, scale_resolution + ) + + reserved = _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=scale_resolution**2) + + assert reserved == expected_blocks * TOKENS_PER_BLOCK + + +def test_a_small_frame_no_longer_reserves_three_hd_blocks() -> None: + """A 448x448 frame is one tile, so a stacked pair is two blocks, not four. + + The size-independent count reserved 264 tokens for this pair. 132 of them + were for patches the model never produces. + """ + frame = _frame((448, 448)) + + assert _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=TILE_448) == 2 * TOKENS_PER_BLOCK + + +def test_one_tile_is_the_whole_slicing_decision() -> None: + """``ceil(w * h / tile)`` capped at 2: a single pixel over the tile slices.""" + + def pair(size: tuple[int, int]) -> int: + frame = _frame(size) + return _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=TILE_448) + + assert pair((448, 448)) == 2 * TOKENS_PER_BLOCK + assert pair((448, 449)) == 4 * TOKENS_PER_BLOCK + + +def test_only_the_first_frame_of_a_pair_is_hd_sliced() -> None: + """``_official_max_slice_nums`` is ``[2, 1]``: the composite is never sliced. + + Sizing the wrong half of the pair would pass every same-size case, so the + two orders have to disagree. + """ + small, large = _frame((448, 448)), _frame((1280, 720)) + + assert _duplex_vision_tokens({"video_frames": [small, large]}, tile_pixels=TILE_448) == 2 * TOKENS_PER_BLOCK + assert _duplex_vision_tokens({"video_frames": [large, small]}, tile_pixels=TILE_448) == 4 * TOKENS_PER_BLOCK + + +def test_a_lone_frame_is_never_hd_sliced() -> None: + """``_official_max_slice_nums(1)`` is ``[1]``: one frame is one tile at any size.""" + assert _duplex_vision_tokens({"video_frames": [_frame((1280, 720))]}, tile_pixels=TILE_448) == TOKENS_PER_BLOCK + + +def test_extra_frames_past_the_pair_each_cost_one_block() -> None: + """The wire caps a frame list at two; the in-process API does not.""" + small, large = _frame((448, 448)), _frame((1280, 720)) + + assert _duplex_vision_tokens({"video_frames": [small] * 3}, tile_pixels=TILE_448) == 3 * TOKENS_PER_BLOCK + assert _duplex_vision_tokens({"video_frames": [large] * 3}, tile_pixels=TILE_448) == 5 * TOKENS_PER_BLOCK + + +def test_an_unknown_tile_size_keeps_the_sliced_reservation() -> None: + """No tile, no shrinking. A checkpoint whose config cannot be read keeps the old number.""" + frame = _frame((224, 224)) + + assert _duplex_vision_tokens({"video_frames": [frame, frame]}) == 4 * TOKENS_PER_BLOCK + assert _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=None) == 4 * TOKENS_PER_BLOCK + + +@pytest.mark.parametrize( + "first_frame", + [ + pytest.param("data:image/jpeg;base64," + _frame((224, 224)), id="data url"), + pytest.param("!!! not base64 !!!", id="corrupt base64"), + pytest.param(base64.b64encode(b"\x00" * 64).decode(), id="base64 of something else"), + pytest.param(base64.b64encode(base64.b64decode(_frame((224, 224)))[:8]).decode(), id="truncated jpeg"), + ], +) +def test_a_frame_whose_header_will_not_parse_keeps_the_sliced_reservation(first_frame: str) -> None: + """The Realtime wire screens most of these; the in-process API does not. + + ``validate_realtime_video_frames`` rejects bad base64 and anything whose + magic bytes are not JPEG or PNG, so on a websocket session only a truncated + but well-headed frame gets this far. ``DuplexOmni.append_audio`` submits + without that screen, so the reservation still has to hold on its own. + """ + good = _frame((224, 224)) + + assert _duplex_vision_tokens({"video_frames": [first_frame, good]}, tile_pixels=TILE_448) == 4 * TOKENS_PER_BLOCK + + +def test_no_camera_track_reserves_nothing() -> None: + for payload in ({}, {"video_frames": []}, {"video_frames": "not a list"}, {"video_frames": [None, ""]}, None): + assert _duplex_vision_tokens(payload, tile_pixels=TILE_448) == 0 + + +def test_the_audio_budget_is_untouched_by_the_camera_track() -> None: + """One second of pcm_f32le is 12 slots; the frames are added on top.""" + one_second = base64.b64encode(b"\x00" * 4 * 16000).decode() + audio_only = {"audio": one_second, "format": "pcm_f32le"} + with_frames = {**audio_only, "video_frames": [_frame((448, 448)), _frame((448, 448))]} + + assert duplex_scheduler_token_budget(audio_only, tile_pixels=TILE_448) == 12 + assert duplex_scheduler_token_budget(with_frames, tile_pixels=TILE_448) == 12 + 2 * TOKENS_PER_BLOCK + + +# ---- where the tile comes from ---- + + +@dataclass(frozen=True) +class _SliceConfig: + scale_resolution: int + + +@dataclass(frozen=True) +class _HFConfig: + slice_config: object = None + image_size: object = None + + +@dataclass(frozen=True) +class _ModelConfig: + hf_config: object = None + + +@pytest.mark.parametrize( + ("hf_config", "expected"), + [ + (_HFConfig(slice_config={"max_slice_nums": 1, "scale_resolution": 448}, image_size=448), TILE_448), + (_HFConfig(slice_config=_SliceConfig(scale_resolution=336), image_size=448), 336 * 336), + (_HFConfig(image_size=560), 560 * 560), + (_HFConfig(), None), + (_HFConfig(slice_config={"scale_resolution": 0}, image_size=0), None), + ], + ids=["released-checkpoint", "attribute", "image-size-fallback", "nothing-to-read", "nonsense"], +) +def test_the_tile_is_read_from_the_checkpoint(hf_config: object, expected: int | None) -> None: + """First row is the released MiniCPM-o 4.5 ``config.json``, where ``slice_config`` is a dict.""" + assert _model_vision_tile_pixels(_ModelConfig(hf_config=hf_config)) == expected + + +def test_no_model_config_means_no_tile() -> None: + assert _model_vision_tile_pixels(None) is None + assert _model_vision_tile_pixels(_ModelConfig()) is None + + +def test_the_tile_reaches_the_budget_through_the_runtime_config() -> None: + """``_apply_default_scheduler_policy`` writes it; ``build_duplex_data_plane_prompt`` reads it back.""" + runtime_config: dict[str, object] = {} + + _apply_default_scheduler_policy( + runtime_config, + config=DuplexSessionConfig(), + tokenizer=None, + model_config=_ModelConfig(hf_config=_HFConfig(slice_config={"scale_resolution": 448})), + ) + + assert runtime_config["duplex_vision_tile_pixels"] == TILE_448 + assert _duplex_vision_tile_pixels(runtime_config) == TILE_448 + + +def test_a_checkpoint_that_cannot_be_read_leaves_the_key_out() -> None: + """No key, no tile, and no tile is the sliced reservation.""" + runtime_config: dict[str, object] = {} + + _apply_default_scheduler_policy(runtime_config, config=DuplexSessionConfig(), tokenizer=None) + + assert "duplex_vision_tile_pixels" not in runtime_config + assert _duplex_vision_tile_pixels(runtime_config) is None + + +def test_a_junk_tile_in_the_runtime_config_is_ignored() -> None: + for value in (0, -1, "448", 448.0, None): + assert _duplex_vision_tile_pixels({"duplex_vision_tile_pixels": value}) is None + + +def test_the_tile_is_server_owned() -> None: + """A client that could set the tile could shrink the reservation under the worker.""" + assert "duplex_vision_tile_pixels" in PRIVATE_RUNTIME_CONFIG_KEYS diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py index 673640461..e57377766 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py @@ -63,6 +63,15 @@ _DUPLEX_VISION_TOKENS_PER_FRAME = 66 # Official stacked pair uses max_slice_nums=[2, 1]: the current frame is HD # sliced (1 source + 2 patches on 960x540) and the composite is not. _DUPLEX_HD_SLICES_PER_BASE_FRAME = 3 +# How many blocks that slicing costs depends on the frame, and at +# ``max_slice_nums=2`` the processor's grid search has only two outcomes. +# ``MiniCPMVImageProcessor.get_sliced_grid`` takes +# ``multiple = min(ceil(w * h / scale_resolution**2), max_slice_nums)`` and +# returns no grid at all for ``multiple <= 1``; at 2 the only candidate split +# is 2, so a frame is either unsliced or a 2-cell grid. One normalization tile +# is therefore the whole decision -- but ``scale_resolution`` is the +# checkpoint's, not a constant, so it is read from the model rather than +# assumed. PRIVATE_RUNTIME_CONFIG_KEYS = frozenset( { @@ -70,6 +79,7 @@ PRIVATE_RUNTIME_CONFIG_KEYS = frozenset( "duplex_stage_max_tokens", "duplex_stage0_max_tokens", "duplex_scheduler_token_id", + "duplex_vision_tile_pixels", "duplex_first_append_context_tokens", "ref_audio_data", "ref_audio_format", @@ -86,29 +96,68 @@ class MiniCPMO45ClientRuntimeConfigError(DuplexRuntimeConfigError): # ---- engine policy helpers: scheduler token budget ---- -def _duplex_frame_count(payload: object) -> int: +def _duplex_frames(payload: object) -> list[str]: if not isinstance(payload, dict): - return 0 + return [] frames = payload.get("video_frames") if not isinstance(frames, list): - return 0 - return sum(1 for frame in frames if isinstance(frame, str) and frame) + return [] + return [frame for frame in frames if isinstance(frame, str) and frame] -def _duplex_vision_tokens(payload: object) -> int: +def _duplex_base_frame_blocks(frame: str, tile_pixels: int | None) -> int: + """Blocks the HD-sliced frame of a stacked pair costs, read from the frame. + + Wire format is the one Stage0 decodes: bare base64 JPEG/PNG, and only the + header is parsed -- the pixels are the worker's job. + + Every uncertainty resolves to the sliced count, because the two directions + are not symmetric. Over-reserving wastes scheduler slots. Under-reserving + hands the worker fewer prompt slots than it has embeddings, and + ``MiniCPMO45OmniModel`` then drops the tail of the unit with a warning + rather than failing, which is silent audio loss. + """ + if tile_pixels is None: + return _DUPLEX_HD_SLICES_PER_BASE_FRAME + + from io import BytesIO + + from PIL import Image + + try: + raw = b64decode(frame, validate=True) + with Image.open(BytesIO(raw)) as image: + width, height = image.size + except (BinasciiError, ValueError, OSError, Image.DecompressionBombError): + return _DUPLEX_HD_SLICES_PER_BASE_FRAME + if width * height <= tile_pixels: + return 1 + return _DUPLEX_HD_SLICES_PER_BASE_FRAME + + +def _duplex_vision_tile_pixels(runtime_config: object) -> int | None: + """Area of the tile this model normalizes a frame to, or ``None`` if unknown.""" + if not isinstance(runtime_config, dict): + return None + value = runtime_config.get("duplex_vision_tile_pixels") + return value if isinstance(value, int) and value > 0 else None + + +def _duplex_vision_tokens(payload: object, *, tile_pixels: int | None = None) -> int: """Scheduler slots for this append's camera track. Audio is never stacked: a unit still carries one second of soundtrack. ``stack_frames`` only adds a second *image*. Official HD on that pair is - ``[2, 1]``, so the base frame reserves three 66-token blocks and every - extra frame reserves one. + ``[2, 1]``, so only the first frame is sliced and every other frame + reserves one 66-token block. """ - count = _duplex_frame_count(payload) - if count <= 0: + frames = _duplex_frames(payload) + if not frames: return 0 - if count >= 2: - return (_DUPLEX_HD_SLICES_PER_BASE_FRAME + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME - return count * _DUPLEX_VISION_TOKENS_PER_FRAME + if len(frames) == 1: + return _DUPLEX_VISION_TOKENS_PER_FRAME + blocks = _duplex_base_frame_blocks(frames[0], tile_pixels) + (len(frames) - 1) + return blocks * _DUPLEX_VISION_TOKENS_PER_FRAME def _duplex_pcm_sample_count(payload: object) -> int | None: @@ -136,8 +185,8 @@ def duplex_first_append_unit_count(payload: object) -> int | None: return max(1, sample_count // _DUPLEX_CHUNK_SAMPLES - 1) -def duplex_scheduler_token_budget(payload: object, *, default: int = 64) -> int: - vision_tokens = _duplex_vision_tokens(payload) +def duplex_scheduler_token_budget(payload: object, *, default: int = 64, tile_pixels: int | None = None) -> int: + vision_tokens = _duplex_vision_tokens(payload, tile_pixels=tile_pixels) sample_count = _duplex_pcm_sample_count(payload) if sample_count is None: return max(1, int(default)) + vision_tokens @@ -185,13 +234,15 @@ def build_duplex_data_plane_prompt( payload: object, final: bool, ) -> dict[str, object]: - token_budget = duplex_scheduler_token_budget(payload) + tile_pixels = _duplex_vision_tile_pixels(runtime_config) + token_budget = duplex_scheduler_token_budget(payload, tile_pixels=tile_pixels) if seq <= 1: context_reserve = duplex_first_append_context_reserve(runtime_config) token_budget += context_reserve first_units = duplex_first_append_unit_count(payload) if first_units is not None: - token_budget = context_reserve + first_units * 12 - 1 + _duplex_vision_tokens(payload) + vision_tokens = _duplex_vision_tokens(payload, tile_pixels=tile_pixels) + token_budget = context_reserve + first_units * 12 - 1 + vision_tokens if seq > 1 and duplex_payload_is_exact_chunks(payload): token_budget += 1 if final and duplex_payload_is_exact_chunks(payload): @@ -466,11 +517,34 @@ def _apply_first_append_context_tokens( runtime_config["duplex_first_append_context_tokens"] = len(prefix_ids) + ref_tokens + len(suffix_ids) +def _model_vision_tile_pixels(model_config: ModelConfig | None) -> int | None: + """Area of one normalization tile, from the checkpoint that will do the slicing. + + ``MiniCPMVImageProcessor`` is built with ``scale_resolution=config.image_size`` + and Stage0 loads the checkpoint's own processor, so this is per-checkpoint + configuration. ``None`` when it cannot be read, which keeps the reservation + at the sliced count. + """ + hf_config = getattr(model_config, "hf_config", None) + if hf_config is None: + return None + slice_config = getattr(hf_config, "slice_config", None) + side = getattr(slice_config, "scale_resolution", None) + if not isinstance(side, int): + side = slice_config.get("scale_resolution") if isinstance(slice_config, dict) else None + if not isinstance(side, int): + side = getattr(hf_config, "image_size", None) + if not isinstance(side, int) or side <= 0: + return None + return side * side + + def _apply_default_scheduler_policy( runtime_config: dict[str, object], *, config: DuplexSessionConfig, tokenizer: PreTrainedTokenizerBase | None, + model_config: ModelConfig | None = None, ) -> None: stage0_max_tokens = config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20 runtime_config["duplex_stage_max_tokens"] = {"0": stage0_max_tokens, "1": 8192} @@ -491,6 +565,9 @@ def _apply_default_scheduler_policy( scheduler_token_id = _scheduler_token_id(tokenizer) if scheduler_token_id is not None: runtime_config["duplex_scheduler_token_id"] = scheduler_token_id + tile_pixels = _model_vision_tile_pixels(model_config) + if tile_pixels is not None: + runtime_config["duplex_vision_tile_pixels"] = tile_pixels class MiniCPMO45DuplexPlugin(DuplexModelPlugin): @@ -653,7 +730,7 @@ class MiniCPMO45DuplexPlugin(DuplexModelPlugin): if isinstance(initial_user_text, str) and initial_user_text: runtime_config["initial_user_text"] = initial_user_text tokenizer = await self._tokenizer_for(model_config) - _apply_default_scheduler_policy(runtime_config, config=config, tokenizer=tokenizer) + _apply_default_scheduler_policy(runtime_config, config=config, tokenizer=tokenizer, model_config=model_config) ref_audio = config.ref_audio extra_ref_audio = extra_body.get("ref_audio")