[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

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,

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,

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()

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,
+    )

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:

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):