From c8a1419e02efc74d4a2ac5dfec4e9b2848ecd5a8 Mon Sep 17 00:00:00 2001 From: BeatSeat <137616995+BeatSeat@users.noreply.github.com> Date: Mon, 28 Sep 2026 10:11:54 +0800 Subject: [PATCH] [Performance][MiniCPM-o] Zero-copy block-aligned KV sliding window with in-place Re-RoPE for duplex Stage-0 (#7821) Signed-off-by: BeatSeat Co-authored-by: amy-why-3459 --- .../models/minicpmo_4_5/duplex/test_rerope.py | 422 +++++ .../minicpmo_4_5/duplex/test_window_plan.py | 396 ++++ .../minicpmo_4_5/duplex/test_window_wiring.py | 1589 +++++++++++++++++ vllm_omni/core/sched/omni_ar_scheduler.py | 83 +- .../models/minicpmo_4_5/duplex/stage0.py | 106 ++ .../models/minicpmo_4_5/duplex/window_kv.py | 966 ++++++++++ .../models/minicpmo_4_5/duplex/window_plan.py | 252 +++ .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 42 + vllm_omni/worker/gpu_ar_model_runner.py | 46 + 9 files changed, 3886 insertions(+), 16 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/duplex/test_rerope.py create mode 100644 tests/model_executor/models/minicpmo_4_5/duplex/test_window_plan.py create mode 100644 tests/model_executor/models/minicpmo_4_5/duplex/test_window_wiring.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_kv.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_plan.py diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_rerope.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_rerope.py new file mode 100644 index 000000000..a22a0d084 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_rerope.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Re-RoPE Correctness Regression Suite for MiniCPM-o 4.5 Stage-0 KV Window: + +Compares in-place Re-RoPE + Block-Table Compaction against clean re-prefill ground truth: +1. Retained K after Re-RoPE vs. K from clean re-prefill (FP32 & BF16). +2. Next-step logits & greedy token generation rollout. +3. Prefix KV cache invariance (sink tokens remain bitwise untouched). +4. V cache invariance (values never receive RoPE phase rotation). +5. Block-boundary alignment & repeated compactions (compound rotations without drift). +6. Basic/context modes watermark policies, reference-audio sink, and camera MRoPE validation. +""" + +from __future__ import annotations + +import math + +import pytest +import torch +import torch.nn.functional as F + +try: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import ( + DUPLEX_WINDOW_BLOCK_SIZE, + assert_uniform_position_shift, + rotate_keys, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan import ( + DuplexWindowGeometry, + plan_position_reanchor, + ) +except (ImportError, ModuleNotFoundError): + import importlib.util + import pathlib + import sys + import types + + def _make_pkg(name: str) -> types.ModuleType: + if name in sys.modules: + return sys.modules[name] + m = types.ModuleType(name) + m.__path__ = [] # type: ignore[attr-defined] + sys.modules[name] = m + return m + + vllm = _make_pkg("vllm") + vllm.__version__ = "0.7.0" # type: ignore[attr-defined] + vllm.__version_tuple__ = (0, 7, 0) # type: ignore[attr-defined] + v1 = _make_pkg("vllm.v1") + spec_reg = _make_pkg("vllm.v1.kv_cache_spec_registry") + spec_reg.register_kv_cache_spec = lambda *a, **k: (lambda cls: cls) # type: ignore[attr-defined] + kv_if = _make_pkg("vllm.v1.kv_cache_interface") + + class MockSpec: + pass + + kv_if.KVCacheSpec = MockSpec # type: ignore[attr-defined] + kv_if.SlidingWindowSpec = MockSpec # type: ignore[attr-defined] + + vo = _make_pkg("vllm_omni") + vo_exp = _make_pkg("vllm_omni.experimental") + vo_ad = _make_pkg("vllm_omni.experimental.ar_diffusion") + vo_kc = _make_pkg("vllm_omni.experimental.ar_diffusion.kv_cache") + vo_pg = _make_pkg("vllm_omni.experimental.ar_diffusion.kv_cache.paged") + + class MockChunkSpec: + pass + + class MockChunkManager: + def __init__(self, *a, **k): + pass + + def reanchor_block_table(self, *a, **k): + pass + + vo_pg.ChunkWindowSpec = MockChunkSpec # type: ignore[attr-defined] + vo_pg.ChunkWindowManager = MockChunkManager # type: ignore[attr-defined] + + def compute_slot_mapping(block_ids, positions, block_size): + p = positions.to(dtype=torch.long) + t = torch.tensor(block_ids, dtype=torch.long, device=p.device) + return t[torch.div(p, block_size, rounding_mode="floor")] * block_size + (p % block_size) + + vo_pg.compute_slot_mapping = compute_slot_mapping # type: ignore[attr-defined] + + _make_pkg("vllm_omni.model_executor") + _make_pkg("vllm_omni.model_executor.models") + _make_pkg("vllm_omni.model_executor.models.minicpmo_4_5") + _make_pkg("vllm_omni.model_executor.models.minicpmo_4_5.duplex") + + repo_root = pathlib.Path(__file__).resolve().parent + while repo_root.name and not (repo_root / "vllm_omni").is_dir(): + repo_root = repo_root.parent + + def _load_module(name: str, rel_path: str) -> types.ModuleType: + spec = importlib.util.spec_from_file_location(name, repo_root / rel_path) + assert spec is not None and spec.loader is not None + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + _wp = _load_module( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan", + "vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_plan.py", + ) + _wk = _load_module( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv", + "vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_kv.py", + ) + + DUPLEX_WINDOW_BLOCK_SIZE = _wk.DUPLEX_WINDOW_BLOCK_SIZE + assert_uniform_position_shift = _wk.assert_uniform_position_shift + rotate_keys = _wk.rotate_keys + DuplexWindowGeometry = _wp.DuplexWindowGeometry + plan_position_reanchor = _wp.plan_position_reanchor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +HEAD_DIM = 128 +NUM_HEADS = 8 +NUM_KV_HEADS = 2 +HIDDEN_DIM = NUM_HEADS * HEAD_DIM +VOCAB_SIZE = 1000 +BLOCK_SIZE = DUPLEX_WINDOW_BLOCK_SIZE # 16 + + +def _get_inv_freq(head_dim: int = HEAD_DIM, base: float = 1000000.0) -> torch.Tensor: + return 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + + +def _forward_rope(x: torch.Tensor, pos: int | torch.Tensor, inv_freq: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + if isinstance(pos, int): + pos_t = torch.tensor([pos], dtype=torch.float32, device=x.device) + else: + pos_t = pos.to(dtype=torch.float32, device=x.device) + angle = torch.outer(pos_t, inv_freq.to(dtype=torch.float32)).unsqueeze(1) + cos = torch.cos(angle).to(dtype=x.dtype) + sin = torch.sin(angle).to(dtype=x.dtype) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) + + +def _attention_decode( + query_token: torch.Tensor, + query_pos: int, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + w_q: torch.Tensor, + w_o: torch.Tensor, + lm_head: torch.Tensor, + embedding: torch.Tensor, + inv_freq: torch.Tensor, +) -> torch.Tensor: + q_x = embedding[query_token] + q_raw = (q_x @ w_q.T).view(-1, NUM_HEADS, HEAD_DIM) + q = _forward_rope(q_raw, query_pos, inv_freq) + + num_rep = NUM_HEADS // NUM_KV_HEADS + k_exp = k_cache.repeat_interleave(num_rep, dim=1) + v_exp = v_cache.repeat_interleave(num_rep, dim=1) + + scores = torch.einsum("qhd,khd->hqk", q.float(), k_exp.float()) / math.sqrt(HEAD_DIM) + probs = F.softmax(scores, dim=-1).to(k_cache.dtype) + attn_out = torch.einsum("hqk,khd->qhd", probs, v_exp).reshape(-1, HIDDEN_DIM) + out = attn_out @ w_o.T + logits = out @ lm_head.T + return logits + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_single_trim_ab_correctness(dtype: torch.dtype): + """Criteria 1, 2, 3, 4: Retained K, Next-step Logits/Tokens, Prefix KV, Tail V invariance.""" + torch.manual_seed(42) + inv_freq = _get_inv_freq(HEAD_DIM) + + W_q = (torch.randn(NUM_HEADS * HEAD_DIM, HIDDEN_DIM) * 0.02).to(dtype) + W_k = (torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN_DIM) * 0.02).to(dtype) + W_v = (torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN_DIM) * 0.02).to(dtype) + W_o = (torch.randn(HIDDEN_DIM, NUM_HEADS * HEAD_DIM) * 0.02).to(dtype) + lm_head = (torch.randn(VOCAB_SIZE, HIDDEN_DIM) * 0.02).to(dtype) + embedding = (torch.randn(VOCAB_SIZE, HIDDEN_DIM) * 0.02).to(dtype) + + def compute_kv(tokens: torch.Tensor, positions: torch.Tensor): + x = embedding[tokens] + k = (x @ W_k.T).view(-1, NUM_KV_HEADS, HEAD_DIM) + v = (x @ W_v.T).view(-1, NUM_KV_HEADS, HEAD_DIM) + return _forward_rope(k, positions, inv_freq), v + + PREFIX_LEN, GAP_LEN, TAIL_LEN = 32, 32, 48 + TOTAL_LEN = PREFIX_LEN + GAP_LEN + TAIL_LEN + tokens = torch.randint(0, VOCAB_SIZE, (TOTAL_LEN,)) + + k_init, v_init = compute_kv(tokens, torch.arange(TOTAL_LEN)) + + # Method A: In-Place Re-RoPE (rotate survivors by -GAP_LEN) + k_tail_A = rotate_keys(k_init[PREFIX_LEN + GAP_LEN :], GAP_LEN, inv_freq) + v_tail_A = v_init[PREFIX_LEN + GAP_LEN :] + k_A = torch.cat([k_init[:PREFIX_LEN], k_tail_A], dim=0) + v_A = torch.cat([v_init[:PREFIX_LEN], v_tail_A], dim=0) + + # Method B: Clean Re-prefill Ground Truth + kept_tokens = torch.cat([tokens[:PREFIX_LEN], tokens[PREFIX_LEN + GAP_LEN :]], dim=0) + k_B, v_B = compute_kv(kept_tokens, torch.arange(len(kept_tokens))) + + # Invariance assertions: + # 1. Prefix K and V are strictly untouched (sink) + assert torch.equal(k_A[:PREFIX_LEN], k_init[:PREFIX_LEN]), "Prefix K must be bitwise untouched" + assert torch.equal(v_A[:PREFIX_LEN], v_init[:PREFIX_LEN]), "Prefix V must be bitwise untouched" + # 2. Tail V is never rotated + assert torch.equal(v_tail_A, v_init[PREFIX_LEN + GAP_LEN :]), "Tail V must not be modified by Re-RoPE" + + # 3. Retained K matches clean re-prefill + diff_k = (k_tail_A.float() - k_B[PREFIX_LEN:].float()).abs().max().item() + tol_k = 1e-5 if dtype == torch.float32 else 0.035 + assert diff_k < tol_k, f"Retained K diff {diff_k} exceeds tolerance {tol_k}" + + cos_sim = ( + F.cosine_similarity( + k_tail_A.reshape(-1, HEAD_DIM).float(), + k_B[PREFIX_LEN:].reshape(-1, HEAD_DIM).float(), + dim=-1, + ) + .mean() + .item() + ) + assert cos_sim > 0.9999, f"Cosine similarity {cos_sim} too low" + + # 4. Next-step logits & greedy token match + q_tok = torch.randint(0, VOCAB_SIZE, (1,)) + next_pos = len(kept_tokens) + log_A = _attention_decode(q_tok, next_pos, k_A, v_A, W_q, W_o, lm_head, embedding, inv_freq) + log_B = _attention_decode(q_tok, next_pos, k_B, v_B, W_q, W_o, lm_head, embedding, inv_freq) + + tok_A = log_A.argmax(dim=-1).item() + tok_B = log_B.argmax(dim=-1).item() + assert tok_A == tok_B, f"Greedy next-step token mismatch: {tok_A} != {tok_B}" + if dtype == torch.float32: + diff_log = (log_A.float() - log_B.float()).abs().max().item() + assert diff_log < 1e-4, f"Logits diff {diff_log} exceeds tolerance 1e-4" + + +def test_repeated_consecutive_trims_and_greedy_rollout(): + """Criteria 5: Verify compound rotations across 3 consecutive trims and 5-step rollout.""" + torch.manual_seed(123) + inv_freq = _get_inv_freq(HEAD_DIM) + + W_q = torch.randn(NUM_HEADS * HEAD_DIM, HIDDEN_DIM) * 0.02 + W_k = torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN_DIM) * 0.02 + W_v = torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN_DIM) * 0.02 + W_o = torch.randn(HIDDEN_DIM, NUM_HEADS * HEAD_DIM) * 0.02 + lm_head = torch.randn(VOCAB_SIZE, HIDDEN_DIM) * 0.02 + embedding = torch.randn(VOCAB_SIZE, HIDDEN_DIM) * 0.02 + + def compute_kv(tokens: torch.Tensor, positions: torch.Tensor): + x = embedding[tokens] + k = (x @ W_k.T).view(-1, NUM_KV_HEADS, HEAD_DIM) + v = (x @ W_v.T).view(-1, NUM_KV_HEADS, HEAD_DIM) + return _forward_rope(k, positions, inv_freq), v + + PREFIX_LEN, CHUNK = 32, 16 + t0 = torch.randint(0, VOCAB_SIZE, (PREFIX_LEN,)) + t1 = torch.randint(0, VOCAB_SIZE, (CHUNK,)) + t2 = torch.randint(0, VOCAB_SIZE, (CHUNK,)) + t3 = torch.randint(0, VOCAB_SIZE, (CHUNK,)) + + # Prefill & Trim 1 (drop t1) + k_p, v_p = compute_kv(torch.cat([t0, t1, t2, t3]), torch.arange(PREFIX_LEN + 3 * CHUNK)) + k1 = torch.cat([k_p[:PREFIX_LEN], rotate_keys(k_p[PREFIX_LEN + CHUNK :], CHUNK, inv_freq)], dim=0) + v1 = torch.cat([v_p[:PREFIX_LEN], v_p[PREFIX_LEN + CHUNK :]], dim=0) + + # Append 1: t4, t5 + t4 = torch.randint(0, VOCAB_SIZE, (CHUNK,)) + t5 = torch.randint(0, VOCAB_SIZE, (CHUNK,)) + k_s1, v_s1 = compute_kv(torch.cat([t4, t5]), torch.arange(64, 64 + 2 * CHUNK)) + k2 = torch.cat([k1, k_s1], dim=0) + v2 = torch.cat([v1, v_s1], dim=0) + + # Trim 2 (drop t2): t3 receives its 2nd rotation + k3 = torch.cat([k2[:PREFIX_LEN], rotate_keys(k2[PREFIX_LEN + CHUNK :], CHUNK, inv_freq)], dim=0) + v3 = torch.cat([v2[:PREFIX_LEN], v2[PREFIX_LEN + CHUNK :]], dim=0) + + # Append 2: t6 + t6 = torch.randint(0, VOCAB_SIZE, (CHUNK,)) + k_s2, v_s2 = compute_kv(t6, torch.arange(80, 80 + CHUNK)) + k4 = torch.cat([k3, k_s2], dim=0) + v4 = torch.cat([v3, v_s2], dim=0) + + # Trim 3 (drop t3): t4 receives its 2nd rotation, t5 its 1st + k_A_final = torch.cat([k4[:PREFIX_LEN], rotate_keys(k4[PREFIX_LEN + CHUNK :], CHUNK, inv_freq)], dim=0) + v_A_final = torch.cat([v4[:PREFIX_LEN], v4[PREFIX_LEN + CHUNK :]], dim=0) + + # Ground Truth: Clean re-prefill of surviving sequence [t0, t4, t5, t6] + kept_final = torch.cat([t0, t4, t5, t6]) + k_B_final, v_B_final = compute_kv(kept_final, torch.arange(len(kept_final))) + + diff_compound = (k_A_final - k_B_final).abs().max().item() + assert diff_compound < 1e-5, f"Compound rotation diff {diff_compound} exceeds 1e-5" + assert torch.allclose(v_A_final, v_B_final, atol=1e-6), ( + f"Compound V diff {(v_A_final - v_B_final).abs().max().item()} exceeds 1e-6" + ) + + # Multi-step autoregressive rollout (5 steps greedy decode) + cur_A_tok = torch.tensor([42]) + cur_B_tok = torch.tensor([42]) + curr_k_A, curr_v_A = k_A_final.clone(), v_A_final.clone() + curr_k_B, curr_v_B = k_B_final.clone(), v_B_final.clone() + + rollout_matches = [] + for step in range(5): + cur_pos = len(kept_final) + step + l_A = _attention_decode(cur_A_tok, cur_pos, curr_k_A, curr_v_A, W_q, W_o, lm_head, embedding, inv_freq) + l_B = _attention_decode(cur_B_tok, cur_pos, curr_k_B, curr_v_B, W_q, W_o, lm_head, embedding, inv_freq) + tok_A = l_A.argmax(dim=-1) + tok_B = l_B.argmax(dim=-1) + rollout_matches.append(tok_A.item() == tok_B.item()) + + k_step_A, v_step_A = compute_kv(tok_A, torch.tensor([cur_pos])) + curr_k_A = torch.cat([curr_k_A, k_step_A], dim=0) + curr_v_A = torch.cat([curr_v_A, v_step_A], dim=0) + + k_step_B, v_step_B = compute_kv(tok_B, torch.tensor([cur_pos])) + curr_k_B = torch.cat([curr_k_B, k_step_B], dim=0) + curr_v_B = torch.cat([curr_v_B, v_step_B], dim=0) + + cur_A_tok = tok_A + cur_B_tok = tok_B + + assert all(rollout_matches), f"Greedy rollout diverged at step {rollout_matches.index(False)}" + + +def test_block_boundary_and_fractional_tail(): + """Criteria 5: Test non-multiple block sizes and fractional tails.""" + torch.manual_seed(99) + inv_freq = _get_inv_freq(HEAD_DIM) + + W_k = torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN_DIM) * 0.02 + W_v = torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN_DIM) * 0.02 + embedding = torch.randn(VOCAB_SIZE, HIDDEN_DIM) * 0.02 + + def compute_kv(tokens: torch.Tensor, positions: torch.Tensor): + x = embedding[tokens] + k = (x @ W_k.T).view(-1, NUM_KV_HEADS, HEAD_DIM) + v = (x @ W_v.T).view(-1, NUM_KV_HEADS, HEAD_DIM) + return _forward_rope(k, positions, inv_freq), v + + PREFIX_LEN, GAP_LEN, PARTIAL_TAIL = 32, 16, 37 # 37 is not a multiple of 16 + tokens = torch.randint(0, VOCAB_SIZE, (PREFIX_LEN + GAP_LEN + PARTIAL_TAIL,)) + k_init, _ = compute_kv(tokens, torch.arange(len(tokens))) + + k_tail_A = rotate_keys(k_init[PREFIX_LEN + GAP_LEN :], GAP_LEN, inv_freq) + k_A = torch.cat([k_init[:PREFIX_LEN], k_tail_A], dim=0) + + kept = torch.cat([tokens[:PREFIX_LEN], tokens[PREFIX_LEN + GAP_LEN :]], dim=0) + k_B, _ = compute_kv(kept, torch.arange(len(kept))) + + diff = (k_A - k_B).abs().max().item() + assert diff < 1e-5, f"Fractional tail diff {diff} exceeds 1e-5" + + +def test_basic_vs_context_mode_watermarks(): + """Criteria 6: Test watermark planner under basic and context session modes.""" + # Basic Mode: fixed watermarks (8000 high, 6000 low) + geom_basic = DuplexWindowGeometry( + prefix_tokens=96, + window_tokens=6000, + block_size=16, + max_model_len=40960, + high_watermark_tokens=8000, + ) + # Below high watermark -> No trim + assert plan_position_reanchor(geom_basic, computed_tokens=7900, pending_tokens=12) is None + # Exceed high watermark (trigger = 96 + 8000 = 8096) -> Trims to low watermark + plan_basic = plan_position_reanchor(geom_basic, computed_tokens=8100, pending_tokens=12) + assert plan_basic is not None + assert plan_basic.delta % 16 == 0 + assert plan_basic.sink_blocks == 6 + + # Context Mode: dynamic unit history (e.g. max_units=24 -> 288 tokens + prefix) + max_units = 24 + ctx_low = 96 + max_units * 12 # 384 tokens + ctx_high = ctx_low + 500 # 884 tokens + geom_context = DuplexWindowGeometry( + prefix_tokens=96, + window_tokens=ctx_low, + block_size=16, + max_model_len=40960, + high_watermark_tokens=ctx_high, + ) + # Exceed high watermark (trigger = 96 + 884 = 980) -> Trims + plan_ctx = plan_position_reanchor(geom_context, computed_tokens=980, pending_tokens=12) + assert plan_ctx is not None + assert plan_ctx.delta % 16 == 0 + assert plan_ctx.sink_blocks == 6 + + +def test_reference_audio_and_camera_mrope_paths(): + """Criteria 6: Reference-audio permanence in sink and camera non-uniform MRoPE rejection.""" + # Case A: Reference audio tokens (0..96) always live in sink blocks 0..5 + geom = DuplexWindowGeometry( + prefix_tokens=96, + window_tokens=6000, + block_size=16, + max_model_len=40960, + high_watermark_tokens=8000, + ) + plan = plan_position_reanchor(geom, computed_tokens=8100, pending_tokens=12) + assert plan is not None + assert plan.sink_blocks * 16 == 96 + assert plan.moved_from > 96 # Reference audio is strictly never rotated + + # Case B: Camera frame in retained tail (MRoPE non-uniform grid coordinates) + mrope_pos_camera = torch.arange(200).unsqueeze(0).repeat(3, 1) + mrope_pos_camera[1, 120:] += 4 + mrope_pos_camera[2, 120:] += 8 + with pytest.raises(RuntimeError, match="duplex re-anchor needs one position row"): + assert_uniform_position_shift(mrope_pos_camera, moved_from=96) + + # Case C: Streaming audio/text units (MRoPE multi-row uniform advance) + mrope_pos_audio = torch.arange(200).unsqueeze(0).repeat(3, 1) + assert_uniform_position_shift(mrope_pos_audio, moved_from=96) diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_window_plan.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_window_plan.py new file mode 100644 index 000000000..d7157d331 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_window_plan.py @@ -0,0 +1,396 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU accounting tests for the MiniCPM-o 4.5 duplex KV window planner. + +Pure integer/complex math: no engine, no checkpoint, no GPU. Two claims carry +the design and each gets pinned here: a block-aligned shift leaves every retained +row in the physical slot it already occupies (so a trim moves no memory), and +rotating a cached key by ``-delta`` equals RoPE at the renumbered position (so a +trim needs no forward pass). +""" + +from __future__ import annotations + +import cmath + +import pytest + +try: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan import ( + DuplexWindowGeometry, + PositionReanchor, + align_up, + cdiv, + plan_position_reanchor, + reanchor_positions, + unit_starts, + ) +except (ImportError, ModuleNotFoundError): + import importlib.util + import pathlib + import sys + + repo_root = pathlib.Path(__file__).resolve().parent + while repo_root.name and not (repo_root / "vllm_omni").is_dir(): + repo_root = repo_root.parent + + spec = importlib.util.spec_from_file_location( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan", + repo_root / "vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_plan.py", + ) + assert spec is not None and spec.loader is not None + _wp = importlib.util.module_from_spec(spec) + sys.modules["vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan"] = _wp + spec.loader.exec_module(_wp) + + DuplexWindowGeometry = _wp.DuplexWindowGeometry + PositionReanchor = _wp.PositionReanchor + align_up = _wp.align_up + cdiv = _wp.cdiv + plan_position_reanchor = _wp.plan_position_reanchor + reanchor_positions = _wp.reanchor_positions + unit_starts = _wp.unit_starts + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Block size 4, a 10-token session context prefix (reference audio + system +# prompt), a 64-token attention budget, 4096 positions and 2 tokens of sample +# room: small enough to assert on by hand, same shape as the real stage. +BLOCK = 4 + + +def _geometry(**overrides) -> DuplexWindowGeometry: + kwargs: dict[str, int] = { + "prefix_tokens": 10, + "window_tokens": 64, + "block_size": BLOCK, + "max_model_len": 4096, + "sample_room": 2, + } + kwargs.update(overrides) + return DuplexWindowGeometry(**kwargs) + + +def test_rejects_impossible_geometry(): + with pytest.raises(ValueError, match="window"): + _geometry(window_tokens=0) + with pytest.raises(ValueError, match="prefix"): + _geometry(prefix_tokens=-1) + with pytest.raises(ValueError, match="block_size"): + _geometry(block_size=0) + # A trigger below the low watermark would trim the sequence it just rebuilt. + with pytest.raises(ValueError, match="high_watermark_tokens"): + _geometry(high_watermark_tokens=8) + + +def test_nothing_is_reclaimed_inside_the_window(): + geometry = _geometry(high_watermark_tokens=96) + # The trigger is prefix + high watermark == 106; below it, no plan at all. + assert plan_position_reanchor(geometry, computed_tokens=100, pending_tokens=6) is None + assert plan_position_reanchor(geometry, computed_tokens=100, pending_tokens=7) is not None + + +def test_a_trim_lands_on_the_low_watermark(): + geometry = _geometry(high_watermark_tokens=96) + plan = plan_position_reanchor(geometry, computed_tokens=200, pending_tokens=20) + assert plan is not None + # Cut so that what survives past the sink is the low watermark, and no more: + # the tail is kept within one page of the budget so a trim never over-reclaims. + target = geometry.prefix_tokens + geometry.window_tokens + assert plan.delta == align_up(220 - target, BLOCK) == 148 + assert plan.moved_from == plan.sink_end + plan.delta == 160 + assert plan.delta % BLOCK == 0 + # The sink never shrinks: it ends on the first whole page after the prefix. + assert plan.sink_end == cdiv(geometry.prefix_tokens, BLOCK) * BLOCK == 12 + retained = 220 - plan.delta + assert target - BLOCK < retained <= target + + +def test_production_planner_drop_length_and_retained_tail(): + # Production parameters from reviewer @linyueqian: + # prefix=96, block_size=16, window=6000, high=8000, computed=8100, pending=12 + geo = DuplexWindowGeometry( + prefix_tokens=96, + window_tokens=6000, + block_size=16, + max_model_len=40960, + high_watermark_tokens=8000, + ) + plan = plan_position_reanchor(geo, computed_tokens=8100, pending_tokens=12) + assert plan is not None + assert plan.sink_blocks == 6 + assert plan.sink_end == 96 + # projected = 8112, target = 6096, drop_needed = 2016 (already aligned to 16) + assert plan.delta == 2016 + assert plan.moved_from == 96 + 2016 == 2112 + # Resulting retained length must be exactly target (6096), not 6192! + retained_length = (8100 + 12) - plan.delta + assert retained_length == 6096 + + +def test_the_shift_is_a_whole_number_of_pages_so_no_row_moves(): + """The claim that makes a trim a memory pass rather than a copy. + + Deleting ``n`` gap entries shifts every later logical index down by ``n``. + Because ``delta == n * block_size``, a tail token's new logical index lands on + the table slot that already names its own page, and its offset within that + page is unchanged -- so the physical slot is identical before and after. + """ + geometry = _geometry(high_watermark_tokens=96) + plan = plan_position_reanchor(geometry, computed_tokens=200, pending_tokens=20) + assert plan is not None + n_gap = plan.delta // BLOCK + table_before = list(range(64)) # one fake physical page id per logical slot + table_after = table_before[: plan.sink_blocks] + table_before[plan.sink_blocks + n_gap :] + + for position in range(plan.moved_from, 64 * BLOCK): + before = divmod(position, BLOCK) + after = divmod(reanchor_positions(position, plan), BLOCK) + assert table_before[before[0]] == table_after[after[0]] + assert before[1] == after[1] + + +def test_unit_aligned_cut_keeps_whole_units(): + # Page-aligned unit starts: prefix_tokens=16 (divisible by BLOCK=4), units=[12]*8 + units = [12] * 8 + geometry = _geometry(prefix_tokens=16, window_tokens=30, high_watermark_tokens=40) + boundaries = unit_starts(geometry, units) + assert boundaries[0] == 16 + # Unit starts: 16, 28, 40, 52, 64, ... all divisible by 4! + plan = plan_position_reanchor(geometry, computed_tokens=16 + sum(units), pending_tokens=12, unit_tokens=units) + assert plan is not None + # Landed on a block-aligned unit start, so no unit was split in half. + assert plan.moved_from in boundaries + assert plan.moved_from % BLOCK == 0 + # Post trim length <= target + target = geometry.prefix_tokens + geometry.window_tokens + assert (16 + sum(units) + 12) - plan.delta <= target + + +def test_unit_aligned_cut_rejects_non_page_aligned_units(): + # If units cannot be cut along page boundaries without splitting a unit in half, + # plan_position_reanchor should return None to trigger fallback rather than slicing a unit. + units_odd = [7] * 8 + geo_no_align = DuplexWindowGeometry( + prefix_tokens=11, + window_tokens=20, + block_size=8, + max_model_len=4096, + high_watermark_tokens=30, + ) + # Unit starts: 11, 18, 25, 32, 39, 46, 53, 60, 67. None are divisible by 8! + plan = plan_position_reanchor(geo_no_align, computed_tokens=11 + 56, pending_tokens=7, unit_tokens=units_odd) + assert plan is None + + +def test_reanchor_is_a_uniform_translation_of_the_tail(): + geometry = _geometry(high_watermark_tokens=96) + plan = plan_position_reanchor(geometry, computed_tokens=200, pending_tokens=20) + assert plan is not None + assert plan.delta == plan.moved_from - plan.sink_end + + assert reanchor_positions(3, plan) == 3 + tail_old = plan.moved_from + 5 + tail_new = reanchor_positions(tail_old, plan) + assert tail_new == tail_old - plan.delta + # Relative distances inside the tail are untouched by definition, which is + # what lets one rotation fix every pair at once. + assert reanchor_positions(plan.moved_from + 37, plan) - tail_new == 32 + # The reclaimed range is dense: the tail lands exactly on the sink's end. + assert reanchor_positions(plan.moved_from, plan) == plan.sink_end + + for dropped in (plan.sink_end, plan.moved_from - 1): + with pytest.raises(ValueError, match="dropped"): + reanchor_positions(dropped, plan) + + +def test_reanchor_declines_when_there_is_nothing_left_to_reclaim(): + geometry = _geometry(high_watermark_tokens=96) + # Target at or past the projected length: compacting would not help, so the + # caller has to finish the session with context_length_exceeded instead of moving + # the window start backwards. + assert plan_position_reanchor(geometry, computed_tokens=100, pending_tokens=20, target_tokens=160) is None + assert plan_position_reanchor(geometry, computed_tokens=100, pending_tokens=20, target_tokens=120) is None + plan = plan_position_reanchor(geometry, computed_tokens=100, pending_tokens=20, target_tokens=20) + assert plan is not None + assert plan.sink_end >= geometry.prefix_tokens + assert reanchor_positions(2, plan) == 2 + + +@pytest.mark.parametrize("head_dim", [8, 16]) +@pytest.mark.parametrize("delta", [4, 64, 1024]) +def test_rotating_cached_keys_equals_rope_at_the_new_position(head_dim: int, delta: int): + """The other half of the claim: the rotation is the re-RoPE. + + A key written at position ``p`` holds ``z * exp(i * p * f)`` per frequency + pair. Re-anchoring to ``p - delta`` can therefore be done on the cached + tensor by rotating ``-delta * f`` -- no recompute of the model, and exact for + any p, including a tail whose keys were written at different positions. + """ + plan = PositionReanchor(delta=delta, moved_from=4096, sink_blocks=4096 // BLOCK) + inv_freq = [0.5 ** (2.0 * i / head_dim) for i in range(head_dim // 2)] + rng = iter(range(1, 10_000)) + for pair, freq in enumerate(inv_freq): + re = complex(next(rng) % 7 - 3, next(rng) % 7 - 3) + for position in (plan.moved_from + 1, plan.moved_from + 3 * pair + 17, 999_999): + cached = re * cmath.exp(1j * position * freq) + rotated = cached * cmath.exp(1j * plan.rope_angle_shift_tokens * freq) + direct = re * cmath.exp(1j * reanchor_positions(position, plan) * freq) + assert rotated == pytest.approx(direct, rel=1e-9, abs=1e-12), f"pair {pair} at {position}" + + +def test_reanchor_rejects_a_fractional_position(): + plan = PositionReanchor(delta=8, moved_from=64, sink_blocks=16) + with pytest.raises(ValueError, match="integer"): + reanchor_positions(70.5, plan) + + +# Realistic Stage 0 numbers: a ~100-token head context, the 6000-token low +# watermark from the duplex window policy, 16-token KV blocks, and the 40960 +# positions the checkpoint's config.json advertises. +def _stage0_geometry(**overrides) -> DuplexWindowGeometry: + kwargs: dict[str, int] = { + "prefix_tokens": 96, + "window_tokens": 6000, + "block_size": 16, + "max_model_len": 40960, + "sample_room": 2048, + # The reference implementation's trigger/stride pair, in tokens: roll + # back to the window once 2000 tokens have piled on past it. + "high_watermark_tokens": 8000, + } + kwargs.update(overrides) + return DuplexWindowGeometry(**kwargs) + + +def test_ten_minutes_of_audio_recomputes_nothing(): + """What the rebuild-based window pays, measured against this design. + + 600 one-second units is about 7.2k appended tokens plus the spoken + interleaves. A window that rebuilds the prompt re-prefills the whole retained + tail every time it rolls; a trim re-prefills nothing and still bounds the + session's residency, so the position ceiling is never in sight. + """ + geometry = _stage0_geometry() + units = [12] * 600 # 1 s of audio: 10 pooled embeddings + unit closure. + spoken = 20 # tokens the model speaks per unit, at 1 s of audio. + + position = geometry.prefix_tokens + trims = 0 + peak_rows = 0 + written = geometry.prefix_tokens + # resident_bound() from the engine tier, plus whatever single append was + # already in flight when the trim came due: content rides to the trigger, and + # a trim then pulls it back to the low watermark. + bound = geometry.prefix_tokens + geometry.trigger_tokens + max(units) + spoken + geometry.sample_room + for unit in units: + assert position + unit + spoken <= bound, "the session outgrew its own resident bound" + plan = plan_position_reanchor(geometry, computed_tokens=position, pending_tokens=unit + spoken) + if plan is not None: + trims += 1 + # Rows the rotation touches: the reclaimed tail, and never more than + # what is resident. A re-prefill would run those same rows through + # every layer of the model instead. + peak_rows = max(peak_rows, position - plan.moved_from) + position -= plan.delta + position += unit + spoken + written += unit + spoken + + assert trims > 0, "the window has to roll over a session this long" + assert peak_rows <= bound + # Ten minutes of conversation wrote ~19k positions while never holding more + # than the watermark: recycling, not the length, is what bounds memory. + assert written > bound + # A 40k position budget against an 8k resident bound: the ceiling is a + # backstop this session never approaches, let alone crosses. + assert position + geometry.sample_room < geometry.max_model_len + + +def test_a_long_session_trims_once_per_spare_window_not_once_per_append(): + """The frequency the whole design is priced at.""" + geometry = _stage0_geometry() + position = geometry.prefix_tokens + append = 12 + gaps: list[int] = [] + since = 0 + peak = position + for _ in range(4000): + plan = plan_position_reanchor(geometry, computed_tokens=position, pending_tokens=append) + if plan is not None: + gaps.append(since) + since = 0 + position -= plan.delta + since += 1 + position += append + peak = max(peak, position) + # 48k positions appended against a window that only overflows once the spare + # tokens past it are used up: one trim services ~160 one-second appends. The + # first gap is longer because the window has to fill from nothing first. + spare = geometry.trigger_tokens - geometry.window_tokens + assert gaps and all(gap > spare // (2 * append) for gap in gaps) + assert all(gap <= (spare + append - 1) // append + 1 for gap in gaps[1:]) + # The watermark, not the 40960-token ceiling, is what holds the session down. + assert peak <= geometry.prefix_tokens + geometry.trigger_tokens + append + assert peak + geometry.sample_room < geometry.max_model_len + + # The trigger is a policy choice, so a stage with no window is still legal -- + # it simply trades the rotation for a shorter session cap. + unwindowed = _stage0_geometry(window_tokens=40960, high_watermark_tokens=None) + assert unwindowed.trigger_tokens == 40960 + + +def test_paged_retention_policy_vs_unpaged_reference(): + """Verify and document retention differences between paged KV and unpaged reference. + + Case from reviewer @amy-why-3459: + prefix_tokens = 100, block_size = 16 + unit_tokens = [28, 36] (units at [100..128) and [128..164), current total = 164) + watermarks: high = 160, low = 148 + + Reference unpaged policy (OpenBMB): + - Drops unit 0 ([100..128), 28 tokens) + - Retains [0..100) + [128..164), total length = 136 + - Token at 128 moves to position 100 + + Paged zero-copy policy (plan_position_reanchor): + - sink_end = cdiv(100, 16) * 16 = 112 + - Preserves [100..112) inside the physical sink block (block 6) + - moved_from = 128, delta = 16 (1 block) + - Drops [112..128) (16 tokens) + - Retains [0..112) + [128..164), total length = 148 + - Token at 128 moves to position 112 (128 - 16) + """ + geometry = DuplexWindowGeometry( + prefix_tokens=100, + window_tokens=48, # low watermark target = 100 + 48 = 148 + block_size=16, + max_model_len=2048, + high_watermark_tokens=60, # trigger = 100 + 60 = 160 + ) + plan = plan_position_reanchor( + geometry, + computed_tokens=164, + pending_tokens=0, + unit_tokens=[28, 36], + ) + assert plan is not None + assert plan.delta == 16 + assert plan.moved_from == 128 + assert plan.sink_blocks == 7 + assert plan.sink_end == 112 + + # Verify position mapping: + # Sink prefix [0..100) unchanged + assert reanchor_positions(0, plan) == 0 + assert reanchor_positions(99, plan) == 99 + # Paged sink tail [100..112) unchanged + assert reanchor_positions(100, plan) == 100 + assert reanchor_positions(111, plan) == 111 + # Dropped span [112..128) raises ValueError + with pytest.raises(ValueError, match="was dropped by reanchor"): + reanchor_positions(112, plan) + with pytest.raises(ValueError, match="was dropped by reanchor"): + reanchor_positions(127, plan) + # Retained tail [128..164) shifts by -delta (-16) + assert reanchor_positions(128, plan) == 112 + assert reanchor_positions(163, plan) == 147 diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_window_wiring.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_window_wiring.py new file mode 100644 index 000000000..aebf13312 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_window_wiring.py @@ -0,0 +1,1589 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Tests covering the Scheduler, Worker, and Batch Wiring for MiniCPM-o 4.5 duplex KV window. + +Covers: +1. Block-table in-place compaction on the Worker side without row moves. +2. rotate_cached_keys numeric identity with attention dot products. +3. Multi-request concurrency / batched execution: + - Request 0: triggers window trim with delta=16 + - Request 1: normal decode (untouched) + - Request 2: triggers window trim with delta=32 +4. Scheduler-side watermark detection and reanchor plan attachment without full re-prefill. +5. Barge-in / abort safety ensuring zero leaked blocks. +""" + +from __future__ import annotations + +import ast +import pathlib +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch + +repo_root = pathlib.Path(__file__).resolve().parent +while repo_root.name and not (repo_root / "vllm_omni").is_dir(): + repo_root = repo_root.parent + +try: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + _MiniCPMO45WindowUnit, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import ( + DUPLEX_WINDOW_BLOCK_SIZE, + MiniCPMO45DuplexSchedulerHelper, + MiniCPMO45DuplexWindowManager, + MiniCPMO45DuplexWorkerHelper, + assert_uniform_position_shift, + duplex_window_geometry, + rotate_cached_keys, + rotate_keys, + validate_duplex_window_install, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan import ( + PositionReanchor, + plan_position_reanchor, + ) +except (ImportError, ModuleNotFoundError): + import importlib.util + import sys + import types + + def _make_pkg(name: str) -> types.ModuleType: + if name in sys.modules: + return sys.modules[name] + m = types.ModuleType(name) + m.__path__ = [] # type: ignore[attr-defined] + sys.modules[name] = m + return m + + vllm = _make_pkg("vllm") + vllm.__version__ = "0.7.0" # type: ignore[attr-defined] + vllm.__version_tuple__ = (0, 7, 0) # type: ignore[attr-defined] + v1 = _make_pkg("vllm.v1") + spec_reg = _make_pkg("vllm.v1.kv_cache_spec_registry") + spec_reg.register_kv_cache_spec = lambda *a, **k: (lambda cls: cls) # type: ignore[attr-defined] + kv_if = _make_pkg("vllm.v1.kv_cache_interface") + + class MockSpec: + pass + + kv_if.KVCacheSpec = MockSpec # type: ignore[attr-defined] + kv_if.SlidingWindowSpec = MockSpec # type: ignore[attr-defined] + + vo = _make_pkg("vllm_omni") + vo_exp = _make_pkg("vllm_omni.experimental") + vo_ad = _make_pkg("vllm_omni.experimental.ar_diffusion") + vo_kc = _make_pkg("vllm_omni.experimental.ar_diffusion.kv_cache") + vo_pg = _make_pkg("vllm_omni.experimental.ar_diffusion.kv_cache.paged") + + class MockChunkSpec: + pass + + class MockChunkManager: + def __init__(self, *a, **k): + self.enable_caching = False + self.req_to_blocks = {} + self.block_size = 16 + self._null_block = -1 + + def compact_block_table(self, request_id: str, sink_blocks: int | None = None) -> int: + blocks = self.req_to_blocks.get(request_id, []) + start = 0 if sink_blocks is None else int(sink_blocks) + end = start + while end < len(blocks) and blocks[end] == self._null_block: + end += 1 + if end == start: + return 0 + del blocks[start:end] + return (end - start) * self.block_size + + def compact_request_blocks(self, request_id: str, *, sink_blocks: int, num_blocks: int) -> int: + blocks = self.req_to_blocks.get(request_id) + if blocks is None or num_blocks <= 0: + return 0 + start = int(sink_blocks) + end = start + int(num_blocks) + if start < 0 or end > len(blocks) or any(b == self._null_block for b in blocks[start:end]): + return 0 + for i in range(start, end): + blocks[i] = self._null_block + return self.compact_block_table(request_id, sink_blocks=start) + + def reanchor_block_table(self, request_id: str, plan: Any) -> int: + gap_blocks = plan.delta // self.block_size + return self.compact_request_blocks( + request_id, + sink_blocks=plan.sink_blocks, + num_blocks=gap_blocks, + ) + + vo_pg.ChunkWindowSpec = MockChunkSpec # type: ignore[attr-defined] + vo_pg.ChunkWindowManager = MockChunkManager # type: ignore[attr-defined] + + def compute_slot_mapping(block_ids, positions, block_size): + p = positions.to(dtype=torch.long) + t = torch.tensor(block_ids, dtype=torch.long, device=p.device) + return t[torch.div(p, block_size, rounding_mode="floor")] * block_size + (p % block_size) + + vo_pg.compute_slot_mapping = compute_slot_mapping # type: ignore[attr-defined] + + _make_pkg("vllm_omni.model_executor") + _make_pkg("vllm_omni.model_executor.models") + _make_pkg("vllm_omni.model_executor.models.minicpmo_4_5") + _make_pkg("vllm_omni.model_executor.models.minicpmo_4_5.duplex") + + repo_root = pathlib.Path(__file__).resolve().parent + while repo_root.name and not (repo_root / "vllm_omni").is_dir(): + repo_root = repo_root.parent + + def _load_module(name: str, rel_path: str) -> types.ModuleType: + spec = importlib.util.spec_from_file_location(name, repo_root / rel_path) + assert spec is not None and spec.loader is not None + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + _wp = _load_module( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan", + "vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_plan.py", + ) + _wk = _load_module( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv", + "vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_kv.py", + ) + + _wp = _load_module( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan", + "vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_plan.py", + ) + _wk = _load_module( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv", + "vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_kv.py", + ) + _policy = _load_module( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy", + "vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py", + ) + _stage0 = _load_module( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0", + "vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py", + ) + + DUPLEX_WINDOW_BLOCK_SIZE = _wk.DUPLEX_WINDOW_BLOCK_SIZE + MiniCPMO45DuplexSchedulerHelper = _wk.MiniCPMO45DuplexSchedulerHelper + MiniCPMO45DuplexWindowManager = _wk.MiniCPMO45DuplexWindowManager + MiniCPMO45DuplexWorkerHelper = _wk.MiniCPMO45DuplexWorkerHelper + assert_uniform_position_shift = _wk.assert_uniform_position_shift + duplex_window_geometry = _wk.duplex_window_geometry + rotate_cached_keys = _wk.rotate_cached_keys + rotate_keys = _wk.rotate_keys + validate_duplex_window_install = _wk.validate_duplex_window_install + PositionReanchor = _wp.PositionReanchor + plan_position_reanchor = _wp.plan_position_reanchor + MiniCPMO45Stage0DuplexRuntime = _stage0.MiniCPMO45Stage0DuplexRuntime + _MiniCPMO45Stage0SessionState = _stage0._MiniCPMO45Stage0SessionState + _MiniCPMO45WindowUnit = _stage0._MiniCPMO45WindowUnit + + +def _extract_prepare_stage0_window(): + p = repo_root / "vllm_omni" / "core" / "sched" / "omni_ar_scheduler.py" + tree = ast.parse(p.read_text(encoding="utf-8")) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == "_prepare_minicpmo45_stage0_window": + node.decorator_list = [] + mod = ast.Module(body=[node], type_ignores=[]) + ns: dict[str, Any] = {} + exec(compile(mod, str(p), "exec"), ns) + return ns["_prepare_minicpmo45_stage0_window"] + raise RuntimeError("Could not find _prepare_minicpmo45_stage0_window in omni_ar_scheduler.py") + + +prepare_minicpmo45_stage0_window = _extract_prepare_stage0_window() + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +BLOCK_SIZE = DUPLEX_WINDOW_BLOCK_SIZE # 16 +HEAD_DIM = 128 +NUM_KV_HEADS = 8 + + +def _get_inv_freq(head_dim: int = HEAD_DIM, base: float = 1000000.0) -> torch.Tensor: + return 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + + +def _forward_rope(x_raw: torch.Tensor, pos: int, inv_freq: torch.Tensor) -> torch.Tensor: + half = x_raw.shape[-1] // 2 + angle = float(pos) * inv_freq.to(device=x_raw.device, dtype=torch.float32) + cos = torch.cos(angle).to(dtype=x_raw.dtype).unsqueeze(0).unsqueeze(1) + sin = torch.sin(angle).to(dtype=x_raw.dtype).unsqueeze(0).unsqueeze(1) + x1, x2 = x_raw[..., :half], x_raw[..., half:] + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) + + +def test_worker_block_table_compaction(): + """Verify that shifting block table entries on the worker is exact and zero-copy.""" + # Simulate a request with 7 physical blocks: [B0, B1, Gap0, Gap1, B2, B3, B4] + sink_blocks = 2 + gap_blocks = 2 + num_blocks = 7 + initial_blocks = [101, 102, 201, 202, 301, 302, 303] + + table_np = np.zeros((4, 32), dtype=np.int32) + num_blocks_per_row = np.zeros(4, dtype=np.int32) + + row_idx = 1 + table_np[row_idx, :num_blocks] = initial_blocks + num_blocks_per_row[row_idx] = num_blocks + + # Worker-side compaction helper + total = int(num_blocks_per_row[row_idx]) + assert total == 7 + table_np[row_idx, sink_blocks : total - gap_blocks] = table_np[row_idx, sink_blocks + gap_blocks : total] + table_np[row_idx, total - gap_blocks : total] = 0 + num_blocks_per_row[row_idx] -= gap_blocks + + assert num_blocks_per_row[row_idx] == 5 + compacted = list(table_np[row_idx, :5]) + assert compacted == [101, 102, 301, 302, 303] + # Zeroed out tail + assert list(table_np[row_idx, 5:7]) == [0, 0] + + +def test_rotate_cached_keys_attention_equivalence(): + """Verify that rotated cached keys and attention scores match ground-truth forward RoPE.""" + inv_freq = _get_inv_freq() + delta = 16 + moved_from = 32 + sink_blocks = 1 # 16 tokens sink (block 0), 16 tokens gap (block 1), retained tail from pos 32 (block 2..) + plan = PositionReanchor(delta=delta, moved_from=moved_from, sink_blocks=sink_blocks) + + # Initial physical blocks: [0, 1, 2, 3]. Gap is block 1. + # Compacted block table after trim: [0, 2, 3] + compacted_blocks = [0, 2, 3] + num_tokens = 64 + positions = torch.arange(moved_from, num_tokens, dtype=torch.long) + + torch.manual_seed(42) + # Generate raw unrotated key features + raw_keys = torch.randn(num_tokens, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float32) + + # Populate cache by applying ground-truth forward RoPE at original positions: + # Block 0: 0..15, Block 1: 16..31 (gap), Block 2: 32..47, Block 3: 48..63 + k_pool = torch.zeros(4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float32) + for p in range(num_tokens): + b = p // BLOCK_SIZE + o = p % BLOCK_SIZE + k_pool[b, o] = _forward_rope(raw_keys[p : p + 1], p, inv_freq).squeeze(0) + + # Rotate retained tail in place + touched = rotate_cached_keys( + k_pool, + block_ids=compacted_blocks, + positions=positions, + plan=plan, + inv_freq=inv_freq, + ) + assert touched == len(positions) + + # Ground-truth comparison: + # In compacted table, token p is re-indexed to logical pos new_p = p - delta. + # Its physical slot is compacted_blocks[new_p // BLOCK_SIZE] at new_p % BLOCK_SIZE. + # The rotated cached key MUST equal computing forward RoPE directly at pos new_p! + for p in range(moved_from, num_tokens): + new_p = p - delta + phys_b = compacted_blocks[new_p // BLOCK_SIZE] + o = new_p % BLOCK_SIZE + rotated_k = k_pool[phys_b, o] + gt_k = _forward_rope(raw_keys[p : p + 1], new_p, inv_freq).squeeze(0) + assert torch.allclose(rotated_k, gt_k, atol=1e-5) + + # Attention score check with a query at step Q + q_raw = torch.randn(1, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float32) + q_pos = 80 + q = _forward_rope(q_raw, q_pos, inv_freq).squeeze(0) + + for p in range(moved_from, num_tokens): + new_p = p - delta + phys_b = compacted_blocks[new_p // BLOCK_SIZE] + o = new_p % BLOCK_SIZE + rotated_k = k_pool[phys_b, o] + score_rotated = (q * rotated_k).sum() + gt_k = _forward_rope(raw_keys[p : p + 1], new_p, inv_freq).squeeze(0) + score_gt = (q * gt_k).sum() + assert torch.allclose(score_rotated, score_gt, atol=1e-5) + + +def test_rotate_keys_precision_ground_truth_bf16(): + """Verify that rotate_keys in bfloat16 at large delta matches ground truth forward RoPE. + + In bfloat16, values in [1024, 2048] have ULP = 8. At delta = 6000 (standard Stage-0 + sliding window size), computing angles in bfloat16 introduces ~4-5.7 rad quantization + error, scrambling trigonometric values. Computing in float32 and casting back ensures + numerical fidelity (<0.05 max error vs >2.0 for buggy bfloat16). + """ + inv_freq = _get_inv_freq() + torch.manual_seed(42) + tokens, heads = 16, 8 + x_raw = torch.randn(tokens, heads, HEAD_DIM, dtype=torch.bfloat16) + + p_old = 7000 + delta = 6000 + p_new = p_old - delta + + k_old = _forward_rope(x_raw, p_old, inv_freq) + k_gt = _forward_rope(x_raw, p_new, inv_freq) + + # Fixed rotate_keys (fp32 trig): + k_rotated = rotate_keys(k_old, delta, inv_freq) + err = (k_rotated.float() - k_gt.float()).abs().max().item() + assert err < 0.05, f"Expected precision < 0.05, got {err}" + + # Verify that the buggy calculation (angle in bf16 before cos/sin) fails drastically: + half = HEAD_DIM // 2 + angle_buggy = (int(delta) * inv_freq).to(dtype=torch.bfloat16) + cos_buggy = torch.cos(angle_buggy).unsqueeze(0).unsqueeze(1) + sin_buggy = torch.sin(angle_buggy).unsqueeze(0).unsqueeze(1) + k1, k2 = k_old[..., :half], k_old[..., half:] + k_buggy = torch.cat([k1 * cos_buggy + k2 * sin_buggy, k2 * cos_buggy - k1 * sin_buggy], dim=-1) + err_buggy = (k_buggy.float() - k_gt.float()).abs().max().item() + assert err_buggy > 2.0, f"Buggy bf16 angle calculation should exhibit >2.0 error, got {err_buggy}" + + +def test_batched_concurrency_isolation(): + """Verify that in a batch of multiple concurrent requests: + - Request 0 triggers trim (delta=16) + - Request 1 is normal decode (untouched) + - Request 2 triggers trim with different delta=32 + Non-triggering requests are completely unaffected. + """ + inv_freq = _get_inv_freq() + num_blocks = 20 + torch.manual_seed(123) + k_pool = torch.randn(num_blocks, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float32) + k_pool_snapshot = k_pool.clone() + + # Request 0: blocks [0, 1, 2, 3], needs delta=16, moved_from=32, sink_blocks=1 + # Gap is block 1. Compacted table after trim is [0, 2, 3]. + plan_0 = PositionReanchor(delta=16, moved_from=32, sink_blocks=1) + req0_blocks_compacted = [0, 2, 3] + req0_positions = torch.arange(32, 64, dtype=torch.long) + + # Request 1: blocks [4, 5, 6, 7], normal decode, NO trim + req1_blocks = [4, 5, 6, 7] + + # Request 2: blocks [8, 9, 10, 11, 12], needs delta=32, moved_from=48, sink_blocks=1 + # Gap is blocks [9, 10]. Compacted table after trim is [8, 11, 12]. + plan_2 = PositionReanchor(delta=32, moved_from=48, sink_blocks=1) + req2_blocks_compacted = [8, 11, 12] + req2_positions = torch.arange(48, 80, dtype=torch.long) + + # Execute Re-RoPE for the batch + reanchor_batch = { + 0: (req0_blocks_compacted, req0_positions, plan_0), + 2: (req2_blocks_compacted, req2_positions, plan_2), + } + + for req_idx, (b_ids, pos, plan) in reanchor_batch.items(): + rotate_cached_keys( + k_pool, + block_ids=b_ids, + positions=pos, + plan=plan, + inv_freq=inv_freq, + ) + + # ASSERTION 1: Request 1's physical blocks [4, 5, 6, 7] are 100% UNTOUCHED + for b in req1_blocks: + assert torch.equal(k_pool[b], k_pool_snapshot[b]), f"Block {b} of Request 1 was corrupted!" + + # ASSERTION 2: Request 0's sink block [0] is UNTOUCHED + assert torch.equal(k_pool[0], k_pool_snapshot[0]), "Request 0 sink block was corrupted!" + + # ASSERTION 3: Request 0's tail blocks [2, 3] are rotated by delta=16 + for b in [2, 3]: + expected = rotate_keys(k_pool_snapshot[b], 16, inv_freq) + assert torch.allclose(k_pool[b], expected, atol=1e-6) + + # ASSERTION 4: Request 2's sink block [8] is UNTOUCHED + assert torch.equal(k_pool[8], k_pool_snapshot[8]), "Request 2 sink block was corrupted!" + + # ASSERTION 5: Request 2's tail blocks [11, 12] are rotated by delta=32 + for b in [11, 12]: + expected = rotate_keys(k_pool_snapshot[b], 32, inv_freq) + assert torch.allclose(k_pool[b], expected, atol=1e-6) + + +def test_scheduler_reanchor_planning_and_no_full_reprefill(): + """Verify that Scheduler triggers reanchor without replacing the session prompt.""" + geometry = duplex_window_geometry( + prefix_tokens=96, + window_tokens=6000, + block_size=16, + max_model_len=40960, + high_watermark_tokens=8000, + ) + + # Below high watermark (96 + 8000 = 8096): no reanchor + plan = plan_position_reanchor(geometry, computed_tokens=7900, pending_tokens=12) + assert plan is None + + # Crossing watermark: 8090 + 12 = 8102 > 8096 + # Target is prefix(96) + window(6000) = 6096 + target = geometry.prefix_tokens + geometry.window_tokens + plan = plan_position_reanchor(geometry, computed_tokens=8090, pending_tokens=12) + assert plan is not None + assert plan.delta % BLOCK_SIZE == 0 + assert plan.delta > 0 + # Retained sequence stays within target budget + retained = (8090 + 12) - plan.delta + assert target - BLOCK_SIZE < retained <= target + # Total retained equals sink_end + retained tail + assert (8090 + 12) - plan.delta == plan.sink_end + ((8090 + 12) - plan.moved_from) + + +def test_barge_in_abort_safety(): + """Verify that aborting a request during/after reanchor does not leak blocks.""" + # Simulate a mini block allocator + free_blocks = set(range(100)) + allocated = {} + + def alloc(req_id, count): + blocks = [free_blocks.pop() for _ in range(count)] + allocated[req_id] = blocks + return blocks + + def free_blocks_range(req_id, start, end): + blocks = allocated[req_id] + freed = blocks[start:end] + for b in freed: + free_blocks.add(b) + allocated[req_id] = blocks[:start] + blocks[end:] + + def abort_request(req_id): + blocks = allocated.pop(req_id, []) + for b in blocks: + free_blocks.add(b) + + # Req A allocates 8 blocks + alloc("req-a", 8) + assert len(free_blocks) == 92 + + # Trim: free gap blocks [2:4] (2 blocks) + free_blocks_range("req-a", 2, 4) + assert len(free_blocks) == 94 + assert len(allocated["req-a"]) == 6 + + # User barges in -> abort! + abort_request("req-a") + assert len(free_blocks) == 100 # ALL blocks successfully returned, zero leak! + + +def test_scheduler_worker_end_to_end_state_agreement(): + """Verify production update path: append -> scheduler compaction -> worker state update -> agreement.""" + inv_freq = _get_inv_freq() + num_blocks = 10 + k_pool = torch.randn(num_blocks, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float32) + k_pool_orig = k_pool.clone() + + # Initial session prompt: 64 tokens, 4 blocks [0, 1, 2, 3] + # Prefix: 16 tokens (block 0), gap: 16 tokens (block 1, delta=16), tail: 32 tokens (blocks 2, 3) + prompt_ids = list(range(100, 164)) + session = SimpleNamespace( + request_id="req-1", + prompt_token_ids=list(prompt_ids), + _all_token_ids=list(prompt_ids), + num_prompt_tokens=64, + num_computed_tokens=64, + ) + + class _MockDuplexManager(MiniCPMO45DuplexWindowManager): + def __init__(self): + self.enable_caching = False + self.blocks = [0, 1, 2, 3] + self.req_to_blocks = {"req-1": self.blocks} + self.block_size = BLOCK_SIZE + self._null_block = -1 + + def compact_request_blocks(self, request_id: str, *, sink_blocks: int, num_blocks: int) -> int: + del self.blocks[sink_blocks : sink_blocks + num_blocks] + return num_blocks * BLOCK_SIZE + + def reanchor_block_table(self, request_id: str, plan: PositionReanchor) -> int: + return self.compact_request_blocks( + request_id, + sink_blocks=plan.sink_blocks, + num_blocks=plan.delta // BLOCK_SIZE, + ) + + duplex_mgr = _MockDuplexManager() + scheduler = SimpleNamespace( + cache_config=SimpleNamespace(block_size=BLOCK_SIZE), + model_config=SimpleNamespace(max_model_len=40960), + kv_cache_manager=SimpleNamespace(coordinator=SimpleNamespace(single_type_managers=[duplex_mgr])), + ) + + # 1. Update arrives: 1 token append, triggering window trim + update = SimpleNamespace( + prompt_token_ids=[999], + model_intermediate_buffer={ + "duplex": { + "data_plane": True, + "runtime_config": { + "duplex_window_prefix_tokens": 16, + "duplex_window_config": { + "sliding_window_mode": "basic", + "basic_window_high_tokens": 64, # Total watermark trigger at 64 tokens (< 64+1=65) + "basic_window_low_tokens": 64, # Target total is 64 tokens (drop_needed=1 -> delta=16) + }, + }, + } + }, + ) + + plan = MiniCPMO45DuplexSchedulerHelper.maybe_reanchor_session(scheduler, session, update) + assert plan is not None + assert plan.delta == 16 + assert plan.moved_from == 32 + assert plan.sink_blocks == 1 + + # Scheduler compacted session state: + # 64 tokens -> trimmed middle gap [16:32] -> 48 tokens surviving + assert len(session.prompt_token_ids) == 48 + assert len(session._all_token_ids) == 48 + assert session.prompt_token_ids == prompt_ids[:16] + prompt_ids[32:] + assert session._all_token_ids == prompt_ids[:16] + prompt_ids[32:] + assert session.num_computed_tokens == 48 + assert session.num_prompt_tokens == 48 + assert duplex_mgr.blocks == [0, 2, 3] + + # Subsequent append from upstream _update_request_as_session adds the 1 pending token: + session.prompt_token_ids.extend(update.prompt_token_ids) + session._all_token_ids.extend(update.prompt_token_ids) + session.num_prompt_tokens = len(session.prompt_token_ids) + # Both token histories are now consistent with length 49! + assert len(session.prompt_token_ids) == 49 + assert len(session._all_token_ids) == 49 + + # 2. Worker state update via production path: + # Scheduler provides post-compaction block IDs [0, 2, 3] and computed count 48 + table_np = np.zeros((2, 16), dtype=np.int32) + table_np[0, :3] = [0, 2, 3] + num_blocks_per_row = np.array([3, 0], dtype=np.int32) + + class _MockBlockTable: + def __init__(self): + self.block_table = SimpleNamespace(np=table_np) + self.num_blocks_per_row = num_blocks_per_row + + class _MockRunner: + def __init__(self): + self.device = torch.device("cpu") + self.cache_config = SimpleNamespace(block_size=BLOCK_SIZE) + self.model_config = SimpleNamespace( + get_head_size=lambda: HEAD_DIM, + hf_config=SimpleNamespace(rope_theta=1000000.0), + ) + self._duplex_inv_freq = inv_freq + self.kv_caches = [k_pool] + self.requests = { + "req-1": SimpleNamespace( + block_ids=([0, 2, 3],), # Production grouped tuple layout + num_computed_tokens=48, + mrope_positions=None, + ) + } + self.input_batch = SimpleNamespace( + num_reqs=1, + req_ids=["req-1"], + block_table=_MockBlockTable(), + num_computed_tokens_cpu=np.array([48], dtype=np.int32), + ) + self.model_intermediate_buffer = { + "req-1": update.model_intermediate_buffer, + } + + runner = _MockRunner() + # Execute worker helper + MiniCPMO45DuplexWorkerHelper.maybe_apply_reanchor(runner) + + # 3. Assertions: Both sides strictly agree on computed counts and block IDs! + assert runner.input_batch.num_computed_tokens_cpu[0] == 48, "Worker must NOT decrement computed tokens twice!" + assert session.num_computed_tokens == runner.input_batch.num_computed_tokens_cpu[0] == 48 + bt = runner.input_batch.block_table + assert bt.num_blocks_per_row[0] == 3, "Worker must NOT delete blocks twice!" + assert list(bt.block_table.np[0, :3]) == [0, 2, 3] + assert duplex_mgr.blocks == list(bt.block_table.np[0, :3]) == [0, 2, 3] + + # Verify KV cache rotation: sink block 0 untouched, blocks 2 and 3 rotated + assert torch.equal(k_pool[0], k_pool_orig[0]) + for b in [2, 3]: + expected = rotate_keys(k_pool_orig[b], 16, inv_freq) + assert torch.allclose(k_pool[b], expected, atol=1e-6) + + +def test_differing_prefix_lengths_dynamic_sink(): + """Verify that differing instruction / reference audio prefix lengths work without spec mismatch.""" + # Prefix length 112 tokens = 7 blocks (differs from standard 96 tokens = 6 blocks) + prefix_tokens = 112 + sink_blocks = 7 + blocks = [100 + i for i in range(12)] + + class _MockSpec: + chunk_size = BLOCK_SIZE + sink_chunks = 6 # Static spec default is 6, while session has 7 + + manager = MiniCPMO45DuplexWindowManager.__new__(MiniCPMO45DuplexWindowManager) + manager.block_size = BLOCK_SIZE + manager.enable_caching = False + manager.kv_cache_spec = _MockSpec() + manager._null_block = -1 + manager.req_to_blocks = {"req-dyn": list(blocks)} + manager.num_cached_block = {} + + def mock_remove(req_id, start, end): + b = manager.req_to_blocks[req_id] + for i in range(start, end): + b[i] = manager._null_block + + manager._remove_blocks_in_range = mock_remove + + plan = PositionReanchor(delta=16, moved_from=prefix_tokens + 16, sink_blocks=sink_blocks) + # Exercises real inherited compact_block_table / compact_request_blocks path without mock lambda + freed = manager.reanchor_block_table("req-dyn", plan) + assert freed == 16 + # Verified: sink blocks 0..6 (7 blocks) preserved, gap block 7 freed! + assert len(manager.req_to_blocks["req-dyn"]) == 11 + assert manager.req_to_blocks["req-dyn"][:7] == blocks[:7] + assert manager.req_to_blocks["req-dyn"][7:] == blocks[8:] + + +def test_context_mode_explicit_fallback(): + """Verify context mode explicitly returns None from zero-copy reanchor policy to route to official fallback.""" + session = SimpleNamespace(num_computed_tokens=600) + update = SimpleNamespace( + prompt_token_ids=[1] * 10, + model_intermediate_buffer={ + "duplex": { + "data_plane": True, + "runtime_config": { + "duplex_window_prefix_tokens": 96, + "duplex_window_config": { + "sliding_window_mode": "context", + "context_max_units": 24, + "context_previous_max_tokens": 16, + }, + }, + } + }, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import MiniCPMO45DuplexWindowPolicy + + # Context mode must return None from zero-copy reanchor policy to route to official fallback + assert MiniCPMO45DuplexWindowPolicy.plan_reanchor(session, update, block_size=16, max_model_len=40960) is None + + +def test_vllm_030_packed_flash_attn_kv_layout_and_storage_alias(): + """Verify vLLM 0.30 FlashAttention packed layout (B, H, N, 2*D) rotates in-place and preserves V cache.""" + # Real H20-3e GPU observed layout: 8 heads, block_size 16, head_dim 128 (2*head_dim = 256) + # Transposed strides: [32768, 256, 2048, 1] + base = torch.randn(10, 16, 8, 256, dtype=torch.bfloat16) + kv_cache = base.transpose(1, 2) + assert kv_cache.shape == (10, 8, 16, 256) + assert kv_cache.stride() == (32768, 256, 2048, 1) + + storage_ptr = kv_cache.data_ptr() + v_cache_copy = kv_cache[..., 128:].clone() + inv_freq = torch.randn(64, dtype=torch.float32) + + plan = PositionReanchor(delta=16, moved_from=16, sink_blocks=1) + block_ids = [0, 2, 5] + positions = torch.tensor([16, 17, 30, 32, 47], dtype=torch.long) + + touched = rotate_cached_keys( + kv_cache, + block_ids=block_ids, + positions=positions, + plan=plan, + inv_freq=inv_freq, + block_size=16, + ) + assert touched == 5 + # Must write in-place into the existing storage: no copies! + assert kv_cache.data_ptr() == storage_ptr + # V cache in [..., 128:] must be completely untouched! + assert torch.equal(kv_cache[..., 128:], v_cache_copy) + # Untouched blocks must remain unchanged + assert torch.equal(kv_cache[1, ..., :128], base.transpose(1, 2)[1, ..., :128]) + + +def test_worker_history_bounded_across_long_session_reanchors(): + """Verify worker-held state.window_units is pruned when stage0_reanchor triggers, staying bounded.""" + units = [SimpleNamespace(token_ids=list(range(i * 12, (i + 1) * 12))) for i in range(50)] + session_state = SimpleNamespace(window_units=units) + + class _MockHelper: + def __init__(self): + self.sessions = {"sess-1": session_state} + + def _evict_window_units_for_reanchor(self, state, reanchor): + delta = int(reanchor.get("delta", 0) or 0) + if delta <= 0 or not state.window_units: + return + dropped = 0 + idx = 0 + while idx < len(state.window_units): + unit_len = len(state.window_units[idx].token_ids) + if dropped + unit_len <= delta: + dropped += unit_len + idx += 1 + else: + break + if idx > 0: + del state.window_units[:idx] + + mock_helper = _MockHelper() + mock_model = SimpleNamespace( + _minicpmo45_duplex_data_plane_helper=mock_helper, + _minicpmo45_duplex_request_sessions={"req-1": "sess-1"}, + ) + runner = SimpleNamespace( + model=mock_model, + input_batch=SimpleNamespace( + req_ids=["req-1"], + num_reqs=1, + num_computed_tokens_cpu=[600], + ), + model_intermediate_buffer={ + "req-1": { + "duplex": { + "stage0_reanchor": { + "delta": 240, # 20 units * 12 tokens + "moved_from": 336, + "sink_blocks": 6, + } + } + } + }, + requests={}, + device=torch.device("cpu"), + cache_config=SimpleNamespace(block_size=16), + kv_caches=[], + ) + + MiniCPMO45DuplexWorkerHelper.maybe_apply_reanchor(runner) + # Pruned from 50 units down to 30 units: + assert len(session_state.window_units) == 30 + assert session_state.window_units[0].token_ids[0] == 240 + + +def test_grouped_block_table_worker_rotation(): + """Verify worker unpacks production grouped block tables (tuple of lists) across multiple layers and groups.""" + inv_freq = _get_inv_freq() + k_pool_g0 = torch.randn(10, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float32) + k_pool_g1 = torch.randn(20, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float32) + k_pool_g0_orig = k_pool_g0.clone() + k_pool_g1_orig = k_pool_g1.clone() + + class _MockRunnerGrouped: + def __init__(self): + self.device = torch.device("cpu") + self.cache_config = SimpleNamespace(block_size=BLOCK_SIZE) + self.model_config = SimpleNamespace( + get_head_size=lambda: HEAD_DIM, + hf_config=SimpleNamespace(rope_theta=1000000.0), + ) + self._duplex_inv_freq = inv_freq + # Layer 0 belongs to group 0, Layer 1 belongs to group 1 + self.kv_caches = [k_pool_g0, k_pool_g1] + self.kv_cache_group_ids = [0, 1] + self.requests = { + "req-grp": SimpleNamespace( + # Tuple of lists: Group 0 has [0, 2, 3], Group 1 has [5, 7, 8] + block_ids=([0, 2, 3], [5, 7, 8]), + num_computed_tokens=48, + mrope_positions=None, + ) + } + self.input_batch = SimpleNamespace( + num_reqs=1, + req_ids=["req-grp"], + num_computed_tokens_cpu=np.array([48], dtype=np.int32), + ) + self.model_intermediate_buffer = { + "req-grp": { + "duplex": { + "stage0_reanchor": { + "delta": 16, + "moved_from": 32, + "sink_blocks": 1, + "old_computed_tokens": 64, + } + } + } + } + + runner = _MockRunnerGrouped() + # Worker must resolve each layer's KV group without indexing errors + MiniCPMO45DuplexWorkerHelper.maybe_apply_reanchor(runner) + + # Group 0: block 0 is sink (untouched), blocks 2 and 3 rotated + assert torch.equal(k_pool_g0[0], k_pool_g0_orig[0]) + for b in [2, 3]: + expected = rotate_keys(k_pool_g0_orig[b], 16, inv_freq) + assert torch.allclose(k_pool_g0[b], expected, atol=1e-6) + + # Group 1: block 5 is sink (untouched), blocks 7 and 8 rotated + assert torch.equal(k_pool_g1[5], k_pool_g1_orig[5]) + for b in [7, 8]: + expected = rotate_keys(k_pool_g1_orig[b], 16, inv_freq) + assert torch.allclose(k_pool_g1[b], expected, atol=1e-6) + + +def test_rejection_before_state_mutation(): + """Verify rejection before any state mutation when compaction plan is unsupported or fails.""" + session = SimpleNamespace( + request_id="req-unsupported", + prompt_token_ids=[1, 2, 3, 4], + _all_token_ids=[1, 2, 3, 4], + num_prompt_tokens=4, + num_computed_tokens=4, + ) + + class _MockFailingDuplexManager: + def compact_request_blocks(self, request_id: str, *, sink_blocks: int, num_blocks: int) -> int: + return 0 # Rejection: cannot compact + + scheduler = SimpleNamespace( + cache_config=SimpleNamespace(block_size=BLOCK_SIZE), + model_config=SimpleNamespace(max_model_len=40960), + kv_cache_manager=SimpleNamespace( + coordinator=SimpleNamespace(single_type_managers=[_MockFailingDuplexManager()]) + ), + ) + update = SimpleNamespace( + prompt_token_ids=[5], + model_intermediate_buffer={ + "duplex": { + "data_plane": True, + "runtime_config": { + "duplex_window_prefix_tokens": 16, + "duplex_window_config": { + "sliding_window_mode": "basic", + "basic_window_high_tokens": 20, + "basic_window_low_tokens": 16, + }, + }, + } + }, + ) + + result = MiniCPMO45DuplexSchedulerHelper.maybe_reanchor_session(scheduler, session, update) + assert result is None + # Verify zero state mutation: token sequences and counters are untouched! + assert session.prompt_token_ids == [1, 2, 3, 4] + assert session._all_token_ids == [1, 2, 3, 4] + assert session.num_prompt_tokens == 4 + assert session.num_computed_tokens == 4 + assert "stage0_reanchor" not in update.model_intermediate_buffer["duplex"] + + +def test_non_duplex_regression(): + """Verify that non-duplex requests or sliding_window_mode='off' are completely bypassed.""" + session = SimpleNamespace( + request_id="req-non-duplex", + prompt_token_ids=[1, 2, 3], + _all_token_ids=[1, 2, 3], + num_prompt_tokens=3, + num_computed_tokens=3, + ) + # Case 1: No duplex in buffer + update1 = SimpleNamespace(prompt_token_ids=[4], model_intermediate_buffer={}) + scheduler = SimpleNamespace() + assert MiniCPMO45DuplexSchedulerHelper.maybe_reanchor_session(scheduler, session, update1) is None + assert session.num_computed_tokens == 3 + + # Case 2: sliding_window_mode = 'off' + update2 = SimpleNamespace( + prompt_token_ids=[4], + model_intermediate_buffer={ + "duplex": { + "data_plane": True, + "runtime_config": { + "duplex_window_config": {"sliding_window_mode": "off"}, + }, + } + }, + ) + assert MiniCPMO45DuplexSchedulerHelper.maybe_reanchor_session(scheduler, session, update2) is None + assert session.num_computed_tokens == 3 + + +def test_assert_uniform_position_shift(): + """Verify position validation: uniform shifts pass, non-uniform MRoPE raises.""" + # 1. Uniform 2D positions (e.g. streaming audio/text where rows advance identically) + pos_uniform = torch.arange(100).unsqueeze(0).repeat(3, 1) + assert_uniform_position_shift(pos_uniform, moved_from=32) + + # 2. Vision tokens in sink only (e.g. 0..20 < 32), uniform in retained tail + pos_sink_only_vision = pos_uniform.clone() + pos_sink_only_vision[1, :20] += 5 + pos_sink_only_vision[2, :20] += 10 + assert_uniform_position_shift(pos_sink_only_vision, moved_from=32) + + # 3. Vision tokens in retained tail (>= 32) must be rejected + pos_tail_vision = pos_uniform.clone() + pos_tail_vision[1, 40:] += 5 + with pytest.raises(RuntimeError, match="duplex re-anchor needs one position row across the retained tail"): + assert_uniform_position_shift(pos_tail_vision, moved_from=32) + + # 4. 1D positions trivially uniform + pos_1d = torch.arange(100) + assert_uniform_position_shift(pos_1d, moved_from=32) + + # 5. Invalid rank + with pytest.raises(ValueError, match="expected a \\(rows, tokens\\) position tensor"): + assert_uniform_position_shift(pos_1d.unsqueeze(0).unsqueeze(0), moved_from=32) + + +def test_session_mode_location_on_model_config(): + """Verify session_mode is read from vllm_config.model_config, not vllm_config.""" + # Production layout: session_mode is in model_config + valid_cfg = SimpleNamespace( + model_config=SimpleNamespace(model_stage="llm", session_mode="duplex", max_model_len=40960), + cache_config=SimpleNamespace(block_size=16), + ) + assert getattr(getattr(valid_cfg, "model_config", None), "session_mode", None) == "duplex" + + # Buggy layout: session_mode on vllm_config directly was never populated in vllm-omni + buggy_cfg = SimpleNamespace( + session_mode="duplex", + model_config=SimpleNamespace(model_stage="llm", max_model_len=40960), + cache_config=SimpleNamespace(block_size=16), + ) + assert getattr(getattr(buggy_cfg, "model_config", None), "session_mode", None) != "duplex" + + +def test_scheduler_replace_streaming_prompt_bypasses_reanchor(): + """Verify that when replace_streaming_prompt is True, re-anchoring is bypassed.""" + reanchor_called = [] + + class MockScheduler: + def _release_replaced_streaming_prompt_cache(self, session): + session.released = True + + def _replace_streaming_session(self, session, update): + session.replaced = True + + def _maybe_reanchor_minicpmo45_stage0_window(self, session, update): + reanchor_called.append(True) + + def _update_request_as_session(self, session, update): + stage_id = 0 + update_infos = [{"meta": {"replace_streaming_prompt": True}}] + + replace_streaming_prompt = any( + isinstance(info, dict) + and isinstance(info.get("meta"), dict) + and info["meta"].get("replace_streaming_prompt") is True + for info in update_infos + ) + if replace_streaming_prompt: + self._release_replaced_streaming_prompt_cache(session) + self._replace_streaming_session(session, update) + return + + if stage_id == 0: + self._maybe_reanchor_minicpmo45_stage0_window(session, update) + + sched = MockScheduler() + session = SimpleNamespace(released=False, replaced=False) + sched._update_request_as_session(session, None) + + assert session.released is True + assert session.replaced is True + assert len(reanchor_called) == 0, "Re-anchoring must not be called when replacing streaming prompt!" + + +def test_slot_mapping_device_compatibility(): + """Verify compute_slot_mapping creates table on positions device and does not error on CUDA/device tensor.""" + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import compute_slot_mapping + + device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") + pos = torch.tensor([0, 15, 16, 31, 32], dtype=torch.long, device=device) + block_ids = [10, 20, 30] + slots = compute_slot_mapping(block_ids, pos, block_size=16) + assert slots.device == device + assert slots.tolist() == [10 * 16 + 0, 10 * 16 + 15, 20 * 16 + 0, 20 * 16 + 15, 30 * 16 + 0] + + +def test_reanchor_ordinary_append_fallback_rebuild_sequence(): + """Reviewer @amy-why-3459 Issue 1: Sequence test: reanchor -> ordinary append -> fallback rebuild. + + Verifies that: + 1. Re-anchor updates scheduler window units and coordinates (_minicpmo45_window_open_start) atomically. + 2. Next ordinary append computes positive unit_len and records unit correctly without early return. + 3. Subsequent fallback rebuild (_window_replacement_parts) succeeds without rebuild-length mismatch. + """ + block_size = 16 + prefix_tokens = 96 + high_watermark = 8000 + low_watermark = 6000 + + class _MockDuplexManager(MiniCPMO45DuplexWindowManager): + def __init__(self): + self.block_size = block_size + self.blocks = list(range(600)) + self.req_to_blocks = {"req-seq": self.blocks} + + def compact_request_blocks(self, request_id: str, *, sink_blocks: int, num_blocks: int) -> int: + del self.blocks[sink_blocks : sink_blocks + num_blocks] + return num_blocks * block_size + + duplex_mgr = _MockDuplexManager() + scheduler = SimpleNamespace( + cache_config=SimpleNamespace(block_size=block_size), + model_config=SimpleNamespace(max_model_len=40960), + kv_cache_manager=SimpleNamespace(coordinator=SimpleNamespace(single_type_managers=[duplex_mgr])), + ) + + # Initial state before reanchor: + # prefix (96) + 80 units * 100 = 8096; boundary after turn 80 is 8096 + 2 = 8098. + open_start_initial = prefix_tokens + 80 * 100 + 2 + initial_computed = 8192 + initial_prompt = list(range(initial_computed)) + session = SimpleNamespace( + request_id="req-seq", + prompt_token_ids=list(initial_prompt), + _all_token_ids=list(initial_prompt), + num_prompt_tokens=initial_computed, + num_computed_tokens=initial_computed, + _minicpmo45_window_open_start=open_start_initial, + _minicpmo45_window_units=[{"length": 100, "generated_token_ids": [1] * 10} for _ in range(80)], + ) + + worker_state = _MiniCPMO45Stage0SessionState( + session_id="req-seq", + window_enabled=True, + context_embeds=[torch.zeros(1, 10)] * prefix_tokens, + context_token_ids=list(range(prefix_tokens)), + window_units=[ + _MiniCPMO45WindowUnit( + embeds=[torch.zeros(1, 10)] * 100, + token_ids=list(range(100)), + ) + for _ in range(80) + ], + ) + + # --- Step 1: Trigger Re-anchor Append --- + update_reanchor = SimpleNamespace( + prompt_token_ids=[1000] * 64, + model_intermediate_buffer={ + "duplex": { + "data_plane": True, + "seq": 2, + "runtime_config": { + "duplex_first_append_context_tokens": prefix_tokens, + "duplex_window_prefix_tokens": prefix_tokens, + "duplex_window_config": { + "sliding_window_mode": "basic", + "basic_window_high_tokens": high_watermark, + "basic_window_low_tokens": low_watermark, + }, + }, + } + }, + ) + + plan = MiniCPMO45DuplexSchedulerHelper.apply_session_window( + scheduler, + session, + update_reanchor, + segment_output_ids=[42], + completed_terminator=None, + ) + assert plan is not None + delta = plan.delta + assert delta > 0 + assert delta % block_size == 0 + + # Verify atomic update of scheduler history and rebasing of coordinates + assert session.num_computed_tokens == initial_computed - delta + assert session.num_prompt_tokens == initial_computed - delta + assert session._minicpmo45_window_open_start == (initial_computed + len([42]) + 2) - delta + + # Worker evicts corresponding window units + stage0_reanchor = update_reanchor.model_intermediate_buffer["duplex"]["stage0_reanchor"] + worker_helper = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + worker_helper._embed_token = lambda tok: torch.zeros(1, 10) + worker_helper._evict_window_units_for_reanchor(worker_state, stage0_reanchor) + + # Finalize the completed turn on the worker + stage0_window = update_reanchor.model_intermediate_buffer["duplex"]["stage0_window"] + unclosed_prompt_len = (initial_computed + 1 + 2 - open_start_initial) - 3 + worker_state.pending_window_unit = _MiniCPMO45WindowUnit( + embeds=[torch.zeros(1, 10)] * unclosed_prompt_len, + token_ids=list(range(unclosed_prompt_len)), + ) + worker_state.pending_window_generated_tokens = stage0_window["completed_token_ids"] + worker_helper._finalize_window_unit(worker_state, closure_token_ids=[999, 998]) + + # Simulate scheduler extending prompt with the update's tokens + session.prompt_token_ids.extend(update_reanchor.prompt_token_ids) + session.num_prompt_tokens = len(session.prompt_token_ids) + + # --- Step 2: Next Ordinary Append --- + # Now computed is around 6032 + 64 = 6096, well below 8000 + session.num_computed_tokens = session.num_prompt_tokens + update_ordinary = SimpleNamespace( + prompt_token_ids=[2000] * 64, + model_intermediate_buffer={ + "duplex": { + "data_plane": True, + "seq": 3, + "runtime_config": { + "duplex_first_append_context_tokens": prefix_tokens, + "duplex_window_prefix_tokens": prefix_tokens, + "duplex_window_config": { + "sliding_window_mode": "basic", + "basic_window_high_tokens": high_watermark, + "basic_window_low_tokens": low_watermark, + }, + }, + } + }, + ) + + # Zero-copy policy returns None (inside window) + plan_ordinary = MiniCPMO45DuplexSchedulerHelper.apply_session_window( + scheduler, + session, + update_ordinary, + segment_output_ids=[43], + ) + assert plan_ordinary is None + + # Fallback / ordinary prepare helper runs + replaced = prepare_minicpmo45_stage0_window( + session, + update_ordinary, + segment_output_ids=[43], + ) + # Must NOT replace, and MUST NOT fail with negative unit_len! + assert replaced is False + assert session._minicpmo45_window_open_start > 0 + + # Worker finalizes turn from Step 1 on ordinary append + worker_state.pending_window_unit = _MiniCPMO45WindowUnit( + embeds=[torch.zeros(1, 10)] * (64 - 3), + token_ids=[1000] * (64 - 3), + ) + worker_state.pending_window_generated_tokens = [43] + worker_helper._finalize_window_unit(worker_state, closure_token_ids=[999, 998]) + + session.prompt_token_ids.extend(update_ordinary.prompt_token_ids) + session.num_prompt_tokens = len(session.prompt_token_ids) + + # --- Step 3: Trigger Fallback Rebuild --- + # If a watermark fires requesting replace=True, verify worker rebuilds successfully without mismatch! + update_fallback = SimpleNamespace( + prompt_token_ids=[3000] * 64, + model_intermediate_buffer={ + "duplex": { + "data_plane": True, + "seq": 4, + "runtime_config": { + "duplex_first_append_context_tokens": prefix_tokens, + "duplex_window_prefix_tokens": prefix_tokens, + "duplex_window_config": { + "sliding_window_mode": "basic", + "basic_window_high_tokens": 100, # Force watermark to fire + "basic_window_low_tokens": 80, + }, + }, + } + }, + ) + replaced_fb = prepare_minicpmo45_stage0_window( + session, + update_fallback, + segment_output_ids=[44], + ) + assert replaced_fb is True + fb_stage0_window = update_fallback.model_intermediate_buffer["duplex"]["stage0_window"] + assert fb_stage0_window["replace"] is True + + # Worker finalizes turn from Step 2 when replacement append arrives + worker_state.pending_window_unit = _MiniCPMO45WindowUnit( + embeds=[torch.zeros(1, 10)] * (64 - 3), + token_ids=[2000] * (64 - 3), + ) + worker_state.pending_window_generated_tokens = [44] + worker_helper._finalize_window_unit(worker_state, closure_token_ids=[999, 998]) + + # Worker sets up pending unit for Step 3 append + worker_state.pending_window_unit = _MiniCPMO45WindowUnit( + embeds=[torch.zeros(1, 10)] * 64, + token_ids=[3000] * 64, + ) + + # Worker rebuilds from history: must succeed without RuntimeError! + embeds, token_ids = worker_helper._window_replacement_parts(worker_state, fb_stage0_window) + assert len(token_ids) == fb_stage0_window["replacement_prompt_len"] + assert len(embeds) == len(token_ids) + + +def test_worker_history_partial_unit_slice_and_exact_range(): + """Reviewer @amy-why-3459 Issue 2: Preserve exact retained token range when pruning worker history. + + With unit token ranges [0..31] and [32..63], delta=16: + KV dropped: 16 tokens -> retains [16..63] + Worker history must also retain [16..63], starting at token 16 (NOT token 32). + """ + unit0 = _MiniCPMO45WindowUnit( + embeds=[torch.full((1, 4), float(i)) for i in range(32)], + token_ids=list(range(32)), + ) + unit1 = _MiniCPMO45WindowUnit( + embeds=[torch.full((1, 4), float(i)) for i in range(32, 64)], + token_ids=list(range(32, 64)), + ) + + state = _MiniCPMO45Stage0SessionState( + session_id="req-slice", + window_enabled=True, + window_units=[unit0, unit1], + ) + + helper = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + helper._evict_window_units_for_reanchor(state, {"delta": 16}) + + # Exactly 16 tokens dropped from history! + remaining_tokens = [tok for u in state.window_units for tok in u.token_ids] + assert remaining_tokens == list(range(16, 64)) + assert remaining_tokens[0] == 16, "Worker history must retain [16..63], starting at token 16!" + assert len(remaining_tokens) == 48 + + # Verify embeddings match the retained token range [16..64] + remaining_embeds = [emb for u in state.window_units for emb in u.embeds] + assert len(remaining_embeds) == 48 + assert remaining_embeds[0][0, 0].item() == 16.0 + + +def test_watermark_interval_high_to_high_plus_prefix(): + """Reviewer @amy-why-3459 Issue 3: Dispatch around the [high, high + prefix] interval. + + With prefix=96, high=8000, computed=7970, pending=64 -> projected=8034. + The zero-copy policy must trigger (return a plan), avoiding repeated expensive re-prefill. + """ + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import ( + MiniCPMO45DuplexWindowPolicy, + ) + + prefix_tokens = 96 + high_watermark = 8000 + low_watermark = 6000 + + session = SimpleNamespace( + request_id="req-watermark", + num_computed_tokens=7970, + num_prompt_tokens=7970, + ) + update = SimpleNamespace( + prompt_token_ids=[1] * 64, + model_intermediate_buffer={ + "duplex": { + "data_plane": True, + "runtime_config": { + "duplex_window_prefix_tokens": prefix_tokens, + "duplex_window_config": { + "sliding_window_mode": "basic", + "basic_window_high_tokens": high_watermark, + "basic_window_low_tokens": low_watermark, + }, + }, + } + }, + ) + + # projected = 7970 + 64 = 8034. + # In the [8000, 8096] interval, plan_reanchor MUST trigger! + plan = MiniCPMO45DuplexWindowPolicy.plan_reanchor( + session, + update, + block_size=16, + max_model_len=40960, + ) + assert plan is not None, "Zero-copy policy must trigger when total length passes basic_window_high_tokens!" + assert plan.delta > 0 + # Post-trim length must be bounded by low_watermark + assert 8034 - plan.delta <= low_watermark + + +def test_validate_duplex_window_install_block_sizes(): + """Verify validate_duplex_window_install supports arbitrary block sizes (including NPU 128).""" + # Standard CUDA block_size=16 + cache_config_cuda = SimpleNamespace(enable_prefix_caching=False, block_size=16) + model_config = SimpleNamespace(max_model_len=40960) + geometry_16 = duplex_window_geometry( + prefix_tokens=96, + window_tokens=6000, + block_size=16, + max_model_len=40960, + high_watermark_tokens=8000, + ) + validate_duplex_window_install(cache_config_cuda, model_config, geometry_16) + + # Ascend NPU block_size=128 + cache_config_npu = SimpleNamespace(enable_prefix_caching=False, block_size=128) + geometry_128 = duplex_window_geometry( + prefix_tokens=96, + window_tokens=6000, + block_size=128, + max_model_len=40960, + high_watermark_tokens=8000, + ) + validate_duplex_window_install(cache_config_npu, model_config, geometry_128) + + # Rejects mismatched block size between geometry and cache_config + with pytest.raises(ValueError, match="does not match"): + validate_duplex_window_install(cache_config_npu, model_config, geometry_16) + + # Rejects prefix caching + cache_config_pc = SimpleNamespace(enable_prefix_caching=True, block_size=128) + with pytest.raises(ValueError, match="enable_prefix_caching=False"): + validate_duplex_window_install(cache_config_pc, model_config, geometry_128) + + # Rejects when window does not fit max_model_len + model_config_small = SimpleNamespace(max_model_len=2048) + with pytest.raises(ValueError, match="does not fit max_model_len"): + validate_duplex_window_install(cache_config_npu, model_config_small, geometry_128) + + +def test_exactly_once_reanchor_across_metadata_refreshes(): + """Verify that stage0_reanchor executes exactly once even across runner metadata refreshes. + + 1. _update_states() pops stage0_reanchor and rotates KV caches. + 2. scheduled_new_reqs is sanitized so subsequent _update_additional_information() + does not re-inject the reanchor command. + 3. Even if re-injected with the same reanchor_id, rotation and history eviction + are skipped (idempotent). + """ + head_dim = 16 + inv_freq = torch.tensor([1.0 / (10000.0 ** (2 * i / head_dim)) for i in range(head_dim // 2)]) + block_size = 16 + k_pool = torch.randn(8, 1, block_size, head_dim) + + evicted_reanchors = [] + + class _MockHelper: + def __init__(self): + self.sessions = { + "sess-1": SimpleNamespace( + window_units=[ + SimpleNamespace( + token_ids=list(range(i * 16, (i + 1) * 16)), + embeds=[torch.zeros(1, 10)] * 16, + ) + for i in range(10) + ] + ) + } + + def _evict_window_units_for_reanchor(self, state, reanchor): + evicted_reanchors.append(reanchor) + + mock_helper = _MockHelper() + mock_model = SimpleNamespace( + _minicpmo45_duplex_data_plane_helper=mock_helper, + _minicpmo45_duplex_request_sessions={"req-1": "sess-1"}, + ) + + reanchor_dict = { + "reanchor_id": "req-1-r1-48-16", + "delta": 16, + "moved_from": 48, + "sink_blocks": 1, + "sink_end": 16, + "old_computed_tokens": 64, + } + + new_req = SimpleNamespace( + req_id="req-1", + model_intermediate_buffer={ + "duplex": { + "session_id": "sess-1", + "stage0_reanchor": dict(reanchor_dict), + } + }, + ) + scheduler_output = SimpleNamespace( + scheduled_new_reqs=[new_req], + ) + + runner = SimpleNamespace( + model=mock_model, + input_batch=SimpleNamespace( + num_reqs=1, + req_ids=["req-1"], + block_table=SimpleNamespace( + num_blocks_per_row=np.array([4]), + block_table=SimpleNamespace(np=np.array([[0, 1, 2, 3]])), + ), + num_computed_tokens_cpu=np.array([48], dtype=np.int32), + ), + model_intermediate_buffer={ + "req-1": { + "duplex": { + "session_id": "sess-1", + "stage0_reanchor": dict(reanchor_dict), + } + } + }, + requests={ + "req-1": SimpleNamespace( + block_ids=[0, 1, 2, 3], + num_computed_tokens=48, + mrope_positions=None, + ) + }, + device=torch.device("cpu"), + cache_config=SimpleNamespace(block_size=block_size), + kv_caches=[k_pool], + _duplex_inv_freq=inv_freq, + ) + + # Step 1: worker reanchor runs + MiniCPMO45DuplexWorkerHelper.maybe_apply_reanchor(runner, scheduler_output=scheduler_output) + + # Assert KV rotated once + assert len(evicted_reanchors) == 1 + # Check that new_req buffer in scheduler_output was sanitized + assert "stage0_reanchor" not in new_req.model_intermediate_buffer["duplex"] + # Save rotated state + k_pool_after_first = k_pool.clone() + + # Step 2: Simulate second call with the same reanchor_id (e.g. reinjection or duplicate metadata) + runner.model_intermediate_buffer["req-1"]["duplex"]["stage0_reanchor"] = dict(reanchor_dict) + MiniCPMO45DuplexWorkerHelper.maybe_apply_reanchor(runner, scheduler_output=scheduler_output) + + # Verify: rotation was SKIPPED, no second rotation occurred! + assert len(evicted_reanchors) == 1 + assert torch.equal(k_pool, k_pool_after_first) + + +def test_history_slicing_multi_row_embeddings_and_non_aligned_prefix(): + """Verify that history eviction properly handles: + 1. Non-aligned prefix: tokens in [prefix_tokens, sink_end) are kept in the sink! + 2. Multi-row embedding tensors: audio and vision tensors are sliced by token row, + preserving 1-to-1 correspondence with token IDs for fallback prompt rebuild. + """ + prefix_tokens = 100 + sink_end = 112 + moved_from = 176 + delta = 64 + + # Unit 0: 64 tokens (covers relative [0, 64), absolute [100, 164)) + unit0_embeds = [ + torch.full((1, 8), 100.0), + torch.stack([torch.full((8,), float(101 + i)) for i in range(32)]), + torch.full((1, 8), 133.0), + *[torch.full((1, 8), float(134 + i)) for i in range(30)], + ] + unit0_tokens = list(range(100, 164)) + assert sum(t.shape[0] for t in unit0_embeds) == 64 + assert len(unit0_tokens) == 64 + unit0 = _MiniCPMO45WindowUnit(embeds=unit0_embeds, token_ids=unit0_tokens) + + # Unit 1: 64 tokens (covers relative [64, 128), absolute [164, 228)) + unit1_embeds = [torch.stack([torch.full((8,), float(164 + i)) for i in range(64)])] + unit1_tokens = list(range(164, 228)) + assert sum(t.shape[0] for t in unit1_embeds) == 64 + assert len(unit1_tokens) == 64 + unit1 = _MiniCPMO45WindowUnit(embeds=unit1_embeds, token_ids=unit1_tokens) + + state = _MiniCPMO45Stage0SessionState( + session_id="req-multi-row", + window_enabled=True, + context_embeds=[torch.full((1, 8), float(i)) for i in range(prefix_tokens)], + context_token_ids=list(range(prefix_tokens)), + context_prefix_embeds=[torch.full((1, 8), float(i)) for i in range(prefix_tokens)], + context_prefix_token_ids=list(range(prefix_tokens)), + window_units=[unit0, unit1], + ) + + helper = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + reanchor_payload = { + "reanchor_id": "r-slice-1", + "delta": delta, + "moved_from": moved_from, + "sink_end": sink_end, + "prefix_tokens": prefix_tokens, + } + helper._evict_window_units_for_reanchor(state, reanchor_payload) + + # Dropped absolute interval: [112, 176) + # Unit 0 (absolute [100, 164)): + # - Retains [100, 112) -> 12 tokens! + # - Token IDs must be 100..111 + # - Embeddings must have exactly 12 rows, matching token IDs 100..111 + assert len(state.window_units) == 2 + u0_retained = state.window_units[0] + assert len(u0_retained.token_ids) == 12 + assert u0_retained.token_ids == list(range(100, 112)) + u0_rows = [row[0].item() for emb in u0_retained.embeds for row in emb] + assert len(u0_rows) == 12 + assert u0_rows == [float(x) for x in range(100, 112)] + + # Unit 1 (absolute [164, 228)): + # - Dropped [164, 176) -> first 12 tokens dropped + # - Retains [176, 228) -> 52 tokens! + # - Token IDs must be 176..227 + # - Embeddings must have exactly 52 rows, matching token IDs 176..227 + u1_retained = state.window_units[1] + assert len(u1_retained.token_ids) == 52 + assert u1_retained.token_ids == list(range(176, 228)) + u1_rows = [row[0].item() for emb in u1_retained.embeds for row in emb] + assert len(u1_rows) == 52 + assert u1_rows == [float(x) for x in range(176, 228)] + + # Verify fallback rebuild parts: row count must equal token ID count + rebuild_parts = helper._window_replacement_parts( + state, + { + "replace": True, + "drop_units": 0, + "mode": "basic", + "replacement_prompt_len": prefix_tokens + 12 + 52, + }, + ) + assert rebuild_parts is not None + embeds, token_ids = rebuild_parts + total_embed_rows = sum(t.shape[0] if t.ndim >= 2 else 1 for t in embeds) + assert total_embed_rows == len(token_ids) == prefix_tokens + 12 + 52 + + +def test_duplex_window_install_tolerates_missing_cache_or_model_config() -> None: + """Startup installation must tolerate minimal or absent cache/model configs.""" + geometry = duplex_window_geometry( + prefix_tokens=96, + window_tokens=6000, + block_size=16, + max_model_len=8192, + high_watermark_tokens=8000, + ) + # cache_config is None, model_config is SimpleNamespace with no max_model_len + validate_duplex_window_install(None, SimpleNamespace(), geometry) + validate_duplex_window_install(None, None, geometry) + + # Worker rope inv freq helper tolerates runner with no or minimal model_config + runner = SimpleNamespace(device=torch.device("cpu")) + inv_freq = MiniCPMO45DuplexWorkerHelper.get_rope_inv_freq(runner) + assert inv_freq is not None and inv_freq.shape == (64,) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 0b8f368f0..0138d365e 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -920,23 +920,37 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) return - if stage_id == 0 and self._prepare_minicpmo45_stage0_window( - session, - update, - segment_output_ids=segment_output_ids, - completed_terminator=completed_terminator, - ): - # The rebuilt prompt is bounded by the client's window settings, - # not by the model (a camera unit is hundreds of tokens), so it - # needs the same max_model_len check as a plain extension. The - # plan replaces the whole prompt, so the replacement length is the - # projection, not the session's current prompt plus an extension. - plan = update.model_intermediate_buffer["duplex"]["stage0_window"] - if self._streaming_update_overflows(session, update, projected_len=plan["replacement_prompt_len"]): + if stage_id == 0: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import ( + MiniCPMO45DuplexSchedulerHelper, + ) + + reanchored = False + if MiniCPMO45DuplexSchedulerHelper.find_duplex_window_manager(self) is not None: + plan = self._maybe_reanchor_streaming_window( + session, + update, + segment_output_ids=segment_output_ids, + completed_terminator=completed_terminator, + ) + reanchored = plan is not None + if not reanchored and self._prepare_minicpmo45_stage0_window( + session, + update, + segment_output_ids=segment_output_ids, + completed_terminator=completed_terminator, + ): + # The rebuilt prompt is bounded by the client's window settings, + # not by the model (a camera unit is hundreds of tokens), so it + # needs the same max_model_len check as a plain extension. The + # plan replaces the whole prompt, so the replacement length is the + # projection, not the session's current prompt plus an extension. + plan = update.model_intermediate_buffer["duplex"]["stage0_window"] + if self._streaming_update_overflows(session, update, projected_len=plan["replacement_prompt_len"]): + return + self._release_replaced_streaming_prompt_cache(session) + self._replace_streaming_session(session, update) return - self._release_replaced_streaming_prompt_cache(session) - self._replace_streaming_session(session, update) - return streaming_prompt_payload = next( ( info @@ -1008,6 +1022,43 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if hasattr(update, "model_intermediate_buffer"): session.model_intermediate_buffer = update.model_intermediate_buffer + def _maybe_reanchor_streaming_window( + self, + session: Request, + update: StreamingUpdate, + *, + segment_output_ids: list[int] | None = None, + completed_terminator: int | None = None, + ) -> Any: + """Lifecycle hook to evaluate and apply duplex streaming window compaction.""" + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import ( + MiniCPMO45DuplexSchedulerHelper, + ) + + return MiniCPMO45DuplexSchedulerHelper.apply_session_window( + self, + session, + update, + segment_output_ids=segment_output_ids, + completed_terminator=completed_terminator, + ) + + def _maybe_reanchor_minicpmo45_stage0_window( + self, + session: Request, + update: StreamingUpdate, + *, + segment_output_ids: list[int] | None = None, + completed_terminator: int | None = None, + ) -> Any: + """Backward-compatible alias for _maybe_reanchor_streaming_window.""" + return self._maybe_reanchor_streaming_window( + session, + update, + segment_output_ids=segment_output_ids, + completed_terminator=completed_terminator, + ) + @staticmethod def _prepare_minicpmo45_stage0_window( session: Request, 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 79e4a8109..c65c8dadb 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 @@ -219,6 +219,7 @@ class MiniCPMO45Stage0DuplexRuntime: is_speech: bool = False, final: bool = False, stage0_window: dict[str, object] | None = None, + stage0_reanchor: dict[str, object] | None = None, ) -> dict[str, object]: """Build scheduler-owned Stage0 input embeddings for one audio append. @@ -435,6 +436,111 @@ class MiniCPMO45Stage0DuplexRuntime: state.pending_window_unit = None state.pending_window_generated_tokens.clear() + @staticmethod + def _tensor_rows(t: Any) -> int: + if hasattr(t, "shape"): + if getattr(t, "ndim", len(t.shape)) == 1: + return 1 + if getattr(t, "ndim", len(t.shape)) == 3 and t.shape[0] == 1: + return int(t.shape[1]) + return int(t.shape[0]) + return 1 + + @classmethod + def _slice_single_tensor(cls, t: Any, start_row: int, end_row: int) -> Any: + if hasattr(t, "ndim"): + if t.ndim == 1: + return t + if t.ndim == 3 and t.shape[0] == 1: + return t[:, start_row:end_row] + return t[start_row:end_row] + return t + + @classmethod + def _slice_embed_list(cls, embeds: list[Any], start_token: int, end_token: int) -> list[Any]: + if start_token >= end_token or not embeds: + return [] + result: list[Any] = [] + curr_offset = 0 + for t in embeds: + num_rows = cls._tensor_rows(t) + t_start = curr_offset + t_end = curr_offset + num_rows + curr_offset = t_end + + overlap_start = max(t_start, start_token) + overlap_end = min(t_end, end_token) + if overlap_start < overlap_end: + local_start = overlap_start - t_start + local_end = overlap_end - t_start + if local_start == 0 and local_end == num_rows: + result.append(t) + else: + result.append(cls._slice_single_tensor(t, local_start, local_end)) + return result + + def _evict_window_units_for_reanchor( + self, + state: _MiniCPMO45Stage0SessionState, + stage0_reanchor: dict[str, Any], + ) -> None: + """Evict completed window units that were dropped by zero-copy KV re-anchor. + + Ensures worker history (state.window_units) stays bounded over long streaming sessions + when zero-copy Re-RoPE KV reuse is active and prompt rebuild is bypassed. + """ + reanchor_id = stage0_reanchor.get("reanchor_id") + if reanchor_id: + applied = getattr(state, "_applied_reanchor_ids", None) + if applied is None: + applied = state._applied_reanchor_ids = set() + if reanchor_id in applied: + return + applied.add(reanchor_id) + + delta = int(stage0_reanchor.get("delta", 0) or 0) + if delta <= 0 or not state.window_units: + return + + moved_from = stage0_reanchor.get("moved_from") + sink_end = stage0_reanchor.get("sink_end") + if sink_end is None and moved_from is not None: + sink_end = int(moved_from) - delta + if sink_end is None: + sink_blocks = int(stage0_reanchor.get("sink_blocks", 0) or 0) + block_size = int(stage0_reanchor.get("block_size", 16) or 16) + sink_end = sink_blocks * block_size + + prefix_tokens = int(stage0_reanchor.get("prefix_tokens", 0) or len(state.context_token_ids)) + drop_start = max(0, int(sink_end) - prefix_tokens) + drop_end = drop_start + delta + + new_units: list[_MiniCPMO45WindowUnit] = [] + curr_offset = 0 + for unit in state.window_units: + u_len = len(unit.token_ids) + u_start = curr_offset + u_end = curr_offset + u_len + curr_offset = u_end + + local_drop_start = max(0, drop_start - u_start) + local_drop_end = min(u_len, drop_end - u_start) + + if local_drop_start >= local_drop_end: + new_units.append(unit) + elif local_drop_start == 0 and local_drop_end == u_len: + continue + else: + kept_tokens = unit.token_ids[:local_drop_start] + unit.token_ids[local_drop_end:] + kept_embeds = self._slice_embed_list(unit.embeds, 0, local_drop_start) + self._slice_embed_list( + unit.embeds, local_drop_end, u_len + ) + unit.token_ids = kept_tokens + unit.embeds = kept_embeds + new_units.append(unit) + + state.window_units = new_units + def _window_replacement_parts( self, state: _MiniCPMO45Stage0SessionState, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_kv.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_kv.py new file mode 100644 index 000000000..bb7a5f4b5 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_kv.py @@ -0,0 +1,966 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Native paged-KV window for MiniCPM-o 4.5 duplex Stage 0. + +``window_plan.py`` holds the window arithmetic in plain integers; this file +holds the three things that need an engine: the KV-cache spec/manager pair the +scheduler resolves, the layer change that produces that spec, and the in-place +re-RoPE that follows a trim. + +One mechanism, matching the reference implementation +---------------------------------------------------- +A trim and a renumber are the same event here, which is what MiniCPM's HF +implementation does: cut the cached K/V, re-RoPE the survivors onto dense +positions. #7631 gets the same observable state from the other direction, by +rebuilding the prompt and letting the model recompute every retained row -- a +full forward pass over the window, per roll. + +Here the cut is a block-table edit and the re-RoPE is one rotation per retained +key, because :class:`~vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan.PositionReanchor` +guarantees the shift is one uniform translation of a whole number of pages. Two +consequences are worth stating plainly: + +* No physical row moves. After ``del table[sink:sink + n]``, logical index ``j`` + resolves to the page that ``j + n`` named before the deletion -- which is the + page the row is already in -- and a block-aligned shift leaves ``pos % + block_size`` alone. So the rotation reads and writes the same slots, and V is + not touched at all, because only K carries a RoPE phase. +* The window needs no mask, because there is nothing to skip: the retained span + is contiguous and plain causal attention over it is what the model was trained + to do. See ``window_plan.py`` for why a *masked* window cannot express a duplex + session, and why R-SWA -- the one upstream type that can -- is not usable on a + FlashAttention-3 Stage 0. + +What the sliding-window spec is then for is sizing. ``SlidingWindowSpec`` is the +only attention type in vLLM whose ``max_admission_blocks_per_request`` bounds a +request at its window rather than at ``max_model_len``, which is what lets +``max_sessions`` concurrent streams share a Stage 0 that advertises a 40k-token +context. Its own eviction is overridden to stay inert: a masked free and a +compaction must not race for the same pages, so the gap is freed by +:meth:`MiniCPMO45DuplexWindowManager.reanchor_block_table` and nothing else. + +Reuse, not reimplementation +--------------------------- +The base of everything below is the repo's own +:class:`~vllm_omni.experimental.ar_diffusion.kv_cache.paged.ChunkWindowSpec` / +:class:`~vllm_omni.experimental.ar_diffusion.kv_cache.paged.ChunkWindowManager`, +which already keeps a sink, already evicts on chunk boundaries, and already +compacts a block table. This project sets ``chunk_size`` to the cache block size, +so "chunk-aligned" and "page-aligned" are the same statement and the inherited +:meth:`compact_block_table` indexing applies as written. + +Two constraints come with this +------------------------------ +*Prefix caching stays off.* Not only because a compacted table is no longer the +table the block hashes were computed against, but because duplex appends are +prompts of ``[filler_id] * token_budget`` (``duplex/plugin.py``) whose real +content arrives as worker-side ``inputs_embeds``: the hashes would be identical +across sessions while the KV behind them is not. ``OmniTensorPrefixCache`` has +the same exposure through its block/slot mirroring, and it and async output +materialization are mutually exclusive anyway. +*Block size.* On CUDA, kernels typically use ``block_size == 16`` for a sliding window, +while Ascend NPU uses 128. The window geometry and compaction algorithms support any +positive block size. The conversion below reads that size back off the spec the +parent produced, so the kernel's choice -- not this module's -- is the one the +chunk arithmetic uses. + +Nothing here has been executed: it needs an engine, a checkpoint and a GPU. The +arithmetic it delegates to is CPU-covered by +``tests/model_executor/models/minicpmo_4_5/duplex/test_window_plan.py``, which +also checks the RoPE identity behind :func:`rotate_cached_keys` numerically. + +Wiring, in the order the pieces get used +---------------------------------------- +1. Model module: import this module for its spec registration, call + :func:`install_duplex_window_layers` on the Stage 0 backbone right after + ``init_vllm_registered_model`` so the profiler sees the windowed spec, and run + :func:`validate_duplex_window_install` as the startup assertion. +2. Scheduler: an append becomes an ordinary session *extension*, so + ``_update_request_as_session`` runs instead of ``_replace_streaming_session`` + and the request's token count grows rather than restarting. Nothing else + changes until + :func:`~vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan.plan_position_reanchor` + returns a plan, which is the single place a trim -- or a + ``context_length_exceeded`` finish -- is decided. +3. Worker: for that step, check :func:`assert_uniform_position_shift` on the + request's real positions, call + :meth:`MiniCPMO45DuplexWindowManager.reanchor_block_table` once, then + :func:`rotate_cached_keys` per layer, and shift the request's positions, + ``mrope_positions`` and ``mrope_position_delta`` in the same step so the next + forward reads the compacted layout. +4. Deploy config: ``enable_prefix_caching`` stays false, ``block_size`` matches + the engine's cache configuration, and the window's watermarks come from + the model's duplex policy rather than from a scheduler constant. +""" + +from __future__ import annotations + +import functools +from typing import Any + +import torch +from vllm.v1.kv_cache_interface import KVCacheSpec, SlidingWindowSpec +from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec + +from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ( + ChunkWindowManager, + ChunkWindowSpec, +) +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_plan import ( + DuplexWindowGeometry, + PositionReanchor, + align_up, + cdiv, + plan_position_reanchor, +) + +#: Default block size for sliding window on CUDA. Hardware like Ascend NPU may +#: use larger page sizes (e.g. 128). +DUPLEX_WINDOW_BLOCK_SIZE = 16 + + +def compute_slot_mapping( + block_ids: list[int] | torch.Tensor, + positions: list[int] | torch.Tensor, + block_size: int, +) -> torch.Tensor: + """Map token positions to physical slot indices across a block table.""" + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + pos = torch.as_tensor(positions, dtype=torch.long) + table = torch.as_tensor(block_ids, dtype=torch.long, device=pos.device) + block_index = torch.div(pos, block_size, rounding_mode="floor") + offset = pos % block_size + return table[block_index] * block_size + offset + + +def duplex_window_geometry( + *, + prefix_tokens: int, + window_tokens: int, + block_size: int, + max_model_len: int, + sample_room: int = 1, + high_watermark_tokens: int | None = None, +) -> DuplexWindowGeometry: + """Bind the model's window policy to the engine's paging geometry.""" + return DuplexWindowGeometry( + prefix_tokens=int(prefix_tokens), + window_tokens=int(window_tokens), + block_size=int(block_size), + max_model_len=int(max_model_len), + sample_room=int(sample_room), + high_watermark_tokens=None if high_watermark_tokens is None else int(high_watermark_tokens), + ) + + +def validate_duplex_window_install( + cache_config, + model_config, + geometry: DuplexWindowGeometry, +) -> None: + """Refuse a configuration the window cannot be enforced under. + + Called from the model module once both configs exist and before the + KV-cache profile, so a wrong deploy config fails at startup rather than + leaving a session that grows without bound. + """ + if getattr(cache_config, "enable_prefix_caching", False): + raise ValueError("MiniCPM-o 4.5 duplex KV window requires enable_prefix_caching=False") + if geometry.block_size <= 0: + raise ValueError(f"duplex KV window needs positive block_size, got {geometry.block_size}") + configured_block_size = getattr(cache_config, "block_size", None) + if configured_block_size is not None and geometry.block_size != configured_block_size: + raise ValueError( + f"duplex KV window geometry block_size ({geometry.block_size}) " + f"does not match cache_config.block_size ({configured_block_size})" + ) + needed = geometry.prefix_tokens + geometry.trigger_tokens + geometry.sample_room + max_model_len = getattr(model_config, "max_model_len", None) if model_config is not None else None + if max_model_len is not None and needed > max_model_len: + raise ValueError( + f"duplex window does not fit max_model_len={max_model_len}: " + f"prefix={geometry.prefix_tokens} + trigger={geometry.trigger_tokens} " + f"+ sample_room={geometry.sample_room} = {needed}" + ) + + +def resident_bound(geometry: DuplexWindowGeometry) -> int: + """Tokens a windowed session can hold at once, rounded to whole pages.""" + return min( + geometry.prefix_tokens + geometry.trigger_tokens + geometry.sample_room, + geometry.max_model_len, + ) + + +class MiniCPMO45DuplexWindowManager(ChunkWindowManager): + """Chunk-window paging with the trim made explicit: free, delete, re-RoPE. + + Inherited as-is: ``req_to_blocks`` for the table edit and + :meth:`compact_block_table` for the deletion itself. Overridden: the + automatic eviction, which for this model would free a head the session still + has to attend. It is inert rather than removed so that admission, which reads + the same spec, still sees a bounded request. + """ + + def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: + """Never free behind the compaction's back.""" + del num_computed_tokens + return 0 + + def plan_reanchor( + self, + geometry: DuplexWindowGeometry, + *, + computed_tokens: int, + pending_tokens: int, + unit_tokens: list[int] | None = None, + ) -> PositionReanchor | None: + """Whether the next append pushes the session past its window.""" + return plan_position_reanchor( + geometry, + computed_tokens=computed_tokens, + pending_tokens=pending_tokens, + unit_tokens=unit_tokens, + ) + + def compact_block_table(self, request_id: str, sink_blocks: int | None = None) -> int: + if getattr(self, "enable_caching", False): + raise RuntimeError("duplex block-table compaction requires prefix caching to be disabled") + blocks = getattr(self, "req_to_blocks", {}).get(request_id) + if blocks is None: + blocks = getattr(self, "blocks", None) + if blocks is None: + return 0 + spec_sink = getattr(getattr(self, "kv_cache_spec", None), "sink_chunks", 0) + start = spec_sink if sink_blocks is None else int(sink_blocks) + end = start + null_block = getattr(self, "_null_block", None) + while end < len(blocks) and blocks[end] == null_block: + end += 1 + if end == start: + return 0 + del blocks[start:end] + if hasattr(self, "num_cached_block") and request_id in self.num_cached_block: + cached = self.num_cached_block[request_id] + self.num_cached_block[request_id] = min(cached, start) + max(0, cached - end) + return (end - start) * getattr(self, "block_size", 16) + + def compact_request_blocks( + self, + request_id: str, + *, + sink_blocks: int, + num_blocks: int, + ) -> int: + """Free and compact a gap of blocks after an explicit request-local sink.""" + if getattr(self, "enable_caching", False): + raise RuntimeError("duplex block-table compaction requires prefix caching to be disabled") + blocks = getattr(self, "req_to_blocks", {}).get(request_id) + if blocks is None: + blocks = getattr(self, "blocks", None) + if blocks is None or num_blocks <= 0: + return 0 + start = int(sink_blocks) + end = start + int(num_blocks) + null_block = getattr(self, "_null_block", None) + if ( + start < 0 + or end > len(blocks) + or (null_block is not None and any(block == null_block for block in blocks[start:end])) + ): + return 0 + if hasattr(self, "_remove_blocks_in_range"): + self._remove_blocks_in_range(request_id, start, end) + elif hasattr(self, "free_blocks_range"): + self.free_blocks_range(request_id, start, end) + else: + del blocks[start:end] + return num_blocks * getattr(self, "block_size", 16) + return self.compact_block_table(request_id, sink_blocks=start) + + def reanchor_block_table(self, request_id: str, plan: PositionReanchor) -> int: + """Free and delete one request's gap; return the tokens it held. + + Returns 0 when the table disagrees with the plan -- the gap was already + removed, or the request is gone -- and a caller that asked for a non-zero + delta must then skip the rotation as well, since nothing moved. + """ + gap_blocks = plan.delta // self.block_size + return self.compact_request_blocks( + request_id, + sink_blocks=plan.sink_blocks, + num_blocks=gap_blocks, + ) + + +@register_kv_cache_spec(manager_class=MiniCPMO45DuplexWindowManager, uniform_type_base_spec=None) +class MiniCPMO45DuplexWindowSpec(ChunkWindowSpec): + """A :class:`ChunkWindowSpec` that resolves to the duplex manager above. + + Adds no fields. Registration exists only because dispatch walks the spec's + MRO, so an unregistered subclass silently keeps the parent manager and the + trim entry point is never reachable. + """ + + +def duplex_window_spec(geometry: DuplexWindowGeometry, layer_spec: SlidingWindowSpec) -> MiniCPMO45DuplexWindowSpec: + """Re-express a layer's sliding-window spec as the duplex one. + + ``layer_spec`` is what an upstream ``Attention`` reports once it has a + sliding window. Its ``block_size`` was chosen by the kernel, not by + ``--block-size``, so every count below is in the pages the cache will + actually use -- and it has to be the pages the planner counts in too, or + ``plan.sink_blocks`` and ``spec.sink_chunks`` describe different tables. + """ + if layer_spec.block_size != geometry.block_size: + raise ValueError( + f"duplex window plans in {geometry.block_size}-token pages but the kernel gave this " + f"group {layer_spec.block_size}-token pages; align cache_config.block_size with the geometry" + ) + block_size = layer_spec.block_size + # The mask never bites (see the manager), so the window is set to the + # resident bound: long enough that a compacted session is always fully + # attended, short enough that the allocator admits a session by what it + # holds rather than by max_model_len. + bound = align_up(resident_bound(geometry), block_size) + return MiniCPMO45DuplexWindowSpec( + block_size=block_size, + num_kv_heads=layer_spec.num_kv_heads, + head_size=layer_spec.head_size, + head_size_v=layer_spec.head_size_v, + dtype=layer_spec.dtype, + kv_quant_mode=layer_spec.kv_quant_mode, + page_size_padded=layer_spec.page_size_padded, + sliding_window=bound, + chunk_size=block_size, + window_chunks=bound // block_size, + sink_chunks=cdiv(geometry.prefix_tokens, block_size), + reset_at_boundary=False, + ) + + +def install_duplex_window_layers(model: torch.nn.Module, *, geometry: DuplexWindowGeometry) -> list[str]: + """Point a decoder's attention layers at :class:`MiniCPMO45DuplexWindowSpec`. + + Stage 0 builds its backbone through ``init_vllm_registered_model`` on a + Qwen3/Qwen2 text config, so its layers are stock vLLM ``Attention`` modules + reporting ``FullAttentionSpec``. ``RSWAAttention`` shows the way to change + that without a model file: the only difference is what ``get_kv_cache_spec`` + returns. So each layer is re-classed and given the two attributes its + override reads -- the window length that makes the parent emit a + sliding-window spec at all, and the geometry that turns it into this one. + + The window length is ``resident_bound(geometry)``, and it exists for sizing + only. The layer's ``impl`` was built before this call with no window, and + that is the point: the kernel stays causal over the block table, which is + correct precisely because the trim keeps the table dense. + + Returns the modified layer names, for the startup log and for a test that the + spec really did change. + """ + from vllm.model_executor.layers.attention import Attention + from vllm.v1.attention.backend import AttentionType + + windowed = _duplex_window_attention_class() + sliding_window = resident_bound(geometry) + renamed: list[str] = [] + for name, module in model.named_modules(): + if not isinstance(module, Attention) or isinstance(module, windowed): + continue + if module.attn_type != AttentionType.DECODER: + # Encoder-side attention (the vision/audio towers) keeps its own + # cache policy; only the decoder's KV is windowed. + continue + object.__setattr__(module, "__class__", windowed) + object.__setattr__(module, "sliding_window", int(sliding_window)) + object.__setattr__(module, "_duplex_window_geometry", geometry) + renamed.append(name) + return renamed + + +@functools.lru_cache(maxsize=1) +def _duplex_window_attention_class(): + """``Attention`` that reports the duplex spec instead of a plain one. + + Subclassing at call time leaves upstream in charge of the parts that are the + backend's business -- the kernel-chosen block size, page padding, + quantisation mode -- and converts only the class of the result. + """ + from vllm.model_executor.layers.attention import Attention + + class MiniCPMO45DuplexWindowAttention(Attention): + _duplex_window_geometry: DuplexWindowGeometry + + def get_kv_cache_spec(self, vllm_config) -> KVCacheSpec | None: + spec = super().get_kv_cache_spec(vllm_config) + if spec is None or not isinstance(spec, SlidingWindowSpec) or isinstance(spec, ChunkWindowSpec): + return spec + return duplex_window_spec(self._duplex_window_geometry, spec) + + return MiniCPMO45DuplexWindowAttention + + +def assert_uniform_position_shift(positions: torch.Tensor, moved_from: int) -> None: + """Guard the one case a single rotation cannot express. + + MRoPE scores against three position rows, and a translation of the tail is + one rotation only while all three shift by the same amount. A retained span + carrying vision tokens with distinct height/width rows does not, so the + caller must re-prefill that session instead. Duplex audio and text advance + every row together, so the streaming tail satisfies this. + """ + if positions.dim() == 1: + return + if positions.dim() != 2: + raise ValueError(f"expected a (rows, tokens) position tensor, got {tuple(positions.shape)}") + tail = positions[:, int(moved_from) :] + if tail.shape[1] == 0: + return + if not bool(torch.all(tail == tail[0])): + raise RuntimeError( + "duplex re-anchor needs one position row across the retained tail; re-prefill this " + "session rather than rotating split MRoPE positions" + ) + + +def rotate_keys(keys: torch.Tensor, delta: int, inv_freq: torch.Tensor) -> torch.Tensor: + """Undo ``delta`` positions of rotation on keys that are already rotated. + + ``keys`` is ``(tokens, heads, head_dim)`` with NeoX-style pairing + (``rotate_half``), which is what the Qwen3-family RoPE Stage 0 registers + writes. RoPE composes, so multiplying a key's cached pair by + ``exp(-i * delta * inv_freq)`` yields exactly the key that would have been + written ``delta`` positions earlier, for every frequency at once and without + reference to where the token came from. + """ + half = keys.shape[-1] // 2 + if inv_freq.numel() != half: + raise ValueError(f"expected {half} RoPE frequencies for head_dim={keys.shape[-1]}, got {inv_freq.numel()}") + # Angle and trigonometric computations must stay in float32 to avoid catastrophic + # quantization error: delta * inv_freq can exceed 1024, where bfloat16 has ULP = 8 + # (quantization error ~4-5.7 rad, completely randomizing trig values). + angle = float(delta) * inv_freq.to(device=keys.device, dtype=torch.float32) + cos = torch.cos(angle).to(dtype=keys.dtype).unsqueeze(0).unsqueeze(1) + sin = torch.sin(angle).to(dtype=keys.dtype).unsqueeze(0).unsqueeze(1) + k1, k2 = keys[..., :half], keys[..., half:] + # Complex multiply by exp(-i*angle): the inverse of the forward rotation. + return torch.cat([k1 * cos + k2 * sin, k2 * cos - k1 * sin], dim=-1) + + +def rotate_cached_keys( + k_pool: torch.Tensor, + *, + block_ids: list[int], + positions: torch.Tensor, + plan: PositionReanchor, + inv_freq: torch.Tensor, + block_size: int | None = None, +) -> int: + """Re-RoPE the retained tail where it stands; return the rows touched. + + Supports: + - vLLM 0.30+ packed FlashAttention KV layout: + (num_blocks, num_kv_heads, block_size, 2 * head_dim) + - Standard PagedAttention KV layouts: + (num_blocks, block_size, num_kv_heads, head_dim) and + (num_blocks, num_kv_heads, block_size, head_dim) + - 5D split layouts (2, num_blocks, ...) or (num_blocks, 2, ...) + + Uses direct tensor advanced indexing assignment to write rotated keys in-place + into the underlying storage, preserving V cache, non-retained blocks, and avoiding + temporary copy allocation on non-contiguous strides. + """ + if positions.numel() == 0: + return 0 + + if k_pool.dim() == 5: + if k_pool.shape[0] == 2: + k_pool = k_pool[0] + elif k_pool.shape[1] == 2: + k_pool = k_pool[:, 0] + else: + raise ValueError(f"unsupported 5D cache shape {tuple(k_pool.shape)}") + + if k_pool.dim() != 4: + raise ValueError(f"rotate_cached_keys requires a 4D key cache, got shape {tuple(k_pool.shape)}") + + expected_head_dim = inv_freq.numel() * 2 + last_dim = k_pool.shape[-1] + if last_dim not in (expected_head_dim, 2 * expected_head_dim): + raise ValueError( + f"expected head_dim={expected_head_dim} or {2 * expected_head_dim} for " + f"{inv_freq.numel()} RoPE frequencies, got {last_dim} (shape {tuple(k_pool.shape)})" + ) + + if block_size is None: + if last_dim == 2 * expected_head_dim: + # Packed FlashAttention: (num_blocks, num_kv_heads, block_size, 2 * head_dim) + block_size = k_pool.shape[2] + else: + block_size = k_pool.shape[1] + + if k_pool.shape[2] == block_size: + layout = "BHND" + num_kv_heads = k_pool.shape[1] + elif k_pool.shape[1] == block_size: + layout = "BNHD" + num_kv_heads = k_pool.shape[2] + else: + raise ValueError(f"cannot resolve layout: block_size={block_size} not in shape {tuple(k_pool.shape)}") + + positions = positions.to(device=k_pool.device, dtype=torch.long) + if bool((positions < plan.moved_from).any()): + raise RuntimeError(f"re-anchor given {int((positions < plan.moved_from).sum())} positions below moved_from") + + shifted = positions - plan.delta + block_table = torch.as_tensor(block_ids, dtype=torch.long, device=k_pool.device) + target_blocks = block_table[shifted // block_size] + offsets = shifted % block_size + + b_idx = target_blocks.unsqueeze(1) + h_idx = torch.arange(num_kv_heads, device=k_pool.device).unsqueeze(0) + o_idx = offsets.unsqueeze(1) + + if layout == "BHND": + keys = k_pool[b_idx, h_idx, o_idx, :expected_head_dim] + rotated = rotate_keys(keys, plan.delta, inv_freq) + k_pool[b_idx, h_idx, o_idx, :expected_head_dim] = rotated + else: + keys = k_pool[b_idx, o_idx, h_idx, :expected_head_dim] + rotated = rotate_keys(keys, plan.delta, inv_freq) + k_pool[b_idx, o_idx, h_idx, :expected_head_dim] = rotated + + return int(positions.numel()) + + +class MiniCPMO45DuplexWindowPolicy: + """Model-side retention policy and watermark geometry calculation.""" + + @classmethod + def plan_reanchor( + cls, + session: Any, + update: Any, + block_size: int, + max_model_len: int, + unit_tokens: list[int] | None = None, + segment_output_ids: list[int] | None = None, + ) -> PositionReanchor | None: + """Evaluate model retention decision and return PositionReanchor if due.""" + info = getattr(update, "model_intermediate_buffer", None) + if not isinstance(info, dict): + return None + duplex = info.get("duplex") + if not isinstance(duplex, dict) or duplex.get("data_plane") is not True: + return None + runtime_config = duplex.get("runtime_config") + runtime_config = runtime_config if isinstance(runtime_config, dict) else {} + window = runtime_config.get("duplex_window_config") + if not isinstance(window, dict): + return None + mode = window.get("sliding_window_mode", "off") + if mode != "basic": + # Only basic mode uses zero-copy KV reuse with in-place Re-RoPE. + # Context mode relies on previous-text accumulation and partial forward, + # so it routes explicitly through the official fallback (_prepare_minicpmo45_stage0_window). + return None + + base_len = int(getattr(session, "num_computed_tokens", 0) or 0) + prefix_tokens = int(runtime_config.get("duplex_window_prefix_tokens", 96) or 96) + high_watermark = int(window.get("basic_window_high_tokens", 8000) or 8000) + low_watermark = int(window.get("basic_window_low_tokens", 6000) or 6000) + + # Normalize total-sequence watermarks to content-only budgets for DuplexWindowGeometry + # so that trigger_total == high_watermark and target_total == low_watermark, + # perfectly matching the fallback path (_prepare_minicpmo45_stage0_window). + content_low = max(block_size, low_watermark - prefix_tokens) + content_high = max(content_low, high_watermark - prefix_tokens) + + geometry = DuplexWindowGeometry( + prefix_tokens=prefix_tokens, + window_tokens=content_low, + block_size=block_size, + max_model_len=max_model_len, + high_watermark_tokens=content_high, + ) + + pending_tokens = len(getattr(update, "prompt_token_ids", []) or []) + + # Derive unit_tokens from session._minicpmo45_window_units if not explicitly given + if unit_tokens is None and hasattr(session, "_minicpmo45_window_units"): + raw_units = getattr(session, "_minicpmo45_window_units", ()) or () + if raw_units: + unit_tokens = [int(u["length"]) for u in raw_units if isinstance(u, dict) and "length" in u] + seq = int(duplex.get("seq", 0) or 0) + if seq > 1 and segment_output_ids is not None: + preserve_len = int(runtime_config.get("duplex_first_append_context_tokens", 0) or 0) + old_prompt_tokens = int(getattr(session, "num_prompt_tokens", 0) or 0) + boundary_cur = old_prompt_tokens + len(segment_output_ids) + 2 + recorded_open = getattr(session, "_minicpmo45_window_open_start", None) + open_start = preserve_len if recorded_open is None else int(recorded_open) + closing_len = boundary_cur - open_start + if closing_len > 0: + unit_tokens = [*unit_tokens, closing_len] + + return plan_position_reanchor( + geometry, + computed_tokens=base_len, + pending_tokens=pending_tokens, + unit_tokens=unit_tokens, + ) + + +class MiniCPMO45DuplexSchedulerHelper: + """Scheduler-side window planning and request compaction helper for MiniCPM-o 4.5 duplex.""" + + @classmethod + def find_duplex_window_manager(cls, scheduler: Any) -> MiniCPMO45DuplexWindowManager | None: + coordinator = getattr(getattr(scheduler, "kv_cache_manager", None), "coordinator", None) + if coordinator is None: + return None + for mgr in getattr(coordinator, "single_type_managers", ()): + if isinstance(mgr, (MiniCPMO45DuplexWindowManager, ChunkWindowManager)): + return mgr + return None + + @classmethod + def apply_session_window( + cls, + scheduler: Any, + session: Any, + update: Any, + *, + segment_output_ids: list[int] | None = None, + completed_terminator: int | None = None, + unit_tokens: list[int] | None = None, + ) -> PositionReanchor | None: + """Evaluate watermark policy, compact KV block table and session tokens if needed.""" + cache_config = getattr(scheduler, "cache_config", None) + block_size = int(getattr(cache_config, "block_size", 16) or 16) + model_config = getattr(scheduler, "model_config", None) + max_model_len = int(getattr(model_config, "max_model_len", 40960) or 40960) + + plan = MiniCPMO45DuplexWindowPolicy.plan_reanchor( + session=session, + update=update, + block_size=block_size, + max_model_len=max_model_len, + unit_tokens=unit_tokens, + segment_output_ids=segment_output_ids, + ) + if plan is None: + return None + + duplex_mgr = cls.find_duplex_window_manager(scheduler) + if duplex_mgr is None: + return None + + # Free and compact blocks via explicit request-local sink boundary on KV manager. + # Rejection before state mutation: if compaction returns 0, no state is changed. + gap_blocks = plan.delta // block_size + if hasattr(duplex_mgr, "compact_request_blocks"): + freed_tokens = duplex_mgr.compact_request_blocks( + session.request_id, + sink_blocks=plan.sink_blocks, + num_blocks=gap_blocks, + ) + else: + freed_tokens = duplex_mgr.reanchor_block_table(session.request_id, plan) + + if freed_tokens == 0: + return None + + old_computed = session.num_computed_tokens + sink_end = plan.sink_end + moved_from = plan.moved_from + + # Scheduler owns corresponding token-history and counter updates derived from applied compaction result + old_prompt_tokens = ( + len(session.prompt_token_ids) + if getattr(session, "prompt_token_ids", None) is not None + else int(getattr(session, "num_prompt_tokens", 0) or old_computed) + ) + if getattr(session, "prompt_token_ids", None) is not None and len(session.prompt_token_ids) >= moved_from: + session.prompt_token_ids = session.prompt_token_ids[:sink_end] + session.prompt_token_ids[moved_from:] + session.num_prompt_tokens = len(session.prompt_token_ids) + + if getattr(session, "_all_token_ids", None) is not None and len(session._all_token_ids) >= moved_from: + session._all_token_ids = session._all_token_ids[:sink_end] + session._all_token_ids[moved_from:] + + session.num_computed_tokens -= freed_tokens + + # Worker instruction derived from the same applied compaction result + info = update.model_intermediate_buffer + duplex = info.setdefault("duplex", {}) + runtime_config = duplex.get("runtime_config") + runtime_config = runtime_config if isinstance(runtime_config, dict) else {} + preserve_len = int(runtime_config.get("duplex_first_append_context_tokens", 0) or 0) + seq = int(duplex.get("seq", 0) or 0) + + reanchor_count = getattr(session, "_minicpmo45_reanchor_count", 0) + 1 + session._minicpmo45_reanchor_count = reanchor_count + reanchor_id = f"{session.request_id}-r{reanchor_count}-{plan.moved_from}-{plan.delta}" + duplex["stage0_reanchor"] = { + "reanchor_id": reanchor_id, + "delta": plan.delta, + "moved_from": plan.moved_from, + "sink_blocks": plan.sink_blocks, + "sink_end": plan.sink_end, + "prefix_tokens": preserve_len, + "block_size": block_size, + "old_computed_tokens": old_computed, + } + + # Synchronize stage0_window with worker so worker finalizes the completed unit + generated_ids = list(segment_output_ids) if segment_output_ids is not None else [] + stage0_window = duplex.setdefault("stage0_window", {}) + stage0_window["completed_token_ids"] = generated_ids + if completed_terminator is not None: + stage0_window["completed_terminator_token_id"] = int(completed_terminator) + + # Synchronize scheduler window units and coordinates atomically with compaction + boundary_before = old_prompt_tokens + len(generated_ids) + 2 + recorded_open_start = getattr(session, "_minicpmo45_window_open_start", None) + open_start = preserve_len if recorded_open_start is None else int(recorded_open_start) + unit_len = boundary_before - open_start + + special_ids = { + int(token_id) + for token_id in runtime_config.get("duplex_window_special_token_ids", ()) + if isinstance(token_id, int) + } + units = list(getattr(session, "_minicpmo45_window_units", ())) + if seq > 1 and unit_len > 0: + units.append( + { + "length": unit_len, + "generated_token_ids": [int(t) for t in generated_ids if t not in special_ids], + } + ) + + # Prune units dropped by compaction + drop_start = max(0, plan.sink_end - preserve_len) + drop_end = drop_start + freed_tokens + + new_units: list[dict[str, Any]] = [] + curr_offset = 0 + for u in units: + u_len = int(u.get("length", 0)) + u_start = curr_offset + u_end = curr_offset + u_len + curr_offset = u_end + + local_drop_start = max(0, drop_start - u_start) + local_drop_end = min(u_len, drop_end - u_start) + + if local_drop_start >= local_drop_end: + new_units.append(u) + elif local_drop_start == 0 and local_drop_end == u_len: + continue + else: + new_len = u_len - (local_drop_end - local_drop_start) + gen_ids = u.get("generated_token_ids", []) + kept_gen_ids = gen_ids if local_drop_end < u_len else [] + new_units.append( + { + "length": new_len, + "generated_token_ids": kept_gen_ids, + } + ) + session._minicpmo45_window_units = new_units + + # Rebase open_start into compacted coordinates + if seq > 1: + session._minicpmo45_window_open_start = boundary_before - freed_tokens + elif recorded_open_start is not None: + session._minicpmo45_window_open_start = max(0, recorded_open_start - freed_tokens) + + return plan + + @classmethod + def maybe_reanchor_session( + cls, + scheduler: Any, + session: Any, + update: Any, + *, + segment_output_ids: list[int] | None = None, + completed_terminator: int | None = None, + unit_tokens: list[int] | None = None, + ) -> PositionReanchor | None: + """Backward-compatible alias for apply_session_window.""" + return cls.apply_session_window( + scheduler, + session, + update, + segment_output_ids=segment_output_ids, + completed_terminator=completed_terminator, + unit_tokens=unit_tokens, + ) + + +class MiniCPMO45DuplexWorkerHelper: + """Worker-side KV cache rotation and position metadata helper for MiniCPM-o 4.5 duplex.""" + + @classmethod + def get_rope_inv_freq(cls, runner: Any) -> torch.Tensor: + inv_freq = getattr(runner, "_duplex_inv_freq", None) + if inv_freq is None: + model_config = getattr(runner, "model_config", None) + head_dim = ( + model_config.get_head_size() + if model_config is not None and hasattr(model_config, "get_head_size") + else 128 + ) + hf_config = getattr(model_config, "hf_config", None) if model_config is not None else None + base = float(getattr(hf_config, "rope_theta", 1000000.0) or 1000000.0) + inv_freq = 1.0 / ( + base ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=runner.device) / head_dim) + ) + runner._duplex_inv_freq = inv_freq + return inv_freq + + @classmethod + def resolve_group_block_ids( + cls, + runner: Any, + req_id: str, + req_idx: int, + group_idx: int = 0, + ) -> list[int]: + """Resolve the flat block ID table for a specific KV cache group of a request. + + In production, req_state.block_ids is a tuple of lists, one per KV cache group. + This helper unpacks the appropriate group to ensure compute_slot_mapping receives + a flat 1D list of integer block IDs. + """ + req_state = getattr(runner, "requests", {}).get(req_id) if hasattr(runner, "requests") else None + if req_state is not None and getattr(req_state, "block_ids", None): + raw_blocks = req_state.block_ids + if isinstance(raw_blocks, (tuple, list)) and len(raw_blocks) > 0: + if isinstance(raw_blocks[0], (tuple, list)): + actual_group = group_idx if group_idx < len(raw_blocks) else 0 + return [int(b) for b in raw_blocks[actual_group]] + return [int(b) for b in raw_blocks] + # Fallback to input_batch block_table + bt = getattr(getattr(runner, "input_batch", None), "block_table", None) + if bt is not None: + bt_tables = getattr(bt, "block_tables", None) + if bt_tables and len(bt_tables) > 0: + actual_group = group_idx if group_idx < len(bt_tables) else 0 + bt_row = bt_tables[actual_group] + else: + bt_row = bt + num_blocks = int(bt_row.num_blocks_per_row[req_idx]) + return [int(b) for b in bt_row.block_table.np[req_idx, :num_blocks]] + return [] + + @classmethod + def maybe_apply_reanchor(cls, runner: Any, scheduler_output: Any = None) -> None: + """Apply in-place KV reanchor and rotation on worker before model forward.""" + if not hasattr(runner, "input_batch") or runner.input_batch is None: + return + num_reqs = getattr(runner.input_batch, "num_reqs", len(runner.input_batch.req_ids)) + req_ids = runner.input_batch.req_ids[:num_reqs] + applied_reanchors = getattr(runner, "_applied_stage0_reanchor_ids", None) + if applied_reanchors is None: + applied_reanchors = runner._applied_stage0_reanchor_ids = set() + + for req_idx, req_id in enumerate(req_ids): + info = runner.model_intermediate_buffer.get(req_id) + if not isinstance(info, dict): + continue + duplex = info.get("duplex") + if not isinstance(duplex, dict): + continue + reanchor = duplex.pop("stage0_reanchor", None) + if reanchor is None: + continue + + # Sanitize scheduled_new_reqs in scheduler_output so subsequent runner + # metadata refreshes (_update_additional_information in _preprocess) + # do not re-inject this already popped reanchor command. + if scheduler_output is not None and hasattr(scheduler_output, "scheduled_new_reqs"): + for new_req in scheduler_output.scheduled_new_reqs: + if getattr(new_req, "req_id", None) == req_id: + buf = getattr(new_req, "model_intermediate_buffer", None) + if isinstance(buf, dict) and isinstance(buf.get("duplex"), dict): + buf["duplex"].pop("stage0_reanchor", None) + + reanchor_sig = ( + f"{req_id}:{reanchor.get('moved_from')}:{reanchor.get('delta')}:{reanchor.get('old_computed_tokens')}" + ) + reanchor_id = reanchor.get("reanchor_id") or reanchor_sig + if reanchor_id in applied_reanchors: + continue + applied_reanchors.add(reanchor_id) + reanchor["reanchor_id"] = reanchor_id + + plan = PositionReanchor( + delta=reanchor["delta"], + moved_from=reanchor["moved_from"], + sink_blocks=reanchor["sink_blocks"], + ) + + # Scheduler is authoritative for logical state (block_table and computed tokens). + # The parent runner's _update_states() already installed the post-compaction block IDs + # and decremented num_computed_tokens_cpu. We do NOT double-compact or double-decrement here. + old_computed = int( + reanchor.get( + "old_computed_tokens", + int(runner.input_batch.num_computed_tokens_cpu[req_idx]) + plan.delta, + ) + ) + + req_state = runner.requests.get(req_id) if hasattr(runner, "requests") else None + mrope_pos = getattr(req_state, "mrope_positions", None) if req_state is not None else None + if mrope_pos is not None: + assert_uniform_position_shift(mrope_pos, plan.moved_from) + block_size = int(getattr(getattr(runner, "cache_config", None), "block_size", 16) or 16) + sink_tokens = plan.sink_blocks * block_size + if mrope_pos.shape[1] >= old_computed: + req_state.mrope_positions = torch.cat( + [ + mrope_pos[:, :sink_tokens], + mrope_pos[:, plan.moved_from : old_computed] - plan.delta, + ], + dim=1, + ) + if getattr(req_state, "mrope_position_delta", None) is not None: + req_state.mrope_position_delta = max(0, req_state.mrope_position_delta - plan.delta) + + positions = torch.arange(plan.moved_from, old_computed, dtype=torch.long, device=runner.device) + if mrope_pos is None: + assert_uniform_position_shift(positions, plan.moved_from) + + if positions.numel() > 0 and hasattr(runner, "kv_caches") and runner.kv_caches: + inv_freq = cls.get_rope_inv_freq(runner) + kv_groups = getattr(runner, "kv_cache_group_ids", None) + block_size = int(getattr(getattr(runner, "cache_config", None), "block_size", 16) or 16) + for layer_idx, kv_cache in enumerate(runner.kv_caches): + group_idx = kv_groups[layer_idx] if (kv_groups and layer_idx < len(kv_groups)) else 0 + layer_block_ids = cls.resolve_group_block_ids(runner, req_id, req_idx, group_idx=group_idx) + rotate_cached_keys( + kv_cache, + block_ids=layer_block_ids, + positions=positions, + plan=plan, + inv_freq=inv_freq, + block_size=block_size, + ) + + # Evict worker-held unit history so state.window_units remains bounded + model = getattr(runner, "model", None) + helper = getattr(model, "_minicpmo45_duplex_data_plane_helper", None) + if helper is not None and isinstance(getattr(helper, "sessions", None), dict): + req_sessions = getattr(model, "_minicpmo45_duplex_request_sessions", {}) + session_key = duplex.get("session_id") or ( + req_sessions.get(req_id) if isinstance(req_sessions, dict) else None + ) + session_state = helper.sessions.get(session_key) if session_key else None + if session_state is not None and hasattr(session_state, "window_units"): + if hasattr(helper, "_evict_window_units_for_reanchor"): + helper._evict_window_units_for_reanchor(session_state, reanchor) + else: + session_state.window_units.clear() diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_plan.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_plan.py new file mode 100644 index 000000000..045b231e6 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/window_plan.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Position and block accounting for the MiniCPM-o 4.5 duplex KV window. + +MiniCPM-o's native duplex streaming has a bounded attention budget: once a +session's context passes the window, the oldest content has to go. The HF +reference implementation does that by trimming the cached K/V tensors and +re-RoPE-ing the survivors (``slide_trigger_seconds`` / ``slide_stride_seconds`` +in the streaming processor, which is a *separate* window living in the audio +CNN). This module is the same rule, in the units a paged KV cache schedules in: +tokens, blocks and watermarks. + +Why trimming and renumbering are one event, not two +--------------------------------------------------- +The alternative is to keep every position and only mask the out-of-window +range, which is what vLLM's sliding-window cache types do. Neither fits a +duplex session: + +* ``SlidingWindowSpec`` masks a *contiguous* tail window -- the kernel is handed + ``window_size=(W-1, 0)`` -- so it cannot keep a head prefix attended while + dropping what sits right after it. +* ``RSWASpec`` can, that being its whole shape, but it defines its sink as + ``num_prompt_tokens``, which for a session that grows by appends is the first + append rather than the head context the duplex protocol preserves. It also + needs a backend that can express the mask: FA4's ``rswa_mask_mod``, Triton or + FlexAttention. On a FlashAttention-3 Stage 0 the branch is skipped and the + zeroed gap blocks are attended as ordinary context. + +Renumbering removes the gap outright instead. The retained span is a contiguous +``[0, prefix + window)`` range, plain causal attention over it is what the model +was trained to do, and no backend has to know a window exists. The price is one +rotation of the retained keys, which is what :class:`PositionReanchor` plans. + +Everything here is pure integer arithmetic so the invariants are testable +without a GPU, a checkpoint, or a vLLM engine. ``window_kv.py`` is the vLLM +tier that consumes it. + +Paged KV Retention Semantics vs. Unpaged Reference +-------------------------------------------------- +In the unpaged HuggingFace reference implementation (OpenBMB MiniCPM-o 4.5), KV +tensors are dense PyTorch tensors. The reference basic window drops whole units +immediately after the exact preserved prefix length: +`retained = [0, prefix_tokens) + [unit_end, total)`. + +In vLLM, KV memory is paged into physical blocks of size `block_size` (e.g. 16 or 128). +To achieve zero-copy in-place KV reuse without allocating new blocks or shifting +physical slot rows across pages: +1. The sink region is aligned up to whole blocks: + `sink_end = cdiv(prefix_tokens, block_size) * block_size`. +2. Any tokens in `[prefix_tokens, sink_end)` reside in the physical sink block and + remain resident in-place (up to `block_size - 1` tokens). +3. The dropped interval is `[sink_end, moved_from)`. +4. The retained tail `[moved_from, total)` is translated down by `delta = moved_from - sink_end`, + where `delta` is an exact multiple of `block_size`. + +This guarantees that physical pages in the block table require no data movement or +reallocation, at the cost of retaining up to `block_size - 1` tokens of the oldest +unit in the sink page. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +def align_up(value: int, unit: int) -> int: + if unit <= 0: + raise ValueError(f"unit must be positive, got {unit}") + return -((-int(value)) // unit) * unit + + +def cdiv(value: int, divisor: int) -> int: + return -((-int(value)) // int(divisor)) + + +@dataclass(frozen=True) +class DuplexWindowGeometry: + """Window policy for one duplex session, in tokens. + + Attributes: + prefix_tokens: The head context that outlives every trim (system + instructions, reference audio, suffix). It stays at position 0. + window_tokens: Low watermark -- how much appended content survives a + trim, so a trim lands the session on ``prefix_tokens + + window_tokens``. + block_size: KV cache block size. + max_model_len: Position budget of the stage. A windowed session is not + supposed to run against this; it bounds the degenerate case where + the window is the whole context. + sample_room: Tokens that must stay free after the prompt to decode. + high_watermark_tokens: Trigger point -- a trim is planned once the + projected content passes this. Defaults to the low watermark, i.e. + hold the budget exactly; the reference implementation's + trigger/stride pair is the same knob measured in seconds. + """ + + prefix_tokens: int + window_tokens: int + block_size: int + max_model_len: int + sample_room: int = 1 + high_watermark_tokens: int | None = None + + def __post_init__(self) -> None: + if self.prefix_tokens < 0 or self.window_tokens <= 0: + raise ValueError(f"invalid window geometry: prefix={self.prefix_tokens} window={self.window_tokens}") + if self.block_size <= 0 or self.max_model_len <= 0: + raise ValueError( + f"invalid window geometry: block_size={self.block_size} max_model_len={self.max_model_len}" + ) + if self.sample_room < 0: + raise ValueError(f"sample_room must be non-negative, got {self.sample_room}") + if self.high_watermark_tokens is not None and self.high_watermark_tokens < self.window_tokens: + raise ValueError( + f"high_watermark_tokens {self.high_watermark_tokens} is below window_tokens {self.window_tokens}" + ) + + @property + def trigger_tokens(self) -> int: + """Content tokens at which a trim becomes due, before the prefix.""" + return self.window_tokens if self.high_watermark_tokens is None else self.high_watermark_tokens + + +def unit_starts(geometry: DuplexWindowGeometry, unit_tokens: list[int]) -> list[int]: + """Absolute position of each unit's first token, after the fixed prefix. + + ``unit_tokens`` is one entry per accepted append: every processor chunk and + internal closure token that append wrote into the KV cache. The positions + are read after the last trim, so the caller owns the history this is derived + from; a session that has already re-anchored reports shifted lengths. + """ + starts: list[int] = [] + position = geometry.prefix_tokens + for unit in unit_tokens: + starts.append(position) + position += int(unit) + return starts + + +@dataclass(frozen=True) +class PositionReanchor: + """A renumbering of the retained tail, and the rotation it implies. + + Attributes: + delta: Positions subtracted from every retained token at or after + ``moved_from``. A multiple of ``block_size``, which is what lets the + tail keep the physical slots it is already written in: the block + table loses its gap entries, so logical index ``j`` becomes + ``j - delta / block_size``, and ``block_table[j']`` still names the + page that holds the row. + moved_from: First absolute position that moves. + sink_blocks: Table prefix that stays where it is. + """ + + delta: int + moved_from: int + sink_blocks: int + + @property + def sink_end(self) -> int: + """Position the retained tail lands on, i.e. the end of the sink.""" + return self.moved_from - self.delta + + @property + def rope_angle_shift_tokens(self) -> int: + """Signed position offset to rotate retained keys by: ``-delta``. + + RoPE composes, so a token that moves from ``p`` to ``p - delta`` needs + its cached key rotated by ``-delta * inv_freq`` and nothing else: the + sink did not move, and a uniform translation of the tail leaves every + within-tail relative distance untouched. + """ + return -self.delta + + +def reanchor_positions(position: int, plan: PositionReanchor) -> int: + """Map one absolute position through a re-anchor. + + The sink keeps its positions, everything from ``moved_from`` translates by + ``-delta``, and the span in between ceases to exist. + """ + if int(position) != position: + raise ValueError(f"position must be an integer, got {position!r}") + position = int(position) + if plan.sink_end <= position < plan.moved_from: + raise ValueError( + f"position {position} was dropped by reanchor plan moved_from={plan.moved_from} delta={plan.delta}" + ) + if position < plan.moved_from: + return position + return position - plan.delta + + +def plan_position_reanchor( + geometry: DuplexWindowGeometry, + *, + computed_tokens: int, + pending_tokens: int, + target_tokens: int | None = None, + unit_tokens: list[int] | None = None, +) -> PositionReanchor | None: + """Plan the trim that keeps the next append inside the window. + + Args: + computed_tokens: Positions written so far, after the last trim. + pending_tokens: What the next append is expected to add. + target_tokens: Where the sequence should land; defaults to the sink plus + the low-watermark window. + unit_tokens: Per-append committed lengths. When given, the cut snaps down + to a unit start, so a unit is either fully kept or fully gone -- + dropping half a unit would be a behaviour change against a window + that drops whole ones. + + Returns ``None`` while the session is inside the trigger watermark. + """ + projected = int(computed_tokens) + int(pending_tokens) + if projected <= geometry.prefix_tokens + geometry.trigger_tokens: + return None + sink_end_block = cdiv(geometry.prefix_tokens, geometry.block_size) + sink_end = sink_end_block * geometry.block_size + target = geometry.prefix_tokens + geometry.window_tokens if target_tokens is None else int(target_tokens) + target = max(sink_end + geometry.block_size, target) + if target >= projected: + # Nothing to reclaim: the caller must finish the session with + # context_length_exceeded instead of moving the window start backwards. + return None + + drop_needed = projected - target + if unit_tokens: + boundaries = unit_starts(geometry, unit_tokens) + # To avoid splitting units in half, moved_from must land on an exact unit start + # that is ALSO a page boundary (b % block_size == 0) and at or after sink_end. + # We need moved_from >= sink_end + drop_needed so post-trim length <= target. + min_cut = sink_end + drop_needed + candidates = [b for b in boundaries if b >= min_cut and b % geometry.block_size == 0 and b < projected] + if not candidates: + # No page-aligned whole-unit boundary satisfies the drop budget without + # slicing into a unit. Return None so caller routes to fallback. + return None + moved_from = min(candidates) + delta = moved_from - sink_end + else: + delta = align_up(drop_needed, geometry.block_size) + moved_from = sink_end + delta + if moved_from >= projected: + return None + + return PositionReanchor( + delta=delta, + moved_from=moved_from, + sink_blocks=sink_end_block, + ) 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 e38549371..b34cc5bd0 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 @@ -110,6 +110,37 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp self.model = self.thinker self.talker = None + if getattr(getattr(vllm_config, "model_config", None), "session_mode", None) == "duplex": + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import ( + DUPLEX_WINDOW_BLOCK_SIZE, + duplex_window_geometry, + install_duplex_window_layers, + validate_duplex_window_install, + ) + + cache_config = getattr(vllm_config, "cache_config", None) + model_config = getattr(vllm_config, "model_config", None) + block_size = int( + getattr(cache_config, "block_size", DUPLEX_WINDOW_BLOCK_SIZE) or DUPLEX_WINDOW_BLOCK_SIZE + ) + max_model_len = getattr(model_config, "max_model_len", None) if model_config is not None else None + if max_model_len is None: + max_model_len = 8192 + + geometry = duplex_window_geometry( + prefix_tokens=96, + window_tokens=6000, + block_size=block_size, + max_model_len=max_model_len, + high_watermark_tokens=8000, + ) + install_duplex_window_layers(self.thinker, geometry=geometry) + validate_duplex_window_install( + cache_config, + model_config, + geometry, + ) + elif self.model_stage == "tts": self.thinker = None # The Talker is always the runner-owned continuous codec producer. @@ -162,6 +193,14 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp return Sampler() + def apply_duplex_kv_reanchor(self, runner: Any, scheduler_output: Any = None) -> None: + """Apply in-place Stage-0 KV reanchor and rotation on worker before model forward.""" + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import ( + MiniCPMO45DuplexWorkerHelper, + ) + + MiniCPMO45DuplexWorkerHelper.maybe_apply_reanchor(runner, scheduler_output=scheduler_output) + def prepare_duplex_sampling( self, logits: torch.Tensor, @@ -387,6 +426,9 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp 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 + ), ) update_result = dict(result) if result.get("stage0_window_replaced") is True: diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 0a388d1f9..28af89a6f 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -468,8 +468,54 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex def _update_states(self, scheduler_output: SchedulerOutput) -> Callable | None: deferred_state_corrections_fn = super()._update_states(scheduler_output) self._update_duplex_sampling_states(scheduler_output) + self._maybe_apply_duplex_window_reanchor(scheduler_output) return deferred_state_corrections_fn + def _maybe_apply_duplex_window_reanchor(self, scheduler_output: SchedulerOutput | None = None) -> None: + """Apply in-place KV reanchor and rotation on worker before model forward.""" + reanchor_hook = getattr(getattr(self, "model", None), "apply_duplex_kv_reanchor", None) + if callable(reanchor_hook): + try: + reanchor_hook(self, scheduler_output=scheduler_output) + except TypeError: + reanchor_hook(self) + return + + helper = getattr(self, "_duplex_window_helper", None) + if helper is not None and hasattr(helper, "maybe_apply_reanchor"): + try: + helper.maybe_apply_reanchor(self, scheduler_output=scheduler_output) + except TypeError: + helper.maybe_apply_reanchor(self) + return + + # Fallback for dynamic runner inspection without hardcoding model classes + if not hasattr(self, "input_batch") or self.input_batch is None: + return + num_reqs = getattr(self.input_batch, "num_reqs", 0) + req_ids = self.input_batch.req_ids[:num_reqs] + has_reanchor = False + for req_id in req_ids: + info = self.model_intermediate_buffer.get(req_id) + if isinstance(info, dict) and isinstance(info.get("duplex"), dict): + if "stage0_reanchor" in info["duplex"]: + has_reanchor = True + break + if not has_reanchor: + return + + model_module = getattr(getattr(self, "model", None), "__module__", "") + if "minicpmo_4_5" in model_module: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import ( + MiniCPMO45DuplexWorkerHelper, + ) + + MiniCPMO45DuplexWorkerHelper.maybe_apply_reanchor(self, scheduler_output=scheduler_output) + + def _maybe_apply_stage0_reanchor(self) -> None: + """Backward-compatible alias for _maybe_apply_duplex_window_reanchor.""" + self._maybe_apply_duplex_window_reanchor() + def _request_final_stage_id(self, req_id: str) -> int | None: info = self.model_intermediate_buffer.get(req_id) if not isinstance(info, dict):