[Bugfix] Add opt-in MiniCPM-o 4.5 Stage-0 sliding window

.buildkite/common/ci_source_file_dependencies.yml

@@ -245,6 +245,8 @@ source_file_dependencies:
     - *minicpmo_4_5_duplex
     - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
     - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
+    - tests/e2e/online_serving/test_minicpmo_4_5_window.py
+    - tests/e2e/online_serving/helpers/minicpmo_window_e2e.py
     - tests/e2e/online_serving/test_duplex_client_live.py
     - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
     - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py

.buildkite/cuda/test-merge.yml

@@ -130,7 +130,7 @@ steps:
         source_file_dependencies: omni_minicpmo_4_5_duplex_function
         timeout_in_minutes: 50
         commands:
-          - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+          - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py tests/e2e/online_serving/test_minicpmo_4_5_window.py -m 'advanced_model and cuda' --run-level 'advanced_model'
         mirror_hardwares: h100_1
 
       - label: "Diffusion · Qwen Image Test"

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -68,6 +68,39 @@ def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate:
     )
 
 
+def _make_minicpm_window_update(
+    *,
+    seq: int,
+    mode: str,
+    high: int = 15,
+    low: int = 10,
+    context_max_units: int = 1,
+) -> StreamingUpdate:
+    update = _make_update([0] * 8)
+    update.model_intermediate_buffer = {
+        "duplex": {
+            "data_plane": True,
+            "seq": seq,
+            "runtime_config": {
+                "duplex_scheduler_token_id": 0,
+                "duplex_first_append_context_tokens": 3,
+                "duplex_window_prefix_tokens": 2,
+                "duplex_window_suffix_token_ids": [3],
+                "duplex_window_previous_marker_token_ids": [70, 71],
+                "duplex_window_special_token_ids": [99],
+                "duplex_window_config": {
+                    "sliding_window_mode": mode,
+                    "basic_window_high_tokens": high,
+                    "basic_window_low_tokens": low,
+                    "context_previous_max_tokens": 4,
+                    "context_max_units": context_max_units,
+                },
+            },
+        }
+    }
+    return update
+
+
 def _make_talker_adapter(
     *,
     max_model_len: int = 100,
@@ -509,6 +542,225 @@ def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder()
     assert sched._new_prompt_len_snapshot[session.request_id] == 2
 
 
+def test_stage0_basic_window_rebuilds_below_low_watermark() -> None:
+    sched = _make_scheduler(stage_id=0)
+    session = _make_request()
+    session.prompt_token_ids = [0] * 9
+    session._all_token_ids.clear()
+    session._all_token_ids.extend(session.prompt_token_ids)
+    session.num_prompt_tokens = 9
+    session.append_output_token_ids([40])
+    session.num_computed_tokens = 10
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+
+    sched._update_request_as_session(session, update)
+
+    assert session.prompt_token_ids == [0] * 9
+    assert session.num_computed_tokens == 0
+    assert update.model_intermediate_buffer["meta"]["replace_streaming_prompt"] is True
+    plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
+    assert plan == {
+        "completed_token_ids": [40],
+        "replace": True,
+        "mode": "basic",
+        "drop_units": 1,
+        "dropped_tokens": 9,
+        "previous_token_ids": [],
+        "previous_marker_token_ids": [70, 71],
+        "replacement_prompt_len": 9,
+    }
+    sched._free_request_blocks.assert_called_once_with(session)
+
+
+def test_stage0_context_window_compacts_dropped_speech() -> None:
+    sched = _make_scheduler(stage_id=0)
+    session = _make_request()
+    session.prompt_token_ids = [0] * 9
+    session._all_token_ids.clear()
+    session._all_token_ids.extend(session.prompt_token_ids)
+    session.num_prompt_tokens = 9
+    session.append_output_token_ids([40])
+    session.num_computed_tokens = 10
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+
+    first = _make_minicpm_window_update(seq=2, mode="context")
+    sched._update_request_as_session(session, first)
+    assert session.num_prompt_tokens == 18
+    assert first.model_intermediate_buffer["duplex"]["stage0_window"] == {"completed_token_ids": [40]}
+
+    session.append_output_token_ids([50])
+    session.num_computed_tokens = 19
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    second = _make_minicpm_window_update(seq=3, mode="context")
+    sched._update_request_as_session(session, second)
+
+    assert session.prompt_token_ids == [0] * 21
+    plan = second.model_intermediate_buffer["duplex"]["stage0_window"]
+    assert plan["completed_token_ids"] == [50]
+    assert plan["drop_units"] == 1
+    assert plan["previous_token_ids"] == [40]
+    # The worker embeds this exact marker, not a re-tokenized one, so the
+    # previous region is len(marker) + len(previous).
+    assert plan["previous_marker_token_ids"] == [70, 71]
+    assert plan["replacement_prompt_len"] == 21
+    assert getattr(session, "_minicpmo45_window_previous_len") == 3
+
+
+@pytest.mark.parametrize("cleared_outputs", [False, True])
+@pytest.mark.parametrize("in_flight", [0, 1])
+def test_stage0_window_uses_confirmed_span_and_terminator(cleared_outputs, in_flight) -> None:
+    sched = _make_scheduler(stage_id=0)
+    session = _make_request()
+    session.prompt_token_ids = [0] * 9
+    session._all_token_ids[:] = session.prompt_token_ids
+    session.num_prompt_tokens = 9
+    session.append_output_token_ids([40, 99])
+    session.num_computed_tokens = 10 + in_flight
+    session.num_output_placeholders = in_flight
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    if cleared_outputs:
+        # A stop without a queued append clears this before the next update.
+        session._output_token_ids.clear()
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+
+    sched._update_request_as_session(session, update)
+
+    plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
+    assert plan["completed_token_ids"] == [40]
+    assert plan["completed_terminator_token_id"] == 99
+    assert plan["replacement_prompt_len"] == 9
+    assert plan["dropped_tokens"] == 9
+
+
+def test_stage0_window_rebuild_that_overflows_max_model_len_finishes_the_session() -> None:
+    """A rebuilt window prompt is bounded by the client's window settings, not
+    by the model: a unit carrying camera frames is hundreds of tokens, and a
+    ``context`` window adds the ``previous`` region on top. The replacement
+    branch must therefore check ``replacement_prompt_len`` against
+    ``max_model_len - sample_room`` and finish the session with
+    ``context_length_exceeded``, exactly like a plain extension."""
+    sched = _make_live_session_scheduler(max_model_len=21)
+    session = _make_request()
+    session.prompt_token_ids = [0] * 9
+    session._all_token_ids[:] = [0] * 9
+    session.num_prompt_tokens = 9
+    session._output_token_ids[:] = [40]
+    session._all_token_ids.append(40)
+    session.num_computed_tokens = 10
+    session.num_output_placeholders = 0
+    _park_session(sched, session)
+    # The parked append state above is the frame before this update: one
+    # confirmed output token, fully computed.
+    session._output_token_ids[:] = [40]
+    session._all_token_ids[:] = [*session.prompt_token_ids, 40]
+    session.num_computed_tokens = 10
+    session.num_in_flight_tokens = 0
+    sched.num_sampled_tokens_per_step = 1
+    # The replacement plan rebuilds 21 tokens (retained context plus the
+    # append) and one step samples one more, so 21 leaves no room.
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+    update.prompt_token_ids = [0] * 20
+    original_prompt = list(session.prompt_token_ids)
+
+    sched._update_request_as_session(session, update)
+
+    assert update.model_intermediate_buffer["duplex"]["stage0_window"]["replacement_prompt_len"] == 21
+    assert session.prompt_token_ids == original_prompt
+    assert session.num_prompt_tokens == 9
+    assert session.status == RequestStatus.FINISHED_ERROR
+    assert session.request_id not in sched.requests
+    assert sched.finished_req_ids == {session.request_id}
+    sched._free_request_blocks.assert_called_once_with(session)
+    client_index, reason = sched._streaming_context_overflow[session.request_id]
+    assert client_index == session.client_index
+    assert reason.startswith("context_length_exceeded: ")
+    assert "21 tokens" in reason and "max_model_len 21" in reason
+
+    engine_core_outputs = _run_idle_step(sched)
+
+    (output,) = engine_core_outputs[session.client_index].outputs
+    assert output.request_id == session.request_id
+    assert output.finish_reason == FinishReason.ERROR
+    assert output.stop_reason == reason
+
+
+def test_stage0_window_rebuild_that_leaves_room_to_sample_replaces_the_prompt() -> None:
+    """One slot above the plan fits: the replacement applies through the normal
+    replacement path."""
+    sched = _make_scheduler(stage_id=0)
+    sched.max_model_len = 25
+    sched.num_sampled_tokens_per_step = 1
+    session = _make_request()
+    session.prompt_token_ids = [0] * 9
+    session._all_token_ids[:] = session.prompt_token_ids
+    session.num_prompt_tokens = 9
+    session.append_output_token_ids([40])
+    session.num_computed_tokens = 10
+    session.num_output_placeholders = 0
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+    update.prompt_token_ids = [0] * 20
+
+    sched._update_request_as_session(session, update)
+
+    assert update.model_intermediate_buffer["meta"]["replace_streaming_prompt"] is True
+    assert update.model_intermediate_buffer["duplex"]["stage0_window"]["replacement_prompt_len"] == 21
+    assert session.prompt_token_ids == [0] * 21
+    assert session.num_prompt_tokens == 21
+    assert session.num_computed_tokens == 0
+    assert session.status == RequestStatus.WAITING
+    sched._free_request_blocks.assert_called_once_with(session)
+    assert not getattr(sched, "_streaming_context_overflow", {})
+
+
+def test_stage0_window_open_start_zero_is_not_replaced_by_the_context_reserve() -> None:
+    """A recorded ``open_start`` of 0 is a legitimate empty context prefix. The
+    fallback must be an explicit ``is None`` check: treating 0 as missing
+    substitutes the context reserve, counts the suffix twice, and the worker's
+    rebuild-length check raises on the first replacement."""
+    session = SimpleNamespace(
+        num_prompt_tokens=9,
+        _all_token_ids=[0] * 9 + [40],
+        num_computed_tokens=10,
+        num_output_placeholders=0,
+        _minicpmo45_window_open_start=0,
+    )
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+    update.prompt_token_ids = [0] * 21
+
+    assert OmniARScheduler._prepare_minicpmo45_stage0_window(
+        session, update, segment_output_ids=[40], completed_terminator=3
+    )
+
+    # open_start stays 0, so the reported unit spans the whole 12-token
+    # boundary instead of the 9 the context reserve would fold in.
+    plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
+    assert plan["dropped_tokens"] == 12
+    assert plan["replacement_prompt_len"] == 19
+
+
+def test_stage0_window_open_start_falls_back_to_the_context_reserve_when_unset() -> None:
+    session = SimpleNamespace(
+        num_prompt_tokens=9,
+        _all_token_ids=[0] * 9 + [40],
+        num_computed_tokens=10,
+        num_output_placeholders=0,
+    )
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+    update.prompt_token_ids = [0] * 21
+
+    assert OmniARScheduler._prepare_minicpmo45_stage0_window(
+        session, update, segment_output_ids=[40], completed_terminator=3
+    )
+
+    # fallback open_start is preserve_len 3 from duplex_first_append_context_tokens,
+    # so the unit is three tokens shorter.
+    plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
+    assert plan["dropped_tokens"] == 9
+    assert plan["replacement_prompt_len"] == 22
+
+
 def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None:
     sched = _make_scheduler(stage_id=1)
     sched.chunk_transfer_adapter = SimpleNamespace(

tests/e2e/online_serving/helpers/minicpmo_window_e2e.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Drive complete native duplex turns while Stage-0 repeatedly rebuilds KV."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+from pathlib import Path
+
+from vllm_omni.clients.duplex import DuplexClient, EventCollector, read_pcm16_wav, wait_for_condition
+from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config
+
+
+async def run_window_turn(
+    *,
+    url: str,
+    model: str,
+    input_wav: Path,
+    mode: str,
+    ref_audio: Path | None = None,
+    repeats: int = 2,
+    video_frames: list[str] | None = None,
+    buffered_flush: bool = False,
+    timeout_s: float = 240.0,
+) -> dict[str, object]:
+    """Use aggressive limits to exercise multiple windows in one audio turn."""
+    pcm = read_pcm16_wav(input_wav) * repeats
+    collector = EventCollector()
+    reference = (
+        "data:audio/wav;base64," + base64.b64encode(ref_audio.read_bytes()).decode("ascii")
+        if ref_audio is not None
+        else None
+    )
+    config = create_duplex_session_config(
+        modalities=("text", "audio") if ref_audio is not None else ("text",),
+        ref_audio=reference,
+        temperature=0.0,
+        extra_body={
+            "sliding_window_mode": mode,
+            "context_max_units": 1,
+            "context_previous_max_tokens": 16,
+            "basic_window_high_tokens": 64,
+            "basic_window_low_tokens": 32,
+        },
+    )
+    client = DuplexClient(
+        url, model=model, config=config, reconnect=None, heartbeat_interval_s=None, handshake_timeout_s=timeout_s
+    )
+    async with client:
+        reader = asyncio.create_task(collector.consume(client))
+        collector.add({"type": "session.created", "session": client.session_info})
+        try:
+            frames_sent = await client.stream_pcm(
+                pcm,
+                chunk_ms=2500 if buffered_flush else 200,
+                realtime=not buffered_flush,
+                is_speech=True,
+                video_frames=video_frames,
+            )
+            await client.commit(final=True)
+            await wait_for_condition(
+                lambda: collector.count("input_audio_buffer.committed") > 0 or bool(collector.errors()),
+                timeout_s=timeout_s,
+                label="final input commit",
+            )
+            await wait_for_condition(
+                lambda: collector.count("response.done") > 0 or bool(collector.errors()),
+                timeout_s=timeout_s,
+                label="completed windowed response",
+            )
+            # Final execution follows the commit acknowledgement asynchronously.
+            await asyncio.sleep(3.0)
+            await client.close(timeout_s=timeout_s)
+            await asyncio.wait_for(reader, timeout=5.0)
+        finally:
+            if not reader.done():
+                reader.cancel()
+                await asyncio.gather(reader, return_exceptions=True)
+    return {
+        "created": collector.count("session.created"),
+        "committed": collector.count("input_audio_buffer.committed"),
+        "done": collector.count("response.done"),
+        "closed": collector.count("session.closed"),
+        "errors": collector.errors(),
+        "audio_bytes": len(collector.audio_bytes()),
+        "frames_sent": frames_sent,
+        "input_seconds": len(pcm) / 32000,
+    }