[Performance][MiniCPM-o] Zero-copy block-aligned KV sliding window with in-place Re-RoPE for duplex Stage-0 (#7821)

Signed-off-by: BeatSeat <wendavid552@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

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)

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