From 7f75d88d3e1673a42d7f3b6665f3da52a122c54c Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 07:02:14 +0000
Subject: [PATCH 01/20] Add engine resources package with KV pool segment
 lifecycle

---
 mstar/engine/resources/__init__.py            |  16 ++
 mstar/engine/resources/base.py                |  59 +++++
 mstar/engine/resources/kv_pool.py             | 119 +++++++++
 test/modular/test_kv_pool.py                  | 241 ++++++++++++++++++
 .../test_page_allocator_thread_safety.py      |  92 +++----
 5 files changed, 482 insertions(+), 45 deletions(-)
 create mode 100644 mstar/engine/resources/__init__.py
 create mode 100644 mstar/engine/resources/base.py
 create mode 100644 mstar/engine/resources/kv_pool.py
 create mode 100644 test/modular/test_kv_pool.py

diff --git a/mstar/engine/resources/__init__.py b/mstar/engine/resources/__init__.py
new file mode 100644
index 000000000..4d2d5da31
--- /dev/null
+++ b/mstar/engine/resources/__init__.py
@@ -0,0 +1,16 @@
+from mstar.engine.resources.base import (
+    PositionPlan,
+    Reservation,
+    Segment,
+    SequenceView,
+)
+from mstar.engine.resources.kv_pool import KVCachePool, PageArena
+
+__all__ = [
+    "KVCachePool",
+    "PageArena",
+    "PositionPlan",
+    "Reservation",
+    "Segment",
+    "SequenceView",
+]
diff --git a/mstar/engine/resources/base.py b/mstar/engine/resources/base.py
new file mode 100644
index 000000000..de3f4f6c3
--- /dev/null
+++ b/mstar/engine/resources/base.py
@@ -0,0 +1,59 @@
+"""Boundary types for engine resources.
+
+A resource (KV cache pool, attention manager, positional embedder) owns one
+piece of per-step machinery. Everything passed between resources is one of
+the immutable values below, valid for a single step.
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from mstar.engine.resources.kv_pool import KVCachePool
+
+
+@dataclass(frozen=True)
+class Segment:
+    """One step's addition to a request's cache stream.
+
+    A request contributes one segment per label active for it in a step;
+    the batch's ordered segment list defines the layout of per-token
+    arrays. ``span`` may be 0: a zero-span segment reads its stream
+    without extending it (admission reserves nothing, commit is a no-op).
+    """
+    request_id: str
+    label: str
+    span: int
+
+
+@dataclass(frozen=True)
+class Reservation:
+    """A KV cache pool's answer to admitting one segment: how much of the
+    span is already resident, how much must actually be computed, and
+    whether residency is still being established asynchronously."""
+    resident: int
+    to_compute: int
+    pending: bool = False
+
+
+@dataclass(frozen=True)
+class SequenceView:
+    """What a pool holds for one segment's stream, after admission: the
+    storage the page table indexes into (by pool reference), the page table
+    itself, and the logical extent it covers. The extent is the positions
+    ``[start, start + length)``, not necessarily a prefix from zero."""
+    pool: "KVCachePool"
+    page_indices: tuple[int, ...]
+    start: int
+    length: int
+
+
+@dataclass(frozen=True)
+class PositionPlan:
+    """A positional embedder's output for one step: position identifiers in
+    segment order (in whatever shape the scheme requires), and the amount by
+    which each segment's position counter advances at commit."""
+    pos_ids: torch.Tensor
+    advance: tuple[int, ...]
diff --git a/mstar/engine/resources/kv_pool.py b/mstar/engine/resources/kv_pool.py
new file mode 100644
index 000000000..415660c54
--- /dev/null
+++ b/mstar/engine/resources/kv_pool.py
@@ -0,0 +1,119 @@
+"""KV cache storage resources.
+
+Storage divides in two. The arena is the physical pool: the backing tensor
+and the page allocator over it. The pool is per-request accounting against
+an arena: which pages a stream holds, its stored length, its position
+counter. Several pools may share one arena; nothing above a pool sees the
+arena.
+"""
+
+import torch
+
+from mstar.engine.kv_store import PageAllocator, PagedAllocationManager
+from mstar.engine.resources.base import Reservation, Segment, SequenceView
+
+
+class PageArena:
+    """Physical page storage: a backing K/V tensor and the allocator over
+    its pages. Pools draw pages from here and return them here."""
+
+    def __init__(
+        self,
+        tensor: torch.Tensor | None,
+        allocator: PageAllocator,
+        page_size: int,
+    ):
+        self.tensor = tensor
+        self.allocator = allocator
+        self.page_size = page_size
+
+    def allocate(self, n: int) -> list[int]:
+        return self.allocator.allocate(n)
+
+    def try_allocate(self, n: int) -> list[int] | None:
+        return self.allocator.try_allocate(n)
+
+    def free(self, pages: list[int]) -> None:
+        self.allocator.free(pages)
+
+    @property
+    def num_free(self) -> int:
+        return self.allocator.num_free
+
+    @property
+    def total_pages(self) -> int:
+        return self.allocator.max_num_pages
+
+
+class KVCachePool:
+    """Per-request cache accounting behind the segment lifecycle.
+
+    ``admit`` reserves capacity for a segment and reports what is already
+    resident; ``view`` describes the stream the segment extends; ``commit``
+    advances the stream's stored length and position counter. The
+    ``PagedAllocationManager`` this pool fronts remains the storage owner
+    (its lock, its request states, its transfer machinery); the pool is the
+    surface planning code goes through, so callers stop reaching into
+    request-state internals.
+    """
+
+    def __init__(self, manager: PagedAllocationManager):
+        self._manager = manager
+        self._arena = PageArena(
+            tensor=manager.kv_cache,
+            allocator=manager.page_allocator,
+            page_size=manager.config.page_size,
+        )
+
+    @property
+    def page_size(self) -> int:
+        return self._arena.page_size
+
+    @property
+    def num_free_pages(self) -> int:
+        return self._arena.num_free
+
+    @property
+    def total_pages(self) -> int:
+        return self._arena.total_pages
+
+    def admit(self, segment: Segment) -> Reservation:
+        """Reserve pages so the segment's stream can hold its history plus
+        this segment's span. Raises ``AllocationFailedError`` when the arena
+        cannot supply the pages; a zero-span segment reserves nothing."""
+        state = self._manager.get_state(segment.request_id, segment.label)
+        resident = state.seq_len
+        self._manager.alloc(
+            segment.request_id, segment.label, resident + segment.span
+        )
+        return Reservation(
+            resident=resident,
+            to_compute=segment.span,
+            pending=state.read_in_progress,
+        )
+
+    def view(self, segment: Segment) -> SequenceView:
+        """The stream as this step's plans must see it: every page backing
+        it and the extent those pages cover once the segment's span lands.
+        Call after ``admit`` for spans that need new pages."""
+        state = self._manager.get_state(segment.request_id, segment.label)
+        return SequenceView(
+            pool=self,
+            page_indices=tuple(state.page_indices),
+            start=0,
+            length=state.seq_len + segment.span,
+        )
+
+    def commit(self, segment: Segment, pos_advance: int | None = None) -> None:
+        """Record that the segment's span was computed: stored length grows
+        by the span, the position counter by ``pos_advance`` (defaults to
+        the span)."""
+        state = self._manager.get_state(segment.request_id, segment.label)
+        state.seq_len += segment.span
+        state.position_id_start += (
+            segment.span if pos_advance is None else pos_advance
+        )
+
+    def positions(self, request_id: str, label: str) -> int:
+        """Current position counter for one stream, read-only."""
+        return self._manager.get_state(request_id, label).position_id_start
diff --git a/test/modular/test_kv_pool.py b/test/modular/test_kv_pool.py
new file mode 100644
index 000000000..5c3201565
--- /dev/null
+++ b/test/modular/test_kv_pool.py
@@ -0,0 +1,241 @@
+"""Unit tests for the KV cache pool and its boundary types.
+
+``KVCachePool`` is the segment lifecycle over per-request cache storage:
+``admit`` reserves pages for a segment, ``view`` describes the stream it
+extends, ``commit`` advances stored length and position counter. The values
+crossing the boundary (``Segment``, ``Reservation``, ``SequenceView``,
+``PositionPlan``) are immutable and valid for one step.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import sys
+import threading
+
+sys.path.insert(0, ".")
+
+import pytest
+import torch
+
+from mstar.engine.kv_store import (
+    AllocationFailedError,
+    KVCacheConfig,
+    PageAllocator,
+    PagedAllocationManager,
+    StoreWritePolicy,
+)
+from mstar.engine.resources import (
+    KVCachePool,
+    PageArena,
+    PositionPlan,
+    Reservation,
+    Segment,
+    SequenceView,
+)
+
+
+def _make_pool(max_num_pages: int = 16, page_size: int = 8) -> tuple[KVCachePool, PagedAllocationManager]:
+    manager = PagedAllocationManager.__new__(PagedAllocationManager)
+    manager.config = KVCacheConfig(
+        num_layers=1,
+        num_kv_heads=1,
+        head_dim=1,
+        max_seq_len=max_num_pages * page_size,
+        max_num_pages=max_num_pages,
+        page_size=page_size,
+    )
+    manager.page_allocator = PageAllocator(max_num_pages)
+    manager.request_states = {}
+    manager.kv_cache = None
+    manager.write_policy = StoreWritePolicy.ALWAYS
+    manager._kv_transfer_engine = None
+    manager._offload_stream = None
+    manager.pending_reads = {}
+    manager._lock = threading.RLock()
+    return KVCachePool(manager), manager
+
+
+class TestAdmit:
+    def test_admit_reserves_whole_pages(self):
+        pool, manager = _make_pool(page_size=8)
+        manager.add_request("r", ["main"])
+
+        reservation = pool.admit(Segment("r", "main", 20))  # 3 pages of 8
+        assert reservation == Reservation(resident=0, to_compute=20, pending=False)
+        assert len(pool.view(Segment("r", "main", 0)).page_indices) == 3
+        assert pool.num_free_pages == 13
+
+    def test_admit_grows_from_resident_length(self):
+        pool, manager = _make_pool(page_size=8)
+        manager.add_request("r", ["main"])
+
+        first = Segment("r", "main", 8)  # exactly one page
+        pool.admit(first)
+        pool.commit(first)
+
+        # 5 more tokens cross into a second page.
+        reservation = pool.admit(Segment("r", "main", 5))
+        assert reservation.resident == 8
+        assert reservation.to_compute == 5
+        assert len(pool.view(Segment("r", "main", 0)).page_indices) == 2
+
+    def test_admit_within_last_page_reserves_nothing(self):
+        pool, manager = _make_pool(page_size=8)
+        manager.add_request("r", ["main"])
+
+        first = Segment("r", "main", 5)
+        pool.admit(first)
+        pool.commit(first)
+        free_before = pool.num_free_pages
+
+        pool.admit(Segment("r", "main", 3))  # fills page to exactly 8
+        assert pool.num_free_pages == free_before
+
+    def test_zero_span_admit_reserves_nothing(self):
+        pool, manager = _make_pool()
+        manager.add_request("r", ["main"])
+        free_before = pool.num_free_pages
+
+        reservation = pool.admit(Segment("r", "main", 0))
+        assert reservation == Reservation(resident=0, to_compute=0, pending=False)
+        assert pool.num_free_pages == free_before
+
+    def test_admit_failure_carries_diagnostics(self):
+        pool, manager = _make_pool(max_num_pages=2, page_size=8)
+        manager.add_request("r", ["main"])
+
+        with pytest.raises(AllocationFailedError) as excinfo:
+            pool.admit(Segment("r", "main", 100))  # needs 13 pages, has 2
+        assert excinfo.value.pages_short == 11
+        assert excinfo.value.request_id == "r"
+        assert excinfo.value.label == "main"
+        # A failed admit must not leak partial reservations.
+        assert pool.num_free_pages == 2
+
+    def test_admit_reports_pending_residency(self):
+        pool, manager = _make_pool()
+        manager.add_request("r", ["main"])
+        manager.get_state("r", "main").read_in_progress = True
+
+        assert pool.admit(Segment("r", "main", 0)).pending is True
+
+
+class TestViewAndCommit:
+    def test_view_extent_covers_resident_plus_span(self):
+        pool, manager = _make_pool(page_size=8)
+        manager.add_request("r", ["main"])
+
+        first = Segment("r", "main", 10)
+        pool.admit(first)
+        view = pool.view(first)
+        assert view.start == 0
+        assert view.length == 10
+        assert view.pool is pool
+
+        pool.commit(first)
+        # Before the next step's admit, a zero-span view sees the committed
+        # stream.
+        assert pool.view(Segment("r", "main", 0)).length == 10
+
+    def test_view_pages_are_an_immutable_copy(self):
+        pool, manager = _make_pool(page_size=8)
+        manager.add_request("r", ["main"])
+
+        segment = Segment("r", "main", 20)
+        pool.admit(segment)
+        view = pool.view(segment)
+        assert isinstance(view.page_indices, tuple)
+        # Mutating the live state afterwards must not change the view.
+        pages_at_plan_time = view.page_indices
+        manager.reset_label("r", "main")
+        assert view.page_indices == pages_at_plan_time
+
+    def test_commit_advances_length_and_positions_together(self):
+        pool, manager = _make_pool()
+        manager.add_request("r", ["main"])
+
+        segment = Segment("r", "main", 7)
+        pool.admit(segment)
+        pool.commit(segment)
+        assert pool.view(Segment("r", "main", 0)).length == 7
+        assert pool.positions("r", "main") == 7
+
+    def test_commit_with_position_override(self):
+        """Steps whose position span differs from their token count (vision
+        grids, single-position image blocks) advance positions by an explicit
+        amount."""
+        pool, manager = _make_pool()
+        manager.add_request("r", ["main"])
+
+        segment = Segment("r", "main", 6)
+        pool.admit(segment)
+        pool.commit(segment, pos_advance=1)
+        assert pool.view(Segment("r", "main", 0)).length == 6
+        assert pool.positions("r", "main") == 1
+
+    def test_streams_advance_independently(self):
+        pool, manager = _make_pool()
+        manager.add_request("r", ["main", "cfg"])
+
+        main_seg = Segment("r", "main", 4)
+        pool.admit(main_seg)
+        pool.commit(main_seg)
+        assert pool.view(Segment("r", "cfg", 0)).length == 0
+        assert pool.positions("r", "cfg") == 0
+
+    def test_write_once_context_flow(self):
+        """The cross-attention shape: one admitted write, position counter
+        pinned, then zero-span reads."""
+        pool, manager = _make_pool(page_size=8)
+        manager.add_request("r", ["ctx"])
+
+        segment = Segment("r", "ctx", 12)
+        pool.admit(segment)
+        pool.commit(segment, pos_advance=0)
+
+        view = pool.view(Segment("r", "ctx", 0))
+        assert view.length == 12
+        assert len(view.page_indices) == 2
+        assert pool.positions("r", "ctx") == 0
+
+
+class TestPageArena:
+    def test_allocate_free_roundtrip(self):
+        arena = PageArena(tensor=None, allocator=PageAllocator(4), page_size=8)
+        assert arena.total_pages == 4
+        pages = arena.allocate(3)
+        assert arena.num_free == 1
+        arena.free(pages)
+        assert arena.num_free == 4
+
+    def test_try_allocate_shortfall_returns_none(self):
+        arena = PageArena(tensor=None, allocator=PageAllocator(2), page_size=8)
+        assert arena.try_allocate(3) is None
+        assert arena.num_free == 2
+
+
+class TestBoundaryValuesAreImmutable:
+    def test_segment_frozen(self):
+        segment = Segment("r", "main", 4)
+        with pytest.raises(dataclasses.FrozenInstanceError):
+            segment.span = 5
+
+    def test_reservation_frozen(self):
+        reservation = Reservation(resident=1, to_compute=2)
+        with pytest.raises(dataclasses.FrozenInstanceError):
+            reservation.resident = 3
+
+    def test_sequence_view_frozen(self):
+        view = SequenceView(pool=None, page_indices=(1, 2), start=0, length=9)
+        with pytest.raises(dataclasses.FrozenInstanceError):
+            view.length = 10
+
+    def test_position_plan_frozen(self):
+        plan = PositionPlan(pos_ids=torch.arange(3), advance=(3,))
+        with pytest.raises(dataclasses.FrozenInstanceError):
+            plan.advance = (4,)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))
diff --git a/test/modular/test_page_allocator_thread_safety.py b/test/modular/test_page_allocator_thread_safety.py
index cbd007e3e..cc97b7e62 100644
--- a/test/modular/test_page_allocator_thread_safety.py
+++ b/test/modular/test_page_allocator_thread_safety.py
@@ -1,23 +1,20 @@
-"""Thread-safety tests for ``PageAllocator`` and ``PagedAllocationManager``.
-
-PR #78 issue #4: under speculative scheduling, the plan thread runs ``alloc`` while
-the GPU thread runs ``reset_label``. The previous implementation had two
-unprotected races:
-
-1. ``PageAllocator.try_allocate`` was ``qsize() < n`` followed by ``n``
-   ``get()`` calls — non-atomic. A concurrent ``free`` could land between
-   the qsize check and the get loop, false-negating ``try_allocate`` (it
-   returns ``None`` even though pages are now available).
-
-2. ``PagedAllocationManager.alloc`` and ``reset_label`` both touch
-   ``request_states[rid][label]``. If they interleaved, a freshly allocated
-   page list could be freed by a concurrent ``reset_label``, or a page
-   list freed by ``reset_label`` could be re-extended by a still-running
-   ``alloc`` on the now-stale state object.
-
-The fix added per-allocator + per-manager locks. These tests exercise the
-race shapes directly with a thread pool and verify page conservation
-across stress runs.
+"""Thread-safety tests for ``PageAllocator`` and ``KVCachePool.admit``.
+
+Under speculative scheduling, the plan thread reserves pages (today via
+``KVCachePool.admit``, which fronts ``PagedAllocationManager.alloc``) while
+the GPU thread runs ``reset_label``. Two race shapes matter:
+
+1. ``PageAllocator.try_allocate`` is qsize-then-get; without the allocator
+   lock a concurrent ``free`` between the two could false-negate the
+   allocation or hand two threads the same page.
+
+2. ``admit`` and ``reset_label`` both touch ``request_states[rid][label]``.
+   Unsynchronized, a freshly reserved page list could be freed by a
+   concurrent ``reset_label``, or a list freed by ``reset_label`` could be
+   re-extended through a stale state reference.
+
+The allocator and manager locks close both; these tests exercise the race
+shapes directly with a thread pool and verify page conservation.
 """
 
 from __future__ import annotations
@@ -30,12 +27,12 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
 from mstar.engine.kv_store import (
-    AllocationStatus,
     KVCacheConfig,
     PageAllocator,
     PagedAllocationManager,
     StoreWritePolicy,
 )
+from mstar.engine.resources import KVCachePool, Segment
 
 
 def _make_test_manager(max_num_pages: int = 32, page_size: int = 8) -> PagedAllocationManager:
@@ -58,7 +55,6 @@ def _make_test_manager(max_num_pages: int = 32, page_size: int = 8) -> PagedAllo
     manager.write_policy = StoreWritePolicy.ALWAYS
     manager._kv_transfer_engine = None
     manager._offload_stream = None
-    manager.alloc_status = AllocationStatus()
     manager.pending_reads = {}
     manager._lock = threading.RLock()
     return manager
@@ -178,9 +174,9 @@ def producer():
         assert alloc.num_free == max_pages
 
 
-class TestPagedAllocationManagerThreadSafety:
-    def test_concurrent_alloc_reset_conserves_pages(self):
-        """Plan-thread ``alloc`` racing GPU-thread ``reset_label`` for
+class TestKVCachePoolThreadSafety:
+    def test_concurrent_admit_reset_conserves_pages(self):
+        """Plan-thread ``admit`` racing GPU-thread ``reset_label`` for
         the same (rid, label) must leave the page pool fully drained
         once both stop. Without the manager lock, the race shape is:
             T1: state = request_states[rid][label]   # old ref
@@ -189,21 +185,22 @@ def test_concurrent_alloc_reset_conserves_pages(self):
             → leaked pages, dict has empty new state.
         """
         manager = _make_test_manager(max_num_pages=64, page_size=8)
+        pool = KVCachePool(manager)
         rid = "rid"
         label = "main"
         manager.add_request(rid, [label])
 
         n_iters = 300
-        # Smaller seq_len so each alloc only takes a couple pages.
-        seq_len_seq = [8, 16, 24, 16, 8]
+        # Small spans so each admit only takes a couple pages.
+        span_seq = [8, 16, 24, 16, 8]
 
-        def alloc_worker():
+        def admit_worker():
             for i in range(n_iters):
                 try:
-                    manager.alloc(rid, label, seq_len_seq[i % len(seq_len_seq)])
+                    pool.admit(Segment(rid, label, span_seq[i % len(span_seq)]))
                 except (KeyError, RuntimeError):
                     # KeyError if reset_label wiped the entry between
-                    # add_request and alloc; RuntimeError if pool empty.
+                    # add_request and admit; RuntimeError if pool empty.
                     pass
 
         def reset_worker():
@@ -215,8 +212,8 @@ def reset_worker():
 
         with ThreadPoolExecutor(max_workers=4) as ex:
             futures = [
-                ex.submit(alloc_worker),
-                ex.submit(alloc_worker),
+                ex.submit(admit_worker),
+                ex.submit(admit_worker),
                 ex.submit(reset_worker),
                 ex.submit(reset_worker),
             ]
@@ -225,16 +222,17 @@ def reset_worker():
 
         # Final reset to drain whatever's still allocated.
         manager.reset_label(rid, label)
-        assert manager.page_allocator.num_free == manager.config.max_num_pages
+        assert pool.num_free_pages == pool.total_pages
         # request_states must still contain a valid (empty) state.
         assert label in manager.request_states[rid]
         assert manager.request_states[rid][label].page_indices == []
 
     def test_concurrent_add_remove_request_conserves_pages(self):
-        """Multiple threads cycling add_request → alloc → remove_request
+        """Multiple threads cycling add_request → admit → remove_request
         must conserve pages. Stresses the request-lifecycle locking.
         """
         manager = _make_test_manager(max_num_pages=128, page_size=8)
+        pool = KVCachePool(manager)
         n_threads = 8
         n_iters = 50
 
@@ -243,7 +241,7 @@ def worker(thread_idx: int):
                 rid = f"rid_{thread_idx}_{i}"
                 manager.add_request(rid, ["main"])
                 try:
-                    manager.alloc(rid, "main", seq_len=16)
+                    pool.admit(Segment(rid, "main", 16))
                 except RuntimeError:
                     pass  # pool exhausted, ok
                 manager.remove_request(rid)
@@ -253,28 +251,32 @@ def worker(thread_idx: int):
             for f in as_completed(futures):
                 f.result()
 
-        assert manager.page_allocator.num_free == manager.config.max_num_pages
+        assert pool.num_free_pages == pool.total_pages
         assert manager.request_states == {}
         assert manager.pending_reads == {}
 
-    def test_alloc_then_reset_releases_correct_pages(self):
-        """Single-threaded sanity: confirm the lock didn't break the
-        normal alloc/free contract.
+    def test_admit_then_reset_releases_correct_pages(self):
+        """Single-threaded sanity: confirm the locking didn't break the
+        normal reserve/release contract.
         """
         manager = _make_test_manager(max_num_pages=16, page_size=8)
+        pool = KVCachePool(manager)
         rid = "rid"
         manager.add_request(rid, ["main"])
 
-        manager.alloc(rid, "main", seq_len=24)  # 3 pages
-        assert len(manager.request_states[rid]["main"].page_indices) == 3
-        assert manager.page_allocator.num_free == 13
+        segment = Segment(rid, "main", 24)  # 3 pages
+        reservation = pool.admit(segment)
+        assert reservation.resident == 0
+        assert reservation.to_compute == 24
+        assert len(pool.view(segment).page_indices) == 3
+        assert pool.num_free_pages == 13
 
         manager.reset_label(rid, "main")
-        assert manager.request_states[rid]["main"].page_indices == []
-        assert manager.page_allocator.num_free == 16
+        assert pool.view(Segment(rid, "main", 0)).page_indices == ()
+        assert pool.num_free_pages == 16
 
         manager.remove_request(rid)
-        assert manager.page_allocator.num_free == 16
+        assert pool.num_free_pages == 16
 
 
 if __name__ == "__main__":

From d81690e29f3e717354556eb4309c635f7bf42a94 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 07:02:14 +0000
Subject: [PATCH 02/20] Route attention planning through KV pool views and
 admits

---
 mstar/engine/cache_manager.py | 103 ++++++++++++++++++++--------------
 1 file changed, 61 insertions(+), 42 deletions(-)

diff --git a/mstar/engine/cache_manager.py b/mstar/engine/cache_manager.py
index 557083e24..fb21953a6 100644
--- a/mstar/engine/cache_manager.py
+++ b/mstar/engine/cache_manager.py
@@ -1,6 +1,7 @@
 import functools
 import logging
 from abc import ABC, abstractmethod
+from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import NamedTuple
 
@@ -12,6 +13,7 @@
     KVRequestState,
     PagedAllocationManager,
 )
+from mstar.engine.resources import KVCachePool, Segment
 from mstar.utils.flashinfer_utils import FlashInferDecodeWrapper, FlashInferPrefillWrapper
 
 logger = logging.getLogger(__name__)
@@ -35,7 +37,7 @@ class PagedIndptrs(NamedTuple):
 
 def build_paged_indptrs(
     q_seq_lens: list[int],
-    page_indices_per_request: list[list[int]],
+    page_indices_per_request: list[Sequence[int]],
     context_lens: list[int],
     page_size: int,
 ) -> PagedIndptrs:
@@ -190,6 +192,17 @@ def __init__(
         # source name -> CrossAttnPool (see KVCacheConfig.cross_attn)
         self.cross_pools = cross_pools or {}
 
+        # Segment-lifecycle surface over the allocation manager: planning
+        # goes through admit/view on these pools instead of reading
+        # KVRequestState fields and allocating mid-plan.
+        self.kv_pool = (
+            KVCachePool(alloc_manager) if alloc_manager is not None else None
+        )
+        self.cross_kv_pools: dict[str, KVCachePool] = {
+            source: KVCachePool(pool.alloc_manager)
+            for source, pool in self.cross_pools.items()
+        }
+
         self.auto_write_store = auto_write_store
 
         # CUDA graph mode: persistent wrappers passed in from CudaGraphRunner.
@@ -693,9 +706,9 @@ def plan_attention(
     ):
         """Pre-compute FlashInfer plan and page positions for a cache label.
 
-        Allocates pages, computes page_indices/page_offsets/token_offsets for
-        vectorized KV writes, builds FlashInfer index tensors, and plans the
-        wrapper. All state is stored in _plan_states[label].
+        Admits one segment per request on the KV pool (reserving pages),
+        builds FlashInfer index tensors from the resulting sequence views,
+        and plans the wrapper. All state is stored in _plan_states[label].
 
         In CUDA graph mode, uses the persistent wrapper from _plan_states
         (pre-built by CudaGraphRunner) and calls its plan() method which
@@ -784,22 +797,18 @@ def _plan_attention_impl(
             kv_cache_locations_list = []
 
             for i, rid in enumerate(self.request_ids):
-                state = self._get_state(rid, effective_label)
-                sl = seq_lens[i]
-                total_len = state.seq_len + sl
-
-                self.alloc_manager.alloc(
-                    rid, label=effective_label, seq_len=total_len
-                )
+                segment = Segment(rid, effective_label, seq_lens[i])
+                self.kv_pool.admit(segment)
+                view = self.kv_pool.view(segment)
 
-                qo_indptr_list.append(qo_indptr_list[-1] + sl)
-                all_page_indices.extend(state.page_indices)
-                kv_indptr_list.append(kv_indptr_list[-1] + len(state.page_indices))
+                qo_indptr_list.append(qo_indptr_list[-1] + segment.span)
+                all_page_indices.extend(view.page_indices)
+                kv_indptr_list.append(kv_indptr_list[-1] + len(view.page_indices))
 
-                last_page_len = total_len % page_size or page_size
+                last_page_len = view.length % page_size or page_size
                 kv_last_page_lens.append(last_page_len)
-                if sl == 1:
-                    kv_cache_locations_list.append([state.page_indices[-1], last_page_len - 1])
+                if segment.span == 1:
+                    kv_cache_locations_list.append([view.page_indices[-1], last_page_len - 1])
         finally:
             if self.enable_nvtx:
                 range_pop(synchronize=False)
@@ -940,21 +949,19 @@ def plan_attention_batched_cfg(
 
         for label in labels:
             for i, rid in enumerate(self.request_ids):
-                state = self._get_state(rid, label)
-                sl = seq_lens[label][i]
-                total_len = state.seq_len + sl
-
-                self.alloc_manager.alloc(rid, label=label, seq_len=total_len)
+                segment = Segment(rid, label, seq_lens[label][i])
+                self.kv_pool.admit(segment)
+                view = self.kv_pool.view(segment)
 
-                qo_indptr_list.append(qo_indptr_list[-1] + sl)
-                all_page_indices.extend(state.page_indices)
+                qo_indptr_list.append(qo_indptr_list[-1] + segment.span)
+                all_page_indices.extend(view.page_indices)
                 kv_indptr_list.append(
-                    kv_indptr_list[-1] + len(state.page_indices)
+                    kv_indptr_list[-1] + len(view.page_indices)
                 )
 
-                last_page_len = total_len % page_size or page_size
+                last_page_len = view.length % page_size or page_size
                 kv_last_page_lens.append(last_page_len)
-                combined_seq_lens.append(sl)
+                combined_seq_lens.append(segment.span)
 
         # CPU tensors — see comment in ``plan_attention`` above. FlashInfer
         # async-H2Ds these inside ``plan()``; passing GPU tensors would
@@ -1110,6 +1117,7 @@ def add_cross_attn_kv(
         allocated on first write (layer 0) and reused for the rest.
         """
         pool = self._get_cross_pool(source)
+        kv_pool = self.cross_kv_pools[source]
         base_label = self._active_base_label(label)
         cross_label = cross_attn_label(base_label, source)
         page_size = pool.alloc_config.page_size
@@ -1122,19 +1130,25 @@ def add_cross_attn_kv(
 
         offset = 0
         for rid, ctx_len in zip(request_ids, seq_lens, strict=True):
-            state = pool.alloc_manager.get_state(rid, cross_label)
-            if state.seq_len == 0:
-                pool.alloc_manager.alloc(rid, label=cross_label, seq_len=ctx_len)
-                state.seq_len = ctx_len
+            resident = kv_pool.view(Segment(rid, cross_label, 0)).length
+            if resident == 0:
+                # First write for this context: reserve its pages and commit
+                # the extent up front; later layers reuse them. The position
+                # counter stays untouched (context positions are baked in at
+                # encode time).
+                segment = Segment(rid, cross_label, ctx_len)
+                kv_pool.admit(segment)
+                kv_pool.commit(segment, pos_advance=0)
             else:
-                assert state.seq_len == ctx_len, (
+                assert resident == ctx_len, (
                     f"cross-attn context for {rid!r}/{source!r} already written "
-                    f"with length {state.seq_len}, got {ctx_len}"
+                    f"with length {resident}, got {ctx_len}"
                 )
+            view = kv_pool.view(Segment(rid, cross_label, 0))
 
             positions = torch.arange(ctx_len, device=self.device)
             page_indices = torch.tensor(
-                state.page_indices, dtype=torch.long, device=self.device,
+                view.page_indices, dtype=torch.long, device=self.device,
             )
             token_to_page = page_indices[
                 torch.div(positions, page_size, rounding_mode="floor")
@@ -1168,6 +1182,7 @@ def plan_cross_attention(
         the full context).
         """
         pool = self._get_cross_pool(source)
+        kv_pool = self.cross_kv_pools[source]
         base_label = self._active_base_label(label)
         cross_label = cross_attn_label(base_label, source)
         cfg = pool.alloc_config
@@ -1176,16 +1191,16 @@ def plan_cross_attention(
         if dtype is None:
             dtype = pool.kv_cache.dtype
 
-        page_indices_per_request: list[list[int]] = []
+        page_indices_per_request: list[tuple[int, ...]] = []
         context_lens: list[int] = []
         for rid in self.request_ids:
-            state = pool.alloc_manager.get_state(rid, cross_label)
-            assert state.seq_len > 0, (
+            view = kv_pool.view(Segment(rid, cross_label, 0))
+            assert view.length > 0, (
                 f"plan_cross_attention before add_cross_attn_kv for {rid!r} "
                 f"(source {source!r})"
             )
-            page_indices_per_request.append(state.page_indices)
-            context_lens.append(state.seq_len)
+            page_indices_per_request.append(view.page_indices)
+            context_lens.append(view.length)
 
         indptrs = build_paged_indptrs(
             q_seq_lens, page_indices_per_request, context_lens, page_size,
@@ -1418,16 +1433,20 @@ def _build_dense_gen_plan(
         max_k = 0
         for label in labels:
             for i, rid in enumerate(self.request_ids):
-                state = self._get_state(rid, label)
-                prefix_len = state.seq_len
+                # The frozen prefix is read without extending the stream (the
+                # generation K/V never enters the pages), so plan it from a
+                # zero-span view.
+                view = self.kv_pool.view(Segment(rid, label, 0))
+                prefix_len = view.length
                 gen_len = seq_lens[label][i]
                 n_pages = (prefix_len + page_size - 1) // page_size
                 idx = torch.tensor(
-                    state.page_indices[:n_pages], dtype=torch.long, device=self.device
+                    view.page_indices[:n_pages], dtype=torch.long, device=self.device
                 )
                 # Carry the persistent KVRequestState so run_attention can cache
                 # the gathered frozen prefix on it across denoise steps (the
                 # manager itself is rebuilt every forward).
+                state = self._get_state(rid, label)
                 segs.append((idx, prefix_len, gen_len, state))
                 cu_q.append(cu_q[-1] + gen_len)
                 cu_k.append(cu_k[-1] + prefix_len + gen_len)

From 854359accbf2aaca6a1285449537ff1f260ab5ef Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 09:31:18 +0000
Subject: [PATCH 03/20] Delegate rope planning to a position embedder and
 commit advances through the pool

---
 mstar/engine/cache_manager.py          | 118 +++++++---------
 mstar/engine/resources/__init__.py     |   2 +
 mstar/engine/resources/positions.py    |  85 ++++++++++++
 test/modular/test_position_embedder.py | 181 +++++++++++++++++++++++++
 4 files changed, 315 insertions(+), 71 deletions(-)
 create mode 100644 mstar/engine/resources/positions.py
 create mode 100644 test/modular/test_position_embedder.py

diff --git a/mstar/engine/cache_manager.py b/mstar/engine/cache_manager.py
index fb21953a6..08d1604de 100644
--- a/mstar/engine/cache_manager.py
+++ b/mstar/engine/cache_manager.py
@@ -13,7 +13,7 @@
     KVRequestState,
     PagedAllocationManager,
 )
-from mstar.engine.resources import KVCachePool, Segment
+from mstar.engine.resources import KVCachePool, RopeEmbedder, Segment
 from mstar.utils.flashinfer_utils import FlashInferDecodeWrapper, FlashInferPrefillWrapper
 
 logger = logging.getLogger(__name__)
@@ -202,6 +202,9 @@ def __init__(
             source: KVCachePool(pool.alloc_manager)
             for source, pool in self.cross_pools.items()
         }
+        # Position semantics live on the embedder; the pool's counters are
+        # read at plan time and advanced only through commit.
+        self.rope_embedder = RopeEmbedder()
 
         self.auto_write_store = auto_write_store
 
@@ -408,21 +411,19 @@ def _plan_rope_impl(
 
         computed_pos_ids = pos_ids
         if computed_pos_ids is None:
-            # CPU-accumulate the position list (1 int per output token). The
-            # old `torch.cat([torch.arange(...) + start for ...])` launched
-            # 2 GPU kernels per request.
+            # The embedder builds the position list on CPU (1 int per output
+            # token); placement is decided here.
             if self.enable_nvtx:
                 range_push("cache.plan_rope.build_pos_ids", synchronize=False)
             try:
-                pos_ids_list: list[int] = []
-                for rid, sl in zip(self.request_ids, seq_lens, strict=True):
-                    start = self._get_state(rid, effective_label).position_id_start
-                    pos_ids_list.extend(range(start, start + sl))
-                computed_pos_ids = torch.tensor(
-                    pos_ids_list,
-                    dtype=torch.long,
-                    device=None if static_copy_from_cpu else self.device,
-                )
+                segments = [
+                    Segment(rid, effective_label, sl)
+                    for rid, sl in zip(self.request_ids, seq_lens, strict=True)
+                ]
+                position_plan = self.rope_embedder.plan(segments, self.kv_pool)
+                computed_pos_ids = position_plan.pos_ids
+                if not static_copy_from_cpu:
+                    computed_pos_ids = computed_pos_ids.to(self.device)
             finally:
                 if self.enable_nvtx:
                     range_pop(synchronize=False)
@@ -482,13 +483,12 @@ def plan_rope_batched_cfg(
             if per_label_pos_ids and label in per_label_pos_ids:
                 parts.append(torch.cat(per_label_pos_ids[label]))
             else:
-                pos_ids_list: list[int] = []
-                for rid, sl in zip(self.request_ids, seq_lens[label], strict=True):
-                    start = self._get_state(rid, label).position_id_start
-                    pos_ids_list.extend(range(start, start + sl))
-                parts.append(torch.tensor(
-                    pos_ids_list, dtype=torch.long, device=self.device,
-                ))
+                segments = [
+                    Segment(rid, label, sl)
+                    for rid, sl in zip(self.request_ids, seq_lens[label], strict=True)
+                ]
+                plan = self.rope_embedder.plan(segments, self.kv_pool)
+                parts.append(plan.pos_ids.to(self.device))
         combined_pos_ids = parts[0] if len(parts) == 1 else torch.cat(parts)
         self._plan_states[combined_label].pos_ids = combined_pos_ids
 
@@ -510,43 +510,15 @@ def apply_rope(
         ps = self._plan_states[label]
         assert ps.pos_ids is not None
 
-        orig_dtype = q.dtype
-
-        if rope_dtype is not None:
-            q, k = q.to(rope_dtype), k.to(rope_dtype)
-        elif torch.is_autocast_enabled():
-            dtype = torch.get_autocast_gpu_dtype()
-            q, k = q.to(dtype), k.to(dtype)
-        elif q.dtype == torch.float32:
-            dtype = torch.bfloat16
-            q, k = q.to(dtype), k.to(dtype)
-
-        llama31_params = {}
-        for key, value in kwargs.items():
-            if key in ['low_freq_factor', 'high_freq_factor', 'old_context_len']:
-                llama31_params[key] = value
-
-        import flashinfer
-
-        if not llama31_params:
-            flashinfer.rope.apply_rope_pos_ids_inplace(
-                q, k, ps.pos_ids,
-                rotary_dim=rotary_dim,
-                interleave=interleave,
-                rope_scale=rope_scale,
-                rope_theta=rope_theta,
-
-            )
-        else:
-            flashinfer.rope.apply_llama31_rope_pos_ids_inplace(
-                q, k, ps.pos_ids,
-                rotary_dim=rotary_dim,
-                interleave=interleave,
-                rope_scale=rope_scale,
-                rope_theta=rope_theta,
-                **llama31_params
-            )
-        return q.to(orig_dtype), k.to(orig_dtype)
+        return self.rope_embedder.apply(
+            q, k, ps.pos_ids,
+            rotary_dim=rotary_dim,
+            interleave=interleave,
+            rope_scale=rope_scale,
+            rope_theta=rope_theta,
+            rope_dtype=rope_dtype,
+            **kwargs,
+        )
 
     @torch.compiler.disable
     def advance_seq_len(self, n: int | None = None, pos_id_n: int | None = None) -> None:
@@ -559,9 +531,11 @@ def advance_seq_len(self, n: int | None = None, pos_id_n: int | None = None) ->
         if n is None:
             return self.advance_seq_lens(pos_id_n)
         for rid in self.request_ids:
-            state = self._get_state(rid)
-            state.seq_len += n
-            state.position_id_start += (pos_id_n if pos_id_n is not None else n)
+            label = self.active_labels.get(rid, "main")
+            self.kv_pool.commit(
+                Segment(rid, label, n),
+                pos_advance=pos_id_n if pos_id_n is not None else n,
+            )
 
     @torch.compiler.disable
     def set_custom_pos_advance(
@@ -612,14 +586,15 @@ def advance_seq_lens(self, pos_id_ns: list[int] | int | None = None) -> None:
             for label, seq_lens in self._batched_cfg_info.per_label_seq_len.items():
                 for i, rid in enumerate(self.request_ids):
                     n = seq_lens[i]
-                    state = self._get_state(rid, label=label)
-                    state.seq_len += n
                     if pos_id_ns is None:
-                        state.position_id_start += n
+                        pos_advance = n
                     elif isinstance(pos_id_ns, int):
-                        state.position_id_start += pos_id_ns
+                        pos_advance = pos_id_ns
                     else:
-                        state.position_id_start += pos_id_ns[i]
+                        pos_advance = pos_id_ns[i]
+                    self.kv_pool.commit(
+                        Segment(rid, label, n), pos_advance=pos_advance
+                    )
         else:
             for i, rid in enumerate(self.request_ids):
                 label = self.active_labels[rid]
@@ -627,17 +602,18 @@ def advance_seq_lens(self, pos_id_ns: list[int] | int | None = None) -> None:
                 if ps.seq_lens is None:
                     continue
                 n = ps.seq_lens[i]
-                state = self._get_state(rid, label=label)
-                state.seq_len += n
                 if pos_id_ns is None:
                     if ps.custom_pos_advance is not None:
-                        state.position_id_start += ps.custom_pos_advance[i]
+                        pos_advance = ps.custom_pos_advance[i]
                     else:
-                        state.position_id_start += n
+                        pos_advance = n
                 elif isinstance(pos_id_ns, int):
-                    state.position_id_start += pos_id_ns
+                    pos_advance = pos_id_ns
                 else:
-                    state.position_id_start += pos_id_ns[i]
+                    pos_advance = pos_id_ns[i]
+                self.kv_pool.commit(
+                    Segment(rid, label, n), pos_advance=pos_advance
+                )
         # Clear the side-channel on every consumer so a stale value can't
         # bleed into a subsequent walk.
         for ps in self._plan_states.values():
diff --git a/mstar/engine/resources/__init__.py b/mstar/engine/resources/__init__.py
index 4d2d5da31..e77a4cf3c 100644
--- a/mstar/engine/resources/__init__.py
+++ b/mstar/engine/resources/__init__.py
@@ -5,12 +5,14 @@
     SequenceView,
 )
 from mstar.engine.resources.kv_pool import KVCachePool, PageArena
+from mstar.engine.resources.positions import RopeEmbedder
 
 __all__ = [
     "KVCachePool",
     "PageArena",
     "PositionPlan",
     "Reservation",
+    "RopeEmbedder",
     "Segment",
     "SequenceView",
 ]
diff --git a/mstar/engine/resources/positions.py b/mstar/engine/resources/positions.py
new file mode 100644
index 000000000..3b43ea9fb
--- /dev/null
+++ b/mstar/engine/resources/positions.py
@@ -0,0 +1,85 @@
+"""Positional embedding resources.
+
+A positional embedder owns position semantics: how position identifiers
+are built for a step, how they are applied to queries and keys, and how
+far each stream's position counter advances when the step commits (the
+plan's ``advance`` field, consumed by the pool's commit). Positions are
+read from the KV cache pool and never mutated outside commit.
+"""
+
+import torch
+
+from mstar.engine.resources.base import PositionPlan, Segment
+from mstar.engine.resources.kv_pool import KVCachePool
+
+
+class RopeEmbedder:
+    """Default 1D rotary scheme: one integer position per token, and the
+    position counter advances by each segment's span."""
+
+    def plan(self, segments: list[Segment], pool: KVCachePool) -> PositionPlan:
+        """Position identifiers for one step, in segment order, read from
+        the pool's counters. The tensor is built on CPU; the caller decides
+        placement (a captured path copies it into a static device buffer,
+        the eager path moves it to the device)."""
+        pos_ids: list[int] = []
+        advance: list[int] = []
+        for segment in segments:
+            start = pool.positions(segment.request_id, segment.label)
+            pos_ids.extend(range(start, start + segment.span))
+            advance.append(segment.span)
+        return PositionPlan(
+            pos_ids=torch.tensor(pos_ids, dtype=torch.long),
+            advance=tuple(advance),
+        )
+
+    def apply(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        pos_ids: torch.Tensor,
+        rotary_dim: int | None = None,
+        interleave: bool = False,
+        rope_scale: float = 1,
+        rope_theta: float = 10000.0,
+        rope_dtype=None,
+        **kwargs,
+    ):
+        """Rotate q and k by ``pos_ids``. Runs on device inside the step's
+        execution, so it allocates nothing beyond the dtype casts."""
+        orig_dtype = q.dtype
+
+        if rope_dtype is not None:
+            q, k = q.to(rope_dtype), k.to(rope_dtype)
+        elif torch.is_autocast_enabled():
+            dtype = torch.get_autocast_gpu_dtype()
+            q, k = q.to(dtype), k.to(dtype)
+        elif q.dtype == torch.float32:
+            dtype = torch.bfloat16
+            q, k = q.to(dtype), k.to(dtype)
+
+        llama31_params = {
+            key: value for key, value in kwargs.items()
+            if key in ("low_freq_factor", "high_freq_factor", "old_context_len")
+        }
+
+        import flashinfer
+
+        if not llama31_params:
+            flashinfer.rope.apply_rope_pos_ids_inplace(
+                q, k, pos_ids,
+                rotary_dim=rotary_dim,
+                interleave=interleave,
+                rope_scale=rope_scale,
+                rope_theta=rope_theta,
+            )
+        else:
+            flashinfer.rope.apply_llama31_rope_pos_ids_inplace(
+                q, k, pos_ids,
+                rotary_dim=rotary_dim,
+                interleave=interleave,
+                rope_scale=rope_scale,
+                rope_theta=rope_theta,
+                **llama31_params,
+            )
+        return q.to(orig_dtype), k.to(orig_dtype)
diff --git a/test/modular/test_position_embedder.py b/test/modular/test_position_embedder.py
new file mode 100644
index 000000000..a7fb56220
--- /dev/null
+++ b/test/modular/test_position_embedder.py
@@ -0,0 +1,181 @@
+"""Unit tests for the positional embedder and pool-committed advances.
+
+``RopeEmbedder.plan`` turns a segment list plus the pool's counters into a
+``PositionPlan``; the cache manager's advance paths resolve each step's
+position delta (default span, explicit ``pos_id_ns``, or the
+``custom_pos_advance`` side channel) and record it only through
+``KVCachePool.commit``.
+"""
+
+from __future__ import annotations
+
+import sys
+import threading
+
+sys.path.insert(0, ".")
+
+import torch
+
+from mstar.engine.cache_manager import BatchedCfgInfo, FlashInferCacheManager
+from mstar.engine.kv_store import (
+    KVCacheConfig,
+    PageAllocator,
+    PagedAllocationManager,
+    StoreWritePolicy,
+)
+from mstar.engine.resources import KVCachePool, RopeEmbedder, Segment
+
+
+def _make_manager(max_num_pages: int = 16, page_size: int = 8) -> PagedAllocationManager:
+    manager = PagedAllocationManager.__new__(PagedAllocationManager)
+    manager.config = KVCacheConfig(
+        num_layers=1,
+        num_kv_heads=1,
+        head_dim=1,
+        max_seq_len=max_num_pages * page_size,
+        max_num_pages=max_num_pages,
+        page_size=page_size,
+    )
+    manager.page_allocator = PageAllocator(max_num_pages)
+    manager.request_states = {}
+    manager.kv_cache = None
+    manager.write_policy = StoreWritePolicy.ALWAYS
+    manager._kv_transfer_engine = None
+    manager._offload_stream = None
+    manager.pending_reads = {}
+    manager._lock = threading.RLock()
+    return manager
+
+
+def _make_cache_manager(request_ids, labels=("main",)):
+    alloc = _make_manager()
+    for rid in request_ids:
+        alloc.add_request(rid, list(labels))
+    cm = FlashInferCacheManager(
+        request_ids=list(request_ids),
+        active_labels_per_request={rid: labels[0] for rid in request_ids},
+        kv_cache=None,
+        alloc_manager=alloc,
+        buffer_manager=None,
+        kv_cache_config=alloc.config,
+        device="cpu",
+    )
+    return cm
+
+
+def _seed_planned_seq_lens(cm, label, seq_lens):
+    """Set a label's planned seq_lens through the pre-planned fast path
+    (the same shortcut the worker's plan thread uses), which records them
+    without building a FlashInfer wrapper."""
+    cm.plan_rope(seq_lens=seq_lens, label=label)
+    cm._pre_planned_labels.add(label)
+    cm.plan_attention(seq_lens=seq_lens, label=label)
+
+
+class TestRopeEmbedderPlan:
+    def test_ids_follow_segment_order_and_counters(self):
+        alloc = _make_manager()
+        alloc.add_request("a", ["main"])
+        alloc.add_request("b", ["main"])
+        pool = KVCachePool(alloc)
+        pool.commit(Segment("a", "main", 5))  # counter at 5
+        embedder = RopeEmbedder()
+
+        plan = embedder.plan(
+            [Segment("a", "main", 3), Segment("b", "main", 2)], pool
+        )
+        assert plan.pos_ids.tolist() == [5, 6, 7, 0, 1]
+        assert plan.advance == (3, 2)
+        assert plan.pos_ids.device.type == "cpu"
+
+    def test_zero_span_segment_contributes_no_ids(self):
+        alloc = _make_manager()
+        alloc.add_request("a", ["main"])
+        pool = KVCachePool(alloc)
+        embedder = RopeEmbedder()
+
+        plan = embedder.plan([Segment("a", "main", 0)], pool)
+        assert plan.pos_ids.numel() == 0
+        assert plan.advance == (0,)
+
+
+class TestAdvanceCommitsThroughPool:
+    def test_plain_decode_advances_by_span(self):
+        cm = _make_cache_manager(["a", "b"])
+        _seed_planned_seq_lens(cm, "main", [1, 1])
+
+        cm.advance_seq_lens()
+        assert cm.kv_pool.positions("a", "main") == 1
+        assert cm.kv_pool.positions("b", "main") == 1
+        assert cm.kv_pool.view(Segment("a", "main", 0)).length == 1
+
+    def test_custom_pos_advance_overrides_span(self):
+        """The vision-prefill shape: position span larger than the token
+        count, delivered through the side channel."""
+        cm = _make_cache_manager(["a"])
+        _seed_planned_seq_lens(cm, "main", [4])
+        cm.set_custom_pos_advance([90], label="main")
+
+        cm.advance_seq_lens()
+        assert cm.kv_pool.view(Segment("a", "main", 0)).length == 4
+        assert cm.kv_pool.positions("a", "main") == 90
+        # The side channel is consumed, not persistent.
+        _seed_planned_seq_lens(cm, "main", [1])
+        cm.advance_seq_lens()
+        assert cm.kv_pool.positions("a", "main") == 91
+
+    def test_pos_id_ns_scalar_override(self):
+        """The image-block shape: many tokens, one position."""
+        cm = _make_cache_manager(["a"])
+        _seed_planned_seq_lens(cm, "main", [16])
+
+        cm.advance_seq_lens(pos_id_ns=1)
+        assert cm.kv_pool.view(Segment("a", "main", 0)).length == 16
+        assert cm.kv_pool.positions("a", "main") == 1
+
+    def test_pos_id_ns_per_request_list(self):
+        cm = _make_cache_manager(["a", "b"])
+        _seed_planned_seq_lens(cm, "main", [2, 2])
+
+        cm.advance_seq_lens(pos_id_ns=[5, 7])
+        assert cm.kv_pool.positions("a", "main") == 5
+        assert cm.kv_pool.positions("b", "main") == 7
+
+    def test_batched_cfg_advances_every_label(self):
+        cm = _make_cache_manager(["a"], labels=("main", "uncond"))
+        cm._batched_cfg_info = BatchedCfgInfo(
+            per_label_seq_len={"main": [3], "uncond": [3]}
+        )
+
+        cm.advance_seq_lens()
+        for label in ("main", "uncond"):
+            assert cm.kv_pool.view(Segment("a", label, 0)).length == 3
+            assert cm.kv_pool.positions("a", label) == 3
+
+    def test_advance_seq_len_singular(self):
+        cm = _make_cache_manager(["a"])
+
+        cm.advance_seq_len(4, pos_id_n=2)
+        assert cm.kv_pool.view(Segment("a", "main", 0)).length == 4
+        assert cm.kv_pool.positions("a", "main") == 2
+
+
+class TestPlanRopeUsesEmbedder:
+    def test_plan_rope_builds_ids_from_pool_counters(self):
+        cm = _make_cache_manager(["a"])
+        cm.kv_pool.commit(Segment("a", "main", 3))  # counter at 3
+
+        cm.plan_rope(seq_lens=[2], label="main")
+        assert cm._plan_states["main"].pos_ids.tolist() == [3, 4]
+
+    def test_explicit_pos_ids_pass_through(self):
+        cm = _make_cache_manager(["a"])
+        explicit = torch.tensor([7, 9], dtype=torch.long)
+
+        cm.plan_rope(seq_lens=[2], pos_ids=explicit, label="main")
+        assert cm._plan_states["main"].pos_ids.tolist() == [7, 9]
+
+
+if __name__ == "__main__":
+    import pytest
+    sys.exit(pytest.main([__file__, "-v"]))

From 2f004f6dd86f18733c1998d5962dcab771fa2f3f Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 10:46:32 +0000
Subject: [PATCH 04/20] Move attention plan machinery into resource managers

---
 mstar/engine/cache_manager.py       | 925 +++++++---------------------
 mstar/engine/resources/__init__.py  |  10 +
 mstar/engine/resources/attention.py | 720 ++++++++++++++++++++++
 mstar/engine/resources/kv_pool.py   |  34 +
 test/modular/test_kv_pool.py        |  90 ++-
 5 files changed, 1063 insertions(+), 716 deletions(-)
 create mode 100644 mstar/engine/resources/attention.py

diff --git a/mstar/engine/cache_manager.py b/mstar/engine/cache_manager.py
index 08d1604de..b0b165040 100644
--- a/mstar/engine/cache_manager.py
+++ b/mstar/engine/cache_manager.py
@@ -1,9 +1,7 @@
 import functools
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import NamedTuple
 
 import torch
 
@@ -14,133 +12,34 @@
     PagedAllocationManager,
 )
 from mstar.engine.resources import KVCachePool, RopeEmbedder, Segment
-from mstar.utils.flashinfer_utils import FlashInferDecodeWrapper, FlashInferPrefillWrapper
-
-logger = logging.getLogger(__name__)
-
-
-def cross_attn_label(label: str, source: str = "default") -> str:
-    """Resolve the cache label under which a cross-attention plan/state for
-    ``source`` is stored, relative to the base (self-attention) label."""
-    return f"{label}::CROSS_ATTN::{source}"
-
-
-class PagedIndptrs(NamedTuple):
-    """The four int32 index tensors a FlashInfer prefill/decode wrapper's
-    ``plan`` consumes, built on CPU (so wrapper.plan's ``.to("cpu")`` is a
-    no-op — see ``_plan_attention_impl``)."""
-    qo_indptr: torch.Tensor
-    paged_kv_indptr: torch.Tensor
-    paged_kv_indices: torch.Tensor
-    paged_kv_last_page_len: torch.Tensor
-
-
-def build_paged_indptrs(
-    q_seq_lens: list[int],
-    page_indices_per_request: list[Sequence[int]],
-    context_lens: list[int],
-    page_size: int,
-) -> PagedIndptrs:
-    """Assemble FlashInfer paged-attention index tensors from per-request
-    query lengths + already-allocated page lists. Shared by the self- and
-    cross-attention plan paths (the difference is only where the pages come
-    from: self grows them per step, cross reads a fixed context)."""
-    qo_indptr = [0]
-    kv_indptr = [0]
-    all_pages: list[int] = []
-    last_page_lens: list[int] = []
-    for q_len, pages, ctx_len in zip(
-        q_seq_lens, page_indices_per_request, context_lens, strict=True,
-    ):
-        qo_indptr.append(qo_indptr[-1] + q_len)
-        all_pages.extend(pages)
-        kv_indptr.append(kv_indptr[-1] + len(pages))
-        last_page_lens.append(ctx_len % page_size or page_size)
-    return PagedIndptrs(
-        qo_indptr=torch.tensor(qo_indptr, dtype=torch.int32),
-        paged_kv_indptr=torch.tensor(kv_indptr, dtype=torch.int32),
-        paged_kv_indices=torch.tensor(all_pages, dtype=torch.int32),
-        paged_kv_last_page_len=torch.tensor(last_page_lens, dtype=torch.int32),
-    )
-
-
-class PlanCacheKey(NamedTuple):
-    """Fingerprint of a wrapper ``plan`` call's inputs. When it is unchanged
-    between steps the re-plan is skippable. Used by cross-attention (context
-    pages are immutable after encode); the mechanism is label-generic, so a
-    fixed-shape self-attention label could reuse it (see the field on
-    ``_PlanState``)."""
-    q_seq_lens: tuple
-    page_indices: tuple
-    last_page_lens: tuple
-    dtype: torch.dtype
-
+from mstar.engine.resources.attention import (
+    CrossAttentionManager,
+    DenseGenAttentionManager,
+    FlashInferAttentionManager,
+    PagedIndptrs,
+    PlanCacheKey,
+    WorkspaceBufferManager,
+    _PlanState,
+    build_paged_indptrs,
+    cross_attn_label,
+)
 
-@dataclass
-class _PlanState:
-    """Pre-computed state from plan_attention/plan_rope for a single cache label.
-
-    Stored per-label so that preprocess can plan for all relevant labels
-    upfront (plan operations are CUDA graph incompatible). During forward,
-    run_attention/apply_rope look up the active label's plan state.
-
-    In CUDA graph mode, wrapper is a persistent FlashInferPrefillWrapper or
-    FlashInferDecodeWrapper created once during capture. plan_attention()
-    calls wrapper.plan() which updates static buffers via .copy_().
-
-    ``custom_pos_advance`` is a generic out-of-band channel for prefill
-    walks whose position-id span differs from the seq_len being prefilled
-    (e.g. Qwen3-Omni's ``prefill_vision``, where the 3D-grid MRoPE span is
-    larger than the number of tokens). The submodule writes a per-request
-    list here via ``BatchedCacheManager.set_custom_pos_advance``;
-    ``advance_seq_lens`` reads it when ``pos_id_ns`` is None and advances
-    ``position_id_start`` by these values instead of by ``seq_len``.
-    Auto-cleared by ``advance_seq_lens`` so it doesn't leak across calls.
-    The CUDA-graph runner's post-replay ``advance_seq_lens()`` call is what
-    actually consumes this — the model's inner ``advance_seq_lens(pos_id_ns=...)``
-    runs at capture time only and is not replayed.
-    """
-    wrapper: FlashInferPrefillWrapper | FlashInferDecodeWrapper | None = None
-    pos_ids: torch.Tensor | None = None
-    seq_lens: list[int] | None = None
-    write_store: bool = True
-    custom_pos_advance: list[int] | None = None
-    # Plan memo: fingerprint of the last wrapper.plan() inputs for this label;
-    # when it matches, the re-plan is skipped. Only the cross-attention path
-    # sets it today (its context pages are immutable after add_cross_attn_kv),
-    # and only where plan states persist across steps (the CUDA-graph runner);
-    # the eager path rebuilds the cache manager per step and still re-plans.
-    #
-    # Future reference — this is a general-purpose tool, not cross-attn only.
-    # A regular (self-attention) label can memo its plan too; the extra care
-    # there is invalidation, since self-attn pages grow every decode step.
-    # The fingerprint would need to include the per-request seq_len (so appending
-    # a token misses the memo and re-plans), and any page-table remap (eviction /
-    # reallocation) must also bust the key. Given that, decode could skip the
-    # re-plan on the common "seq_len += 1, same pages" step. Deferred until a
-    # model needs it; the eager-path persistence noted above is the prerequisite.
-    plan_cache_key: "PlanCacheKey | None" = None
-    # Set when DenseGenCacheManager planned this label dense: the per-segment
-    # gather indices + varlen cu_seqlens needed to attend each generation
-    # segment over its contiguous frozen prefix. None on paged plans, which
-    # keep the FlashInfer path. See DenseGenCacheManager._build_dense_gen_plan.
-    dense_gen: dict | None = None
-
-
-class WorkspaceBufferManager:
-    def __init__(
-        self, size, device
-    ):
-        self.size = size
-        self.device = device
-        self.buffers = {}
+__all__ = [
+    "ATTENTION_BACKENDS",
+    "BatchedCacheManager",
+    "BatchedCfgInfo",
+    "DenseGenCacheManager",
+    "FlashInferCacheManager",
+    "PagedIndptrs",
+    "PlanCacheKey",
+    "WorkspaceBufferManager",
+    "_PlanState",
+    "build_paged_indptrs",
+    "create_cache_manager",
+    "cross_attn_label",
+]
 
-    def get(self, label: str="main"):
-        if label not in self.buffers:
-            self.buffers[label] = torch.empty(
-                self.size, dtype=torch.uint8, device=self.device
-            )
-        return self.buffers[label]
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -149,23 +48,30 @@ class BatchedCfgInfo:
 
 
 class BatchedCacheManager(ABC):
-    """Attention/KV-cache backend interface for batched multi-request forwards.
-
-    Owns the backend-agnostic machinery: per-label plan state, active-label
-    switching, RoPE position planning/application, sequence-length stepping,
-    KV snapshots, store flushes, and the qo_indptr accessor. Concrete backends
-    implement the attention ops (``plan_attention``,
-    ``plan_attention_batched_cfg``, ``run_attention``); ``ATTENTION_BACKENDS``
-    maps ``KVCacheConfig.attention_backend`` names to backend classes and
+    """Model-facing facade over the engine's per-step resources.
+
+    Holds the step's addressing (which requests, which label each is on)
+    and dispatches every call into the resources behind it: the KV cache
+    pool (admission, views, commits, forks), the attention manager (plans,
+    wrappers, workspaces), the positional embedder, and one cross-attention
+    manager per declared source. The domain state lives on those resources;
+    what stays here is per-step bookkeeping the model or the graph runner
+    drives directly (active labels, the batched-CFG advance info, the
+    pre-plan short-circuit set).
+
+    Concrete backends pick the attention manager kind via
+    ``ATTENTION_MANAGER_CLS``; ``ATTENTION_BACKENDS`` maps
+    ``KVCacheConfig.attention_backend`` names to backend classes and
     ``create_cache_manager`` instantiates the configured one.
 
-    Replaces per-request CacheHandle for decode and simple prefill batches where
-    all requests use the same graph_walk. Constructed per batch: one manager
-    serves the whole batch with a single attention call per layer instead of N
-    per-request calls. Complex paths like image_gen (3-pass CFG with label
-    switching) continue using per-request CacheHandle.
+    Constructed per batch: one facade serves the whole batch with a single
+    attention call per layer instead of N per-request calls. Complex paths
+    like image_gen (3-pass CFG with label switching) continue using
+    per-request construction.
     """
 
+    ATTENTION_MANAGER_CLS: type[FlashInferAttentionManager] = FlashInferAttentionManager
+
     def __init__(
         self,
         request_ids: list[str],
@@ -179,6 +85,9 @@ def __init__(
         auto_write_store: bool=False,
         enable_nvtx: bool=False,
         cross_pools: dict[str, CrossAttnPool] | None = None,
+        kv_pool: KVCachePool | None = None,
+        cross_kv_pools: dict[str, KVCachePool] | None = None,
+        rope_embedder: RopeEmbedder | None = None,
     ):
         self.request_ids = request_ids
         self.active_labels = active_labels_per_request  # {req_id: label}
@@ -192,37 +101,56 @@ def __init__(
         # source name -> CrossAttnPool (see KVCacheConfig.cross_attn)
         self.cross_pools = cross_pools or {}
 
-        # Segment-lifecycle surface over the allocation manager: planning
-        # goes through admit/view on these pools instead of reading
-        # KVRequestState fields and allocating mid-plan.
+        # The resources this facade dispatches into. The engine passes its
+        # persistent pools and embedder; callers that construct the facade
+        # standalone (the graph runner, tests) get fresh fronts over the
+        # same allocation manager, which hold no state of their own.
         self.kv_pool = (
-            KVCachePool(alloc_manager) if alloc_manager is not None else None
+            kv_pool if kv_pool is not None
+            else KVCachePool(alloc_manager) if alloc_manager is not None
+            else None
+        )
+        self.cross_kv_pools: dict[str, KVCachePool] = (
+            cross_kv_pools if cross_kv_pools is not None
+            else {
+                source: KVCachePool(pool.alloc_manager)
+                for source, pool in self.cross_pools.items()
+            }
         )
-        self.cross_kv_pools: dict[str, KVCachePool] = {
-            source: KVCachePool(pool.alloc_manager)
-            for source, pool in self.cross_pools.items()
-        }
         # Position semantics live on the embedder; the pool's counters are
         # read at plan time and advanced only through commit.
-        self.rope_embedder = RopeEmbedder()
+        self.rope_embedder = rope_embedder if rope_embedder is not None else RopeEmbedder()
 
         self.auto_write_store = auto_write_store
 
-        # CUDA graph mode: persistent wrappers passed in from CudaGraphRunner.
-        # When set, plan_attention() uses the persistent wrapper's plan()
-        # method instead of creating a new wrapper each call.
+        # CUDA graph mode: persistent plan states passed in from
+        # CudaGraphRunner. The attention manager plans onto them, so a
+        # persistent wrapper's plan() updates static buffers via .copy_()
+        # instead of creating a new wrapper each call.
         self._cuda_graph_mode = cuda_graph_plan_states is not None
-
-        # Per-label plan state: plan_attention/plan_rope store results here,
-        # run_attention/apply_rope look up by active label.
-        if cuda_graph_plan_states is not None:
-            self._plan_states: dict[str, _PlanState] = cuda_graph_plan_states
-        else:
-            self._plan_states: dict[str, _PlanState] = {}
-
-        self.base_pos_ids = torch.arange(
-            kv_cache_config.max_seq_len, dtype=torch.long, device=device
+        self.attention = self.ATTENTION_MANAGER_CLS(
+            kv_cache=kv_cache,
+            kv_cache_config=kv_cache_config,
+            buffer_manager=buffer_manager,
+            device=device,
+            states=cuda_graph_plan_states,
+            cuda_graph_mode=self._cuda_graph_mode,
+            enable_nvtx=enable_nvtx,
         )
+        # One manager per cross-attention source, sharing the per-step plan
+        # store (cross plans live under the resolved cross label).
+        self.cross_attention: dict[str, CrossAttentionManager] = {
+            source: CrossAttentionManager(
+                source=source,
+                pool=pool,
+                kv_pool=self.cross_kv_pools[source],
+                buffer_manager=buffer_manager,
+                device=device,
+                states=self.attention.states,
+                enable_nvtx=enable_nvtx,
+            )
+            for source, pool in self.cross_pools.items()
+        }
 
         # Labels the Worker's plan_executor has pre-planned for the
         # next batch. Each entry causes the matching plan_attention(label=L)
@@ -250,6 +178,12 @@ def __init__(
 
         self._batched_cfg_info: BatchedCfgInfo | None = None
 
+    @property
+    def _plan_states(self) -> dict[str, _PlanState]:
+        """The per-label plan store, owned by the attention manager (the
+        graph runner's per-slot store when one was passed in)."""
+        return self.attention.states
+
     @torch.compiler.disable
     def _get_state(self, request_id: str, label: str | None = None) -> KVRequestState:
         label = label or self.active_labels.get(request_id, "main")
@@ -628,26 +562,7 @@ def snapshot_all(
     ) -> None:
         """Snapshot KV cache for all requests in batch."""
         for rid in self.request_ids:
-            from_state = self._get_state(rid, from_label)
-
-            if realloc:
-                self.alloc_manager.reset_label(rid, to_label)
-
-            to_state = self._get_state(rid, to_label)
-            start_pos =  to_state.seq_len // self.kv_cache_config.page_size
-            self.alloc_manager.alloc(
-                rid, to_label, seq_len=from_state.seq_len
-            )
-
-            to_state.seq_len = from_state.seq_len
-            to_state.position_id_start = from_state.position_id_start
-
-            for src_page, dst_page in zip(
-                from_state.page_indices[start_pos:],
-                to_state.page_indices[start_pos:],
-                strict=True
-            ):
-                self.kv_cache[:, dst_page] = self.kv_cache[:, src_page]
+            self.kv_pool.fork(rid, from_label, to_label, realloc=realloc)
             if write_store:
                 self.alloc_manager.flush_to_store(
                     rid, label=to_label
@@ -664,13 +579,16 @@ def flush_to_store(self):
 
 
 class FlashInferCacheManager(BatchedCacheManager):
-    """Paged FlashInfer attention backend (the default).
+    """Facade over the paged FlashInfer attention manager (the default).
 
-    Constructs batch-level FlashInfer index tensors (qo_indptr, paged_kv_indptr,
-    paged_kv_indices) and issues a single FlashInfer call per layer instead of
-    N separate calls. K/V for every planned token is written to the paged cache.
+    Admits one segment per request on the KV pool, hands the resulting
+    sequence views to the attention manager, and issues a single FlashInfer
+    call per layer instead of N per-request calls. K/V for every planned
+    token is written to the paged cache.
     """
 
+    ATTENTION_MANAGER_CLS = FlashInferAttentionManager
+
     def plan_attention(
         self,
         seq_lens: list[int] | None = None,
@@ -682,14 +600,14 @@ def plan_attention(
     ):
         """Pre-compute FlashInfer plan and page positions for a cache label.
 
-        Admits one segment per request on the KV pool (reserving pages),
-        builds FlashInfer index tensors from the resulting sequence views,
-        and plans the wrapper. All state is stored in _plan_states[label].
+        Admits one segment per request on the KV pool (reserving pages) and
+        has the attention manager plan against the resulting views. All
+        state is stored in the per-label plan store.
 
-        In CUDA graph mode, uses the persistent wrapper from _plan_states
-        (pre-built by CudaGraphRunner) and calls its plan() method which
-        updates static buffers via .copy_(). In eager mode, creates a new
-        wrapper each call.
+        In CUDA graph mode, the manager plans onto the persistent wrapper
+        from the store (pre-built by CudaGraphRunner), whose plan() updates
+        static buffers via .copy_(). In eager mode, it creates a new wrapper
+        each call.
 
         Planning hints for other backends arriving via **kwargs are ignored.
         """
@@ -716,171 +634,25 @@ def plan_attention(
                     range_push("cache.plan_attention.skipped_pre_planned", synchronize=False)
                     range_pop(synchronize=False)
                 return
-            self._plan_attention_impl(
-                seq_lens=seq_lens,
-                dtype=dtype,
-                is_causal=is_causal,
-                write_store=write_store,
-                label=label,
-            )
-        finally:
-            if self.enable_nvtx:
-                range_pop(synchronize=False)
-
-    def _plan_attention_impl(
-        self,
-        seq_lens: list[int] | None = None,
-        dtype: torch.dtype | None = None,
-        is_causal=True,
-        write_store: bool=True,
-        label: str | None = None,
-    ):
-        from mstar.utils.profiler import range_pop, range_push
-
-        assert self.kv_cache is not None
-
-        # Default the FlashInfer wrapper's dtype to whatever dtype the KV
-        # cache tensor was actually allocated in. Hardcoding bf16 here breaks
-        # any model that runs in fp32 (the wrapper would try to write
-        # bf16-cast K/V into an fp32 cache and torch raises a dtype mismatch
-        # in flashinfer_utils.set_kv_cache).
-        if dtype is None:
-            dtype = self.kv_cache.dtype
-
-        effective_label = label if label is not None else self._active_label()
-
-        cfg = self.kv_cache_config
-        page_size = cfg.page_size
-        num_kv_heads = cfg.num_kv_heads
-        head_dim = cfg.head_dim
-        num_qo_heads = cfg.num_qo_heads
-        device = self.device
-
-        # CPU-side accumulation. The old implementation launched 4-5 tiny GPU
-        # kernels per request (arange, tensor(state.page_indices), indexing,
-        # mod) to build page_indices/page_offsets/token_offsets — all of which
-        # turn out to be unused bookkeeping (grep: no reader in the codebase).
-        # We only need the four int32 tensors the FlashInfer wrapper consumes,
-        # so do the arithmetic in pure Python and send them over in one H2D
-        # each.
-        if self.enable_nvtx:
-            range_push("cache.plan_attention.build_lists", synchronize=False)
-        try:
-            qo_indptr_list = [0]
-            kv_indptr_list = [0]
-            all_page_indices = []
-            kv_last_page_lens = []
-            kv_cache_locations_list = []
-
+            assert self.kv_cache is not None
+            segments = []
+            views = []
             for i, rid in enumerate(self.request_ids):
                 segment = Segment(rid, effective_label, seq_lens[i])
                 self.kv_pool.admit(segment)
-                view = self.kv_pool.view(segment)
-
-                qo_indptr_list.append(qo_indptr_list[-1] + segment.span)
-                all_page_indices.extend(view.page_indices)
-                kv_indptr_list.append(kv_indptr_list[-1] + len(view.page_indices))
-
-                last_page_len = view.length % page_size or page_size
-                kv_last_page_lens.append(last_page_len)
-                if segment.span == 1:
-                    kv_cache_locations_list.append([view.page_indices[-1], last_page_len - 1])
-        finally:
-            if self.enable_nvtx:
-                range_pop(synchronize=False)
-
-        # Build batched FlashInfer index tensors on CPU so wrapper.plan()
-        # doesn't trigger a synchronous D→H inside its body. FlashInfer
-        # calls ``indptr.to("cpu")`` / ``last_page_len.to("cpu")`` near the
-        # top of ``plan()`` to get host views of those metadata tensors;
-        # if we hand them GPU tensors that ``.to("cpu")`` becomes a
-        # synchronous default-stream sync that waits for the entire
-        # outstanding stream — including the speculatively-queued next
-        # decode step. By creating these on CPU directly, ``.to("cpu")``
-        # is a no-op. FlashInfer later copies the tiny int32 metadata to
-        # the device when it needs it; the source is pageable CPU memory, so
-        # ``non_blocking=True`` does not make that H2D copy asynchronous, but
-        # the tensors are batch-size length and the cost is inconsequential.
-        if self.enable_nvtx:
-            range_push("cache.plan_attention.make_tensors", synchronize=False)
-        try:
-            qo_indptr = torch.tensor(qo_indptr_list, dtype=torch.int32)
-            paged_kv_indptr = torch.tensor(kv_indptr_list, dtype=torch.int32)
-            paged_kv_indices = torch.tensor(all_page_indices, dtype=torch.int32)
-            paged_kv_last_page_len = torch.tensor(kv_last_page_lens, dtype=torch.int32)
-            kv_cache_locations = (
-                torch.tensor(kv_cache_locations_list, dtype=torch.long)
-                if len(kv_cache_locations_list) == len(self.request_ids)
-                else None
-            )
-        finally:
-            if self.enable_nvtx:
-                range_pop(synchronize=False)
-
-
-        is_decode = all([sl == 1 for sl in seq_lens])
-        ps = self._plan_states.get(effective_label)
-        if ps is not None and ps.wrapper is not None:
-            wrapper = ps.wrapper
-        elif is_decode:
-            wrapper = FlashInferDecodeWrapper(
-                workspace_buffer=self.buffer_manager.get(effective_label),
-                num_qo_heads=num_qo_heads,
-                num_kv_heads=num_kv_heads,
-                head_dim=head_dim,
-                page_size=page_size,
-                device=self.device,
-                enable_nvtx=self.enable_nvtx,
-                backend=cfg.flashinfer_backend,
-            )
-            ps = _PlanState(wrapper=wrapper)
-            self._plan_states[effective_label] = ps
-        else:
-            wrapper = FlashInferPrefillWrapper(
-                workspace_buffer=self.buffer_manager.get(effective_label),
-                num_qo_heads=num_qo_heads,
-                num_kv_heads=num_kv_heads,
-                head_dim=head_dim,
-                page_size=page_size,
-                device=self.device,
-                enable_nvtx=self.enable_nvtx,
-                backend=cfg.flashinfer_backend,
+                segments.append(segment)
+                views.append(self.kv_pool.view(segment))
+            self.attention.plan(
+                label=effective_label,
+                segments=segments,
+                views=views,
+                dtype=dtype,
+                is_causal=is_causal,
+                write_store=write_store,
             )
-            ps = _PlanState(wrapper=wrapper)
-            self._plan_states[effective_label] = ps
-
-        if self.enable_nvtx:
-            range_push("cache.plan_attention.wrapper_plan", synchronize=False)
-        try:
-            if isinstance(wrapper, FlashInferDecodeWrapper):
-                wrapper.plan(
-                    paged_kv_indptr=paged_kv_indptr,
-                    paged_kv_indices=paged_kv_indices,
-                    paged_kv_last_page_len=paged_kv_last_page_len,
-                    kv_cache_locations=kv_cache_locations,
-                    dtype=dtype,
-                )
-            else:
-                wrapper.plan(
-                    qo_indptr=qo_indptr,
-                    paged_kv_indptr=paged_kv_indptr,
-                    paged_kv_indices=paged_kv_indices,
-                    paged_kv_last_page_len=paged_kv_last_page_len,
-                    causal=is_causal,
-                    dtype=dtype,
-                )
         finally:
             if self.enable_nvtx:
                 range_pop(synchronize=False)
-        # seq_lens is read by the flush_to_store path; write_store by
-        # run_attention. The page_indices / page_offsets / token_offsets /
-        # per_req_page_indices fields were legacy bookkeeping and had no
-        # reader — dropped along with their per-rid GPU construction above.
-        ps.seq_lens = seq_lens
-        ps.write_store = write_store
-        # A paged plan clears any prior dense plan (set by DenseGenCacheManager)
-        # so run_attention routes this label back through the wrapper.
-        ps.dense_gen = None
 
     @torch.compiler.disable
     def plan_attention_batched_cfg(
@@ -909,98 +681,25 @@ def plan_attention_batched_cfg(
             per_label_seq_len=seq_lens
         )
 
-        cfg = self.kv_cache_config
-        page_size = cfg.page_size
-        num_kv_heads = cfg.num_kv_heads
-        head_dim = cfg.head_dim
-        num_qo_heads = cfg.num_qo_heads
-        device = self.device
-
-        # CPU-side accumulation (see plan_attention for the same pattern).
-        qo_indptr_list = [0]
-        kv_indptr_list = [0]
-        all_page_indices = []
-        kv_last_page_lens = []
-        combined_seq_lens = []
-
+        # Label-major segment order: every (label, request) pair in batch
+        # order, the packed layout the CFG forward uses.
+        segments = []
+        views = []
         for label in labels:
             for i, rid in enumerate(self.request_ids):
                 segment = Segment(rid, label, seq_lens[label][i])
                 self.kv_pool.admit(segment)
-                view = self.kv_pool.view(segment)
+                segments.append(segment)
+                views.append(self.kv_pool.view(segment))
 
-                qo_indptr_list.append(qo_indptr_list[-1] + segment.span)
-                all_page_indices.extend(view.page_indices)
-                kv_indptr_list.append(
-                    kv_indptr_list[-1] + len(view.page_indices)
-                )
-
-                last_page_len = view.length % page_size or page_size
-                kv_last_page_lens.append(last_page_len)
-                combined_seq_lens.append(segment.span)
-
-        # CPU tensors — see comment in ``plan_attention`` above. FlashInfer
-        # async-H2Ds these inside ``plan()``; passing GPU tensors would
-        # trigger a synchronous default-stream sync via the internal
-        # ``.to("cpu")`` call.
-        qo_indptr = torch.tensor(qo_indptr_list, dtype=torch.int32)
-        paged_kv_indptr = torch.tensor(kv_indptr_list, dtype=torch.int32)
-        paged_kv_indices = torch.tensor(all_page_indices, dtype=torch.int32)
-        paged_kv_last_page_len = torch.tensor(kv_last_page_lens, dtype=torch.int32)
-
-        ps = self._plan_states.get(combined_label)
-        if self._cuda_graph_mode and ps is not None and ps.wrapper is not None:
-            # CUDA-graph mode: reuse the persistent wrapper across denoise steps.
-            # plan() updates its static buffers via .copy_() so the captured
-            # kernel picks up each step's page table without reallocating.
-            wrapper = ps.wrapper
-        elif self._cuda_graph_mode:
-            # First call under capture: build the persistent wrapper sized for the
-            # fixed batch (labels x requests) and token budget.
-            wrapper = FlashInferPrefillWrapper(
-                workspace_buffer=self.buffer_manager.get(combined_label),
-                num_qo_heads=num_qo_heads,
-                num_kv_heads=num_kv_heads,
-                head_dim=head_dim,
-                page_size=page_size,
-                batch_size=len(labels) * len(self.request_ids),
-                max_total_tokens=sum(combined_seq_lens),
-                max_num_pages=cfg.max_num_pages,
-                device=self.device,
-                use_cuda_graph=True,
-                enable_nvtx=self.enable_nvtx,
-                backend=cfg.flashinfer_backend,
-            )
-            ps = _PlanState(wrapper=wrapper)
-            self._plan_states[combined_label] = ps
-        else:
-            # Eager mode: a fresh wrapper each call (the cache manager is rebuilt
-            # per forward, so there is nothing persistent to reuse).
-            wrapper = FlashInferPrefillWrapper(
-                workspace_buffer=self.buffer_manager.get(combined_label),
-                num_qo_heads=num_qo_heads,
-                num_kv_heads=num_kv_heads,
-                head_dim=head_dim,
-                page_size=page_size,
-                enable_nvtx=self.enable_nvtx,
-                backend=cfg.flashinfer_backend,
-            )
-            ps = _PlanState(wrapper=wrapper)
-            self._plan_states[combined_label] = ps
-
-        wrapper.plan(
-            qo_indptr=qo_indptr,
-            paged_kv_indptr=paged_kv_indptr,
-            paged_kv_indices=paged_kv_indices,
-            paged_kv_last_page_len=paged_kv_last_page_len,
-            causal=is_causal,
+        self.attention.plan_batched_cfg(
+            combined_label=combined_label,
+            segments=segments,
+            views=views,
             dtype=dtype,
+            is_causal=is_causal,
+            write_store=write_store,
         )
-        ps.seq_lens = combined_seq_lens
-        ps.write_store = write_store
-        # A paged plan clears any prior dense plan (set by DenseGenCacheManager)
-        # so run_attention routes this label back through the wrapper.
-        ps.dense_gen = None
 
     @torch.compiler.disable
     def run_attention(
@@ -1016,12 +715,10 @@ def run_attention(
         call). Writes K and V to the paged KV cache at pre-computed page
         positions, then runs the FlashInfer wrapper for batched attention.
 
-        In CUDA graph mode, uses wrapper.set_kv_cache() + wrapper.run()
-        which operates on pre-computed token_to_page/token_to_cache or
-        kv_cache_locations tensors (static GPU addresses).
-
-        In eager mode, uses direct fancy indexing for KV writes and
-        the raw FlashInfer wrapper's run().
+        In CUDA graph mode, the wrapper's set_kv_cache() + run() operate on
+        pre-computed token_to_page/token_to_cache or kv_cache_locations
+        tensors (static GPU addresses). In eager mode, direct fancy indexing
+        for KV writes and the raw FlashInfer wrapper's run().
         """
         if layer_idx is None:
             layer_idx = self.layer_idx
@@ -1033,7 +730,7 @@ def run_attention(
 
         assert self.kv_cache is not None and ps.wrapper is not None
 
-        ps.wrapper.set_kv_cache(self.kv_cache[layer_idx], k, v)
+        self.attention.write_kv(k, v, layer_idx, label)
 
         if self.auto_write_store and ps.write_store:
             for req_id in self.request_ids:
@@ -1041,30 +738,23 @@ def run_attention(
                     req_id, label=label, layers=layer_idx
                 )
 
-        return ps.wrapper.run(q, self.kv_cache[layer_idx]).to(orig_dtype)
+        return self.attention.run(q, layer_idx, label).to(orig_dtype)
 
     # ------------------------------------------------------------------
-    # Cross-attention (issue #160)
-    #
-    # Cross-attention is non-causal attention over a separate, fixed
-    # encoder-context KV: written once at encode time (add_cross_attn_kv),
-    # planned per step against the decoder's query lengths
-    # (plan_cross_attention), and executed per layer (run_cross_attn). The
-    # context KV lives in per-source pools (KVCacheConfig.cross_attn) whose
-    # head config may differ from the decoder's self-attention; plan/run
-    # state rides the existing per-label _PlanState machinery under the
-    # resolved ``{label}::CROSS_ATTN::{source}`` label. Cross labels never
-    # plan RoPE — context positions, if any, are baked in at encode time.
+    # Cross-attention: dispatched to the per-source CrossAttentionManager
+    # (a read-only context pool plus its own wrapper machinery). The label
+    # namespace under which cross plans are stored is that manager's
+    # concern; callers keep addressing by base label and source.
     # ------------------------------------------------------------------
 
-    def _get_cross_pool(self, source: str) -> CrossAttnPool:
-        pool = self.cross_pools.get(source)
-        if pool is None:
+    def _get_cross_manager(self, source: str) -> CrossAttentionManager:
+        manager = self.cross_attention.get(source)
+        if manager is None:
             raise KeyError(
                 f"No cross-attention pool for source {source!r}; declare it in "
-                f"KVCacheConfig.cross_attn (available: {list(self.cross_pools)})"
+                f"KVCacheConfig.cross_attn (available: {list(self.cross_attention)})"
             )
-        return pool
+        return manager
 
     def _active_base_label(self, label: str | None) -> str:
         if label is not None:
@@ -1086,58 +776,16 @@ def add_cross_attn_kv(
     ) -> None:
         """Write encoder-context K/V for one layer into ``source``'s pool.
 
-        Called once per request per layer at encode time. ``k``/``v`` are
-        packed ``(total_context_tokens, num_kv_heads, head_dim)`` across
-        ``request_ids``; ``seq_lens`` gives per-request context lengths
-        (defaults to a single request owning the whole tensor). Pages are
-        allocated on first write (layer 0) and reused for the rest.
+        See ``CrossAttentionManager.write_context``.
         """
-        pool = self._get_cross_pool(source)
-        kv_pool = self.cross_kv_pools[source]
-        base_label = self._active_base_label(label)
-        cross_label = cross_attn_label(base_label, source)
-        page_size = pool.alloc_config.page_size
-
-        if seq_lens is None:
-            assert len(request_ids) == 1, (
-                "add_cross_attn_kv needs seq_lens for multi-request batches"
-            )
-            seq_lens = [k.shape[0]]
-
-        offset = 0
-        for rid, ctx_len in zip(request_ids, seq_lens, strict=True):
-            resident = kv_pool.view(Segment(rid, cross_label, 0)).length
-            if resident == 0:
-                # First write for this context: reserve its pages and commit
-                # the extent up front; later layers reuse them. The position
-                # counter stays untouched (context positions are baked in at
-                # encode time).
-                segment = Segment(rid, cross_label, ctx_len)
-                kv_pool.admit(segment)
-                kv_pool.commit(segment, pos_advance=0)
-            else:
-                assert resident == ctx_len, (
-                    f"cross-attn context for {rid!r}/{source!r} already written "
-                    f"with length {resident}, got {ctx_len}"
-                )
-            view = kv_pool.view(Segment(rid, cross_label, 0))
-
-            positions = torch.arange(ctx_len, device=self.device)
-            page_indices = torch.tensor(
-                view.page_indices, dtype=torch.long, device=self.device,
-            )
-            token_to_page = page_indices[
-                torch.div(positions, page_size, rounding_mode="floor")
-            ]
-            token_to_cache = positions % page_size
-
-            layer_cache = pool.kv_cache[layer_idx]
-            dtype = pool.kv_cache.dtype
-            layer_cache[token_to_page, 0, token_to_cache] = \
-                k[offset:offset + ctx_len].to(dtype)
-            layer_cache[token_to_page, 1, token_to_cache] = \
-                v[offset:offset + ctx_len].to(dtype)
-            offset += ctx_len
+        self._get_cross_manager(source).write_context(
+            request_ids=request_ids,
+            k=k,
+            v=v,
+            layer_idx=layer_idx,
+            seq_lens=seq_lens,
+            base_label=self._active_base_label(label),
+        )
 
     def plan_cross_attention(
         self,
@@ -1148,80 +796,17 @@ def plan_cross_attention(
     ) -> None:
         """Plan the cross-attention wrapper for this step's decoder queries.
 
-        ``q_seq_lens`` is the number of decoder query tokens per request
-        (matches the self-attention plan's seq_lens). The context side is
-        read from the pool state written by ``add_cross_attn_kv`` — pages
-        are fixed, so unlike ``plan_attention`` nothing is allocated here.
-        ``label`` is the base (self-attention) label; the plan is stored
-        under the resolved cross label. Always uses the prefill wrapper
-        with ``causal=False`` (decode-style single queries still attend to
-        the full context).
+        ``q_seq_lens`` matches the self-attention plan's seq_lens; ``label``
+        is the base (self-attention) label. See
+        ``CrossAttentionManager.plan``.
         """
-        pool = self._get_cross_pool(source)
-        kv_pool = self.cross_kv_pools[source]
-        base_label = self._active_base_label(label)
-        cross_label = cross_attn_label(base_label, source)
-        cfg = pool.alloc_config
-        page_size = cfg.page_size
-
-        if dtype is None:
-            dtype = pool.kv_cache.dtype
-
-        page_indices_per_request: list[tuple[int, ...]] = []
-        context_lens: list[int] = []
-        for rid in self.request_ids:
-            view = kv_pool.view(Segment(rid, cross_label, 0))
-            assert view.length > 0, (
-                f"plan_cross_attention before add_cross_attn_kv for {rid!r} "
-                f"(source {source!r})"
-            )
-            page_indices_per_request.append(view.page_indices)
-            context_lens.append(view.length)
-
-        indptrs = build_paged_indptrs(
-            q_seq_lens, page_indices_per_request, context_lens, page_size,
-        )
-
-        # Skip the re-plan when its inputs are unchanged (common decode case).
-        # Keyed on the page indices so a reused request id with a new context
-        # can't alias a stale plan.
-        plan_key = PlanCacheKey(
-            q_seq_lens=tuple(q_seq_lens),
-            page_indices=tuple(indptrs.paged_kv_indices.tolist()),
-            last_page_lens=tuple(indptrs.paged_kv_last_page_len.tolist()),
+        self._get_cross_manager(source).plan(
+            request_ids=self.request_ids,
+            q_seq_lens=q_seq_lens,
             dtype=dtype,
+            base_label=self._active_base_label(label),
         )
 
-        ps = self._plan_states.get(cross_label)
-        if ps is not None and ps.wrapper is not None and ps.plan_cache_key == plan_key:
-            # plan_key carries q_seq_lens, so a match means ps.seq_lens (set at
-            # plan time below) already equals the current q_seq_lens.
-            return
-        if ps is None or ps.wrapper is None:
-            wrapper = FlashInferPrefillWrapper(
-                workspace_buffer=self.buffer_manager.get(cross_label),
-                num_qo_heads=cfg.num_qo_heads,
-                num_kv_heads=cfg.num_kv_heads,
-                head_dim=cfg.head_dim,
-                page_size=page_size,
-                device=self.device,
-                enable_nvtx=self.enable_nvtx,
-            )
-            ps = _PlanState(wrapper=wrapper)
-            self._plan_states[cross_label] = ps
-
-        ps.wrapper.plan(
-            qo_indptr=indptrs.qo_indptr,
-            paged_kv_indptr=indptrs.paged_kv_indptr,
-            paged_kv_indices=indptrs.paged_kv_indices,
-            paged_kv_last_page_len=indptrs.paged_kv_last_page_len,
-            causal=False,
-            dtype=dtype,
-        )
-        ps.seq_lens = q_seq_lens
-        ps.write_store = False
-        ps.plan_cache_key = plan_key
-
     # Like run_attention: kept out of the compiled decoder — its query's leading
     # dim varies per step, which would blow dynamo's recompile limit (#160).
     @torch.compiler.disable
@@ -1243,32 +828,25 @@ def run_cross_attn(
         """
         if layer_idx is None:
             layer_idx = self.layer_idx
-        pool = self._get_cross_pool(source)
-        base_label = self._active_base_label(None)
-        cross_label = cross_attn_label(base_label, source)
 
         orig_dtype = q.dtype
-        ps = self._plan_states.get(cross_label)
-        assert ps is not None and ps.wrapper is not None, (
-            f"run_cross_attn before plan_cross_attention (label {cross_label!r})"
-        )
-        return ps.wrapper.run(q, pool.kv_cache[layer_idx]).to(orig_dtype)
+        return self._get_cross_manager(source).run(
+            q, layer_idx, self._active_base_label(None),
+        ).to(orig_dtype)
 
 class DenseGenCacheManager(FlashInferCacheManager):
-    """FlashInfer backend with a dense generation-attention fast path.
+    """Facade over the dense generation-attention manager.
 
     Runs non-causal generation attention (planned with ``dense_gen=True``) as a
     dense FlashAttention-3 pass over a contiguous [frozen-prefix | fresh]
-    sequence instead of the paged FlashInfer prefill. Diffusion recomputes every
-    generation K/V each step (only the tiny text prefix is reused), so the paged
-    path's per-step full-buffer K/V write is pure overhead here; a dense pass
-    gathers the small prefix, concatenates it with the freshly projected K/V,
-    and runs one varlen kernel — which is also the faster attention kernel at
-    these shapes. Eager-only and single-request only (see ``_dense_gen_applies``);
-    everything else — prefill, captured graphs, multi-request batches — falls
-    through to the inherited paged FlashInfer path.
+    sequence instead of the paged FlashInfer prefill. Eager-only and
+    single-request only (see ``_dense_gen_applies``); everything else —
+    prefill, captured graphs, multi-request batches — falls through to the
+    inherited paged FlashInfer path.
     """
 
+    ATTENTION_MANAGER_CLS = DenseGenAttentionManager
+
     def _dense_gen_applies(self) -> bool:
         """Dense generation attention is eager-only (captured paths keep the
         persistent paged wrapper the graph was planned with) and single-request
@@ -1277,6 +855,22 @@ def _dense_gen_applies(self) -> bool:
         on the paged path)."""
         return not self._cuda_graph_mode and len(self.request_ids) == 1
 
+    def _dense_entries(
+        self, labels: list[str], seq_lens: dict[str, list[int]],
+    ) -> list[tuple]:
+        """The dense plan's per-segment inputs, in the same (label, request)
+        batch order the generation tokens are packed in. Each frozen prefix
+        is read without extending its stream (the generation K/V never
+        enters the pages), so it is planned from a zero-span view; the
+        persistent request state rides along so the manager can cache the
+        gathered prefix across denoise steps."""
+        entries = []
+        for label in labels:
+            for i, rid in enumerate(self.request_ids):
+                view = self.kv_pool.view(Segment(rid, label, 0))
+                entries.append((view, seq_lens[label][i], self._get_state(rid, label)))
+        return entries
+
     def plan_attention(
         self,
         seq_lens: list[int] | None = None,
@@ -1309,14 +903,17 @@ def plan_attention(
             # written to pages and attention is one varlen FA3 over the
             # frozen prefix + fresh gen, so the whole paged-FlashInfer plan
             # (per-step page alloc, index tensors, and wrapper.plan()'s
-            # radix-sort/fills) is dead work — _run_dense_gen reads none of
-            # it. Build only the dense gather/varlen plan.
+            # radix-sort/fills) is dead work — run_dense reads none of it.
+            # Build only the dense gather/varlen plan.
             effective_label = label if label is not None else self._active_label()
-            ps = self._plan_states.get(effective_label) or _PlanState()
-            self._plan_states[effective_label] = ps
-            ps.seq_lens = seq_lens
-            ps.write_store = write_store
-            ps.dense_gen = self._build_dense_gen_plan([effective_label], seq_lens)
+            self.attention.plan_dense(
+                label=effective_label,
+                entries=self._dense_entries(
+                    [effective_label], {effective_label: seq_lens}
+                ),
+                seq_lens=seq_lens,
+                write_store=write_store,
+            )
         finally:
             if self.enable_nvtx:
                 range_pop(synchronize=False)
@@ -1355,14 +952,14 @@ def plan_attention_batched_cfg(
         )
 
         # Lean dense generation-attention path (see plan_attention): skip the
-        # paged-FlashInfer plan (per-label page alloc, index tensors, and
-        # wrapper.plan()'s radix-sort/fills) — _run_dense_gen reads none of
-        # it. Build only the per-segment dense gather/varlen plan.
-        ps = self._plan_states.get(combined_label) or _PlanState()
-        self._plan_states[combined_label] = ps
-        ps.seq_lens = seq_lens
-        ps.write_store = write_store
-        ps.dense_gen = self._build_dense_gen_plan(labels, seq_lens)
+        # paged-FlashInfer plan entirely and build only the per-segment dense
+        # gather/varlen plan.
+        self.attention.plan_dense(
+            label=combined_label,
+            entries=self._dense_entries(labels, seq_lens),
+            seq_lens=seq_lens,
+            write_store=write_store,
+        )
 
     @torch.compiler.disable
     def run_attention(
@@ -1379,109 +976,9 @@ def run_attention(
         if ps.dense_gen is not None:
             if layer_idx is None:
                 layer_idx = self.layer_idx
-            return self._run_dense_gen(q, k, v, layer_idx, ps.dense_gen).to(q.dtype)
+            return self.attention.run_dense(q, k, v, layer_idx, ps.dense_gen).to(q.dtype)
         return super().run_attention(q, k, v, layer_idx=layer_idx)
 
-    def _build_dense_gen_plan(
-        self, labels: list[str],
-        seq_lens: list[int] | dict[str, list[int]]
-    ) -> dict:
-        """Pre-compute the per-segment gather + varlen layout for the dense
-        generation-attention path, in the same (label, request) batch order the
-        generation tokens are packed in. Each segment attends its fresh
-        generation tokens over its frozen text prefix; the prefix lives in the
-        pages written at prefill, so we record the page indices to gather it from
-        (the same across all layers) and the cumulative-sequence-length tensors a
-        single varlen kernel needs. Built once per denoise step, reused by every
-        layer's run_attention."""
-
-        if isinstance(seq_lens, list):
-            seq_lens = {
-                key: seq_lens for key in labels
-            }
-
-        cfg = self.kv_cache_config
-        page_size = cfg.page_size
-        segs = []  # (prefix_page_indices, prefix_len, gen_len)
-        cu_q = [0]
-        cu_k = [0]
-        max_q = 0
-        max_k = 0
-        for label in labels:
-            for i, rid in enumerate(self.request_ids):
-                # The frozen prefix is read without extending the stream (the
-                # generation K/V never enters the pages), so plan it from a
-                # zero-span view.
-                view = self.kv_pool.view(Segment(rid, label, 0))
-                prefix_len = view.length
-                gen_len = seq_lens[label][i]
-                n_pages = (prefix_len + page_size - 1) // page_size
-                idx = torch.tensor(
-                    view.page_indices[:n_pages], dtype=torch.long, device=self.device
-                )
-                # Carry the persistent KVRequestState so run_attention can cache
-                # the gathered frozen prefix on it across denoise steps (the
-                # manager itself is rebuilt every forward).
-                state = self._get_state(rid, label)
-                segs.append((idx, prefix_len, gen_len, state))
-                cu_q.append(cu_q[-1] + gen_len)
-                cu_k.append(cu_k[-1] + prefix_len + gen_len)
-                max_q = max(max_q, gen_len)
-                max_k = max(max_k, prefix_len + gen_len)
-        return {
-            "segs": segs,
-            "cu_q": torch.tensor(cu_q, dtype=torch.int32, device=self.device),
-            "cu_k": torch.tensor(cu_k, dtype=torch.int32, device=self.device),
-            "max_q": max_q,
-            "max_k": max_k,
-        }
-
-    @torch.compiler.disable
-    def _run_dense_gen(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_idx: int, dg: dict
-    ) -> torch.Tensor:
-        """Dense generation attention: per segment, take the frozen text-prefix
-        K/V, concatenate it with this segment's fresh K/V, and attend
-        non-causally with one FlashAttention-3 varlen kernel. Bypasses the paged
-        write entirely (the generation K/V is recomputed every step, so
-        persisting it is wasted work). The frozen prefix is gathered from the
-        paged cache once per layer and cached on the request state, then reused
-        across denoise steps (it never changes during denoise)."""
-        from fa3_fwd_interface import flash_attn_varlen_func
-
-        cfg = self.kv_cache_config
-        num_kv_heads, head_dim = cfg.num_kv_heads, cfg.head_dim
-        kv_layer = self.kv_cache[layer_idx]  # [max_pages, 2, page_size, num_kv_heads, head_dim]
-
-        k_parts, v_parts = [], []
-        offset = 0
-        for idx, prefix_len, gen_len, state in dg["segs"]:
-            prefix_cache = state.dense_prefix_kv
-            if prefix_cache is None:
-                prefix_cache = state.dense_prefix_kv = {}
-            cached = prefix_cache.get(layer_idx)
-            if cached is None:
-                sub = kv_layer[idx]  # [n_pages, 2, page_size, num_kv_heads, head_dim]
-                k_pref = sub[:, 0].reshape(-1, num_kv_heads, head_dim)[:prefix_len].clone()
-                v_pref = sub[:, 1].reshape(-1, num_kv_heads, head_dim)[:prefix_len].clone()
-                prefix_cache[layer_idx] = (k_pref, v_pref)
-            else:
-                k_pref, v_pref = cached
-            k_parts.append(k_pref)
-            k_parts.append(k[offset:offset + gen_len])
-            v_parts.append(v_pref)
-            v_parts.append(v[offset:offset + gen_len])
-            offset += gen_len
-        key = torch.cat(k_parts, dim=0)
-        val = torch.cat(v_parts, dim=0)
-        if q.dtype != key.dtype:
-            q = q.to(key.dtype)
-
-        out = flash_attn_varlen_func(
-            q, key, val, dg["cu_q"], dg["cu_k"], dg["max_q"], dg["max_k"], causal=False,
-        )
-        return out[0] if isinstance(out, tuple) else out
-
 
 # Backend registry: KVCacheConfig.attention_backend names one of these.
 ATTENTION_BACKENDS: dict[str, type[BatchedCacheManager]] = {
diff --git a/mstar/engine/resources/__init__.py b/mstar/engine/resources/__init__.py
index e77a4cf3c..af7837984 100644
--- a/mstar/engine/resources/__init__.py
+++ b/mstar/engine/resources/__init__.py
@@ -1,3 +1,9 @@
+from mstar.engine.resources.attention import (
+    CrossAttentionManager,
+    DenseGenAttentionManager,
+    FlashInferAttentionManager,
+    WorkspaceBufferManager,
+)
 from mstar.engine.resources.base import (
     PositionPlan,
     Reservation,
@@ -8,6 +14,9 @@
 from mstar.engine.resources.positions import RopeEmbedder
 
 __all__ = [
+    "CrossAttentionManager",
+    "DenseGenAttentionManager",
+    "FlashInferAttentionManager",
     "KVCachePool",
     "PageArena",
     "PositionPlan",
@@ -15,4 +24,5 @@
     "RopeEmbedder",
     "Segment",
     "SequenceView",
+    "WorkspaceBufferManager",
 ]
diff --git a/mstar/engine/resources/attention.py b/mstar/engine/resources/attention.py
new file mode 100644
index 000000000..b71b907e9
--- /dev/null
+++ b/mstar/engine/resources/attention.py
@@ -0,0 +1,720 @@
+"""Attention manager resources.
+
+An attention manager owns the attention execution plan for its set of
+layers and the static buffers behind it: the FlashInfer wrappers, the
+workspace buffers, and the per-label plan state. It plans against the
+sequence views handed to it rather than reading cache internals, and it
+exposes the device-side write and run calls the model's forward uses.
+Cross-attention gets its own manager over a read-only context pool.
+"""
+
+import logging
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import torch
+
+from mstar.engine.kv_store import CrossAttnPool, KVCacheConfig, KVRequestState
+from mstar.engine.resources.base import Segment, SequenceView
+from mstar.engine.resources.kv_pool import KVCachePool
+from mstar.utils.flashinfer_utils import FlashInferDecodeWrapper, FlashInferPrefillWrapper
+
+logger = logging.getLogger(__name__)
+
+
+def cross_attn_label(label: str, source: str = "default") -> str:
+    """Resolve the cache label under which a cross-attention plan/state for
+    ``source`` is stored, relative to the base (self-attention) label."""
+    return f"{label}::CROSS_ATTN::{source}"
+
+
+class PagedIndptrs(NamedTuple):
+    """The four int32 index tensors a FlashInfer prefill/decode wrapper's
+    ``plan`` consumes, built on CPU (so wrapper.plan's ``.to("cpu")`` is a
+    no-op — see ``FlashInferAttentionManager.plan``)."""
+    qo_indptr: torch.Tensor
+    paged_kv_indptr: torch.Tensor
+    paged_kv_indices: torch.Tensor
+    paged_kv_last_page_len: torch.Tensor
+
+
+def build_paged_indptrs(
+    q_seq_lens: list[int],
+    page_indices_per_request: list[Sequence[int]],
+    context_lens: list[int],
+    page_size: int,
+) -> PagedIndptrs:
+    """Assemble FlashInfer paged-attention index tensors from per-request
+    query lengths + already-allocated page lists. Shared by the self- and
+    cross-attention plan paths (the difference is only where the pages come
+    from: self grows them per step, cross reads a fixed context)."""
+    qo_indptr = [0]
+    kv_indptr = [0]
+    all_pages: list[int] = []
+    last_page_lens: list[int] = []
+    for q_len, pages, ctx_len in zip(
+        q_seq_lens, page_indices_per_request, context_lens, strict=True,
+    ):
+        qo_indptr.append(qo_indptr[-1] + q_len)
+        all_pages.extend(pages)
+        kv_indptr.append(kv_indptr[-1] + len(pages))
+        last_page_lens.append(ctx_len % page_size or page_size)
+    return PagedIndptrs(
+        qo_indptr=torch.tensor(qo_indptr, dtype=torch.int32),
+        paged_kv_indptr=torch.tensor(kv_indptr, dtype=torch.int32),
+        paged_kv_indices=torch.tensor(all_pages, dtype=torch.int32),
+        paged_kv_last_page_len=torch.tensor(last_page_lens, dtype=torch.int32),
+    )
+
+
+class PlanCacheKey(NamedTuple):
+    """Fingerprint of a wrapper ``plan`` call's inputs. When it is unchanged
+    between steps the re-plan is skippable. Used by cross-attention (context
+    pages are immutable after encode); the mechanism is label-generic, so a
+    fixed-shape self-attention label could reuse it (see the field on
+    ``_PlanState``)."""
+    q_seq_lens: tuple
+    page_indices: tuple
+    last_page_lens: tuple
+    dtype: torch.dtype
+
+
+@dataclass
+class _PlanState:
+    """Pre-computed state from plan_attention/plan_rope for a single cache label.
+
+    Stored per-label so that preprocess can plan for all relevant labels
+    upfront (plan operations are CUDA graph incompatible). During forward,
+    run_attention/apply_rope look up the active label's plan state.
+
+    In CUDA graph mode, wrapper is a persistent FlashInferPrefillWrapper or
+    FlashInferDecodeWrapper created once during capture. plan_attention()
+    calls wrapper.plan() which updates static buffers via .copy_().
+
+    ``custom_pos_advance`` is a generic out-of-band channel for prefill
+    walks whose position-id span differs from the seq_len being prefilled
+    (e.g. Qwen3-Omni's ``prefill_vision``, where the 3D-grid MRoPE span is
+    larger than the number of tokens). The submodule writes a per-request
+    list here via ``BatchedCacheManager.set_custom_pos_advance``;
+    ``advance_seq_lens`` reads it when ``pos_id_ns`` is None and advances
+    ``position_id_start`` by these values instead of by ``seq_len``.
+    Auto-cleared by ``advance_seq_lens`` so it doesn't leak across calls.
+    The CUDA-graph runner's post-replay ``advance_seq_lens()`` call is what
+    actually consumes this — the model's inner ``advance_seq_lens(pos_id_ns=...)``
+    runs at capture time only and is not replayed.
+    """
+    wrapper: FlashInferPrefillWrapper | FlashInferDecodeWrapper | None = None
+    pos_ids: torch.Tensor | None = None
+    seq_lens: list[int] | None = None
+    write_store: bool = True
+    custom_pos_advance: list[int] | None = None
+    # Plan memo: fingerprint of the last wrapper.plan() inputs for this label;
+    # when it matches, the re-plan is skipped. Only the cross-attention path
+    # sets it today (its context pages are immutable after add_cross_attn_kv),
+    # and only where plan states persist across steps (the CUDA-graph runner);
+    # the eager path rebuilds the plan surface per step and still re-plans.
+    #
+    # Future reference — this is a general-purpose tool, not cross-attn only.
+    # A regular (self-attention) label can memo its plan too; the extra care
+    # there is invalidation, since self-attn pages grow every decode step.
+    # The fingerprint would need to include the per-request seq_len (so appending
+    # a token misses the memo and re-plans), and any page-table remap (eviction /
+    # reallocation) must also bust the key. Given that, decode could skip the
+    # re-plan on the common "seq_len += 1, same pages" step. Deferred until a
+    # model needs it; the eager-path persistence noted above is the prerequisite.
+    plan_cache_key: "PlanCacheKey | None" = None
+    # Set when the dense-gen manager planned this label dense: the per-segment
+    # gather indices + varlen cu_seqlens needed to attend each generation
+    # segment over its contiguous frozen prefix. None on paged plans, which
+    # keep the FlashInfer path. See DenseGenAttentionManager.build_dense_plan.
+    dense_gen: dict | None = None
+
+
+class WorkspaceBufferManager:
+    def __init__(
+        self, size, device
+    ):
+        self.size = size
+        self.device = device
+        self.buffers = {}
+
+    def get(self, label: str="main"):
+        if label not in self.buffers:
+            self.buffers[label] = torch.empty(
+                self.size, dtype=torch.uint8, device=self.device
+            )
+        return self.buffers[label]
+
+
+class FlashInferAttentionManager:
+    """Paged FlashInfer attention (the default backend).
+
+    Builds batch-level FlashInfer index tensors from the sequence views it
+    is handed, plans the label's wrapper, and runs one FlashInfer call per
+    layer. The per-label ``_PlanState`` store belongs here; in CUDA-graph
+    mode the graph runner passes in its own per-slot store so the persistent
+    wrappers survive across steps.
+    """
+
+    def __init__(
+        self,
+        kv_cache: torch.Tensor | None,
+        kv_cache_config: KVCacheConfig,
+        buffer_manager: WorkspaceBufferManager | None,
+        device,
+        states: dict[str, _PlanState] | None = None,
+        cuda_graph_mode: bool = False,
+        enable_nvtx: bool = False,
+    ):
+        self.kv_cache = kv_cache
+        self.kv_cache_config = kv_cache_config
+        self.buffer_manager = buffer_manager
+        self.device = device
+        self.states = states if states is not None else {}
+        self.cuda_graph_mode = cuda_graph_mode
+        self.enable_nvtx = enable_nvtx
+
+    def plan(
+        self,
+        label: str,
+        segments: list[Segment],
+        views: list[SequenceView],
+        dtype: torch.dtype | None = None,
+        is_causal: bool = True,
+        write_store: bool = True,
+    ) -> None:
+        """Plan one label's attention over the given segment/view pairs."""
+        from mstar.utils.profiler import range_pop, range_push
+
+        assert self.kv_cache is not None
+
+        # Default the FlashInfer wrapper's dtype to whatever dtype the KV
+        # cache tensor was actually allocated in. Hardcoding bf16 here breaks
+        # any model that runs in fp32 (the wrapper would try to write
+        # bf16-cast K/V into an fp32 cache and torch raises a dtype mismatch
+        # in flashinfer_utils.set_kv_cache).
+        if dtype is None:
+            dtype = self.kv_cache.dtype
+
+        cfg = self.kv_cache_config
+        page_size = cfg.page_size
+
+        # CPU-side accumulation: only the four int32 tensors the FlashInfer
+        # wrapper consumes are built, in pure Python, one H2D each.
+        if self.enable_nvtx:
+            range_push("cache.plan_attention.build_lists", synchronize=False)
+        try:
+            qo_indptr_list = [0]
+            kv_indptr_list = [0]
+            all_page_indices = []
+            kv_last_page_lens = []
+            kv_cache_locations_list = []
+
+            for segment, view in zip(segments, views, strict=True):
+                qo_indptr_list.append(qo_indptr_list[-1] + segment.span)
+                all_page_indices.extend(view.page_indices)
+                kv_indptr_list.append(kv_indptr_list[-1] + len(view.page_indices))
+
+                last_page_len = view.length % page_size or page_size
+                kv_last_page_lens.append(last_page_len)
+                if segment.span == 1:
+                    kv_cache_locations_list.append([view.page_indices[-1], last_page_len - 1])
+        finally:
+            if self.enable_nvtx:
+                range_pop(synchronize=False)
+
+        # Build batched FlashInfer index tensors on CPU so wrapper.plan()
+        # doesn't trigger a synchronous D→H inside its body. FlashInfer
+        # calls ``indptr.to("cpu")`` / ``last_page_len.to("cpu")`` near the
+        # top of ``plan()`` to get host views of those metadata tensors;
+        # if we hand them GPU tensors that ``.to("cpu")`` becomes a
+        # synchronous default-stream sync that waits for the entire
+        # outstanding stream — including the speculatively-queued next
+        # decode step. By creating these on CPU directly, ``.to("cpu")``
+        # is a no-op. FlashInfer later copies the tiny int32 metadata to
+        # the device when it needs it; the source is pageable CPU memory, so
+        # ``non_blocking=True`` does not make that H2D copy asynchronous, but
+        # the tensors are batch-size length and the cost is inconsequential.
+        if self.enable_nvtx:
+            range_push("cache.plan_attention.make_tensors", synchronize=False)
+        try:
+            qo_indptr = torch.tensor(qo_indptr_list, dtype=torch.int32)
+            paged_kv_indptr = torch.tensor(kv_indptr_list, dtype=torch.int32)
+            paged_kv_indices = torch.tensor(all_page_indices, dtype=torch.int32)
+            paged_kv_last_page_len = torch.tensor(kv_last_page_lens, dtype=torch.int32)
+            kv_cache_locations = (
+                torch.tensor(kv_cache_locations_list, dtype=torch.long)
+                if len(kv_cache_locations_list) == len(segments)
+                else None
+            )
+        finally:
+            if self.enable_nvtx:
+                range_pop(synchronize=False)
+
+        seq_lens = [segment.span for segment in segments]
+        is_decode = all([sl == 1 for sl in seq_lens])
+        ps = self.states.get(label)
+        if ps is not None and ps.wrapper is not None:
+            wrapper = ps.wrapper
+        elif is_decode:
+            wrapper = FlashInferDecodeWrapper(
+                workspace_buffer=self.buffer_manager.get(label),
+                num_qo_heads=cfg.num_qo_heads,
+                num_kv_heads=cfg.num_kv_heads,
+                head_dim=cfg.head_dim,
+                page_size=page_size,
+                device=self.device,
+                enable_nvtx=self.enable_nvtx,
+                backend=cfg.flashinfer_backend,
+            )
+            ps = _PlanState(wrapper=wrapper)
+            self.states[label] = ps
+        else:
+            wrapper = FlashInferPrefillWrapper(
+                workspace_buffer=self.buffer_manager.get(label),
+                num_qo_heads=cfg.num_qo_heads,
+                num_kv_heads=cfg.num_kv_heads,
+                head_dim=cfg.head_dim,
+                page_size=page_size,
+                device=self.device,
+                enable_nvtx=self.enable_nvtx,
+                backend=cfg.flashinfer_backend,
+            )
+            ps = _PlanState(wrapper=wrapper)
+            self.states[label] = ps
+
+        if self.enable_nvtx:
+            range_push("cache.plan_attention.wrapper_plan", synchronize=False)
+        try:
+            if isinstance(wrapper, FlashInferDecodeWrapper):
+                wrapper.plan(
+                    paged_kv_indptr=paged_kv_indptr,
+                    paged_kv_indices=paged_kv_indices,
+                    paged_kv_last_page_len=paged_kv_last_page_len,
+                    kv_cache_locations=kv_cache_locations,
+                    dtype=dtype,
+                )
+            else:
+                wrapper.plan(
+                    qo_indptr=qo_indptr,
+                    paged_kv_indptr=paged_kv_indptr,
+                    paged_kv_indices=paged_kv_indices,
+                    paged_kv_last_page_len=paged_kv_last_page_len,
+                    causal=is_causal,
+                    dtype=dtype,
+                )
+        finally:
+            if self.enable_nvtx:
+                range_pop(synchronize=False)
+        # seq_lens is read by the flush_to_store path; write_store by
+        # run_attention.
+        ps.seq_lens = seq_lens
+        ps.write_store = write_store
+        # A paged plan clears any prior dense plan (set by the dense-gen
+        # manager) so run routes this label back through the wrapper.
+        ps.dense_gen = None
+
+    def plan_batched_cfg(
+        self,
+        combined_label: str,
+        segments: list[Segment],
+        views: list[SequenceView],
+        dtype=torch.bfloat16,
+        is_causal: bool = False,
+        write_store: bool = False,
+    ) -> None:
+        """Plan a single FlashInfer batch across multiple cache labels.
+
+        ``segments``/``views`` arrive label-major: every (label, request)
+        pair in batch order, exactly the packed layout the CFG forward uses.
+        """
+        assert self.kv_cache is not None
+
+        cfg = self.kv_cache_config
+        page_size = cfg.page_size
+
+        # CPU-side accumulation (see plan for the same pattern).
+        qo_indptr_list = [0]
+        kv_indptr_list = [0]
+        all_page_indices = []
+        kv_last_page_lens = []
+        combined_seq_lens = []
+
+        for segment, view in zip(segments, views, strict=True):
+            qo_indptr_list.append(qo_indptr_list[-1] + segment.span)
+            all_page_indices.extend(view.page_indices)
+            kv_indptr_list.append(
+                kv_indptr_list[-1] + len(view.page_indices)
+            )
+
+            last_page_len = view.length % page_size or page_size
+            kv_last_page_lens.append(last_page_len)
+            combined_seq_lens.append(segment.span)
+
+        # CPU tensors — see comment in ``plan`` above. FlashInfer
+        # async-H2Ds these inside ``plan()``; passing GPU tensors would
+        # trigger a synchronous default-stream sync via the internal
+        # ``.to("cpu")`` call.
+        qo_indptr = torch.tensor(qo_indptr_list, dtype=torch.int32)
+        paged_kv_indptr = torch.tensor(kv_indptr_list, dtype=torch.int32)
+        paged_kv_indices = torch.tensor(all_page_indices, dtype=torch.int32)
+        paged_kv_last_page_len = torch.tensor(kv_last_page_lens, dtype=torch.int32)
+
+        ps = self.states.get(combined_label)
+        if self.cuda_graph_mode and ps is not None and ps.wrapper is not None:
+            # CUDA-graph mode: reuse the persistent wrapper across denoise steps.
+            # plan() updates its static buffers via .copy_() so the captured
+            # kernel picks up each step's page table without reallocating.
+            wrapper = ps.wrapper
+        elif self.cuda_graph_mode:
+            # First call under capture: build the persistent wrapper sized for the
+            # fixed batch (labels x requests) and token budget.
+            wrapper = FlashInferPrefillWrapper(
+                workspace_buffer=self.buffer_manager.get(combined_label),
+                num_qo_heads=cfg.num_qo_heads,
+                num_kv_heads=cfg.num_kv_heads,
+                head_dim=cfg.head_dim,
+                page_size=page_size,
+                batch_size=len(segments),
+                max_total_tokens=sum(combined_seq_lens),
+                max_num_pages=cfg.max_num_pages,
+                device=self.device,
+                use_cuda_graph=True,
+                enable_nvtx=self.enable_nvtx,
+                backend=cfg.flashinfer_backend,
+            )
+            ps = _PlanState(wrapper=wrapper)
+            self.states[combined_label] = ps
+        else:
+            # Eager mode: a fresh wrapper each call (the plan surface is rebuilt
+            # per forward, so there is nothing persistent to reuse).
+            wrapper = FlashInferPrefillWrapper(
+                workspace_buffer=self.buffer_manager.get(combined_label),
+                num_qo_heads=cfg.num_qo_heads,
+                num_kv_heads=cfg.num_kv_heads,
+                head_dim=cfg.head_dim,
+                page_size=page_size,
+                enable_nvtx=self.enable_nvtx,
+                backend=cfg.flashinfer_backend,
+            )
+            ps = _PlanState(wrapper=wrapper)
+            self.states[combined_label] = ps
+
+        wrapper.plan(
+            qo_indptr=qo_indptr,
+            paged_kv_indptr=paged_kv_indptr,
+            paged_kv_indices=paged_kv_indices,
+            paged_kv_last_page_len=paged_kv_last_page_len,
+            causal=is_causal,
+            dtype=dtype,
+        )
+        ps.seq_lens = combined_seq_lens
+        ps.write_store = write_store
+        ps.dense_gen = None
+
+    def write_kv(self, k: torch.Tensor, v: torch.Tensor, layer_idx: int, label: str) -> None:
+        """Write this step's K/V into the paged cache at the label's planned
+        positions. Separate from ``run`` so strategies that attend without
+        storing (or store without attending) stay expressible."""
+        self.states[label].wrapper.set_kv_cache(self.kv_cache[layer_idx], k, v)
+
+    def run(self, q: torch.Tensor, layer_idx: int, label: str) -> torch.Tensor:
+        """Run the label's pre-planned paged attention for one layer."""
+        return self.states[label].wrapper.run(q, self.kv_cache[layer_idx])
+
+
+class DenseGenAttentionManager(FlashInferAttentionManager):
+    """FlashInfer manager with a dense generation-attention fast path.
+
+    Runs non-causal generation attention as a dense FlashAttention-3 pass
+    over a contiguous [frozen-prefix | fresh] sequence instead of the paged
+    FlashInfer prefill. Diffusion recomputes every generation K/V each step
+    (only the tiny text prefix is reused), so the paged path's per-step
+    full-buffer K/V write is pure overhead here; a dense pass gathers the
+    small prefix, concatenates it with the freshly projected K/V, and runs
+    one varlen kernel.
+    """
+
+    def plan_dense(
+        self,
+        label: str,
+        entries: list[tuple[SequenceView, int, KVRequestState]],
+        seq_lens,
+        write_store: bool,
+    ) -> None:
+        """Build the dense gather/varlen plan for ``label`` and store it on
+        the label's plan state. ``entries`` are (frozen-prefix view, fresh
+        generation length, request state) per batch element, in the packed
+        batch order; the request state carries the cross-step cache slot for
+        the gathered prefix."""
+        ps = self.states.get(label) or _PlanState()
+        self.states[label] = ps
+        ps.seq_lens = seq_lens
+        ps.write_store = write_store
+        ps.dense_gen = self._build_dense_plan(entries)
+
+    def _build_dense_plan(
+        self, entries: list[tuple[SequenceView, int, KVRequestState]]
+    ) -> dict:
+        """Pre-compute the per-segment gather + varlen layout for the dense
+        generation-attention path. Each segment attends its fresh generation
+        tokens over its frozen text prefix; the prefix lives in the pages
+        written at prefill, so we record the page indices to gather it from
+        (the same across all layers) and the cumulative-sequence-length
+        tensors a single varlen kernel needs. Built once per denoise step,
+        reused by every layer's run."""
+        page_size = self.kv_cache_config.page_size
+        segs = []  # (prefix_page_indices, prefix_len, gen_len, state)
+        cu_q = [0]
+        cu_k = [0]
+        max_q = 0
+        max_k = 0
+        for view, gen_len, state in entries:
+            prefix_len = view.length
+            n_pages = (prefix_len + page_size - 1) // page_size
+            idx = torch.tensor(
+                view.page_indices[:n_pages], dtype=torch.long, device=self.device
+            )
+            segs.append((idx, prefix_len, gen_len, state))
+            cu_q.append(cu_q[-1] + gen_len)
+            cu_k.append(cu_k[-1] + prefix_len + gen_len)
+            max_q = max(max_q, gen_len)
+            max_k = max(max_k, prefix_len + gen_len)
+        return {
+            "segs": segs,
+            "cu_q": torch.tensor(cu_q, dtype=torch.int32, device=self.device),
+            "cu_k": torch.tensor(cu_k, dtype=torch.int32, device=self.device),
+            "max_q": max_q,
+            "max_k": max_k,
+        }
+
+    @torch.compiler.disable
+    def run_dense(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_idx: int, dg: dict
+    ) -> torch.Tensor:
+        """Dense generation attention: per segment, take the frozen text-prefix
+        K/V, concatenate it with this segment's fresh K/V, and attend
+        non-causally with one FlashAttention-3 varlen kernel. Bypasses the paged
+        write entirely (the generation K/V is recomputed every step, so
+        persisting it is wasted work). The frozen prefix is gathered from the
+        paged cache once per layer and cached on the request state, then reused
+        across denoise steps (it never changes during denoise)."""
+        from fa3_fwd_interface import flash_attn_varlen_func
+
+        cfg = self.kv_cache_config
+        num_kv_heads, head_dim = cfg.num_kv_heads, cfg.head_dim
+        kv_layer = self.kv_cache[layer_idx]  # [max_pages, 2, page_size, num_kv_heads, head_dim]
+
+        k_parts, v_parts = [], []
+        offset = 0
+        for idx, prefix_len, gen_len, state in dg["segs"]:
+            prefix_cache = state.dense_prefix_kv
+            if prefix_cache is None:
+                prefix_cache = state.dense_prefix_kv = {}
+            cached = prefix_cache.get(layer_idx)
+            if cached is None:
+                sub = kv_layer[idx]  # [n_pages, 2, page_size, num_kv_heads, head_dim]
+                k_pref = sub[:, 0].reshape(-1, num_kv_heads, head_dim)[:prefix_len].clone()
+                v_pref = sub[:, 1].reshape(-1, num_kv_heads, head_dim)[:prefix_len].clone()
+                prefix_cache[layer_idx] = (k_pref, v_pref)
+            else:
+                k_pref, v_pref = cached
+            k_parts.append(k_pref)
+            k_parts.append(k[offset:offset + gen_len])
+            v_parts.append(v_pref)
+            v_parts.append(v[offset:offset + gen_len])
+            offset += gen_len
+        key = torch.cat(k_parts, dim=0)
+        val = torch.cat(v_parts, dim=0)
+        if q.dtype != key.dtype:
+            q = q.to(key.dtype)
+
+        out = flash_attn_varlen_func(
+            q, key, val, dg["cu_q"], dg["cu_k"], dg["max_q"], dg["max_k"], causal=False,
+        )
+        return out[0] if isinstance(out, tuple) else out
+
+
+class CrossAttentionManager:
+    """Cross-attention over a fixed encoder context (issue #160).
+
+    Cross-attention is non-causal attention over a separate, read-only
+    context KV: written once at encode time (``write_context``), planned per
+    step against the decoder's query lengths (``plan``), and executed per
+    layer (``run``). The context lives in a per-source pool whose head
+    config may differ from the decoder's self-attention; plan state rides
+    the shared per-label store under the resolved
+    ``{label}::CROSS_ATTN::{source}`` label. Cross labels never plan RoPE —
+    context positions, if any, are baked in at encode time.
+    """
+
+    def __init__(
+        self,
+        source: str,
+        pool: CrossAttnPool,
+        kv_pool: KVCachePool,
+        buffer_manager: WorkspaceBufferManager | None,
+        device,
+        states: dict[str, _PlanState],
+        enable_nvtx: bool = False,
+    ):
+        self.source = source
+        self.pool = pool
+        self.kv_pool = kv_pool
+        self.buffer_manager = buffer_manager
+        self.device = device
+        self.states = states
+        self.enable_nvtx = enable_nvtx
+
+    def write_context(
+        self,
+        request_ids: list[str],
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_idx: int,
+        seq_lens: list[int] | None,
+        base_label: str,
+    ) -> None:
+        """Write encoder-context K/V for one layer into the source's pool.
+
+        Called once per request per layer at encode time. ``k``/``v`` are
+        packed ``(total_context_tokens, num_kv_heads, head_dim)`` across
+        ``request_ids``; ``seq_lens`` gives per-request context lengths
+        (defaults to a single request owning the whole tensor). Pages are
+        allocated on first write (layer 0) and reused for the rest.
+        """
+        cross_label = cross_attn_label(base_label, self.source)
+        page_size = self.pool.alloc_config.page_size
+
+        if seq_lens is None:
+            assert len(request_ids) == 1, (
+                "add_cross_attn_kv needs seq_lens for multi-request batches"
+            )
+            seq_lens = [k.shape[0]]
+
+        offset = 0
+        for rid, ctx_len in zip(request_ids, seq_lens, strict=True):
+            resident = self.kv_pool.view(Segment(rid, cross_label, 0)).length
+            if resident == 0:
+                # First write for this context: reserve its pages and commit
+                # the extent up front; later layers reuse them. The position
+                # counter stays untouched (context positions are baked in at
+                # encode time).
+                segment = Segment(rid, cross_label, ctx_len)
+                self.kv_pool.admit(segment)
+                self.kv_pool.commit(segment, pos_advance=0)
+            else:
+                assert resident == ctx_len, (
+                    f"cross-attn context for {rid!r}/{self.source!r} already "
+                    f"written with length {resident}, got {ctx_len}"
+                )
+            view = self.kv_pool.view(Segment(rid, cross_label, 0))
+
+            positions = torch.arange(ctx_len, device=self.device)
+            page_indices = torch.tensor(
+                view.page_indices, dtype=torch.long, device=self.device,
+            )
+            token_to_page = page_indices[
+                torch.div(positions, page_size, rounding_mode="floor")
+            ]
+            token_to_cache = positions % page_size
+
+            layer_cache = self.pool.kv_cache[layer_idx]
+            dtype = self.pool.kv_cache.dtype
+            layer_cache[token_to_page, 0, token_to_cache] = \
+                k[offset:offset + ctx_len].to(dtype)
+            layer_cache[token_to_page, 1, token_to_cache] = \
+                v[offset:offset + ctx_len].to(dtype)
+            offset += ctx_len
+
+    def plan(
+        self,
+        request_ids: list[str],
+        q_seq_lens: list[int],
+        dtype: torch.dtype | None,
+        base_label: str,
+    ) -> None:
+        """Plan the cross-attention wrapper for this step's decoder queries.
+
+        ``q_seq_lens`` is the number of decoder query tokens per request
+        (matches the self-attention plan's seq_lens). The context side is
+        read from the pool state written by ``write_context`` — pages are
+        fixed, so nothing is allocated here. Always uses the prefill wrapper
+        with ``causal=False`` (decode-style single queries still attend to
+        the full context).
+        """
+        cross_label = cross_attn_label(base_label, self.source)
+        cfg = self.pool.alloc_config
+        page_size = cfg.page_size
+
+        if dtype is None:
+            dtype = self.pool.kv_cache.dtype
+
+        page_indices_per_request: list[tuple[int, ...]] = []
+        context_lens: list[int] = []
+        for rid in request_ids:
+            view = self.kv_pool.view(Segment(rid, cross_label, 0))
+            assert view.length > 0, (
+                f"plan_cross_attention before add_cross_attn_kv for {rid!r} "
+                f"(source {self.source!r})"
+            )
+            page_indices_per_request.append(view.page_indices)
+            context_lens.append(view.length)
+
+        indptrs = build_paged_indptrs(
+            q_seq_lens, page_indices_per_request, context_lens, page_size,
+        )
+
+        # Skip the re-plan when its inputs are unchanged (common decode case).
+        # Keyed on the page indices so a reused request id with a new context
+        # can't alias a stale plan.
+        plan_key = PlanCacheKey(
+            q_seq_lens=tuple(q_seq_lens),
+            page_indices=tuple(indptrs.paged_kv_indices.tolist()),
+            last_page_lens=tuple(indptrs.paged_kv_last_page_len.tolist()),
+            dtype=dtype,
+        )
+
+        ps = self.states.get(cross_label)
+        if ps is not None and ps.wrapper is not None and ps.plan_cache_key == plan_key:
+            # plan_key carries q_seq_lens, so a match means ps.seq_lens (set at
+            # plan time below) already equals the current q_seq_lens.
+            return
+        if ps is None or ps.wrapper is None:
+            wrapper = FlashInferPrefillWrapper(
+                workspace_buffer=self.buffer_manager.get(cross_label),
+                num_qo_heads=cfg.num_qo_heads,
+                num_kv_heads=cfg.num_kv_heads,
+                head_dim=cfg.head_dim,
+                page_size=page_size,
+                device=self.device,
+                enable_nvtx=self.enable_nvtx,
+            )
+            ps = _PlanState(wrapper=wrapper)
+            self.states[cross_label] = ps
+
+        ps.wrapper.plan(
+            qo_indptr=indptrs.qo_indptr,
+            paged_kv_indptr=indptrs.paged_kv_indptr,
+            paged_kv_indices=indptrs.paged_kv_indices,
+            paged_kv_last_page_len=indptrs.paged_kv_last_page_len,
+            causal=False,
+            dtype=dtype,
+        )
+        ps.seq_lens = q_seq_lens
+        ps.write_store = False
+        ps.plan_cache_key = plan_key
+
+    def run(self, q: torch.Tensor, layer_idx: int, base_label: str) -> torch.Tensor:
+        """Run pre-planned cross-attention against the source's context pool.
+
+        The context K/V were written once by ``write_context``; nothing is
+        written here.
+        """
+        cross_label = cross_attn_label(base_label, self.source)
+        ps = self.states.get(cross_label)
+        assert ps is not None and ps.wrapper is not None, (
+            f"run_cross_attn before plan_cross_attention (label {cross_label!r})"
+        )
+        return ps.wrapper.run(q, self.pool.kv_cache[layer_idx])
diff --git a/mstar/engine/resources/kv_pool.py b/mstar/engine/resources/kv_pool.py
index 415660c54..878274f0b 100644
--- a/mstar/engine/resources/kv_pool.py
+++ b/mstar/engine/resources/kv_pool.py
@@ -117,3 +117,37 @@ def commit(self, segment: Segment, pos_advance: int | None = None) -> None:
     def positions(self, request_id: str, label: str) -> int:
         """Current position counter for one stream, read-only."""
         return self._manager.get_state(request_id, label).position_id_start
+
+    def fork(
+        self,
+        request_id: str,
+        from_label: str,
+        to_label: str,
+        realloc: bool = False,
+    ) -> None:
+        """Make one stream's content the state of another stream of the same
+        request: reserve the destination's pages, mirror the stored length
+        and position counter, and copy the page data. With ``realloc`` the
+        destination is reset first; otherwise only the pages past its
+        current length are copied (an incremental top-up). Allocates, so it
+        can fail like any admit."""
+        manager = self._manager
+        from_state = manager.get_state(request_id, from_label)
+
+        if realloc:
+            manager.reset_label(request_id, to_label)
+
+        to_state = manager.get_state(request_id, to_label)
+        start_pos = to_state.seq_len // self.page_size
+        manager.alloc(request_id, to_label, seq_len=from_state.seq_len)
+
+        to_state.seq_len = from_state.seq_len
+        to_state.position_id_start = from_state.position_id_start
+
+        tensor = self._arena.tensor
+        for src_page, dst_page in zip(
+            from_state.page_indices[start_pos:],
+            to_state.page_indices[start_pos:],
+            strict=True,
+        ):
+            tensor[:, dst_page] = tensor[:, src_page]
diff --git a/test/modular/test_kv_pool.py b/test/modular/test_kv_pool.py
index 5c3201565..e0912e705 100644
--- a/test/modular/test_kv_pool.py
+++ b/test/modular/test_kv_pool.py
@@ -35,7 +35,11 @@
 )
 
 
-def _make_pool(max_num_pages: int = 16, page_size: int = 8) -> tuple[KVCachePool, PagedAllocationManager]:
+def _make_pool(
+    max_num_pages: int = 16,
+    page_size: int = 8,
+    with_tensor: bool = False,
+) -> tuple[KVCachePool, PagedAllocationManager]:
     manager = PagedAllocationManager.__new__(PagedAllocationManager)
     manager.config = KVCacheConfig(
         num_layers=1,
@@ -47,7 +51,9 @@ def _make_pool(max_num_pages: int = 16, page_size: int = 8) -> tuple[KVCachePool
     )
     manager.page_allocator = PageAllocator(max_num_pages)
     manager.request_states = {}
-    manager.kv_cache = None
+    manager.kv_cache = (
+        torch.zeros(1, max_num_pages, 2, page_size, 1, 1) if with_tensor else None
+    )
     manager.write_policy = StoreWritePolicy.ALWAYS
     manager._kv_transfer_engine = None
     manager._offload_stream = None
@@ -215,6 +221,86 @@ def test_try_allocate_shortfall_returns_none(self):
         assert arena.num_free == 2
 
 
+class TestFork:
+    def _fill_stream(self, pool, manager, rid, label, tokens):
+        segment = Segment(rid, label, tokens)
+        pool.admit(segment)
+        pool.commit(segment)
+        # Stamp each page with its own index so copies are checkable.
+        for page in manager.get_state(rid, label).page_indices:
+            manager.kv_cache[:, page] = float(page + 1)
+
+    def test_fork_mirrors_accounting_and_copies_pages(self):
+        pool, manager = _make_pool(page_size=8, with_tensor=True)
+        manager.add_request("r", ["main", "snap"])
+        self._fill_stream(pool, manager, "r", "main", 12)
+        pool.commit(Segment("r", "main", 0), pos_advance=30)  # counter != length
+
+        pool.fork("r", "main", "snap")
+
+        main_state = manager.get_state("r", "main")
+        snap_state = manager.get_state("r", "snap")
+        assert snap_state.seq_len == 12
+        assert snap_state.position_id_start == 42
+        assert snap_state.page_indices != main_state.page_indices
+        for src, dst in zip(
+            main_state.page_indices, snap_state.page_indices, strict=True
+        ):
+            assert torch.equal(manager.kv_cache[:, dst], manager.kv_cache[:, src])
+
+    def test_fork_tops_up_only_past_the_destination_length(self):
+        pool, manager = _make_pool(page_size=8, with_tensor=True)
+        manager.add_request("r", ["main", "snap"])
+        self._fill_stream(pool, manager, "r", "main", 16)  # 2 full pages
+
+        # The destination already holds one committed page of its own data.
+        snap_seg = Segment("r", "snap", 8)
+        pool.admit(snap_seg)
+        pool.commit(snap_seg)
+        first_snap_page = manager.get_state("r", "snap").page_indices[0]
+        manager.kv_cache[:, first_snap_page] = -1.0
+
+        pool.fork("r", "main", "snap")
+
+        snap_state = manager.get_state("r", "snap")
+        assert snap_state.seq_len == 16
+        # Page 0 was already resident on the destination and is not re-copied.
+        assert torch.all(manager.kv_cache[:, snap_state.page_indices[0]] == -1.0)
+        src_page_1 = manager.get_state("r", "main").page_indices[1]
+        assert torch.equal(
+            manager.kv_cache[:, snap_state.page_indices[1]],
+            manager.kv_cache[:, src_page_1],
+        )
+
+    def test_fork_with_realloc_replaces_the_destination(self):
+        pool, manager = _make_pool(page_size=8, with_tensor=True)
+        manager.add_request("r", ["main", "snap"])
+        self._fill_stream(pool, manager, "r", "main", 8)
+        self._fill_stream(pool, manager, "r", "snap", 16)
+        free_before = pool.num_free_pages
+
+        pool.fork("r", "main", "snap", realloc=True)
+
+        snap_state = manager.get_state("r", "snap")
+        assert snap_state.seq_len == 8
+        assert len(snap_state.page_indices) == 1
+        # The two old destination pages went back, one new page came out.
+        assert pool.num_free_pages == free_before + 1
+        src_page = manager.get_state("r", "main").page_indices[0]
+        assert torch.equal(
+            manager.kv_cache[:, snap_state.page_indices[0]],
+            manager.kv_cache[:, src_page],
+        )
+
+    def test_fork_can_fail_like_any_admit(self):
+        pool, manager = _make_pool(max_num_pages=2, page_size=8, with_tensor=True)
+        manager.add_request("r", ["main", "snap"])
+        self._fill_stream(pool, manager, "r", "main", 16)  # both pages
+
+        with pytest.raises(AllocationFailedError):
+            pool.fork("r", "main", "snap")
+
+
 class TestBoundaryValuesAreImmutable:
     def test_segment_frozen(self):
         segment = Segment("r", "main", 4)

From 82c76ee852a60970f0bfa5425223847295af9c0b Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 10:47:00 +0000
Subject: [PATCH 05/20] Drive the engine step lifecycle through the pool and a
 step runner

---
 mstar/engine/kv_cache_engine.py    | 145 ++++++++++++++++--------
 mstar/engine/kv_store.py           |  18 +--
 mstar/engine/resources/__init__.py |   3 +
 mstar/engine/resources/kv_pool.py  |  40 ++++++-
 mstar/engine/resources/step.py     | 100 ++++++++++++++++
 test/modular/test_kv_pool.py       |  95 +++++++++++++++-
 test/modular/test_step_runner.py   | 176 +++++++++++++++++++++++++++++
 7 files changed, 512 insertions(+), 65 deletions(-)
 create mode 100644 mstar/engine/resources/step.py
 create mode 100644 test/modular/test_step_runner.py

diff --git a/mstar/engine/kv_cache_engine.py b/mstar/engine/kv_cache_engine.py
index df85289b5..349ea7bd7 100644
--- a/mstar/engine/kv_cache_engine.py
+++ b/mstar/engine/kv_cache_engine.py
@@ -36,6 +36,8 @@
     StoreWritePolicy,
     TransferEngineInfo,
 )
+from mstar.engine.resources import KVCachePool, RopeEmbedder
+from mstar.engine.resources.step import StepPlan, StepRunner
 from mstar.model.submodule_base import ARNodeInputs, ARNodeSubmodule, ModelInputsFromEngine
 from mstar.utils.profiler import range_pop, range_push
 from mstar.utils.sampling import MultiSampler, MultiSamplingConfig
@@ -57,6 +59,13 @@ class KVManagement:
     # sources); precomputed at build time since it can't change after
     # startup, so per-request add/remove doesn't re-walk cross_pools.
     cross_alloc_managers: list[PagedAllocationManager] = field(default_factory=list)
+    # Persistent resource fronts over the managers above: the pool is the
+    # engine's surface for admission, retrieval, position reads, and
+    # publishing; the embedder owns position semantics. Built once at
+    # load_model and shared by every step's facade.
+    kv_pool: KVCachePool | None = None
+    cross_kv_pools: dict[str, KVCachePool] = field(default_factory=dict)
+    rope_embedder: RopeEmbedder | None = None
 
 
 def _build_cross_pools(
@@ -155,6 +164,10 @@ def __init__(
         self.kv_management: dict[str, KVManagement] = {}
         self.submodule_management: dict[str, SubmoduleManagement] = {}
 
+        # Sequences each batch's resource lifecycle: admit before prepare,
+        # plan surface before forward, publish after.
+        self.step_runner = StepRunner()
+
         self.device = None
         self.autocast_dtype = autocast_dtype
 
@@ -256,14 +269,15 @@ def load_model(
                 if all(pool.alloc_manager is not m for m in cross_alloc_managers):
                     cross_alloc_managers.append(pool.alloc_manager)
 
+            alloc_manager = PagedAllocationManager(
+                config=cfg,
+                kv_cache=kv_cache,
+                transfer_engine_info=transfer_engine_info
+            )
             kv_mgmt = KVManagement(
                 kv_cache_config=cfg,
                 kv_cache=kv_cache,
-                alloc_manager=PagedAllocationManager(
-                    config=cfg,
-                    kv_cache=kv_cache,
-                    transfer_engine_info=transfer_engine_info
-                ),
+                alloc_manager=alloc_manager,
                 cpu_page_pool=cpu_page_pool,
                 buffer_manager = WorkspaceBufferManager(
                     int(os.environ.get("MSTAR_WORKSPACE_BUFFER_MB", "512")) * 1024 * 1024,
@@ -271,6 +285,12 @@ def load_model(
                 ),
                 cross_pools=cross_pools,
                 cross_alloc_managers=cross_alloc_managers,
+                kv_pool=KVCachePool(alloc_manager),
+                cross_kv_pools={
+                    source: KVCachePool(pool.alloc_manager)
+                    for source, pool in cross_pools.items()
+                },
+                rope_embedder=RopeEmbedder(),
             )
             self.kv_management[cfg.get_node_str()] = kv_mgmt
 
@@ -316,6 +336,9 @@ def _create_cache_manager(
             auto_write_store=autowrite,
             enable_nvtx=self.enable_nvtx,
             cross_pools=cache_mgmt.cross_pools,
+            kv_pool=cache_mgmt.kv_pool,
+            cross_kv_pools=cache_mgmt.cross_kv_pools,
+            rope_embedder=cache_mgmt.rope_embedder,
         )
 
     def _compile_submodules(self) -> None:
@@ -517,11 +540,9 @@ def _sample_decode_outputs(
     def _execute_batched(
         self, batch: NodeBatch, submodule: ARNodeSubmodule,
         inputs: list[ARNodeInputs], sampler: MultiSampler,
+        cache_manager: BatchedCacheManager,
     ) -> NodeOutput:
         """Execute batch with BatchedCacheManager for true vectorized batching."""
-        cache_manager = self._create_cache_manager(
-            batch.request_ids, batch.node_name
-        )
         engine_inputs = ModelInputsFromEngine(
             request_ids=batch.request_ids,
             per_request_info=batch.per_request_info,
@@ -639,13 +660,14 @@ def _execute_sequential(
         submodule: ARNodeSubmodule,
         inputs: list[ARNodeInputs],
         sampler: MultiSampler,
+        cache_managers: list[BatchedCacheManager],
     ) -> NodeOutput:
         """Original per-request execution with CacheHandle."""
         per_request_outputs = {}
 
-        for rid, node_inputs in zip(batch.request_ids, inputs, strict=True):
-            cache_manager = self._create_cache_manager([rid], batch.node_name)
-            inputs = batch.per_request_input_tensors.get(rid, {})
+        for rid, node_inputs, cache_manager in zip(
+            batch.request_ids, inputs, cache_managers, strict=True
+        ):
             engine_inputs = ModelInputsFromEngine(
                 request_ids=[rid],
                 per_request_info={
@@ -869,30 +891,27 @@ def execute_batch(self, batch: NodeBatch) -> NodeOutput:
                 range_pop()
 
     def finalize_batch(self, batch: NodeBatch) -> None:
-        """Mirror this engine's per-request KV seq_info back onto
+        """Publish each request's durable pool state onto
         ``batch.per_request_info`` so the next iter / conductor sees the
         updated page indices, seq_len, and position_id_start.
 
         Safe to call after a successful forward, after an allocation
-        failure, or after an unrelated exception — the writeback reads
-        the alloc manager's current state, which always reflects whatever
-        progress this batch made.
+        failure, or after an unrelated exception — publish describes the
+        pool's current state, which always reflects whatever progress this
+        batch made.
         """
         if batch.node_name not in self.submodule_management:
             return
         submod_mgmt = self.submodule_management[batch.node_name]
         cache_mgmt = submod_mgmt.kv_management
-        kv_cache_string = cache_mgmt.kv_cache_config.get_node_str()
-        for req_id in batch.request_ids:
-            info = batch.per_request_info.get(req_id)
-            if info is None:
-                continue
-            info.per_label_seq_info.add(
-                kv_cache_string,
-                submod_mgmt.tp_group.rank,
-                submod_mgmt.tp_group.world_size,
-                cache_mgmt.alloc_manager.get_per_label_seq_info(req_id),
-            )
+        self.step_runner.publish(
+            request_ids=batch.request_ids,
+            per_request_info=batch.per_request_info,
+            pool=cache_mgmt.kv_pool,
+            kv_cache_string=cache_mgmt.kv_cache_config.get_node_str(),
+            tp_rank=submod_mgmt.tp_group.rank,
+            tp_world_size=submod_mgmt.tp_group.world_size,
+        )
 
     def prepare_batch(self, batch: NodeBatch) -> PreparedBatch:
         """KV sync retrieve, per-request sampler config, then per-rid
@@ -912,18 +931,14 @@ def prepare_batch(self, batch: NodeBatch) -> PreparedBatch:
 
         if self.enable_nvtx:
             range_push("kv_cache.kv_sync_retrieve", synchronize=False)
-        world_size = submod_mgmt.tp_group.world_size
-        for req_id, info in batch.per_request_info.items():
-            if info.per_label_seq_info.world_size.get(kv_cache_string, world_size) != world_size:
-                raise RuntimeError(
-                    "KV cache transfer across TP world size is currently disallowed"
-                ) # TODO: figure out fanin/fanout for KV cache transfer
-            for label, seq_info in info.per_label_seq_info.get(
-                kv_cache_string, submod_mgmt.tp_group.rank
-            ).items():
-                if needed_labels is not None and label not in needed_labels:
-                    continue
-                cache_mgmt.alloc_manager.sync_retrieve(req_id, label, seq_info)
+        self.step_runner.admit(
+            per_request_info=batch.per_request_info,
+            pool=cache_mgmt.kv_pool,
+            kv_cache_string=kv_cache_string,
+            tp_rank=submod_mgmt.tp_group.rank,
+            tp_world_size=submod_mgmt.tp_group.world_size,
+            needed_labels=needed_labels,
+        )
         if self.enable_nvtx:
             range_pop(synchronize=False)
 
@@ -949,11 +964,9 @@ def prepare_batch(self, batch: NodeBatch) -> PreparedBatch:
             range_push("kv_cache.prepare_inputs")
         for rid in batch.request_ids:
             try:
-                labels = cache_mgmt.alloc_manager.get_labels(rid)
                 pos_info = {
-                    label: cache_mgmt.alloc_manager.get_state(
-                        rid, label
-                    ).get_pos_info() for label in labels
+                    label: cache_mgmt.kv_pool.pos_info(rid, label)
+                    for label in cache_mgmt.kv_pool.labels(rid)
                 }
                 req_inputs = submodule.prepare_inputs(
                     graph_walk=batch.graph_walk,
@@ -998,12 +1011,43 @@ def prepare_batch(self, batch: NodeBatch) -> PreparedBatch:
             failed_requests=failed
         )
 
-    def execute_forward(self, planned: PlannedBatch) -> NodeOutput:
-        """Dispatch CUDA-graph / batched / sequential.
+    def plan_batch(self, prepared: PreparedBatch) -> PlannedBatch:
+        """Choose the execution path and build the step's plan surface.
 
         Priority: CUDA graph (largest single launch) > batched (single
-        FlashInfer plan + forward) > sequential (per-rid fallback).
+        FlashInfer plan + forward) > sequential (per-rid fallback). The
+        graph path keeps its captured surface inside the runner; the eager
+        paths get their cache-manager facades built here, before the
+        forward, so the model's plan calls in preprocess land on a surface
+        whose pools were admitted first.
         """
+        batch = prepared.batch
+        submodule = prepared.submodule
+        node_inputs = prepared.node_inputs
+
+        if not batch.request_ids:
+            return PlannedBatch(
+                prepared=prepared, metadata={"step": StepPlan(mode="sequential")}
+            )
+
+        if self._can_use_cuda_graph(batch, node_inputs):
+            mode = "graph"
+        elif submodule.can_batch(batch, node_inputs):
+            mode = "batched"
+        else:
+            mode = "sequential"
+
+        step = self.step_runner.plan(
+            mode=mode,
+            request_ids=batch.request_ids,
+            build_manager=lambda rids: self._create_cache_manager(
+                rids, batch.node_name
+            ),
+        )
+        return PlannedBatch(prepared=prepared, metadata={"step": step})
+
+    def execute_forward(self, planned: PlannedBatch) -> NodeOutput:
+        """Run the batch on the path ``plan_batch`` chose."""
         batch = planned.batch
         submodule = planned.submodule
         node_inputs = planned.node_inputs
@@ -1015,7 +1059,8 @@ def execute_forward(self, planned: PlannedBatch) -> NodeOutput:
 
         submod_mgmt.tp_group.barrier()
 
-        if self._can_use_cuda_graph(batch, node_inputs):
+        step: StepPlan = planned.metadata["step"]
+        if step.mode == "graph":
             if self.enable_nvtx:
                 range_push("kv_cache.cuda_graph_path", synchronize=False)
             try:
@@ -1023,12 +1068,13 @@ def execute_forward(self, planned: PlannedBatch) -> NodeOutput:
             finally:
                 if self.enable_nvtx:
                     range_pop(synchronize=False)
-        elif submodule.can_batch(batch, node_inputs):
+        elif step.mode == "batched":
             if self.enable_nvtx:
                 range_push("kv_cache.batched_path", synchronize=False)
             try:
                 output = self._execute_batched(
-                    batch, submodule, node_inputs, sampler=sampler
+                    batch, submodule, node_inputs, sampler=sampler,
+                    cache_manager=step.cache_manager,
                 )
             finally:
                 if self.enable_nvtx:
@@ -1038,7 +1084,8 @@ def execute_forward(self, planned: PlannedBatch) -> NodeOutput:
                 range_push("kv_cache.sequential_path", synchronize=False)
             try:
                 output = self._execute_sequential(
-                    batch, submodule, node_inputs, sampler=sampler
+                    batch, submodule, node_inputs, sampler=sampler,
+                    cache_managers=step.per_request_managers,
                 )
             finally:
                 if self.enable_nvtx:
diff --git a/mstar/engine/kv_store.py b/mstar/engine/kv_store.py
index 3db7606ad..4528757b7 100644
--- a/mstar/engine/kv_store.py
+++ b/mstar/engine/kv_store.py
@@ -634,20 +634,10 @@ def start_async_retrieve(
         state.position_id_start = seq_info.pos_id
         state.read_in_progress = future is not None
 
-    def get_per_label_seq_info(self, request_id: str):
-        per_label_seq_info: dict[str, SequenceInfo] = {}
-        transfer_info = self._kv_transfer_engine.get_kv_transfer_info()
-        for label, state in self.request_states.get(request_id, {}).items():
-            self.wait_for_retrieves(request_id, label)
-
-            state = self.get_state(request_id, label)
-            per_label_seq_info[label] = SequenceInfo(
-                seq_len = state.seq_len,
-                pos_id = state.position_id_start,
-                latest_kv_transfer_info=transfer_info,
-                page_indices=state.page_indices
-            )
-        return per_label_seq_info
+    def get_kv_transfer_info(self):
+        """Descriptor another process needs to read this cache remotely.
+        ``KVCachePool.publish`` stamps it onto every ``SequenceInfo``."""
+        return self._kv_transfer_engine.get_kv_transfer_info()
 
     def get_labels(self, request_id: str):
         return list(self.request_states[request_id].keys())
diff --git a/mstar/engine/resources/__init__.py b/mstar/engine/resources/__init__.py
index af7837984..0b58550dd 100644
--- a/mstar/engine/resources/__init__.py
+++ b/mstar/engine/resources/__init__.py
@@ -12,6 +12,7 @@
 )
 from mstar.engine.resources.kv_pool import KVCachePool, PageArena
 from mstar.engine.resources.positions import RopeEmbedder
+from mstar.engine.resources.step import StepPlan, StepRunner
 
 __all__ = [
     "CrossAttentionManager",
@@ -24,5 +25,7 @@
     "RopeEmbedder",
     "Segment",
     "SequenceView",
+    "StepPlan",
+    "StepRunner",
     "WorkspaceBufferManager",
 ]
diff --git a/mstar/engine/resources/kv_pool.py b/mstar/engine/resources/kv_pool.py
index 878274f0b..d8d0b7886 100644
--- a/mstar/engine/resources/kv_pool.py
+++ b/mstar/engine/resources/kv_pool.py
@@ -9,7 +9,12 @@
 
 import torch
 
-from mstar.engine.kv_store import PageAllocator, PagedAllocationManager
+from mstar.conductor.request_info import SequenceInfo
+from mstar.engine.kv_store import (
+    PageAllocator,
+    PagedAllocationManager,
+    PositionInfo,
+)
 from mstar.engine.resources.base import Reservation, Segment, SequenceView
 
 
@@ -118,6 +123,14 @@ def positions(self, request_id: str, label: str) -> int:
         """Current position counter for one stream, read-only."""
         return self._manager.get_state(request_id, label).position_id_start
 
+    def pos_info(self, request_id: str, label: str) -> PositionInfo:
+        """One stream's stored length and position counter, read-only."""
+        return self._manager.get_state(request_id, label).get_pos_info()
+
+    def labels(self, request_id: str) -> list[str]:
+        """The cache streams currently existing for one request."""
+        return self._manager.get_labels(request_id)
+
     def fork(
         self,
         request_id: str,
@@ -151,3 +164,28 @@ def fork(
             strict=True,
         ):
             tensor[:, dst_page] = tensor[:, src_page]
+
+    def retrieve(self, request_id: str, label: str, seq_info: SequenceInfo) -> None:
+        """Bring one stream's published state into residency, synchronously.
+        Allocates the pages the published length needs, so it can fail like
+        any admit."""
+        self._manager.sync_retrieve(request_id, label, seq_info)
+
+    def publish(self, request_id: str) -> dict[str, SequenceInfo]:
+        """Describe the request's durable streams to another process, one
+        ``SequenceInfo`` per label. Waits out any in-flight retrieves first
+        so the description never covers half-arrived pages."""
+        manager = self._manager
+        transfer_info = manager.get_kv_transfer_info()
+        per_label_seq_info: dict[str, SequenceInfo] = {}
+        for label in list(manager.request_states.get(request_id, {})):
+            manager.wait_for_retrieves(request_id, label)
+
+            state = manager.get_state(request_id, label)
+            per_label_seq_info[label] = SequenceInfo(
+                seq_len=state.seq_len,
+                pos_id=state.position_id_start,
+                latest_kv_transfer_info=transfer_info,
+                page_indices=state.page_indices,
+            )
+        return per_label_seq_info
diff --git a/mstar/engine/resources/step.py b/mstar/engine/resources/step.py
new file mode 100644
index 000000000..eb4ee564d
--- /dev/null
+++ b/mstar/engine/resources/step.py
@@ -0,0 +1,100 @@
+"""Step sequencing for engines with KV-cache resources.
+
+The runner owns which lifecycle stage runs when. It holds no domain state
+of its own: residency is admitted through the pool before inputs are
+prepared, the step's plan surface is built before the forward runs, and
+durable state is published through the pool once the step is done. What
+each stage does lives on the resources; the runner only drives them in
+order.
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from mstar.engine.resources.kv_pool import KVCachePool
+
+if TYPE_CHECKING:
+    from mstar.engine.cache_manager import BatchedCacheManager
+
+
+@dataclass
+class StepPlan:
+    """One batch's planned execution shape: which path runs it and the plan
+    surface each forward call uses. The graph path builds no surface here
+    (the graph runner owns its captured one); the batched path shares one
+    across the batch; the sequential path gets one per request."""
+    mode: str
+    cache_manager: "BatchedCacheManager | None" = None
+    per_request_managers: list["BatchedCacheManager"] = field(default_factory=list)
+
+
+class StepRunner:
+    """Drives the per-batch resource lifecycle in dependency order."""
+
+    def admit(
+        self,
+        per_request_info: dict[str, Any],
+        pool: KVCachePool,
+        kv_cache_string: str,
+        tp_rank: int,
+        tp_world_size: int,
+        needed_labels: set[str] | None,
+    ) -> None:
+        """Establish residency for every stream this step reads: any state
+        published by another process is retrieved into the pool before
+        anything plans against it."""
+        for req_id, info in per_request_info.items():
+            if info.per_label_seq_info.world_size.get(
+                kv_cache_string, tp_world_size
+            ) != tp_world_size:
+                raise RuntimeError(
+                    "KV cache transfer across TP world size is currently disallowed"
+                ) # TODO: figure out fanin/fanout for KV cache transfer
+            for label, seq_info in info.per_label_seq_info.get(
+                kv_cache_string, tp_rank
+            ).items():
+                if needed_labels is not None and label not in needed_labels:
+                    continue
+                pool.retrieve(req_id, label, seq_info)
+
+    def plan(
+        self,
+        mode: str,
+        request_ids: list[str],
+        build_manager: Callable[[list[str]], "BatchedCacheManager"],
+    ) -> StepPlan:
+        """Build the step's plan surface for the chosen execution path. The
+        model's plan calls run against this surface during its preprocess;
+        the pools behind it were admitted before anything got here."""
+        if mode == "graph":
+            return StepPlan(mode=mode)
+        if mode == "batched":
+            return StepPlan(mode=mode, cache_manager=build_manager(request_ids))
+        return StepPlan(
+            mode=mode,
+            per_request_managers=[build_manager([rid]) for rid in request_ids],
+        )
+
+    def publish(
+        self,
+        request_ids: list[str],
+        per_request_info: dict[str, Any],
+        pool: KVCachePool,
+        kv_cache_string: str,
+        tp_rank: int,
+        tp_world_size: int,
+    ) -> None:
+        """Describe each request's durable pool state outward, after the
+        step: the published descriptors are what another process retrieves
+        from, and the next pass's routing reads them too."""
+        for req_id in request_ids:
+            info = per_request_info.get(req_id)
+            if info is None:
+                continue
+            info.per_label_seq_info.add(
+                kv_cache_string,
+                tp_rank,
+                tp_world_size,
+                pool.publish(req_id),
+            )
diff --git a/test/modular/test_kv_pool.py b/test/modular/test_kv_pool.py
index e0912e705..ca6659369 100644
--- a/test/modular/test_kv_pool.py
+++ b/test/modular/test_kv_pool.py
@@ -18,6 +18,7 @@
 import pytest
 import torch
 
+from mstar.conductor.request_info import SequenceInfo
 from mstar.engine.kv_store import (
     AllocationFailedError,
     KVCacheConfig,
@@ -35,6 +36,20 @@
 )
 
 
+class _StubTransferEngine:
+    """Transfer engine that never actually moves bytes: reads complete
+    immediately and the transfer descriptor is a sentinel."""
+
+    def __init__(self):
+        self.transfer_info = object()
+
+    def read_batched_async(self, remote_kv_info, read_info):
+        return None
+
+    def get_kv_transfer_info(self):
+        return self.transfer_info
+
+
 def _make_pool(
     max_num_pages: int = 16,
     page_size: int = 8,
@@ -55,7 +70,7 @@ def _make_pool(
         torch.zeros(1, max_num_pages, 2, page_size, 1, 1) if with_tensor else None
     )
     manager.write_policy = StoreWritePolicy.ALWAYS
-    manager._kv_transfer_engine = None
+    manager._kv_transfer_engine = _StubTransferEngine()
     manager._offload_stream = None
     manager.pending_reads = {}
     manager._lock = threading.RLock()
@@ -301,6 +316,84 @@ def test_fork_can_fail_like_any_admit(self):
             pool.fork("r", "main", "snap")
 
 
+class TestRetrieveAndPublish:
+    def test_retrieve_installs_published_state(self):
+        pool, manager = _make_pool(page_size=8)
+        manager.add_request("r", ["main"])
+        seq_info = SequenceInfo(
+            seq_len=12,
+            pos_id=34,
+            latest_kv_transfer_info=object(),
+            page_indices=[5, 6],
+        )
+
+        pool.retrieve("r", "main", seq_info)
+
+        state = manager.get_state("r", "main")
+        assert state.seq_len == 12
+        assert state.position_id_start == 34
+        assert len(state.page_indices) == 2
+        assert state.read_in_progress is False
+
+    def test_publish_describes_every_stream(self):
+        pool, manager = _make_pool(page_size=8)
+        manager.add_request("r", ["main", "cfg"])
+        for label, span, pos in (("main", 12, 12), ("cfg", 3, 0)):
+            segment = Segment("r", label, span)
+            pool.admit(segment)
+            pool.commit(segment, pos_advance=pos)
+
+        published = pool.publish("r")
+
+        assert set(published) == {"main", "cfg"}
+        assert published["main"].seq_len == 12
+        assert published["main"].pos_id == 12
+        assert published["cfg"].seq_len == 3
+        assert published["cfg"].pos_id == 0
+        transfer_info = manager._kv_transfer_engine.transfer_info
+        for label, info in published.items():
+            assert info.latest_kv_transfer_info is transfer_info
+            assert info.page_indices == manager.get_state("r", label).page_indices
+
+    def test_publish_roundtrips_through_retrieve(self):
+        """What one pool publishes, another can retrieve: the consumer ends
+        up with the producer's stored length and position counter."""
+        producer, producer_manager = _make_pool(page_size=8)
+        producer_manager.add_request("r", ["main"])
+        segment = Segment("r", "main", 20)
+        producer.admit(segment)
+        producer.commit(segment)
+
+        consumer, consumer_manager = _make_pool(page_size=8)
+        consumer_manager.add_request("r", ["main"])
+        consumer.retrieve("r", "main", producer.publish("r")["main"])
+
+        assert consumer.view(Segment("r", "main", 0)).length == 20
+        assert consumer.positions("r", "main") == 20
+
+    def test_publish_unknown_request_is_empty(self):
+        pool, _ = _make_pool()
+        assert pool.publish("ghost") == {}
+
+
+class TestPosInfoAndLabels:
+    def test_pos_info_reads_one_stream(self):
+        pool, manager = _make_pool()
+        manager.add_request("r", ["main"])
+        segment = Segment("r", "main", 9)
+        pool.admit(segment)
+        pool.commit(segment, pos_advance=4)
+
+        info = pool.pos_info("r", "main")
+        assert info.full_seq_len == 9
+        assert info.position_id_start == 4
+
+    def test_labels_lists_the_request_streams(self):
+        pool, manager = _make_pool()
+        manager.add_request("r", ["main", "cfg"])
+        assert pool.labels("r") == ["main", "cfg"]
+
+
 class TestBoundaryValuesAreImmutable:
     def test_segment_frozen(self):
         segment = Segment("r", "main", 4)
diff --git a/test/modular/test_step_runner.py b/test/modular/test_step_runner.py
new file mode 100644
index 000000000..a1f83cef7
--- /dev/null
+++ b/test/modular/test_step_runner.py
@@ -0,0 +1,176 @@
+"""Unit tests for the step runner's lifecycle sequencing.
+
+``StepRunner`` drives the per-batch resource lifecycle: ``admit`` retrieves
+published stream state into the pool before anything plans, ``plan`` builds
+the step's plan surface for the chosen execution path, and ``publish``
+describes each request's durable pool state outward after the step.
+"""
+
+from __future__ import annotations
+
+import sys
+
+sys.path.insert(0, ".")
+
+import pytest
+
+from mstar.conductor.request_info import PerLabelSeqInfo, SequenceInfo
+from mstar.engine.resources import StepPlan, StepRunner
+
+
+class _Info:
+    """Just the slice of CurrentForwardPassInfo the runner touches."""
+
+    def __init__(self, per_label_seq_info=None):
+        self.per_label_seq_info = per_label_seq_info or PerLabelSeqInfo()
+
+
+class _RecordingPool:
+    def __init__(self):
+        self.retrieved = []
+        self.published = []
+
+    def retrieve(self, request_id, label, seq_info):
+        self.retrieved.append((request_id, label, seq_info))
+
+    def publish(self, request_id):
+        self.published.append(request_id)
+        return {"main": SequenceInfo(seq_len=7, pos_id=7, latest_kv_transfer_info=None)}
+
+
+def _seq_info(seq_len=4):
+    return SequenceInfo(seq_len=seq_len, pos_id=seq_len, latest_kv_transfer_info=None)
+
+
+class TestAdmit:
+    def test_retrieves_every_published_stream(self):
+        runner, pool = StepRunner(), _RecordingPool()
+        info = _Info()
+        info.per_label_seq_info.add("kv", 0, 1, {"main": _seq_info(), "cfg": _seq_info()})
+
+        runner.admit(
+            per_request_info={"r": info},
+            pool=pool,
+            kv_cache_string="kv",
+            tp_rank=0,
+            tp_world_size=1,
+            needed_labels=None,
+        )
+        assert [(rid, label) for rid, label, _ in pool.retrieved] == [
+            ("r", "main"), ("r", "cfg"),
+        ]
+
+    def test_needed_labels_filter(self):
+        runner, pool = StepRunner(), _RecordingPool()
+        info = _Info()
+        info.per_label_seq_info.add("kv", 0, 1, {"main": _seq_info(), "cfg": _seq_info()})
+
+        runner.admit(
+            per_request_info={"r": info},
+            pool=pool,
+            kv_cache_string="kv",
+            tp_rank=0,
+            tp_world_size=1,
+            needed_labels={"main"},
+        )
+        assert [(rid, label) for rid, label, _ in pool.retrieved] == [("r", "main")]
+
+    def test_world_size_mismatch_is_rejected(self):
+        runner, pool = StepRunner(), _RecordingPool()
+        info = _Info()
+        info.per_label_seq_info.add("kv", 0, 2, {"main": _seq_info()})
+
+        with pytest.raises(RuntimeError, match="TP world size"):
+            runner.admit(
+                per_request_info={"r": info},
+                pool=pool,
+                kv_cache_string="kv",
+                tp_rank=0,
+                tp_world_size=1,
+                needed_labels=None,
+            )
+
+    def test_unpublished_requests_retrieve_nothing(self):
+        runner, pool = StepRunner(), _RecordingPool()
+        runner.admit(
+            per_request_info={"r": _Info()},
+            pool=pool,
+            kv_cache_string="kv",
+            tp_rank=0,
+            tp_world_size=1,
+            needed_labels=None,
+        )
+        assert pool.retrieved == []
+
+
+class TestPlan:
+    def test_graph_mode_builds_no_surface(self):
+        built = []
+        step = StepRunner().plan(
+            mode="graph", request_ids=["a", "b"], build_manager=built.append,
+        )
+        assert step == StepPlan(mode="graph")
+        assert built == []
+
+    def test_batched_mode_shares_one_surface(self):
+        built = []
+
+        def build(rids):
+            built.append(rids)
+            return f"manager{len(built)}"
+
+        step = StepRunner().plan(
+            mode="batched", request_ids=["a", "b"], build_manager=build,
+        )
+        assert built == [["a", "b"]]
+        assert step.cache_manager == "manager1"
+        assert step.per_request_managers == []
+
+    def test_sequential_mode_builds_one_surface_per_request(self):
+        built = []
+
+        def build(rids):
+            built.append(rids)
+            return f"manager{len(built)}"
+
+        step = StepRunner().plan(
+            mode="sequential", request_ids=["a", "b"], build_manager=build,
+        )
+        assert built == [["a"], ["b"]]
+        assert step.cache_manager is None
+        assert step.per_request_managers == ["manager1", "manager2"]
+
+
+class TestPublish:
+    def test_publishes_onto_each_request_info(self):
+        runner, pool = StepRunner(), _RecordingPool()
+        infos = {"a": _Info(), "b": _Info()}
+
+        runner.publish(
+            request_ids=["a", "b"],
+            per_request_info=infos,
+            pool=pool,
+            kv_cache_string="kv",
+            tp_rank=0,
+            tp_world_size=2,
+        )
+        assert pool.published == ["a", "b"]
+        for info in infos.values():
+            assert info.per_label_seq_info.get("kv", 0)["main"].seq_len == 7
+            assert info.per_label_seq_info.world_size["kv"] == 2
+
+    def test_requests_without_info_are_skipped(self):
+        runner, pool = StepRunner(), _RecordingPool()
+        runner.publish(
+            request_ids=["a"],
+            per_request_info={},
+            pool=pool,
+            kv_cache_string="kv",
+            tp_rank=0,
+            tp_world_size=1,
+        )
+        assert pool.published == []
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

From c8dcdf71b5d0402b30e5f7d50f7fa271309dbcd8 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 17:50:37 +0000
Subject: [PATCH 06/20] Address decode replay slots at real request ids instead
 of aliasing states

---
 mstar/engine/cuda_graph_runner.py             | 192 ++++++++++++------
 .../test_cuda_graph_step_addressing.py        | 170 ++++++++++++++++
 2 files changed, 301 insertions(+), 61 deletions(-)
 create mode 100644 test/modular/test_cuda_graph_step_addressing.py

diff --git a/mstar/engine/cuda_graph_runner.py b/mstar/engine/cuda_graph_runner.py
index a9252e3b6..1803ffb3d 100644
--- a/mstar/engine/cuda_graph_runner.py
+++ b/mstar/engine/cuda_graph_runner.py
@@ -156,11 +156,13 @@ class CudaGraphRunner:
 
     Runtime flow:
     1. Look up graph by (graph_walk, requires_cfg, padded_batch_size)
-    2. Re-plan persistent wrappers with real page tables (outside graph)
-    3. Copy real input embeddings to static buffers
-    4. graph.replay()
-    5. advance_seq_lens on real request states (Python-only, post-replay)
-    6. Clone and remap outputs from dummy to real request IDs
+    2. Address the slot's cache manager at the real request ids (the
+       padding tail keeps the slot's capture-time ids)
+    3. Re-plan persistent wrappers with real page tables (outside graph)
+    4. Copy real input embeddings to static buffers
+    5. graph.replay()
+    6. advance_seq_lens on the real request states (Python-only, post-replay)
+    7. Clone and remap outputs from capture-time dummy IDs to real request IDs
     """
 
     CAPTURE_BATCH_SIZES = DEFAULT_AR_CAPTURE_BATCH_SIZES
@@ -1127,13 +1129,13 @@ def pre_plan_for_batch(
         plan_stream = self._get_or_make_plan_stream()
         plan_done_event: torch.cuda.Event | None = None
 
-        # Temporarily alias real rids onto this slot's cache_manager so
-        # plan_attention reads real request states. The slot's static_cm
-        # will be aliased to the same real rids again at replay time
-        # (Step 1 swap_states in _run_basic_batched), so the wrapper
-        # buffers we write here stay valid for the matching replay.
-        saved_request_ids = static_cm.request_ids
-        saved_active_labels = static_cm.active_labels
+        # Address the slot at the step's request ids, through the same entry
+        # the replay path uses, so plan_attention reads and grows real
+        # request state directly. The matching replay re-addresses with the
+        # same ids, so the wrapper buffers written here stay valid for it; a
+        # dropped speculation restores the addressing via
+        # ``reset_pre_plan_state_for_slot``.
+        #
         # Pre-plan every label this captured graph's preprocess will ask for.
         # Multi-label captures (e.g. BAGEL CFG decode, labels=["main",
         # "cfg_img"]) inline-plan once per label; we cover all of them so
@@ -1149,12 +1151,12 @@ def pre_plan_for_batch(
         # captures (e.g. tree-spec) would silently mis-plan with [1]*bs.
         per_req_seq_len = config.single_request_inputs.input_seq_len
         try:
-            static_cm.request_ids = list(request_ids) + saved_request_ids[len(request_ids):]
-            seq_lens = [per_req_seq_len] * len(saved_request_ids)
+            step_ids = self._slot_step_ids(slot_data, request_ids)
+            self._address_slot(slot_data, step_ids)
+            seq_lens = [per_req_seq_len] * len(step_ids)
             if plan_stream is not None:
                 with torch.cuda.stream(plan_stream):
                     for label_name in config_labels:
-                        static_cm.active_labels = {rid: label_name for rid in request_ids}
                         static_cm.plan_attention(
                             seq_lens=seq_lens,
                             dtype=self.autocast_dtype,
@@ -1164,7 +1166,6 @@ def pre_plan_for_batch(
                 plan_done_event.record(plan_stream)
             else:
                 for label_name in config_labels:
-                    static_cm.active_labels = {rid: label_name for rid in request_ids}
                     static_cm.plan_attention(
                         seq_lens=seq_lens,
                         dtype=self.autocast_dtype,
@@ -1173,8 +1174,6 @@ def pre_plan_for_batch(
             static_cm._pre_planned_labels = set(config_labels)
             static_cm._plan_done_event = plan_done_event
         finally:
-            static_cm.request_ids = saved_request_ids
-            static_cm.active_labels = saved_active_labels
             if self.enable_nvtx:
                 range_pop(synchronize=False)
         return True
@@ -1218,11 +1217,14 @@ def reset_pre_plan_state_for_slot(
         cm._pre_planned_labels.clear()
         cm._plan_done_event = None
 
-        # pre_plan_for_batch allocates pages on the slot's tail dummy rids
-        # (positions [real_bs..padded_bs) in static_cm.request_ids). If
-        # replay never runs to free them via _restore_dummy_states, the
-        # pages leak. Free every dummy rid's pages — no-op for positions
-        # pre-plan didn't touch.
+        # pre_plan_for_batch addressed the slot at the dropped batch's ids
+        # and allocated pages on the padding tail's capture ids. If replay
+        # never runs to release the slot, the tail pages leak and the
+        # addressing keeps referencing the dropped requests. Restore the
+        # capture addressing and free every capture id's pages (a no-op
+        # for positions pre-plan didn't touch). Real requests keep the pages
+        # their streams grew by; the next plan for them finds them resident.
+        self._reset_slot_addressing(slot_data)
         for rid in slot_data.static_inputs.get("dummy_rids", []):
             for label in graph_data.config.labels:
                 self.alloc_manager.reset_label(rid, label, free=True)
@@ -1324,8 +1326,11 @@ def _run_basic_batched(
         into the static buffers before replay.
 
         ``slot_data`` is the chosen double-buffer slot (graph + persistent
-        wrappers + cache_manager). Same logic as before — we just look up the
-        slot's graph/cm instead of reading flat fields off ``graph_data``.
+        wrappers + cache_manager). The slot's cache manager is addressed at
+        the real request ids for the step, so every plan, advance, and page
+        allocation lands on real request state directly; the padding tail
+        keeps the slot's capture-time ids and its pages are freed after the
+        step.
         """
         real_bs = len(request_ids)
         padded_bs = key.bs
@@ -1341,34 +1346,23 @@ def _run_basic_batched(
         capture_template = static["capture_template"]
         config_labels = graph_data.config.labels
 
-        # Swap-and-restore must be paired: if any step between swap and restore
-        # raises (e.g., submodule.preprocess hitting an insufficient-KV alloc
-        # failure), the dummy slots are still aliased to real RequestState
-        # objects. The finally below un-aliases them; flush_writes stays False
-        # on failure since the captured forward never replayed.
-        swapped = False
+        # Address-and-release must be paired: if any step in between raises
+        # (e.g., submodule.preprocess hitting an insufficient-KV alloc
+        # failure), the finally below restores the capture addressing and
+        # frees the padding tail's pages; flush_writes stays False on failure
+        # since the captured forward never replayed.
+        addressed = False
         success = False
         try:
-            # --- Step 1: Swap real request states onto dummy slots ---
+            # --- Step 1: Address the slot at the step's request ids ---
             if self.enable_nvtx:
                 mark("gpu_thread.preprocess_start")
                 range_push("gpu_thread.preprocess", synchronize=False)
             if self.enable_nvtx:
-                range_push("cg.swap_states", synchronize=False)
-            for i, rid in enumerate(request_ids):
-                dummy_rid = dummy_rids[i]
-                for label in config_labels:
-                    real_state = self.alloc_manager.get_state(rid, label)
-                    # makes state if it doesn't exist
-                    self.alloc_manager.get_state(dummy_rid, label)
-                    self.alloc_manager.request_states[dummy_rid][label] = real_state
-
-            # For padding slots (i >= real_bs), ensure dummy states exist
-            for i in range(real_bs, padded_bs):
-                dummy_rid = dummy_rids[i]
-                for label in config_labels:
-                    self.alloc_manager.get_state(dummy_rid, label)
-            swapped = True
+                range_push("cg.address_slot", synchronize=False)
+            step_ids = self._slot_step_ids(slot_data, request_ids)
+            self._address_slot(slot_data, step_ids)
+            addressed = True
             if self.enable_nvtx:
                 range_pop(synchronize=False)
 
@@ -1387,9 +1381,8 @@ def _run_basic_batched(
 
             if self.enable_nvtx:
                 range_push("cg.preprocess_replan.metadata", synchronize=False)
-            real_metadata = self._build_replay_metadata(
-                dummy_rids, request_ids, real_bs,
-                per_request_info, static["dummy_metadata"],
+            step_metadata = self._build_step_metadata(
+                step_ids, real_bs, per_request_info, static["dummy_metadata"],
             )
             # Stage the live seen-token masks into master before the gather so
             # the per-step buffer reflects the request's accumulated tokens for
@@ -1399,8 +1392,8 @@ def _run_basic_batched(
                     request_ids, self.sampler,
                 )
             engine_inputs = ModelInputsFromEngine(
-                request_ids=dummy_rids,
-                per_request_info=real_metadata,
+                request_ids=step_ids,
+                per_request_info=step_metadata,
                 cache_manager=static_cm,
                 sampler=self._get_sampler(
                     request_ids=request_ids,
@@ -1469,9 +1462,9 @@ def _run_basic_batched(
             # masters (GPU-only; keyed by slot, so batch-position invariant).
             self.sampler_buffer.scatter_offsets()
 
-            # --- Step 5: Advance seq_lens on REAL request states (Python-only) ---
-            # advance_seq_lens is not captured in the graph; we call it manually so
-            # the real states (aliased onto dummy slots) move forward.
+            # --- Step 5: Advance seq_lens on the real request states (Python-only) ---
+            # advance_seq_lens is not captured in the graph; we call it manually,
+            # and the slot addressing routes it to the real states.
             if self.enable_nvtx:
                 mark("gpu_thread.postprocess_start")
                 range_push("gpu_thread.postprocess", synchronize=False)
@@ -1519,14 +1512,13 @@ def _run_basic_batched(
             success = True
             return outputs
         finally:
-            # --- Step 7: Restore dummy states (always — un-aliases dummy slots) ---
-            if swapped:
-                self._restore_dummy_states(
-                    dummy_rids=dummy_rids,
+            # --- Step 7: Release the slot (capture addressing back, padding freed) ---
+            if addressed:
+                self._release_slot_step(
+                    slot_data=slot_data,
                     request_ids=request_ids,
                     real_bs=real_bs,
                     config_labels=config_labels,
-                    static_cm=static_cm,
                     flush_writes=success,
                 )
             if self.enable_nvtx:
@@ -1762,6 +1754,52 @@ def _run_flashinfer_packed(
                 range_pop(synchronize=False)
                 mark("gpu_thread.postprocess_end")
 
+    def _slot_step_ids(
+        self, slot_data: CudaGraphSlot, request_ids: list[str],
+    ) -> list[str]:
+        """The padded addressing for one step on this slot: the real request
+        ids first, the slot's capture-time ids for the padding tail. Plans,
+        advances, and commits address real request state by its own id; only
+        the padding rows run against capture-time state."""
+        dummy_rids = slot_data.static_inputs["dummy_rids"]
+        return list(request_ids) + dummy_rids[len(request_ids):]
+
+    @staticmethod
+    def _address_slot(slot_data: CudaGraphSlot, step_ids: list[str]) -> None:
+        """Point the slot's cache manager at the step's request ids. Shared
+        by pre-plan and replay so both plan the same real segment list into
+        the slot's static buffers."""
+        static_cm = slot_data.static_cache_manager
+        static_cm.request_ids = step_ids
+        static_cm.active_labels = {rid: "main" for rid in step_ids}
+
+    @staticmethod
+    def _reset_slot_addressing(slot_data: CudaGraphSlot) -> None:
+        """Return the slot's cache manager to its capture-time addressing so
+        a slot at rest references no live request. No-op for stub slots
+        without recorded capture ids."""
+        dummy_rids = slot_data.static_inputs.get("dummy_rids")
+        if not dummy_rids:
+            return
+        static_cm = slot_data.static_cache_manager
+        static_cm.request_ids = list(dummy_rids)
+        static_cm.active_labels = {rid: "main" for rid in dummy_rids}
+
+    def _build_step_metadata(
+        self,
+        step_ids: list[str],
+        real_bs: int,
+        per_request_info: dict[str, CurrentForwardPassInfo],
+        dummy_metadata: dict[str, CurrentForwardPassInfo],
+    ) -> dict[str, CurrentForwardPassInfo]:
+        """Per-request info keyed by the step's addressing: the real info for
+        real ids, the capture-time info for the padding tail. Used by both
+        replay paths."""
+        out = {}
+        for i, rid in enumerate(step_ids):
+            out[rid] = per_request_info[rid] if i < real_bs else dummy_metadata[rid]
+        return out
+
     def _build_replay_metadata(
         self,
         dummy_rids: list[str],
@@ -1771,7 +1809,7 @@ def _build_replay_metadata(
         dummy_metadata: dict[str, CurrentForwardPassInfo],
     ) -> dict[str, CurrentForwardPassInfo]:
         """Map dummy_rid → real per_request_info for [:real_bs], dummy_metadata
-        from capture for [real_bs:]. Used by both replay paths."""
+        from capture for [real_bs:]. Used by the packed replay path."""
         out = {}
         for i, dummy_rid in enumerate(dummy_rids):
             if i < real_bs:
@@ -1832,6 +1870,38 @@ def _zero_seq_dim(tensor: torch.Tensor, seq_len: int) -> torch.Tensor:
                 return tensor[tuple(slicer)]
         return tensor[:0]
 
+    def _release_slot_step(
+        self,
+        slot_data: CudaGraphSlot,
+        request_ids: list[str],
+        real_bs: int,
+        config_labels: list[str],
+        flush_writes: bool = True,
+    ) -> None:
+        """Return the slot to its capture-time state after a step: restore
+        the capture addressing, free the pages the padding tail's plan
+        allocated, and (on the success path) flush real-request KV writes
+        for any label planned with write_store. ``flush_writes=False`` on a
+        failure path skips the flush, since the captured forward never
+        executed. Real request state was addressed by its own id throughout,
+        so there is nothing to restore on it."""
+        if self.enable_nvtx:
+            range_push("cg.release_slot", synchronize=False)
+        static_cm = slot_data.static_cache_manager
+        dummy_rids = slot_data.static_inputs["dummy_rids"]
+        self._reset_slot_addressing(slot_data)
+        for i in range(real_bs, len(dummy_rids)):
+            for label in config_labels:
+                self.alloc_manager.reset_label(dummy_rids[i], label, free=True)
+        if flush_writes:
+            for rid in request_ids:
+                for label in config_labels:
+                    ps = static_cm._plan_states.get(label)
+                    if ps is not None and ps.write_store:
+                        self.alloc_manager.flush_to_store(rid, label)
+        if self.enable_nvtx:
+            range_pop(synchronize=False)
+
     def _restore_dummy_states(
         self,
         dummy_rids: list[str],
diff --git a/test/modular/test_cuda_graph_step_addressing.py b/test/modular/test_cuda_graph_step_addressing.py
new file mode 100644
index 000000000..af9880e72
--- /dev/null
+++ b/test/modular/test_cuda_graph_step_addressing.py
@@ -0,0 +1,170 @@
+"""Tests for the runner's slot step addressing.
+
+Replay and pre-plan point a slot's cache manager at the step's request ids
+(real ids first, the slot's capture-time ids for the padding tail) instead
+of aliasing live request state onto dummy slots. These tests drive the
+addressing helpers and the release paths directly with stub slots,
+verifying that:
+
+  - step ids compose real ids with the capture tail;
+  - addressing and its reset write the cache manager's request_ids and
+    active_labels and nothing else;
+  - step metadata keys real info by real id and capture info by capture id;
+  - the pre-plan drop path restores capture addressing and frees only
+    capture-id pages.
+"""
+
+from __future__ import annotations
+
+import sys
+import types
+
+sys.path.insert(0, ".")
+
+from mstar.engine.cuda_graph_runner import (
+    CudaGraphData,
+    CudaGraphKey,
+    CudaGraphRunner,
+    CudaGraphSlot,
+)
+
+
+def _make_stub_cm() -> types.SimpleNamespace:
+    cm = types.SimpleNamespace()
+    cm.request_ids = []
+    cm.active_labels = {}
+    cm._pre_planned_labels = set()
+    cm._plan_done_event = None
+    return cm
+
+
+def _make_slot(dummy_rids: list[str]) -> CudaGraphSlot:
+    cm = _make_stub_cm()
+    cm.request_ids = list(dummy_rids)
+    cm.active_labels = {rid: "main" for rid in dummy_rids}
+    return CudaGraphSlot(
+        graph=object(),
+        static_inputs={"dummy_rids": list(dummy_rids)},
+        static_outputs={},
+        static_cache_manager=cm,
+    )
+
+
+def _make_runner() -> CudaGraphRunner:
+    runner = CudaGraphRunner.__new__(CudaGraphRunner)
+    runner.enable_nvtx = False
+    return runner
+
+
+DUMMIES = ["__cg_d0__", "__cg_d1__", "__cg_d2__", "__cg_d3__"]
+
+
+class TestSlotStepIds:
+    def test_partial_batch_keeps_capture_tail(self):
+        runner = _make_runner()
+        slot = _make_slot(DUMMIES)
+        step_ids = runner._slot_step_ids(slot, ["r0", "r1"])
+        assert step_ids == ["r0", "r1", "__cg_d2__", "__cg_d3__"]
+
+    def test_full_batch_has_no_tail(self):
+        runner = _make_runner()
+        slot = _make_slot(DUMMIES)
+        step_ids = runner._slot_step_ids(slot, ["r0", "r1", "r2", "r3"])
+        assert step_ids == ["r0", "r1", "r2", "r3"]
+
+    def test_input_list_not_mutated(self):
+        runner = _make_runner()
+        slot = _make_slot(DUMMIES)
+        real = ["r0"]
+        runner._slot_step_ids(slot, real)
+        assert real == ["r0"]
+
+
+class TestAddressing:
+    def test_address_sets_ids_and_labels(self):
+        slot = _make_slot(DUMMIES)
+        step_ids = ["r0", "r1", "__cg_d2__", "__cg_d3__"]
+        CudaGraphRunner._address_slot(slot, step_ids)
+        cm = slot.static_cache_manager
+        assert cm.request_ids == step_ids
+        assert cm.active_labels == {rid: "main" for rid in step_ids}
+
+    def test_reset_restores_capture_addressing(self):
+        slot = _make_slot(DUMMIES)
+        CudaGraphRunner._address_slot(slot, ["r0", "r1", "__cg_d2__", "__cg_d3__"])
+        CudaGraphRunner._reset_slot_addressing(slot)
+        cm = slot.static_cache_manager
+        assert cm.request_ids == DUMMIES
+        assert cm.active_labels == {rid: "main" for rid in DUMMIES}
+
+    def test_reset_noops_without_capture_ids(self):
+        slot = CudaGraphSlot(
+            graph=object(),
+            static_inputs={},
+            static_outputs={},
+            static_cache_manager=_make_stub_cm(),
+        )
+        slot.static_cache_manager.request_ids = ["r0"]
+        CudaGraphRunner._reset_slot_addressing(slot)
+        assert slot.static_cache_manager.request_ids == ["r0"]
+
+
+class TestStepMetadata:
+    def test_real_head_and_capture_tail(self):
+        runner = _make_runner()
+        real_info = {"r0": "info0", "r1": "info1"}
+        dummy_info = {d: f"cap_{d}" for d in DUMMIES}
+        out = runner._build_step_metadata(
+            ["r0", "r1", "__cg_d2__", "__cg_d3__"], 2, real_info, dummy_info,
+        )
+        assert out == {
+            "r0": "info0",
+            "r1": "info1",
+            "__cg_d2__": "cap___cg_d2__",
+            "__cg_d3__": "cap___cg_d3__",
+        }
+
+
+class _RecordingAllocManager:
+    def __init__(self):
+        self.resets: list[tuple[str, str, bool]] = []
+
+    def reset_label(self, rid: str, label: str, free: bool = True) -> None:
+        self.resets.append((rid, label, free))
+
+
+class TestResetPrePlanRestoresAddressing:
+    def _runner_with_slot(self):
+        runner = _make_runner()
+        runner.alloc_manager = _RecordingAllocManager()
+        key = CudaGraphKey(graph_walk="decode", requires_cfg=False, bs=4, num_tokens=4)
+        slot = _make_slot(DUMMIES)
+        config = types.SimpleNamespace(labels=["main"])
+        runner.graphs = {
+            key: CudaGraphData(config=config, bs=4, index=0, slots=[slot]),
+        }
+        runner._get_basic_batched_key_for = lambda **kw: key
+        return runner, slot
+
+    def test_drop_path_restores_addressing_and_frees_capture_pages(self):
+        runner, slot = self._runner_with_slot()
+        CudaGraphRunner._address_slot(
+            slot, ["r0", "r1", "__cg_d2__", "__cg_d3__"],
+        )
+        slot.static_cache_manager._pre_planned_labels = {"main"}
+        runner.reset_pre_plan_state_for_slot(
+            graph_walk="decode", requires_cfg=False, batch_size=4, slot=0,
+        )
+        cm = slot.static_cache_manager
+        assert cm._pre_planned_labels == set()
+        assert cm._plan_done_event is None
+        assert cm.request_ids == DUMMIES
+        assert cm.active_labels == {rid: "main" for rid in DUMMIES}
+        assert runner.alloc_manager.resets == [
+            (rid, "main", True) for rid in DUMMIES
+        ]
+
+
+if __name__ == "__main__":
+    import pytest
+    sys.exit(pytest.main([__file__, "-v"]))

From ef10a886375803b19bf48256eb1cd0309df945ec Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 18:10:35 +0000
Subject: [PATCH 07/20] Address packed prefill replay slots the same way and
 drop the state swap

---
 mstar/engine/cuda_graph_runner.py | 101 +++++++-----------------------
 mstar/model/bagel/submodules.py   |  16 ++---
 2 files changed, 29 insertions(+), 88 deletions(-)

diff --git a/mstar/engine/cuda_graph_runner.py b/mstar/engine/cuda_graph_runner.py
index 1803ffb3d..e46b0d5a8 100644
--- a/mstar/engine/cuda_graph_runner.py
+++ b/mstar/engine/cuda_graph_runner.py
@@ -1545,7 +1545,8 @@ def _run_flashinfer_packed(
         which FlashInfer's attention path actually walks. Trailing static-buffer
         slots [real_num_tokens : padded_num_tokens] keep their capture-time
         contents; non-attention compute over them is wasted work, not a correctness
-        issue. State swap / advance_seq_lens / output remap mirror _run_basic_matched.
+        issue. Slot addressing / advance_seq_lens / output remap mirror
+        _run_basic_batched.
 
         ``slot_data`` selects one of the captured double-buffer slots.
         Prefill paths don't speculate or pre-plan today, so slot alternation
@@ -1565,31 +1566,22 @@ def _run_flashinfer_packed(
         static_input_keys = static["static_input_keys"]
         config_labels = graph_data.config.labels
 
-        # Swap-and-restore must be paired (see _run_basic_batched). On a
-        # submodule.preprocess failure mid-flight, the dummy slots are still
-        # aliased to real RequestState objects; the finally below un-aliases
-        # them and skips the store flush (no captured forward ran).
-        swapped = False
+        # Address-and-release must be paired (see _run_basic_batched). On a
+        # submodule.preprocess failure mid-flight, the finally below restores
+        # the capture addressing, frees the padding tail's pages, and skips
+        # the store flush (no captured forward ran).
+        addressed = False
         success = False
         try:
-            # --- Step 1: Swap real request states onto dummy slots ---
+            # --- Step 1: Address the slot at the step's request ids ---
             if self.enable_nvtx:
                 mark("gpu_thread.preprocess_start")
                 range_push("gpu_thread.preprocess", synchronize=False)
             if self.enable_nvtx:
-                range_push("cg.swap_states", synchronize=False)
-            for i, rid in enumerate(request_ids):
-                dummy_rid = dummy_rids[i]
-                for label in config_labels:
-                    real_state = self.alloc_manager.get_state(rid, label)
-                    self.alloc_manager.get_state(dummy_rid, label)
-                    self.alloc_manager.request_states[dummy_rid][label] = real_state
-
-            for i in range(real_bs, padded_bs):
-                dummy_rid = dummy_rids[i]
-                for label in config_labels:
-                    self.alloc_manager.get_state(dummy_rid, label)
-            swapped = True
+                range_push("cg.address_slot", synchronize=False)
+            step_ids = self._slot_step_ids(slot_data, request_ids)
+            self._address_slot(slot_data, step_ids)
+            addressed = True
             if self.enable_nvtx:
                 range_pop(synchronize=False)
 
@@ -1615,9 +1607,8 @@ def _run_flashinfer_packed(
 
             if self.enable_nvtx:
                 range_push("cg.preprocess_replan.metadata", synchronize=False)
-            real_metadata = self._build_replay_metadata(
-                dummy_rids, request_ids, real_bs,
-                per_request_info, static["dummy_metadata"],
+            step_metadata = self._build_step_metadata(
+                step_ids, real_bs, per_request_info, static["dummy_metadata"],
             )
             # Stage the live seen-token masks into master before the gather so
             # the per-step buffer reflects the request's accumulated tokens for
@@ -1627,8 +1618,8 @@ def _run_flashinfer_packed(
                     request_ids, self.sampler,
                 )
             engine_inputs = ModelInputsFromEngine(
-                request_ids=dummy_rids,
-                per_request_info=real_metadata,
+                request_ids=step_ids,
+                per_request_info=step_metadata,
                 cache_manager=static_cm,
                 sampler=self._get_sampler(
                     request_ids=request_ids,
@@ -1693,7 +1684,7 @@ def _run_flashinfer_packed(
             # Persist the in-graph-advanced RNG offsets back to their slot masters.
             self.sampler_buffer.scatter_offsets()
 
-            # --- Step 5: Advance seq_lens on REAL request states (Python-only) ---
+            # --- Step 5: Advance seq_lens on the real request states (Python-only) ---
             if self.enable_nvtx:
                 mark("gpu_thread.postprocess_start")
                 range_push("gpu_thread.postprocess", synchronize=False)
@@ -1740,14 +1731,13 @@ def _run_flashinfer_packed(
             success = True
             return outputs
         finally:
-            # --- Step 7: Restore dummy states (always — un-aliases dummy slots) ---
-            if swapped:
-                self._restore_dummy_states(
-                    dummy_rids=dummy_rids,
+            # --- Step 7: Release the slot (capture addressing back, padding freed) ---
+            if addressed:
+                self._release_slot_step(
+                    slot_data=slot_data,
                     request_ids=request_ids,
                     real_bs=real_bs,
                     config_labels=config_labels,
-                    static_cm=static_cm,
                     flush_writes=success,
                 )
             if self.enable_nvtx:
@@ -1800,24 +1790,6 @@ def _build_step_metadata(
             out[rid] = per_request_info[rid] if i < real_bs else dummy_metadata[rid]
         return out
 
-    def _build_replay_metadata(
-        self,
-        dummy_rids: list[str],
-        request_ids: list[str],
-        real_bs: int,
-        per_request_info: dict[str, CurrentForwardPassInfo],
-        dummy_metadata: dict[str, CurrentForwardPassInfo],
-    ) -> dict[str, CurrentForwardPassInfo]:
-        """Map dummy_rid → real per_request_info for [:real_bs], dummy_metadata
-        from capture for [real_bs:]. Used by the packed replay path."""
-        out = {}
-        for i, dummy_rid in enumerate(dummy_rids):
-            if i < real_bs:
-                out[dummy_rid] = per_request_info[request_ids[i]]
-            else:
-                out[dummy_rid] = dummy_metadata[dummy_rid]
-        return out
-
     def _zero_padding_input(self, template: ARNodeInputs) -> ARNodeInputs:
         """Synthetic zero-length ARNodeInputs for prefill padding slots.
 
@@ -1902,37 +1874,6 @@ def _release_slot_step(
         if self.enable_nvtx:
             range_pop(synchronize=False)
 
-    def _restore_dummy_states(
-        self,
-        dummy_rids: list[str],
-        request_ids: list[str],
-        real_bs: int,
-        config_labels: list[str],
-        static_cm: BatchedCacheManager,
-        flush_writes: bool = True,
-    ) -> None:
-        """Reset every dummy slot's per-label state and (on the success path)
-        flush real-request KV writes to the store for any label whose plan_state
-        had write_store enabled. ``flush_writes=False`` on a failure path skips
-        the flush — the captured forward never executed, so there's nothing to
-        commit, and flushing partial state would publish stale pages.
-        """
-        if self.enable_nvtx:
-            range_push("cg.restore_states", synchronize=False)
-        for i, rid in enumerate(dummy_rids):
-            for label in config_labels:
-                self.alloc_manager.reset_label(
-                    rid, label, free=i >= real_bs,
-                )
-        if flush_writes:
-            for rid in request_ids:
-                for label in config_labels:
-                    ps = static_cm._plan_states.get(label)
-                    if ps is not None and ps.write_store:
-                        self.alloc_manager.flush_to_store(rid, label)
-        if self.enable_nvtx:
-            range_pop(synchronize=False)
-
     def _sample_and_remap(
         self,
         request_ids: list[str],
diff --git a/mstar/model/bagel/submodules.py b/mstar/model/bagel/submodules.py
index a27ad48e3..4ca8f3b2b 100644
--- a/mstar/model/bagel/submodules.py
+++ b/mstar/model/bagel/submodules.py
@@ -500,14 +500,14 @@ def get_cuda_graph_configs(
 
         cfg-on prefill_text is intentionally NOT captured. BAGEL's
         ``preprocess`` for prefill_text+cfg calls
-        ``cache_handle.snapshot_all("main", "cfg_text")`` which writes to
-        the cache_manager's ``request_ids`` (= ``dummy_rids`` at replay).
-        ``cfg_text`` is not in ``config.labels`` (only ``main`` + ``cfg_img``
-        get FlashInfer wrappers), so the runner's state-swap doesn't alias
-        it onto the real request and the snapshot lands on the dummy slot.
-        cfg-on prefill_text continues to use the eager path; downstream
-        image_gen / decode+cfg captures are unaffected (they don't depend
-        on this capture's snapshot semantics).
+        ``cache_handle.snapshot_all("main", "cfg_text")``, an allocating
+        fork of the main stream. Replay addresses the cache manager at the
+        real request ids, so the snapshot would land on real state now,
+        but a captured cfg-on prefill has never been exercised and stays
+        off until verified on its own. cfg-on prefill_text continues to
+        use the eager path; downstream image_gen / decode+cfg captures are
+        unaffected (they don't depend on this capture's snapshot
+        semantics).
         """
         dummy = ARNodeInputs(
             input_ids=torch.zeros(1, dtype=torch.long, device=device),

From a010de92c96993682ee011e0254660f2b27d319a Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 18:19:47 +0000
Subject: [PATCH 08/20] Address piecewise replay at real request ids and
 release padding pages

---
 mstar/engine/cuda_graph_runner.py | 107 ++++++++++++++++--------------
 1 file changed, 59 insertions(+), 48 deletions(-)

diff --git a/mstar/engine/cuda_graph_runner.py b/mstar/engine/cuda_graph_runner.py
index e46b0d5a8..1d710ac1a 100644
--- a/mstar/engine/cuda_graph_runner.py
+++ b/mstar/engine/cuda_graph_runner.py
@@ -2434,7 +2434,9 @@ class PiecewiseCudaGraphRunner:
     - FlashInfer wrappers are PERSISTENT (created once per bucket at capture).
     - plan_attention is called OUTSIDE the graph before each replay.
     - advance_seq_lens is called OUTSIDE the graph after each replay.
-    - KV state is swapped onto dummy slots before replay and restored after.
+    - The cache manager is addressed at the real request ids before each
+      replay; padding rows keep the capture-time ids and their pages are
+      freed after the step.
     """
 
     DEFAULT_CAPTURE_BATCH_SIZES = DEFAULT_AR_CAPTURE_BATCH_SIZES
@@ -2757,10 +2759,11 @@ def run(
         Steps (mirroring CudaGraphRunner._run_basic_batched):
           1. Copy each real input tensor into the runner-owned static buffer of
              the same name (zeroing any padded tail).
-          2. Swap real KV states onto dummy slots + plan_attention (if KV).
+          2. Address the cache manager at the real request ids + plan_attention
+             (if KV); padding rows keep the capture-time ids.
           3. graph.replay().
           4. advance_seq_lens (Python-only, outside graph).
-          5. Restore dummy states.
+          5. Release the slot: capture addressing back, padding pages freed.
           6. Return a ``PiecewiseOutput`` over the captured output buffers.
 
         ``real_bs`` is inferred from ``request_ids`` or ``seq_lens`` when not
@@ -2800,55 +2803,63 @@ def run(
             if n < buf.shape[0]:
                 buf[n:].zero_()
 
-        # --- 2: KV state swap + plan_attention ---
+        # --- 2: address the slot at the step's request ids + plan_attention ---
         static_cm = data.static_cache_manager
-        if static_cm is not None and request_ids is not None:
-            for i, rid in enumerate(request_ids):
-                dummy_rid = data.dummy_rids[i]
-                for label in self.cache_labels:
-                    real_state = self.alloc_manager.get_state(rid, label)
-                    self.alloc_manager.get_state(dummy_rid, label)  # ensure slot exists
-                    self.alloc_manager.request_states[dummy_rid][label] = real_state
-            self._plan(
-                static_cm,
-                data.shape,
-                seq_lens=self._replay_seq_lens(data.shape, seq_lens, real_bs),
-            )
-
-        if self.sampler_buffer is not None and request_ids is not None:
-            # TODO: add "gather_seen_tokens" as an explicit flag in the piecewise cuda
-            # graph config so we don't do unnecessary work here
-            self.sampler_buffer.gather_for_request_ids(
-                request_ids=request_ids, padded_bs=data.shape.bs,
-                gather_seen_tokens=True,
-            )
-
-        # --- 3: replay ---
-        data.graph.replay()
-
-        if self.sampler_buffer is not None and request_ids is not None:
-            # Persist the in-graph-advanced RNG offsets back to their slot
-            # masters (mirrors the gather above; GPU-only, real rows only).
-            self.sampler_buffer.scatter_offsets()
+        addressed = False
+        try:
+            if static_cm is not None and request_ids is not None:
+                step_ids = list(request_ids) + data.dummy_rids[len(request_ids):]
+                static_cm.request_ids = step_ids
+                static_cm.active_labels = {
+                    rid: self.cache_labels[0] for rid in step_ids
+                }
+                addressed = True
+                self._plan(
+                    static_cm,
+                    data.shape,
+                    seq_lens=self._replay_seq_lens(data.shape, seq_lens, real_bs),
+                )
 
-        # --- 4: advance seq_lens (Python-only, post-replay) ---
-        # Uses the per-request lengths planned in step 2, so this is correct for
-        # both uniform (BATCHED) and variable (PACKED) sequences. Opt out via
-        # config.advance_seq_lens=False when the caller advances the cache itself.
-        if (
-            self.config.advance_seq_lens
-            and static_cm is not None
-            and request_ids is not None
-        ):
-            for label in self.cache_labels:
-                static_cm.set_active_label(label)
-                static_cm.advance_seq_lens()
+            if self.sampler_buffer is not None and request_ids is not None:
+                # TODO: add "gather_seen_tokens" as an explicit flag in the piecewise cuda
+                # graph config so we don't do unnecessary work here
+                self.sampler_buffer.gather_for_request_ids(
+                    request_ids=request_ids, padded_bs=data.shape.bs,
+                    gather_seen_tokens=True,
+                )
 
-        # --- 5: restore dummy states ---
-        if static_cm is not None and request_ids is not None:
-            for i, dummy_rid in enumerate(data.dummy_rids):
+            # --- 3: replay ---
+            data.graph.replay()
+
+            if self.sampler_buffer is not None and request_ids is not None:
+                # Persist the in-graph-advanced RNG offsets back to their slot
+                # masters (mirrors the gather above; GPU-only, real rows only).
+                self.sampler_buffer.scatter_offsets()
+
+            # --- 4: advance seq_lens (Python-only, post-replay) ---
+            # Uses the per-request lengths planned in step 2, so this is correct for
+            # both uniform (BATCHED) and variable (PACKED) sequences. Opt out via
+            # config.advance_seq_lens=False when the caller advances the cache itself.
+            if (
+                self.config.advance_seq_lens
+                and static_cm is not None
+                and request_ids is not None
+            ):
                 for label in self.cache_labels:
-                    self.alloc_manager.reset_label(dummy_rid, label, free=i >= real_bs)
+                    static_cm.set_active_label(label)
+                    static_cm.advance_seq_lens()
+        finally:
+            # --- 5: release the slot (capture addressing back, padding freed) ---
+            if addressed:
+                static_cm.request_ids = list(data.dummy_rids)
+                static_cm.active_labels = {
+                    rid: self.cache_labels[0] for rid in data.dummy_rids
+                }
+                for i in range(real_bs, len(data.dummy_rids)):
+                    for label in self.cache_labels:
+                        self.alloc_manager.reset_label(
+                            data.dummy_rids[i], label, free=True,
+                        )
 
         # --- 6: return output view ---
         real_len = real_total_tokens if is_packed else real_bs

From 47404f47a4ae9dd45b2b424e4cf39d9ed2530a60 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 22:02:11 +0000
Subject: [PATCH 09/20] Build per-node resource dicts from model-declared specs

---
 mstar/engine/base.py                |  7 +++
 mstar/engine/kv_cache_engine.py     | 36 +++++++++++-
 mstar/engine/resources/__init__.py  |  6 +-
 mstar/engine/resources/kv_pool.py   | 11 ++++
 mstar/engine/resources/spec.py      | 36 ++++++++++++
 mstar/model/base.py                 | 15 +++++
 mstar/model/submodule_base.py       |  8 +++
 mstar/worker/engine_manager.py      |  2 +
 test/modular/test_node_resources.py | 90 +++++++++++++++++++++++++++++
 9 files changed, 208 insertions(+), 3 deletions(-)
 create mode 100644 mstar/engine/resources/spec.py
 create mode 100644 test/modular/test_node_resources.py

diff --git a/mstar/engine/base.py b/mstar/engine/base.py
index 8a879f4fd..da65a9c54 100644
--- a/mstar/engine/base.py
+++ b/mstar/engine/base.py
@@ -320,6 +320,13 @@ def reset_pre_plan_for_batch(self, batch: NodeBatch) -> None:
 
     capabilities = EngineCapabilities()
 
+    def node_resources(self, node_name: str) -> dict[str, Any]:
+        """The named resources this engine built for ``node_name`` (KV
+        cache pool, embedder, cross pools, scratch caches). Engines that
+        build none return an empty dict.
+        """
+        return {}
+
     def lru_tracked_nodes(self) -> list[str]:
         """Nodes for which the worker should LRU-track per-request activity
         (used to pick CPU-offload victims). Default: no nodes — stateless
diff --git a/mstar/engine/kv_cache_engine.py b/mstar/engine/kv_cache_engine.py
index 349ea7bd7..9376ccae2 100644
--- a/mstar/engine/kv_cache_engine.py
+++ b/mstar/engine/kv_cache_engine.py
@@ -36,7 +36,8 @@
     StoreWritePolicy,
     TransferEngineInfo,
 )
-from mstar.engine.resources import KVCachePool, RopeEmbedder
+from mstar.engine.resources import KVCachePool, RopeEmbedder, ScratchKVPool
+from mstar.engine.resources.spec import NodeResourceSpec
 from mstar.engine.resources.step import StepPlan, StepRunner
 from mstar.model.submodule_base import ARNodeInputs, ARNodeSubmodule, ModelInputsFromEngine
 from mstar.utils.profiler import range_pop, range_push
@@ -163,6 +164,9 @@ def __init__(
 
         self.kv_management: dict[str, KVManagement] = {}
         self.submodule_management: dict[str, SubmoduleManagement] = {}
+        # node name -> the named resources built for it at load_model
+        # (shared between nodes on the same KV cache group).
+        self._node_resources: dict[str, dict[str, object]] = {}
 
         # Sequences each batch's resource lifecycle: admit before prepare,
         # plan surface before forward, publish after.
@@ -184,6 +188,9 @@ def __init__(
         supports_cpu_offload=True,
     )
 
+    def node_resources(self, node_name: str) -> dict[str, object]:
+        return self._node_resources.get(node_name, {})
+
     def engine_type(self) -> EngineType:
         return EngineType.KV_CACHE
 
@@ -196,13 +203,21 @@ def load_model(
         transfer_engine_info: TransferEngineInfo,
         default_sampling_config: dict[str, MultiSamplingConfig],
         kv_cache_type=None,
+        node_resources: list[NodeResourceSpec] | None = None,
     ) -> None:
         self.device = device
         if kv_cache_type is None:
             kv_cache_type = self.autocast_dtype
 
+        # Callers that pass only raw configs get the default declaration,
+        # the same one Model.get_node_resources derives.
+        specs = node_resources or [
+            NodeResourceSpec(kv_cache_config=cfg) for cfg in kv_cache_config
+        ]
+
         node_to_kv_mgmt = {}
-        for cfg in kv_cache_config:
+        for spec in specs:
+            cfg = spec.kv_cache_config
             num_layers = cfg.num_layers
             max_num_pages = cfg.max_num_pages
             page_size = cfg.page_size
@@ -294,11 +309,28 @@ def load_model(
             )
             self.kv_management[cfg.get_node_str()] = kv_mgmt
 
+            resources: dict[str, object] = {
+                "kv": kv_mgmt.kv_pool,
+                "rope": kv_mgmt.rope_embedder,
+            }
+            for source, cross_kv_pool in kv_mgmt.cross_kv_pools.items():
+                resources[f"cross_kv:{source}"] = cross_kv_pool
+            for key, scratch_spec in spec.scratch.items():
+                resources[key] = ScratchKVPool(torch.zeros(
+                    scratch_spec.shape,
+                    dtype=scratch_spec.dtype or kv_cache_type,
+                    device=device,
+                ))
+
             for node_name in nodes:
                 node_to_kv_mgmt[node_name] = kv_mgmt
+                self._node_resources[node_name] = resources
 
         for node_name, submodule in submodules.items():
             tp_group = parallel_groups.get_tp_config_for_node(node_name)
+            submodule.bind_node_resources(
+                self._node_resources.get(node_name, {})
+            )
             sampl_cfg = default_sampling_config.get(
                 node_name, MultiSamplingConfig()
             )
diff --git a/mstar/engine/resources/__init__.py b/mstar/engine/resources/__init__.py
index 0b58550dd..4604448af 100644
--- a/mstar/engine/resources/__init__.py
+++ b/mstar/engine/resources/__init__.py
@@ -10,8 +10,9 @@
     Segment,
     SequenceView,
 )
-from mstar.engine.resources.kv_pool import KVCachePool, PageArena
+from mstar.engine.resources.kv_pool import KVCachePool, PageArena, ScratchKVPool
 from mstar.engine.resources.positions import RopeEmbedder
+from mstar.engine.resources.spec import NodeResourceSpec, ScratchKVSpec
 from mstar.engine.resources.step import StepPlan, StepRunner
 
 __all__ = [
@@ -19,10 +20,13 @@
     "DenseGenAttentionManager",
     "FlashInferAttentionManager",
     "KVCachePool",
+    "NodeResourceSpec",
     "PageArena",
     "PositionPlan",
     "Reservation",
     "RopeEmbedder",
+    "ScratchKVPool",
+    "ScratchKVSpec",
     "Segment",
     "SequenceView",
     "StepPlan",
diff --git a/mstar/engine/resources/kv_pool.py b/mstar/engine/resources/kv_pool.py
index d8d0b7886..f9b806495 100644
--- a/mstar/engine/resources/kv_pool.py
+++ b/mstar/engine/resources/kv_pool.py
@@ -50,6 +50,17 @@ def total_pages(self) -> int:
         return self.allocator.max_num_pages
 
 
+class ScratchKVPool:
+    """Fixed-shape scratch KV storage with a trivial lifecycle: no admit,
+    no publish, no per-request lifetime. The tensor is overwritten every
+    step and slot-indexed by batch position. Models that need a depth
+    loop's per-step cache take it from here instead of allocating their
+    own tensors on the side."""
+
+    def __init__(self, tensor: torch.Tensor):
+        self.tensor = tensor
+
+
 class KVCachePool:
     """Per-request cache accounting behind the segment lifecycle.
 
diff --git a/mstar/engine/resources/spec.py b/mstar/engine/resources/spec.py
new file mode 100644
index 000000000..e7fa696ca
--- /dev/null
+++ b/mstar/engine/resources/spec.py
@@ -0,0 +1,36 @@
+"""Resource declarations models hand to the engine.
+
+The engine builds each node's resources once, at load time, from these
+specs. A spec names what to build and its parameters; the model declares,
+the engine constructs. The default declaration wraps a model's KV cache
+configs unchanged, so a model only overrides it to add resources beyond
+what those configs already describe.
+"""
+
+from dataclasses import dataclass, field
+
+import torch
+
+from mstar.engine.kv_store import KVCacheConfig
+
+
+@dataclass(frozen=True)
+class ScratchKVSpec:
+    """A fixed-shape scratch cache: overwritten every step, slot-indexed
+    by batch position, no per-request lifetime. A ``dtype`` of None means
+    the engine's KV cache dtype."""
+    shape: tuple[int, ...]
+    dtype: "torch.dtype | None" = None
+
+
+@dataclass
+class NodeResourceSpec:
+    """One KV cache group's resource declaration.
+
+    The cache config derives the self-attention pool, the attention
+    backend, the rope embedder, and the cross-attention pools, exactly as
+    it always has. ``scratch`` adds keyed fixed-shape caches built
+    alongside them (resource key to spec).
+    """
+    kv_cache_config: KVCacheConfig
+    scratch: dict[str, ScratchKVSpec] = field(default_factory=dict)
diff --git a/mstar/model/base.py b/mstar/model/base.py
index e8bbd0ea0..dc987d837 100644
--- a/mstar/model/base.py
+++ b/mstar/model/base.py
@@ -15,6 +15,7 @@
 from mstar.distributed.base import ShardingConfig, ShardingGroup
 from mstar.engine.base import EngineType
 from mstar.engine.kv_store import KVCacheConfig
+from mstar.engine.resources.spec import NodeResourceSpec
 from mstar.graph.base import (
     GraphEdge,
     GraphNode,
@@ -360,6 +361,20 @@ def get_kv_cache_config(self) -> list[KVCacheConfig]:
         """
         pass
 
+    def get_node_resources(
+        self, kv_cache_config: list[KVCacheConfig],
+    ) -> list[NodeResourceSpec]:
+        """Declare the resources the engine builds per node group.
+
+        ``kv_cache_config`` is this model's config list after any
+        deployment overrides were applied. The default wraps each config
+        unchanged, so models that only define ``get_kv_cache_config``
+        need no change here. Models with resources those configs cannot
+        describe (e.g. a fixed-shape scratch cache) override this and
+        extend the returned specs.
+        """
+        return [NodeResourceSpec(kv_cache_config=cfg) for cfg in kv_cache_config]
+
     def get_sampling_config(
         self, node_name: str,
         model_kwargs: dict | None = None,
diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py
index 871928b71..dd9f74664 100644
--- a/mstar/model/submodule_base.py
+++ b/mstar/model/submodule_base.py
@@ -236,6 +236,14 @@ def __init__(self):
         # of the same objects. The engine removes a request's entry via
         # ``cleanup_request`` when the request is removed.
         self.request_states: dict[str, PerRequestState] = {}
+        # Engine-built resources for this submodule's node (KV cache pool,
+        # embedder, scratch caches), bound once at load. Empty until then
+        # and on engines that build none.
+        self.node_resources: dict[str, Any] = {}
+
+    def bind_node_resources(self, resources: dict[str, Any]) -> None:
+        """Receive the engine-built resources for this submodule's node."""
+        self.node_resources = resources
 
     def request_state(self, request_id: str) -> PerRequestState:
         """The request's state, created on first access."""
diff --git a/mstar/worker/engine_manager.py b/mstar/worker/engine_manager.py
index 82ba0648e..1a9248db8 100644
--- a/mstar/worker/engine_manager.py
+++ b/mstar/worker/engine_manager.py
@@ -87,6 +87,7 @@ def build(
         real computation.
         """
         node_to_engine_type = model.get_node_engine_types()
+        node_resource_specs = model.get_node_resources(kv_config)
 
         # Resolve autocast dtype: explicit YAML config wins; otherwise we
         # fall back to the Model's own preference (so models that need to
@@ -163,6 +164,7 @@ def build(
                 device=device,
                 transfer_engine_info=transfer_engine_info,
                 kv_cache_type=autocast_dtype,
+                node_resources=node_resource_specs,
                 default_sampling_config={
                     node: model.resolve_sampling_configs(node, {}) \
                         for node in submodules
diff --git a/test/modular/test_node_resources.py b/test/modular/test_node_resources.py
new file mode 100644
index 000000000..0c2d28268
--- /dev/null
+++ b/test/modular/test_node_resources.py
@@ -0,0 +1,90 @@
+"""Unit tests for node resource declaration and construction.
+
+Models declare per-node resources as ``NodeResourceSpec`` lists; the
+engine builds them at load time and exposes them per node through
+``engine.node_resources``. Submodules receive their node's dict once,
+at bind time.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import sys
+
+sys.path.insert(0, ".")
+
+import pytest
+import torch
+
+from mstar.engine.kv_store import KVCacheConfig
+from mstar.engine.resources import NodeResourceSpec, ScratchKVPool, ScratchKVSpec
+from mstar.model.base import Model
+from mstar.model.submodule_base import NodeInputs, NodeSubmodule
+
+
+def _config(nodes: list[str]) -> KVCacheConfig:
+    return KVCacheConfig(
+        num_layers=2,
+        num_kv_heads=1,
+        head_dim=4,
+        max_seq_len=64,
+        max_num_pages=8,
+        page_size=8,
+        nodes=nodes,
+    )
+
+
+class TestDefaultDeclaration:
+    def test_default_wraps_each_config_unchanged(self):
+        configs = [_config(["A"]), _config(["B", "C"])]
+        specs = Model.get_node_resources(None, configs)
+
+        assert [s.kv_cache_config for s in specs] == configs
+        assert all(s.scratch == {} for s in specs)
+
+    def test_override_can_extend_a_spec(self):
+        class _Declares(Model):
+            def get_node_resources(self, kv_cache_config):
+                specs = super().get_node_resources(kv_cache_config)
+                specs[0].scratch["scratch_kv"] = ScratchKVSpec(shape=(2, 4))
+                return specs
+
+        _Declares.__abstractmethods__ = frozenset()
+        specs = _Declares().get_node_resources([_config(["A"])])
+        assert specs[0].scratch["scratch_kv"].shape == (2, 4)
+        assert specs[0].scratch["scratch_kv"].dtype is None
+
+
+class TestSpecTypes:
+    def test_scratch_spec_is_immutable(self):
+        spec = ScratchKVSpec(shape=(1, 2, 3))
+        with pytest.raises(dataclasses.FrozenInstanceError):
+            spec.shape = (4,)
+
+    def test_scratch_pool_holds_its_tensor(self):
+        tensor = torch.zeros(2, 3)
+        pool = ScratchKVPool(tensor)
+        assert pool.tensor is tensor
+
+
+class _StubSubmodule(NodeSubmodule):
+    def prepare_inputs(self, graph_walk, fwd_info, inputs, **kwargs):
+        return NodeInputs()
+
+    def forward(self, graph_walk, engine_inputs, **kwargs):
+        return {}
+
+
+class TestSubmoduleBinding:
+    def test_unbound_submodule_has_no_resources(self):
+        assert _StubSubmodule().node_resources == {}
+
+    def test_bind_replaces_the_dict(self):
+        submodule = _StubSubmodule()
+        resources = {"scratch_kv": ScratchKVPool(torch.zeros(1))}
+        submodule.bind_node_resources(resources)
+        assert submodule.node_resources is resources
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

From 708501f3c73b23dbfbd8f4bcca1c686fbad67bca Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 22:09:07 +0000
Subject: [PATCH 10/20] Move offload and LRU surfaces onto the KV cache pool

---
 mstar/engine/base.py                         |  66 ++--------
 mstar/engine/kv_cache_engine.py              |  89 ++-----------
 mstar/engine/resources/kv_pool.py            |  66 +++++++++-
 mstar/worker/engine_manager.py               |  25 ++--
 mstar/worker/worker.py                       |  18 ++-
 test/modular/test_kv_cache_engine_cleanup.py |   3 +-
 test/modular/test_kv_pool.py                 | 128 ++++++++++++++++++-
 7 files changed, 241 insertions(+), 154 deletions(-)

diff --git a/mstar/engine/base.py b/mstar/engine/base.py
index da65a9c54..a64b6ba45 100644
--- a/mstar/engine/base.py
+++ b/mstar/engine/base.py
@@ -9,7 +9,7 @@
 from mstar.communication.tensors import NameToTensorList
 from mstar.conductor.request_info import CurrentForwardPassInfo
 from mstar.distributed.communication import WorkerParallelGroups
-from mstar.engine.kv_store import KVCacheConfig, StoreWritePolicy
+from mstar.engine.kv_store import KVCacheConfig
 from mstar.profile.worker import ExecTimings
 
 
@@ -23,14 +23,12 @@ class EngineCapabilities:
     """Static declaration of optional surfaces an engine implements.
 
     The worker and ``EngineManager`` consult these flags instead of
-    ``isinstance`` / ``hasattr`` probes to decide whether to iterate or
-    dispatch into engine-specific code paths (e.g. CPU-offload victim
-    selection, KV-cache LRU tracking, store write-policy push). The
-    default ``EngineCapabilities()`` declares an engine that needs none
-    of the optional surfaces — stateless engines leave it untouched.
+    ``isinstance`` / ``hasattr`` probes. The default declares an engine
+    that needs none of the optional surfaces; stateless engines leave it
+    untouched. Per-node resource questions (offload tiers, LRU tracking,
+    write policy) go through ``node_resources`` instead of flags here.
     """
     requires_kv_cache: bool = False
-    supports_cpu_offload: bool = False
 
 
 @dataclass
@@ -308,15 +306,14 @@ def reset_pre_plan_for_batch(self, batch: NodeBatch) -> None:
         """
         return
 
-    # ── Capabilities + optional surfaces ────────────────────────────────
+    # ── Capabilities + per-node resources ───────────────────────────────
     #
     # ``capabilities`` is a class-level declaration of which optional
-    # surfaces this engine class implements. Worker / EngineManager check
-    # it instead of ``isinstance`` / ``hasattr`` probes. The methods below
-    # are the corresponding surfaces — all safe no-op defaults so engines
-    # that don't opt in can still be called uniformly. KVCacheEngine
-    # overrides both the capability flags and the methods; stateless
-    # engines leave them at default.
+    # surfaces this engine class implements; Worker / EngineManager check
+    # it instead of ``isinstance`` / ``hasattr`` probes. ``node_resources``
+    # is the per-node counterpart: the worker reaches an engine's KV cache
+    # pool (offload, LRU eligibility, write policy) through it rather than
+    # through per-question engine methods.
 
     capabilities = EngineCapabilities()
 
@@ -327,47 +324,6 @@ def node_resources(self, node_name: str) -> dict[str, Any]:
         """
         return {}
 
-    def lru_tracked_nodes(self) -> list[str]:
-        """Nodes for which the worker should LRU-track per-request activity
-        (used to pick CPU-offload victims). Default: no nodes — stateless
-        engines have no KV state to age out.
-        """
-        return []
-
-    def set_alloc_write_policy(self, policy: StoreWritePolicy) -> None:
-        """Apply a store write policy. Default: no-op — engines without an
-        alloc manager have nothing to set.
-        """
-        return
-
-    def offload_candidates(self, node_name: str) -> list[tuple[str, int]]:
-        """Return ``(request_id, gpu_pages_held)`` for every request with
-        GPU pages on ``node_name``. The worker partitions the result into
-        in-batch vs external candidates and picks an eviction victim.
-        Default: empty list — no offloadable state.
-        """
-        return []
-
-    def offload_request(self, node_name: str, request_id: str) -> int:
-        """Offload ``request_id``'s KV pages on ``node_name`` to CPU.
-        Returns the number of GPU pages freed (0 if nothing was freed or
-        the engine doesn't support offload).
-        """
-        return 0
-
-    def reload_request(self, node_name: str, request_id: str) -> bool:
-        """Reload an offloaded request back to GPU on ``node_name``.
-        Returns True on success; False if the request isn't offloaded, GPU
-        pages are insufficient, or the engine doesn't support offload.
-        """
-        return False
-
-    def is_offloaded(self, node_name: str, request_id: str) -> bool:
-        """Whether ``request_id`` is currently CPU-offloaded on ``node_name``.
-        Default: False.
-        """
-        return False
-
     def execute_with_max_batch_size(self, batch: NodeBatch) -> NodeOutput:
         if self.enable_profile:
             batch.exec_timings.start = time.perf_counter()
diff --git a/mstar/engine/kv_cache_engine.py b/mstar/engine/kv_cache_engine.py
index 9376ccae2..dc1c8624f 100644
--- a/mstar/engine/kv_cache_engine.py
+++ b/mstar/engine/kv_cache_engine.py
@@ -52,7 +52,6 @@ class KVManagement:
     kv_cache_config: KVCacheConfig
     kv_cache: torch.Tensor
     alloc_manager: PagedAllocationManager
-    cpu_page_pool: CPUPagePool | None
     buffer_manager: WorkspaceBufferManager
     # source name -> cross-attention context pool (see KVCacheConfig.cross_attn)
     cross_pools: dict[str, CrossAttnPool] = field(default_factory=dict)
@@ -61,9 +60,9 @@ class KVManagement:
     # startup, so per-request add/remove doesn't re-walk cross_pools.
     cross_alloc_managers: list[PagedAllocationManager] = field(default_factory=list)
     # Persistent resource fronts over the managers above: the pool is the
-    # engine's surface for admission, retrieval, position reads, and
-    # publishing; the embedder owns position semantics. Built once at
-    # load_model and shared by every step's facade.
+    # engine's surface for admission, retrieval, offload tiers, position
+    # reads, and publishing; the embedder owns position semantics. Built
+    # once at load_model and shared by every step's facade.
     kv_pool: KVCachePool | None = None
     cross_kv_pools: dict[str, KVCachePool] = field(default_factory=dict)
     rope_embedder: RopeEmbedder | None = None
@@ -183,10 +182,7 @@ def __init__(
         # warnings — logged at most once per (node, graph walk).
         self._logged_missing_rid_outputs: set[tuple[str, str]] = set()
 
-    capabilities = EngineCapabilities(
-        requires_kv_cache=True,
-        supports_cpu_offload=True,
-    )
+    capabilities = EngineCapabilities(requires_kv_cache=True)
 
     def node_resources(self, node_name: str) -> dict[str, object]:
         return self._node_resources.get(node_name, {})
@@ -293,14 +289,13 @@ def load_model(
                 kv_cache_config=cfg,
                 kv_cache=kv_cache,
                 alloc_manager=alloc_manager,
-                cpu_page_pool=cpu_page_pool,
                 buffer_manager = WorkspaceBufferManager(
                     int(os.environ.get("MSTAR_WORKSPACE_BUFFER_MB", "512")) * 1024 * 1024,
                     device=device,
                 ),
                 cross_pools=cross_pools,
                 cross_alloc_managers=cross_alloc_managers,
-                kv_pool=KVCachePool(alloc_manager),
+                kv_pool=KVCachePool(alloc_manager, cpu_pool=cpu_page_pool),
                 cross_kv_pools={
                     source: KVCachePool(pool.alloc_manager)
                     for source, pool in cross_pools.items()
@@ -354,7 +349,6 @@ def _create_cache_manager(
         submod_mgmt = self.submodule_management[node_name]
         cache_mgmt = submod_mgmt.kv_management
 
-        from mstar.engine.kv_store import StoreWritePolicy
         autowrite = (cache_mgmt.alloc_manager.write_policy == StoreWritePolicy.ALWAYS)
 
         return create_cache_manager(
@@ -1166,12 +1160,10 @@ def check_ready(
         submod_mgmt = self.submodule_management[node_name]
         cache_mgmt = submod_mgmt.kv_management
         # If this request was offloaded to CPU, try reloading first
-        if cache_mgmt.cpu_page_pool is not None and cache_mgmt.cpu_page_pool.is_offloaded(request_id):
-            try:
-                cache_mgmt.alloc_manager.reload_request(request_id, cache_mgmt.cpu_page_pool)
-                logger.info("Reloaded offloaded request %s from CPU", request_id)
-            except RuntimeError:
+        if cache_mgmt.kv_pool.is_offloaded(request_id):
+            if not cache_mgmt.kv_pool.reload(request_id):
                 return False  # can't reload yet, not ready
+            logger.info("Reloaded offloaded request %s from CPU", request_id)
 
         needed_labels = self._get_needed_labels(
             node_name, request_info.graph_walk, {
@@ -1329,7 +1321,7 @@ def add_request(
         self, request_id: str, cache_labels: list[str] | None = None,
     ) -> None:
         for submodule_mgmt in self.submodule_management.values():
-            submodule_mgmt.kv_management.alloc_manager.add_request(request_id, cache_labels or ["main"])
+            submodule_mgmt.kv_management.kv_pool.add_request(request_id, cache_labels or ["main"])
             for cross_mgr in submodule_mgmt.kv_management.cross_alloc_managers:
                 cross_mgr.add_request(request_id)
             submodule_mgmt.sampler.add_request(request_id)
@@ -1342,9 +1334,7 @@ def add_request(
     def remove_request(self, request_id: str) -> None:
         for submodule_mgmt in self.submodule_management.values():
             cache_mgmt = submodule_mgmt.kv_management
-            if cache_mgmt.cpu_page_pool is not None:
-                cache_mgmt.cpu_page_pool.remove_request(request_id)
-            cache_mgmt.alloc_manager.remove_request(request_id)
+            cache_mgmt.kv_pool.remove_request(request_id)
             for cross_mgr in cache_mgmt.cross_alloc_managers:
                 cross_mgr.remove_request(request_id)
             submodule_mgmt.sampler.remove_request(request_id)
@@ -1368,65 +1358,6 @@ def resume_request(
             cache_mgmt = submodule_mgmt.kv_management
             cache_mgmt.alloc_manager.get_state(request_id, cache_label).is_paused = False
 
-    # ── Optional surfaces declared via ``capabilities`` ─────────────────
-
-    def lru_tracked_nodes(self) -> list[str]:
-        return list(self.submodule_management.keys())
-
-    def set_alloc_write_policy(self, policy: StoreWritePolicy) -> None:
-        for submod_mgmt in self.submodule_management.values():
-            submod_mgmt.kv_management.alloc_manager.write_policy = policy
-
-    def offload_candidates(self, node_name: str) -> list[tuple[str, int]]:
-        submod_mgmt = self.submodule_management.get(node_name)
-        if submod_mgmt is None or submod_mgmt.kv_management.cpu_page_pool is None:
-            return []
-        alloc = submod_mgmt.kv_management.alloc_manager
-        out: list[tuple[str, int]] = []
-        for rid, labels in alloc.request_states.items():
-            total_pages = sum(len(s.page_indices) for s in labels.values())
-            if total_pages > 0:
-                out.append((rid, total_pages))
-        return out
-
-    def offload_request(self, node_name: str, request_id: str) -> int:
-        submod_mgmt = self.submodule_management.get(node_name)
-        if submod_mgmt is None:
-            return 0
-        cache_mgmt = submod_mgmt.kv_management
-        if cache_mgmt.cpu_page_pool is None:
-            return 0
-        return cache_mgmt.alloc_manager.offload_request(
-            request_id, cache_mgmt.cpu_page_pool,
-        )
-
-    def reload_request(self, node_name: str, request_id: str) -> bool:
-        submod_mgmt = self.submodule_management.get(node_name)
-        if submod_mgmt is None:
-            return False
-        cache_mgmt = submod_mgmt.kv_management
-        if cache_mgmt.cpu_page_pool is None:
-            return False
-        if not cache_mgmt.cpu_page_pool.is_offloaded(request_id):
-            return False
-        try:
-            cache_mgmt.alloc_manager.reload_request(
-                request_id, cache_mgmt.cpu_page_pool,
-            )
-            return True
-        except RuntimeError:
-            # Not enough GPU pages to reload — caller will retry later.
-            return False
-
-    def is_offloaded(self, node_name: str, request_id: str) -> bool:
-        submod_mgmt = self.submodule_management.get(node_name)
-        if submod_mgmt is None:
-            return False
-        cache_mgmt = submod_mgmt.kv_management
-        if cache_mgmt.cpu_page_pool is None:
-            return False
-        return cache_mgmt.cpu_page_pool.is_offloaded(request_id)
-
     def shutdown(self) -> None:
         for submodule_mgmt in self.submodule_management.values():
             cache_mgmt = submodule_mgmt.kv_management
diff --git a/mstar/engine/resources/kv_pool.py b/mstar/engine/resources/kv_pool.py
index f9b806495..83921fdb0 100644
--- a/mstar/engine/resources/kv_pool.py
+++ b/mstar/engine/resources/kv_pool.py
@@ -10,10 +10,12 @@
 import torch
 
 from mstar.conductor.request_info import SequenceInfo
+from mstar.engine.cpu_page_pool import CPUPagePool
 from mstar.engine.kv_store import (
     PageAllocator,
     PagedAllocationManager,
     PositionInfo,
+    StoreWritePolicy,
 )
 from mstar.engine.resources.base import Reservation, Segment, SequenceView
 
@@ -71,10 +73,19 @@ class KVCachePool:
     (its lock, its request states, its transfer machinery); the pool is the
     surface planning code goes through, so callers stop reaching into
     request-state internals.
+
+    ``cpu_pool`` is the pool's optional CPU tier: with one attached, whole
+    requests can be offloaded there and reloaded later, and the pool
+    answers the offload questions eviction policy asks.
     """
 
-    def __init__(self, manager: PagedAllocationManager):
+    def __init__(
+        self,
+        manager: PagedAllocationManager,
+        cpu_pool: CPUPagePool | None = None,
+    ):
         self._manager = manager
+        self._cpu_pool = cpu_pool
         self._arena = PageArena(
             tensor=manager.kv_cache,
             allocator=manager.page_allocator,
@@ -142,6 +153,59 @@ def labels(self, request_id: str) -> list[str]:
         """The cache streams currently existing for one request."""
         return self._manager.get_labels(request_id)
 
+    def add_request(self, request_id: str, labels: list[str]) -> None:
+        """Open per-request accounting with the given initial streams."""
+        self._manager.add_request(request_id, labels)
+
+    def remove_request(self, request_id: str) -> None:
+        """Drop a request's accounting on every tier and free its pages."""
+        if self._cpu_pool is not None:
+            self._cpu_pool.remove_request(request_id)
+        self._manager.remove_request(request_id)
+
+    def set_write_policy(self, policy: StoreWritePolicy) -> None:
+        """Set whether committed pages are pushed to the distributed store."""
+        self._manager.write_policy = policy
+
+    @property
+    def supports_offload(self) -> bool:
+        """Whether a CPU tier is attached to offload requests into."""
+        return self._cpu_pool is not None
+
+    def offload_candidates(self) -> list[tuple[str, int]]:
+        """``(request_id, gpu_pages_held)`` for every request holding GPU
+        pages, for eviction-victim selection. Empty without a CPU tier."""
+        if self._cpu_pool is None:
+            return []
+        out: list[tuple[str, int]] = []
+        for request_id, states in self._manager.request_states.items():
+            total_pages = sum(len(s.page_indices) for s in states.values())
+            if total_pages > 0:
+                out.append((request_id, total_pages))
+        return out
+
+    def offload(self, request_id: str) -> int:
+        """Move all of one request's streams to the CPU tier. Returns the
+        number of GPU pages freed (0 without a CPU tier)."""
+        if self._cpu_pool is None:
+            return 0
+        return self._manager.offload_request(request_id, self._cpu_pool)
+
+    def reload(self, request_id: str) -> bool:
+        """Bring an offloaded request back to GPU. False when the request
+        is not offloaded or GPU pages are insufficient right now."""
+        if self._cpu_pool is None or not self._cpu_pool.is_offloaded(request_id):
+            return False
+        try:
+            self._manager.reload_request(request_id, self._cpu_pool)
+            return True
+        except RuntimeError:
+            return False
+
+    def is_offloaded(self, request_id: str) -> bool:
+        """Whether the request currently lives on the CPU tier."""
+        return self._cpu_pool is not None and self._cpu_pool.is_offloaded(request_id)
+
     def fork(
         self,
         request_id: str,
diff --git a/mstar/worker/engine_manager.py b/mstar/worker/engine_manager.py
index 1a9248db8..71cfce14d 100644
--- a/mstar/worker/engine_manager.py
+++ b/mstar/worker/engine_manager.py
@@ -211,18 +211,25 @@ def remove_request(self, request_id: str) -> None:
                 engine.remove_request(request_id)
 
     def set_alloc_write_policies(self, policy):
-        for engine in self._unique_engines():
-            engine.set_alloc_write_policy(policy)
+        """Apply a store write policy to every distinct KV cache pool."""
+        seen: set[int] = set()
+        for node_name, engine in self.node_to_engine.items():
+            pool = engine.node_resources(node_name).get("kv")
+            if pool is None or id(pool) in seen:
+                continue
+            seen.add(id(pool))
+            pool.set_write_policy(policy)
 
     def lru_tracked_nodes(self) -> list[str]:
-        """Aggregate ``engine.lru_tracked_nodes()`` across unique engines.
-        The worker uses this to seed / clean up the per-request LRU
-        timestamps it needs for offload-victim selection.
+        """Nodes whose engine holds a KV cache pool for them. The worker
+        uses this to seed / clean up the per-request LRU timestamps it
+        needs for offload-victim selection.
         """
-        out: list[str] = []
-        for engine in self._unique_engines():
-            out.extend(engine.lru_tracked_nodes())
-        return out
+        return [
+            node_name
+            for node_name, engine in self.node_to_engine.items()
+            if "kv" in engine.node_resources(node_name)
+        ]
 
     def _unique_engines(self) -> list[BaseEngine]:
         seen = set()
diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py
index 97be3ed33..59ddd9e28 100644
--- a/mstar/worker/worker.py
+++ b/mstar/worker/worker.py
@@ -797,11 +797,13 @@ def _try_offload_cold_request(
 
         Returns the victim request_id, or None if offloading wasn't possible.
         """
-        engine = self.engine_manager.get_engine(node_name)
-        if not engine.capabilities.supports_cpu_offload:
+        pool = self.engine_manager.get_engine(node_name).node_resources(
+            node_name
+        ).get("kv")
+        if pool is None or not pool.supports_offload:
             return None
 
-        candidates_raw = engine.offload_candidates(node_name)
+        candidates_raw = pool.offload_candidates()
         if not candidates_raw:
             return None
 
@@ -820,7 +822,7 @@ def _try_offload_cold_request(
             return None
 
         victim_id = self._select_eviction_victim(node_name, candidates)
-        freed = engine.offload_request(node_name, victim_id)
+        freed = pool.offload(victim_id)
         logger.info(
             "Offloaded request %s to CPU (%d GPU pages freed, "
             "policy=%s, in_batch=%s)",
@@ -851,10 +853,12 @@ def _select_eviction_victim(
 
     def _try_reload_request(self, node_name: str, request_id: str) -> bool:
         """Reload an offloaded request back to GPU. Returns True if reloaded."""
-        engine = self.engine_manager.get_engine(node_name)
-        if not engine.is_offloaded(node_name, request_id):
+        pool = self.engine_manager.get_engine(node_name).node_resources(
+            node_name
+        ).get("kv")
+        if pool is None or not pool.is_offloaded(request_id):
             return False
-        if engine.reload_request(node_name, request_id):
+        if pool.reload(request_id):
             logger.info("Reloaded request %s from CPU to GPU", request_id)
             return True
         logger.debug(
diff --git a/test/modular/test_kv_cache_engine_cleanup.py b/test/modular/test_kv_cache_engine_cleanup.py
index 941d57857..dba4b41d4 100644
--- a/test/modular/test_kv_cache_engine_cleanup.py
+++ b/test/modular/test_kv_cache_engine_cleanup.py
@@ -24,7 +24,6 @@ def _mgmt(submodule):
     """A SubmoduleManagement with mocked infra and the given submodule. Uses the
     real dataclass so a field rename breaks the test loudly."""
     kv = MagicMock(name="kv_management")
-    kv.cpu_page_pool = None  # exercise the None branch (skip cpu pool teardown)
     return SubmoduleManagement(
         submodule=submodule,
         kv_management=kv,
@@ -53,7 +52,7 @@ def test_remove_request_invokes_submodule_cleanup():
     b.cleanup_request.assert_called_once_with("req-1")
     # the teardown that already worked must still fire
     engine.submodule_management["n0"].sampler.remove_request.assert_called_once_with("req-1")
-    engine.submodule_management["n0"].kv_management.alloc_manager.remove_request.assert_called_once_with("req-1")
+    engine.submodule_management["n0"].kv_management.kv_pool.remove_request.assert_called_once_with("req-1")
 
 
 class _PooledSubmodule:
diff --git a/test/modular/test_kv_pool.py b/test/modular/test_kv_pool.py
index ca6659369..8016d1063 100644
--- a/test/modular/test_kv_pool.py
+++ b/test/modular/test_kv_pool.py
@@ -19,6 +19,7 @@
 import torch
 
 from mstar.conductor.request_info import SequenceInfo
+from mstar.engine.cpu_page_pool import OffloadedState
 from mstar.engine.kv_store import (
     AllocationFailedError,
     KVCacheConfig,
@@ -50,10 +51,51 @@ def get_kv_transfer_info(self):
         return self.transfer_info
 
 
+class _FakeCpuPool:
+    """CPU tier double: tracks offloaded page bookkeeping without pinned
+    memory or GPU copies."""
+
+    def __init__(self, max_pages: int = 16):
+        self.page_allocator = PageAllocator(max_pages)
+        self.offloaded: dict[str, dict[str, OffloadedState]] = {}
+
+    def is_offloaded(self, request_id: str) -> bool:
+        return bool(self.offloaded.get(request_id))
+
+    def offload_pages(
+        self, request_id, label, gpu_kv_cache, gpu_page_indices,
+        seq_len, position_id_start,
+    ) -> None:
+        cpu_pages = self.page_allocator.try_allocate(len(gpu_page_indices))
+        if cpu_pages is None:
+            return
+        self.offloaded.setdefault(request_id, {})[label] = OffloadedState(
+            cpu_page_indices=cpu_pages,
+            seq_len=seq_len,
+            position_id_start=position_id_start,
+        )
+
+    def reload_pages(self, request_id, label, gpu_kv_cache, gpu_page_indices):
+        state = self.offloaded[request_id][label]
+        self.page_allocator.free(state.cpu_page_indices)
+        del self.offloaded[request_id][label]
+        if not self.offloaded[request_id]:
+            del self.offloaded[request_id]
+        return state.seq_len, state.position_id_start
+
+    def sync(self) -> None:
+        return
+
+    def remove_request(self, request_id: str) -> None:
+        for state in self.offloaded.pop(request_id, {}).values():
+            self.page_allocator.free(state.cpu_page_indices)
+
+
 def _make_pool(
     max_num_pages: int = 16,
     page_size: int = 8,
     with_tensor: bool = False,
+    cpu_pool: _FakeCpuPool | None = None,
 ) -> tuple[KVCachePool, PagedAllocationManager]:
     manager = PagedAllocationManager.__new__(PagedAllocationManager)
     manager.config = KVCacheConfig(
@@ -74,7 +116,7 @@ def _make_pool(
     manager._offload_stream = None
     manager.pending_reads = {}
     manager._lock = threading.RLock()
-    return KVCachePool(manager), manager
+    return KVCachePool(manager, cpu_pool=cpu_pool), manager
 
 
 class TestAdmit:
@@ -376,6 +418,90 @@ def test_publish_unknown_request_is_empty(self):
         assert pool.publish("ghost") == {}
 
 
+class TestOffloadTier:
+    def _fill(self, pool, tokens=20):
+        segment = Segment("r", "main", tokens)
+        pool.admit(segment)
+        pool.commit(segment)
+
+    def test_without_tier_the_pool_declines(self):
+        pool, manager = _make_pool()
+        manager.add_request("r", ["main"])
+        self._fill(pool)
+
+        assert pool.supports_offload is False
+        assert pool.offload_candidates() == []
+        assert pool.offload("r") == 0
+        assert pool.is_offloaded("r") is False
+        assert pool.reload("r") is False
+
+    def test_candidates_list_page_holding_requests(self):
+        pool, manager = _make_pool(cpu_pool=_FakeCpuPool())
+        manager.add_request("r", ["main"])
+        manager.add_request("empty", ["main"])
+        self._fill(pool, tokens=20)  # 3 pages for "r", none for "empty"
+
+        assert dict(pool.offload_candidates()) == {"r": 3}
+
+    def test_offload_frees_gpu_pages_and_reload_restores(self):
+        pool, manager = _make_pool(with_tensor=True, cpu_pool=_FakeCpuPool())
+        manager.add_request("r", ["main"])
+        self._fill(pool, tokens=20)  # 3 pages
+        free_before = pool.num_free_pages
+
+        freed = pool.offload("r")
+        assert freed == 3
+        assert pool.num_free_pages == free_before + 3
+        assert pool.is_offloaded("r") is True
+        assert manager.get_state("r", "main").page_indices == []
+
+        assert pool.reload("r") is True
+        assert pool.is_offloaded("r") is False
+        state = manager.get_state("r", "main")
+        assert state.seq_len == 20
+        assert len(state.page_indices) == 3
+        assert pool.num_free_pages == free_before
+
+    def test_reload_declines_when_gpu_pages_are_short(self):
+        pool, manager = _make_pool(
+            max_num_pages=3, with_tensor=True, cpu_pool=_FakeCpuPool(),
+        )
+        manager.add_request("r", ["main"])
+        self._fill(pool, tokens=20)  # all 3 pages
+        assert pool.offload("r") == 3
+
+        # Another request takes the freed pages; the reload cannot fit.
+        manager.add_request("greedy", ["main"])
+        greedy = Segment("greedy", "main", 24)
+        pool.admit(greedy)
+        pool.commit(greedy)
+
+        assert pool.reload("r") is False
+        assert pool.is_offloaded("r") is True
+
+    def test_remove_request_clears_the_cpu_tier(self):
+        cpu_pool = _FakeCpuPool(max_pages=4)
+        pool, manager = _make_pool(with_tensor=True, cpu_pool=cpu_pool)
+        manager.add_request("r", ["main"])
+        self._fill(pool, tokens=20)
+        pool.offload("r")
+        assert cpu_pool.page_allocator.num_free == 1
+
+        pool.remove_request("r")
+        assert cpu_pool.page_allocator.num_free == 4
+        assert "r" not in manager.request_states
+
+    def test_write_policy_reaches_the_manager(self):
+        pool, manager = _make_pool()
+        pool.set_write_policy(StoreWritePolicy.NEVER)
+        assert manager.write_policy is StoreWritePolicy.NEVER
+
+    def test_add_request_opens_streams(self):
+        pool, manager = _make_pool()
+        pool.add_request("r", ["main", "cfg"])
+        assert pool.labels("r") == ["main", "cfg"]
+
+
 class TestPosInfoAndLabels:
     def test_pos_info_reads_one_stream(self):
         pool, manager = _make_pool()

From bb24a499ea7722fadd0e9466a03600d09aa86c05 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 22:12:14 +0000
Subject: [PATCH 11/20] Read admit outcomes in check_ready via the pool

---
 mstar/engine/kv_cache_engine.py   | 13 +++----
 mstar/engine/resources/kv_pool.py | 16 +++++++++
 test/modular/test_kv_pool.py      | 57 +++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 9 deletions(-)

diff --git a/mstar/engine/kv_cache_engine.py b/mstar/engine/kv_cache_engine.py
index dc1c8624f..7dd51a289 100644
--- a/mstar/engine/kv_cache_engine.py
+++ b/mstar/engine/kv_cache_engine.py
@@ -1171,7 +1171,7 @@ def check_ready(
             }
         )
 
-        labels_to_check = []
+        outcomes = []
         try:
             kv_cache_string = cache_mgmt.kv_cache_config.get_node_str()
             world_size = submod_mgmt.tp_group.world_size
@@ -1184,19 +1184,14 @@ def check_ready(
             ).items():
                 if needed_labels is not None and label not in needed_labels:
                     continue
-                cache_mgmt.alloc_manager.start_async_retrieve(
+                outcomes.append(cache_mgmt.kv_pool.admit_retrieve(
                     request_id, label, seq_info
-                )
-                labels_to_check.append(label)
+                ))
         except RuntimeError:
             # Not enough pages to allocate for retrieval — not ready
             return False
 
-        ar_ready = all([
-            cache_mgmt.alloc_manager.check_retrieve_ready(request_id, label)
-            for label in labels_to_check
-        ])
-        if not ar_ready:
+        if any(outcome.pending for outcome in outcomes):
             return False
         return super().check_ready(node_name, request_id, request_info)
 
diff --git a/mstar/engine/resources/kv_pool.py b/mstar/engine/resources/kv_pool.py
index 83921fdb0..2adec40d3 100644
--- a/mstar/engine/resources/kv_pool.py
+++ b/mstar/engine/resources/kv_pool.py
@@ -246,6 +246,22 @@ def retrieve(self, request_id: str, label: str, seq_info: SequenceInfo) -> None:
         any admit."""
         self._manager.sync_retrieve(request_id, label, seq_info)
 
+    def admit_retrieve(
+        self, request_id: str, label: str, seq_info: SequenceInfo,
+    ) -> Reservation:
+        """Start bringing one stream's published state into residency and
+        report the admit outcome: ``pending`` stays True while pages are
+        in flight, and polling again re-checks without restarting the
+        transfer. Allocates on first call, so it can fail like any
+        admit."""
+        self._manager.start_async_retrieve(request_id, label, seq_info)
+        ready = self._manager.check_retrieve_ready(request_id, label)
+        return Reservation(
+            resident=self._manager.get_state(request_id, label).seq_len,
+            to_compute=0,
+            pending=not ready,
+        )
+
     def publish(self, request_id: str) -> dict[str, SequenceInfo]:
         """Describe the request's durable streams to another process, one
         ``SequenceInfo`` per label. Waits out any in-flight retrieves first
diff --git a/test/modular/test_kv_pool.py b/test/modular/test_kv_pool.py
index 8016d1063..84fa9a5cb 100644
--- a/test/modular/test_kv_pool.py
+++ b/test/modular/test_kv_pool.py
@@ -502,6 +502,63 @@ def test_add_request_opens_streams(self):
         assert pool.labels("r") == ["main", "cfg"]
 
 
+class _FakeFuture:
+    def __init__(self):
+        self._done = False
+
+    def done(self) -> bool:
+        return self._done
+
+    def result(self):
+        return None
+
+
+class TestAdmitRetrieve:
+    def _seq_info(self, seq_len=12, pos_id=34):
+        return SequenceInfo(
+            seq_len=seq_len,
+            pos_id=pos_id,
+            latest_kv_transfer_info=object(),
+            page_indices=[5, 6],
+        )
+
+    def test_instant_transfer_is_not_pending(self):
+        pool, manager = _make_pool(page_size=8)
+        manager.add_request("r", ["main"])
+
+        outcome = pool.admit_retrieve("r", "main", self._seq_info())
+
+        assert outcome.pending is False
+        assert outcome.resident == 12
+        assert manager.get_state("r", "main").position_id_start == 34
+
+    def test_pending_until_the_transfer_lands(self):
+        pool, manager = _make_pool(page_size=8)
+        manager.add_request("r", ["main"])
+        future = _FakeFuture()
+        manager._kv_transfer_engine.read_batched_async = (
+            lambda remote_kv_info, read_info: future
+        )
+
+        first = pool.admit_retrieve("r", "main", self._seq_info())
+        assert first.pending is True
+
+        # Polling again neither restarts the transfer nor clears it early.
+        second = pool.admit_retrieve("r", "main", self._seq_info())
+        assert second.pending is True
+
+        future._done = True
+        third = pool.admit_retrieve("r", "main", self._seq_info())
+        assert third.pending is False
+
+    def test_admit_retrieve_can_fail_like_any_admit(self):
+        pool, manager = _make_pool(max_num_pages=1, page_size=8)
+        manager.add_request("r", ["main"])
+
+        with pytest.raises(AllocationFailedError):
+            pool.admit_retrieve("r", "main", self._seq_info(seq_len=100))
+
+
 class TestPosInfoAndLabels:
     def test_pos_info_reads_one_stream(self):
         pool, manager = _make_pool()

From e6abef22b23852b8cad50aa6dea06236bf510789 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 22:18:53 +0000
Subject: [PATCH 12/20] Serve the code predictor scratch cache from the engine

---
 mstar/model/qwen3_omni/qwen3_omni_model.py | 28 +++++++++++
 mstar/model/qwen3_omni/submodules.py       | 17 ++-----
 mstar/model/qwen3_tts/qwen3_tts_model.py   | 28 +++++++++++
 mstar/model/qwen3_tts/submodules.py        | 31 +++---------
 test/modular/test_node_resources.py        | 57 ++++++++++++++++++++++
 test/modular/test_qwen3_tts_model.py       | 15 ++++++
 6 files changed, 138 insertions(+), 38 deletions(-)

diff --git a/mstar/model/qwen3_omni/qwen3_omni_model.py b/mstar/model/qwen3_omni/qwen3_omni_model.py
index 1cddfc7fa..72cbbf894 100644
--- a/mstar/model/qwen3_omni/qwen3_omni_model.py
+++ b/mstar/model/qwen3_omni/qwen3_omni_model.py
@@ -42,6 +42,7 @@
 )
 from mstar.engine.base import EngineType
 from mstar.engine.kv_store import KVCacheConfig
+from mstar.engine.resources.spec import NodeResourceSpec, ScratchKVSpec
 from mstar.graph.base import GraphEdge, GraphNode, Loop, Sequential, TensorPointerInfo
 from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION
 from mstar.model.base import MAX_OUTPUT_TOKENS, ForwardPassArgs, Model, TensorAndMetadata
@@ -171,6 +172,33 @@ def get_kv_cache_config(self) -> list[KVCacheConfig]:
         )
         return [thinker_cfg, talker_cfg]
 
+    def get_node_resources(
+        self, kv_cache_config: list[KVCacheConfig],
+    ) -> list[NodeResourceSpec]:
+        """The Talker adds the code predictor's fixed-shape scratch cache.
+
+        CodePredictor attention is local to one 16-group frame, so its
+        cache is overwritten every Talker step rather than paged across
+        steps; a maximum-batch allocation gives the captured decode graph
+        stable addresses.
+        """
+        from mstar.model.qwen3_omni.submodules import TalkerSubmodule
+
+        specs = super().get_node_resources(kv_cache_config)
+        cp = self.config.code_predictor
+        for spec in specs:
+            nodes = spec.kv_cache_config.nodes or []
+            if "Talker" in nodes:
+                spec.scratch["code_predictor"] = ScratchKVSpec(shape=(
+                    cp.num_hidden_layers,
+                    TalkerSubmodule.MAX_BATCH_SIZE,
+                    2,
+                    cp.num_code_groups,
+                    cp.num_key_value_heads,
+                    cp.head_dim,
+                ))
+        return specs
+
     # -----------------------------------------------------------------------
     # Model ABC: node engine types
     # -----------------------------------------------------------------------
diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py
index a41289859..84e70bf37 100644
--- a/mstar/model/qwen3_omni/submodules.py
+++ b/mstar/model/qwen3_omni/submodules.py
@@ -1239,21 +1239,10 @@ def __init__(
         # inject tts_eos_embed for ONE step before falling back to pad.
         self._eos_embed_sent: set[str] = set()
 
-        # TODO: this is hacky; when we have time, refactor it to make this
-        # come from the engine
-        self._cp_kv_cache: torch.Tensor | None = None
-
     def _get_cp_kv_cache(self):
-        if self._cp_kv_cache is None:
-            self._cp_kv_cache = torch.zeros((
-                    self.cp_cfg.num_hidden_layers,
-                    self.MAX_BATCH_SIZE, 2, self.num_codes,
-                    self.cp_cfg.num_key_value_heads,
-                    self.cp_cfg.head_dim
-                ), dtype=self.talker_code_emb.weight.dtype,
-                device=self.get_device(),
-            )
-        return self._cp_kv_cache
+        """The engine-built CodePredictor scratch cache, overwritten every
+        Talker step."""
+        return self.node_resources["code_predictor"].tensor
 
     def init_tts_embeds(self, thinker_embed_tokens: nn.Embedding) -> None:
         """Pre-compute TTS pad/bos/eos hidden states using the Thinker's
diff --git a/mstar/model/qwen3_tts/qwen3_tts_model.py b/mstar/model/qwen3_tts/qwen3_tts_model.py
index 92b7a98a1..31aec6346 100644
--- a/mstar/model/qwen3_tts/qwen3_tts_model.py
+++ b/mstar/model/qwen3_tts/qwen3_tts_model.py
@@ -40,6 +40,7 @@
 )
 from mstar.engine.base import EngineType
 from mstar.engine.kv_cache_engine import KVCacheConfig
+from mstar.engine.resources.spec import NodeResourceSpec, ScratchKVSpec
 from mstar.graph.base import (
     GraphEdge,
     GraphNode,
@@ -233,6 +234,33 @@ def get_kv_cache_config(self) -> list[KVCacheConfig]:
             flashinfer_backend="auto",
         )]
 
+    def get_node_resources(
+        self, kv_cache_config: list[KVCacheConfig],
+    ) -> list[NodeResourceSpec]:
+        """The Talker adds the CodePredictor's fixed-shape scratch cache.
+
+        CodePredictor attention is local to one 16-group frame, so its
+        cache is overwritten every Talker step rather than paged across
+        steps; a maximum-batch allocation gives the captured decode graph
+        stable addresses.
+        """
+        from mstar.model.qwen3_tts.submodules import TalkerSubmodule
+
+        specs = super().get_node_resources(kv_cache_config)
+        cp = self.config.talker.code_predictor
+        for spec in specs:
+            nodes = spec.kv_cache_config.nodes or []
+            if "Talker" in nodes:
+                spec.scratch["code_predictor"] = ScratchKVSpec(shape=(
+                    cp.num_hidden_layers,
+                    TalkerSubmodule.MAX_BATCH_SIZE,
+                    2,
+                    self.config.talker.num_code_groups,
+                    cp.num_key_value_heads,
+                    cp.head_dim,
+                ))
+        return specs
+
     def get_node_engine_types(self) -> dict[str, EngineType]:
         """Talker keeps cross-step KV state; Codec is a pure frame decoder."""
         return {
diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py
index 27af3a399..8ab4c0516 100644
--- a/mstar/model/qwen3_tts/submodules.py
+++ b/mstar/model/qwen3_tts/submodules.py
@@ -97,7 +97,6 @@ def __init__(
         self.cp_config = config.talker.code_predictor
         self.num_codes = config.talker.num_code_groups
         self._suppress_mask: torch.Tensor | None = None
-        self._cp_kv_cache: torch.Tensor | None = None
 
     def _get_suppress_mask(self) -> torch.Tensor:
         """Cache the checkpoint's static invalid-token mask on the worker GPU."""
@@ -134,28 +133,11 @@ def _get_batch_suppress_mask(
         return mask
 
     def _get_cp_kv_cache(self, batch_size: int) -> torch.Tensor:
-        """Return the fixed CodePredictor scratch cache for this micro-batch.
-
-        CodePredictor attention is local to one 16-group frame, so this cache
-        does not belong to the engine's cross-step paged KV cache. A maximum
-        batch allocation is reused and overwritten for every Talker step,
-        which also gives the captured decode graph stable addresses.
+        """The engine-built CodePredictor scratch cache, sliced to this
+        micro-batch. The fixed maximum-batch tensor is overwritten every
+        Talker step and gives the captured decode graph stable addresses.
         """
-        expected = (
-            self.cp_config.num_hidden_layers,
-            self.MAX_BATCH_SIZE,
-            2,
-            self.num_codes,
-            self.cp_config.num_key_value_heads,
-            self.cp_config.head_dim,
-        )
-        if self._cp_kv_cache is None:
-            self._cp_kv_cache = torch.empty(
-                expected,
-                dtype=self.model.model.codec_embedding.weight.dtype,
-                device=self.get_device(),
-            )
-        return self._cp_kv_cache[:, :batch_size]
+        return self.node_resources["code_predictor"].tensor[:, :batch_size]
 
     def _project_text(self, token_ids: torch.Tensor) -> torch.Tensor:
         """Map tokenizer embeddings into the Talker hidden width."""
@@ -659,8 +641,9 @@ def get_piecewise_cuda_graph_configs(
 
         The whole-walk decode graph already covers this loop; this runner serves
         the paths it can't — chiefly ``talker_prefill``, which stays eager
-        because it is variable-length and runs once per request. It has no
-        engine KV cache: its frame-local cache is a static tensor owned here.
+        because it is variable-length and runs once per request. It uses no
+        paged KV cache: its frame-local cache is the engine's fixed-shape
+        scratch pool.
         """
         del tp_world_size
         hidden_size = self.talker_config.hidden_size
diff --git a/test/modular/test_node_resources.py b/test/modular/test_node_resources.py
index 0c2d28268..b1684b4d2 100644
--- a/test/modular/test_node_resources.py
+++ b/test/modular/test_node_resources.py
@@ -10,6 +10,7 @@
 
 import dataclasses
 import sys
+from types import SimpleNamespace
 
 sys.path.insert(0, ".")
 
@@ -67,6 +68,62 @@ def test_scratch_pool_holds_its_tensor(self):
         assert pool.tensor is tensor
 
 
+def _model_stub(model_cls, config):
+    """An instance carrying only ``config``, skipping the heavy loader
+    the real constructor runs."""
+    class _Stub(model_cls):
+        def __init__(self):
+            self.config = config
+
+    _Stub.__abstractmethods__ = frozenset()
+    return _Stub()
+
+
+class TestModelDeclarations:
+    def test_qwen3_tts_declares_the_code_predictor_scratch(self):
+        from mstar.model.qwen3_tts.qwen3_tts_model import Qwen3TTSModel
+
+        stub = _model_stub(Qwen3TTSModel, SimpleNamespace(
+            talker=SimpleNamespace(
+                num_hidden_layers=2, num_key_value_heads=1, head_dim=4,
+                max_position_embeddings=64, num_attention_heads=2,
+                num_code_groups=16,
+                code_predictor=SimpleNamespace(
+                    num_hidden_layers=3, num_key_value_heads=2, head_dim=8,
+                ),
+            ),
+        ))
+        specs = stub.get_node_resources(stub.get_kv_cache_config())
+        scratch = specs[0].scratch["code_predictor"]
+        assert scratch.shape == (3, 32, 2, 16, 2, 8)
+        assert scratch.dtype is None
+
+    def test_qwen3_omni_declares_it_on_the_talker_only(self):
+        from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel
+
+        stub = _model_stub(Qwen3OmniModel, SimpleNamespace(
+            thinker_text=SimpleNamespace(
+                num_hidden_layers=2, num_key_value_heads=1,
+                max_position_embeddings=64, num_attention_heads=2,
+            ),
+            thinker_head_dim=4,
+            talker_text=SimpleNamespace(
+                num_hidden_layers=2, num_key_value_heads=1,
+                num_attention_heads=2,
+            ),
+            talker_head_dim=4,
+            code_predictor=SimpleNamespace(
+                num_hidden_layers=5, num_code_groups=16,
+                num_key_value_heads=2, head_dim=8,
+            ),
+        ))
+        specs = stub.get_node_resources(stub.get_kv_cache_config())
+        by_nodes = {tuple(s.kv_cache_config.nodes): s for s in specs}
+        assert by_nodes[("Thinker",)].scratch == {}
+        scratch = by_nodes[("Talker",)].scratch["code_predictor"]
+        assert scratch.shape == (5, 32, 2, 16, 2, 8)
+
+
 class _StubSubmodule(NodeSubmodule):
     def prepare_inputs(self, graph_walk, fwd_info, inputs, **kwargs):
         return NodeInputs()
diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py
index 8775975ed..ab28a37ad 100644
--- a/test/modular/test_qwen3_tts_model.py
+++ b/test/modular/test_qwen3_tts_model.py
@@ -11,6 +11,7 @@
 
 from mstar.conductor.request_info import CurrentForwardConductorMetadata
 from mstar.engine.base import EngineType, NodeBatch
+from mstar.engine.resources import ScratchKVPool
 from mstar.model.qwen3_tts.components.talker import (
     Qwen3TTSCodePredictor,
     Qwen3TTSTalkerModel,
@@ -768,6 +769,20 @@ def test_qwen3_tts_depth_loop_piecewise_captures_its_own_sampling():
     ).to(device=dev, dtype=torch.bfloat16).eval()
     submodule.code_predictor.consolidate_stacked_weights()
     submodule.DECODE_CAPTURE_BATCH_SIZES = [1, 2]
+    # The depth loop's scratch cache is engine-built in serving; bind the
+    # same declared shape here since no engine is constructed.
+    class _DeclaringModel(Qwen3TTSModel):
+        def __init__(self):
+            self.config = config
+
+    _DeclaringModel.__abstractmethods__ = frozenset()
+    declaring = _DeclaringModel()
+    scratch = declaring.get_node_resources(
+        declaring.get_kv_cache_config()
+    )[0].scratch["code_predictor"]
+    submodule.bind_node_resources({"code_predictor": ScratchKVPool(
+        torch.zeros(scratch.shape, dtype=torch.bfloat16, device=dev)
+    )})
 
     rids = ["a", "b"]
     multi = MultiSamplingConfig(

From 6e7f0823ad965cbe31a5b5c0599eb97bb8b656da Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 9 Aug 2026 23:54:55 +0000
Subject: [PATCH 13/20] Add step declarations the runner drives for adopting
 models

---
 mstar/engine/cache_manager.py         |  40 ++--
 mstar/engine/cuda_graph_runner.py     |  52 ++++-
 mstar/engine/kv_cache_engine.py       |  18 ++
 mstar/engine/resources/__init__.py    |   3 +
 mstar/engine/resources/declare.py     |  52 +++++
 mstar/engine/resources/step.py        |  78 ++++++++
 mstar/model/submodule_base.py         |  16 ++
 test/modular/test_step_declaration.py | 270 ++++++++++++++++++++++++++
 8 files changed, 511 insertions(+), 18 deletions(-)
 create mode 100644 mstar/engine/resources/declare.py
 create mode 100644 test/modular/test_step_declaration.py

diff --git a/mstar/engine/cache_manager.py b/mstar/engine/cache_manager.py
index b0b165040..3eca3d928 100644
--- a/mstar/engine/cache_manager.py
+++ b/mstar/engine/cache_manager.py
@@ -184,6 +184,12 @@ def _plan_states(self) -> dict[str, _PlanState]:
         graph runner's per-slot store when one was passed in)."""
         return self.attention.states
 
+    @property
+    def is_captured(self) -> bool:
+        """Whether this step surface fronts a captured graph's persistent
+        per-slot plan store rather than a per-step eager one."""
+        return self._cuda_graph_mode
+
     @torch.compiler.disable
     def _get_state(self, request_id: str, label: str | None = None) -> KVRequestState:
         label = label or self.active_labels.get(request_id, "main")
@@ -281,14 +287,17 @@ def run_attention(
         k: torch.Tensor,
         v: torch.Tensor,
         layer_idx: int | None=None,
+        label: str | None = None,
     ) -> torch.Tensor:
-        """Run the pre-planned attention for the active label.
+        """Run the pre-planned attention for a cache label.
 
         Args:
             q: [total_tokens, num_q_heads, head_dim]
             k: [total_tokens, num_kv_heads, head_dim]
             v: [total_tokens, num_kv_heads, head_dim]
             layer_idx: transformer layer index
+            label: plan key to run (a cache label, or a combined key such
+                as ``_cfg_batched``). If None, uses the active label.
         Returns:
             output: [total_tokens, num_q_heads, head_dim]
         """
@@ -436,10 +445,13 @@ def apply_rope(
         rope_scale: float = 1,
         rope_theta: float = 10000.0,
         rope_dtype=None,
+        label: str | None = None,
         **kwargs
     ):
-        """Apply RoPE using the active label's pre-computed position IDs."""
-        label = self._active_label()
+        """Apply RoPE using a label's pre-computed position IDs (the
+        active label when ``label`` is None)."""
+        if label is None:
+            label = self._active_label()
 
         ps = self._plan_states[label]
         assert ps.pos_ids is not None
@@ -708,12 +720,14 @@ def run_attention(
         k: torch.Tensor,
         v: torch.Tensor,
         layer_idx: int | None=None,
+        label: str | None = None,
     ) -> torch.Tensor:
         """Run pre-planned FlashInfer attention with KV cache write.
 
-        Uses the active label's plan state (set up by a prior plan_attention
-        call). Writes K and V to the paged KV cache at pre-computed page
-        positions, then runs the FlashInfer wrapper for batched attention.
+        Uses the given label's plan state (the active label when ``label``
+        is None; set up by a prior plan_attention call). Writes K and V to
+        the paged KV cache at pre-computed page positions, then runs the
+        FlashInfer wrapper for batched attention.
 
         In CUDA graph mode, the wrapper's set_kv_cache() + run() operate on
         pre-computed token_to_page/token_to_cache or kv_cache_locations
@@ -725,7 +739,8 @@ def run_attention(
 
         orig_dtype = q.dtype
 
-        label = self._active_label()
+        if label is None:
+            label = self._active_label()
         ps = self._plan_states[label]
 
         assert self.kv_cache is not None and ps.wrapper is not None
@@ -968,16 +983,19 @@ def run_attention(
         k: torch.Tensor,
         v: torch.Tensor,
         layer_idx: int | None=None,
+        label: str | None = None,
     ) -> torch.Tensor:
-        """Route the active label to its dense plan when one was built, else
-        run the inherited paged FlashInfer attention."""
-        label = self._active_label()
+        """Route the label (the active one when None) to its dense plan
+        when one was built, else run the inherited paged FlashInfer
+        attention."""
+        if label is None:
+            label = self._active_label()
         ps = self._plan_states[label]
         if ps.dense_gen is not None:
             if layer_idx is None:
                 layer_idx = self.layer_idx
             return self.attention.run_dense(q, k, v, layer_idx, ps.dense_gen).to(q.dtype)
-        return super().run_attention(q, k, v, layer_idx=layer_idx)
+        return super().run_attention(q, k, v, layer_idx=layer_idx, label=label)
 
 
 # Backend registry: KVCacheConfig.attention_backend names one of these.
diff --git a/mstar/engine/cuda_graph_runner.py b/mstar/engine/cuda_graph_runner.py
index 1d710ac1a..6fae14ae4 100644
--- a/mstar/engine/cuda_graph_runner.py
+++ b/mstar/engine/cuda_graph_runner.py
@@ -41,6 +41,7 @@
     PiecewiseCudaGraphConfig,
 )
 from mstar.engine.kv_store import KVCacheConfig, PagedAllocationManager
+from mstar.engine.resources.step import StepRunner
 from mstar.model.submodule_base import ARNodeInputs, ARNodeSubmodule, ModelInputsFromEngine, NodeSubmodule
 from mstar.profile.worker import ExecTimings
 from mstar.utils.profiler import mark, range_pop, range_push
@@ -207,6 +208,10 @@ def __init__(
         self.default_sampling_config = default_sampling_config or MultiSamplingConfig()
         self.enable_nvtx = False  # set by KVCacheEngine after construction
 
+        # Drives declared steps (submodules that return a StepDeclaration)
+        # at capture and replay, in place of plan calls in preprocess.
+        self.step_runner = StepRunner()
+
         self.graphs: dict[CudaGraphKey, CudaGraphData] = {}
 
         self.memory_pool = None
@@ -806,9 +811,18 @@ def prepare_slot(slot_idx: int) -> _SlotCaptureSpec:
                 dummy_rids=dummy_rids, plan_states=plan_states, config=config,
             )
 
-            # Preprocess (plans attention+rope outside graph) and intern
-            # the resulting tensors into the shared static-buffer pool so
-            # both slots' captures read from the same GPU addresses.
+            # Declared steps plan through the runner; preprocess then only
+            # marshals data. Undeclared submodules keep planning inside
+            # preprocess. Either way the resulting tensors are interned
+            # into the shared static-buffer pool so both slots' captures
+            # read from the same GPU addresses.
+            declaration = submodule.declare_step(
+                graph_walk=config.capture_graph_walk,
+                engine_inputs=engine_inputs,
+                inputs=dummy_inputs,
+            )
+            if declaration is not None:
+                self.step_runner.drive(declaration, engine_inputs.cache_manager)
             preprocessed = submodule.preprocess(
                 graph_walk=config.capture_graph_walk,
                 engine_inputs=engine_inputs,
@@ -825,6 +839,9 @@ def prepare_slot(slot_idx: int) -> _SlotCaptureSpec:
             ]
 
             def re_prepare() -> None:
+                if declaration is not None:
+                    self.step_runner.drive(declaration, engine_inputs.cache_manager)
+                    return
                 submodule.preprocess(
                     graph_walk=config.capture_graph_walk,
                     engine_inputs=engine_inputs,
@@ -1404,6 +1421,13 @@ def _run_basic_batched(
             if self.enable_nvtx:
                 range_pop(synchronize=False)
                 range_push("cg.preprocess_replan.submodule_preprocess", synchronize=False)
+            declaration = submodule.declare_step(
+                graph_walk=key.graph_walk,
+                engine_inputs=engine_inputs,
+                inputs=real_inputs,
+            )
+            if declaration is not None:
+                self.step_runner.drive(declaration, static_cm)
             real_inputs = submodule.preprocess(
                 graph_walk=key.graph_walk,
                 engine_inputs=engine_inputs,
@@ -1472,8 +1496,12 @@ def _run_basic_batched(
                 range_push("cg.advance_seq_lens", synchronize=False)
             # Frozen-prefix denoise walks re-read a fixed prefix and overwrite the
             # same tail pages every step, so they opt out of the advance (it would
-            # grow the prefix across steps and corrupt attention).
-            if graph_data.config.advance_seq_lens:
+            # grow the prefix across steps and corrupt attention). Declared
+            # steps carry their own commit spans (including the padding
+            # tail, whose inputs the declaration covered).
+            if declaration is not None:
+                self.step_runner.commit(declaration, static_cm)
+            elif graph_data.config.advance_seq_lens:
                 for label in config_labels:
                     static_cm.set_active_label(label)
                     static_cm.advance_seq_lens()
@@ -1630,6 +1658,13 @@ def _run_flashinfer_packed(
             if self.enable_nvtx:
                 range_pop(synchronize=False)
                 range_push("cg.preprocess_replan.submodule_preprocess", synchronize=False)
+            declaration = submodule.declare_step(
+                graph_walk=key.graph_walk,
+                engine_inputs=engine_inputs,
+                inputs=padded_inputs,
+            )
+            if declaration is not None:
+                self.step_runner.drive(declaration, static_cm)
             real_packed = submodule.preprocess(
                 graph_walk=key.graph_walk,
                 engine_inputs=engine_inputs,
@@ -1692,8 +1727,11 @@ def _run_flashinfer_packed(
                 range_push("cg.advance_seq_lens", synchronize=False)
             # Frozen-prefix denoise walks re-read a fixed prefix and overwrite the
             # same tail pages every step, so they opt out of the advance (it would
-            # grow the prefix across steps and corrupt attention).
-            if graph_data.config.advance_seq_lens:
+            # grow the prefix across steps and corrupt attention). Declared
+            # steps carry their own commit spans (padding tail included).
+            if declaration is not None:
+                self.step_runner.commit(declaration, static_cm)
+            elif graph_data.config.advance_seq_lens:
                 if graph_data.config.batched_cfg:
                     # _batched_cfg_info (set by the preprocess plan above) makes a
                     # single advance_seq_lens walk every label's state; looping
diff --git a/mstar/engine/kv_cache_engine.py b/mstar/engine/kv_cache_engine.py
index 7dd51a289..f917adfce 100644
--- a/mstar/engine/kv_cache_engine.py
+++ b/mstar/engine/kv_cache_engine.py
@@ -581,6 +581,13 @@ def _execute_batched(
         )
         if self.enable_nvtx:
             range_push("ar.batched.preprocess", synchronize=False)
+        declaration = submodule.declare_step(
+            graph_walk=batch.graph_walk,
+            engine_inputs=engine_inputs,
+            inputs=inputs,
+        )
+        if declaration is not None:
+            self.step_runner.drive(declaration, cache_manager)
         preprocessed = submodule.preprocess(
             graph_walk=batch.graph_walk,
             engine_inputs=engine_inputs,
@@ -607,6 +614,8 @@ def _execute_batched(
         if self.enable_nvtx:
             range_pop()
 
+        if declaration is not None:
+            self.step_runner.commit(declaration, cache_manager)
         cache_manager.flush_to_store()
 
         # `__batched_logits__` is the stacked [B, V] logits the submodule
@@ -707,6 +716,13 @@ def _execute_sequential(
 
             if self.enable_nvtx:
                 range_push("ar.seq.preprocess", synchronize=False)
+            declaration = submodule.declare_step(
+                graph_walk=batch.graph_walk,
+                engine_inputs=engine_inputs,
+                inputs=[node_inputs],
+            )
+            if declaration is not None:
+                self.step_runner.drive(declaration, cache_manager)
             preprocessed = submodule.preprocess(
                 batch.graph_walk,
                 engine_inputs=engine_inputs,
@@ -733,6 +749,8 @@ def _execute_sequential(
             if self.enable_nvtx:
                 range_pop()
 
+            if declaration is not None:
+                self.step_runner.commit(declaration, cache_manager)
             cache_manager.flush_to_store()
             per_request_outputs[rid] = output
 
diff --git a/mstar/engine/resources/__init__.py b/mstar/engine/resources/__init__.py
index 4604448af..661875f5c 100644
--- a/mstar/engine/resources/__init__.py
+++ b/mstar/engine/resources/__init__.py
@@ -10,6 +10,7 @@
     Segment,
     SequenceView,
 )
+from mstar.engine.resources.declare import PlanSpec, StepDeclaration
 from mstar.engine.resources.kv_pool import KVCachePool, PageArena, ScratchKVPool
 from mstar.engine.resources.positions import RopeEmbedder
 from mstar.engine.resources.spec import NodeResourceSpec, ScratchKVSpec
@@ -22,6 +23,7 @@
     "KVCachePool",
     "NodeResourceSpec",
     "PageArena",
+    "PlanSpec",
     "PositionPlan",
     "Reservation",
     "RopeEmbedder",
@@ -29,6 +31,7 @@
     "ScratchKVSpec",
     "Segment",
     "SequenceView",
+    "StepDeclaration",
     "StepPlan",
     "StepRunner",
     "WorkspaceBufferManager",
diff --git a/mstar/engine/resources/declare.py b/mstar/engine/resources/declare.py
new file mode 100644
index 000000000..ef1a5defa
--- /dev/null
+++ b/mstar/engine/resources/declare.py
@@ -0,0 +1,52 @@
+"""Step declarations for runner-driven models.
+
+A migrated model does not sequence lifecycle stages. It declares what its
+step is: which cache streams it touches, what spans they grow by, which
+attention and rope plans back it, which streams fork, and what commits
+when the step lands. The runner drives the declaration against the step
+surface; the values below are immutable and valid for one step.
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True)
+class PlanSpec:
+    """One attention (and optional rope) plan of a declared step.
+
+    ``labels`` with a single entry plans that label alone; more than one
+    entry plans a single combined batch across them, label-major, under
+    ``combined_key``. ``spans`` gives each label's per-request token
+    counts in batch order. ``rope_pos_ids`` optionally supplies explicit
+    position ids (a tensor for a single label, a per-label list dict for
+    a combined plan); None lets the embedder build them from the pool's
+    counters. ``commit`` says whether the step's spans become part of the
+    streams' history; ``pos_advance`` overrides the per-request position
+    advance when it differs from the span.
+    """
+    labels: tuple[str, ...]
+    spans: dict[str, tuple[int, ...]]
+    is_causal: bool = True
+    write_store: bool = True
+    dense_gen: bool = False
+    rope: bool = False
+    rope_pos_ids: "dict[str, list[torch.Tensor]] | torch.Tensor | None" = None
+    combined_key: str = "_cfg_batched"
+    commit: bool = True
+    pos_advance: tuple[int, ...] | None = None
+
+
+@dataclass(frozen=True)
+class StepDeclaration:
+    """A model's declaration of one step.
+
+    ``plans`` run in order. ``pre_forks`` are (from_label, to_label)
+    pairs applied before anything plans; ``post_forks`` are applied at
+    commit, after the committed spans have landed (a fork copies the
+    source stream as it stands, so a post fork sees the step's writes).
+    """
+    plans: tuple[PlanSpec, ...]
+    pre_forks: tuple[tuple[str, str], ...] = ()
+    post_forks: tuple[tuple[str, str], ...] = ()
diff --git a/mstar/engine/resources/step.py b/mstar/engine/resources/step.py
index eb4ee564d..cf7a23f25 100644
--- a/mstar/engine/resources/step.py
+++ b/mstar/engine/resources/step.py
@@ -12,6 +12,8 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
+from mstar.engine.resources.base import Segment
+from mstar.engine.resources.declare import StepDeclaration
 from mstar.engine.resources.kv_pool import KVCachePool
 
 if TYPE_CHECKING:
@@ -76,6 +78,82 @@ def plan(
             per_request_managers=[build_manager([rid]) for rid in request_ids],
         )
 
+    def drive(
+        self,
+        declaration: StepDeclaration,
+        cache_manager: "BatchedCacheManager",
+    ) -> None:
+        """Execute a declared step's admission and planning: forks first,
+        then each plan in declaration order. Every plan goes through the
+        step surface's own plan calls, so a declared step admits, plans,
+        and short-circuits pre-planned labels exactly as a facade-driven
+        model would."""
+        for from_label, to_label in declaration.pre_forks:
+            cache_manager.snapshot_all(from_label, to_label)
+        for plan in declaration.plans:
+            if len(plan.labels) > 1:
+                seq_lens = {
+                    label: list(plan.spans[label]) for label in plan.labels
+                }
+                cache_manager.plan_attention_batched_cfg(
+                    labels=list(plan.labels),
+                    seq_lens=seq_lens,
+                    is_causal=plan.is_causal,
+                    write_store=plan.write_store,
+                    combined_label=plan.combined_key,
+                    dense_gen=plan.dense_gen,
+                )
+                if plan.rope:
+                    cache_manager.plan_rope_batched_cfg(
+                        labels=list(plan.labels),
+                        seq_lens=seq_lens,
+                        per_label_pos_ids=plan.rope_pos_ids,
+                        combined_label=plan.combined_key,
+                    )
+            else:
+                label = plan.labels[0]
+                spans = list(plan.spans[label])
+                cache_manager.plan_attention(
+                    seq_lens=spans,
+                    is_causal=plan.is_causal,
+                    write_store=plan.write_store,
+                    label=label,
+                    dense_gen=plan.dense_gen,
+                )
+                if plan.rope:
+                    cache_manager.plan_rope(
+                        seq_lens=spans,
+                        pos_ids=plan.rope_pos_ids,
+                        label=label,
+                    )
+
+    def commit(
+        self,
+        declaration: StepDeclaration,
+        cache_manager: "BatchedCacheManager",
+    ) -> None:
+        """Commit a declared step once its forward is launched: advance
+        each committing plan's streams by their planned spans (and any
+        declared position advances) straight on the pool, then apply the
+        declared post-forward forks."""
+        pool = cache_manager.kv_pool
+        request_ids = cache_manager.request_ids
+        for plan in declaration.plans:
+            if not plan.commit:
+                continue
+            for label in plan.labels:
+                spans = plan.spans[label]
+                for i, rid in enumerate(request_ids):
+                    pool.commit(
+                        Segment(rid, label, spans[i]),
+                        pos_advance=(
+                            None if plan.pos_advance is None
+                            else plan.pos_advance[i]
+                        ),
+                    )
+        for from_label, to_label in declaration.post_forks:
+            cache_manager.snapshot_all(from_label, to_label)
+
     def publish(
         self,
         request_ids: list[str],
diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py
index dd9f74664..52811f0aa 100644
--- a/mstar/model/submodule_base.py
+++ b/mstar/model/submodule_base.py
@@ -18,6 +18,7 @@
 if TYPE_CHECKING:
     from mstar.engine.cuda_graph_config import CudaGraphConfig, PiecewiseCudaGraphConfig
     from mstar.engine.cuda_graph_runner import PiecewiseCudaGraphRunner
+    from mstar.engine.resources.declare import StepDeclaration
 
 
 @dataclass
@@ -280,6 +281,21 @@ def preprocess(
             **inputs[0].kwargs
         }
 
+    def declare_step(
+        self,
+        graph_walk: str,
+        engine_inputs: ModelInputsFromEngine,
+        inputs: list[NodeInputs],
+    ) -> "StepDeclaration | None":
+        """Declare this batch's step for the runner to drive: which cache
+        streams it touches, what spans they grow by, which plans back it,
+        which streams fork, and what commits when it lands. The runner
+        drives the declaration before ``preprocess`` and commits it after
+        the forward, so a declaring submodule keeps no plan or advance
+        calls of its own. None means the submodule still plans and
+        advances through the facade itself."""
+        return None
+
     @abstractmethod
     def forward(
         self,
diff --git a/test/modular/test_step_declaration.py b/test/modular/test_step_declaration.py
new file mode 100644
index 000000000..9e2903197
--- /dev/null
+++ b/test/modular/test_step_declaration.py
@@ -0,0 +1,270 @@
+"""Unit tests for step declarations and the runner driving them.
+
+A declaring submodule hands the runner a ``StepDeclaration``; the runner's
+``drive`` must make exactly the plan calls a facade-driven model would, in
+declaration order, and ``commit`` must advance exactly the declared spans
+on the pool (with declared position advances), then apply post forks.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import sys
+import threading
+
+sys.path.insert(0, ".")
+
+import pytest
+
+from mstar.engine.kv_store import (
+    KVCacheConfig,
+    PageAllocator,
+    PagedAllocationManager,
+    StoreWritePolicy,
+)
+from mstar.engine.resources import KVCachePool, PlanSpec, StepDeclaration
+from mstar.engine.resources.step import StepRunner
+
+
+def _make_pool(max_num_pages: int = 16, page_size: int = 8):
+    manager = PagedAllocationManager.__new__(PagedAllocationManager)
+    manager.config = KVCacheConfig(
+        num_layers=1,
+        num_kv_heads=1,
+        head_dim=1,
+        max_seq_len=max_num_pages * page_size,
+        max_num_pages=max_num_pages,
+        page_size=page_size,
+    )
+    manager.page_allocator = PageAllocator(max_num_pages)
+    manager.request_states = {}
+    manager.kv_cache = None
+    manager.write_policy = StoreWritePolicy.ALWAYS
+    manager._kv_transfer_engine = None
+    manager._offload_stream = None
+    manager.pending_reads = {}
+    manager._lock = threading.RLock()
+    return KVCachePool(manager), manager
+
+
+class _RecordingManager:
+    """Step-surface double that records every call with its kwargs."""
+
+    def __init__(self, request_ids):
+        self.request_ids = request_ids
+        self.calls = []
+
+    def snapshot_all(self, from_label, to_label):
+        self.calls.append(("snapshot_all", from_label, to_label))
+
+    def plan_attention(self, **kwargs):
+        self.calls.append(("plan_attention", kwargs))
+
+    def plan_attention_batched_cfg(self, **kwargs):
+        self.calls.append(("plan_attention_batched_cfg", kwargs))
+
+    def plan_rope(self, **kwargs):
+        self.calls.append(("plan_rope", kwargs))
+
+    def plan_rope_batched_cfg(self, **kwargs):
+        self.calls.append(("plan_rope_batched_cfg", kwargs))
+
+
+class _PoolManager:
+    """Commit-side double: a real pool plus the step addressing commit
+    reads, recording fork calls so ordering is checkable."""
+
+    def __init__(self, pool, request_ids):
+        self.kv_pool = pool
+        self.request_ids = request_ids
+        self.forks = []
+
+    def snapshot_all(self, from_label, to_label):
+        self.forks.append((from_label, to_label))
+
+
+class TestDeclarationTypes:
+    def test_plan_spec_is_frozen(self):
+        plan = PlanSpec(labels=("main",), spans={"main": (4,)})
+        with pytest.raises(dataclasses.FrozenInstanceError):
+            plan.is_causal = False
+
+    def test_declaration_is_frozen(self):
+        decl = StepDeclaration(
+            plans=(PlanSpec(labels=("main",), spans={"main": (1,)}),)
+        )
+        with pytest.raises(dataclasses.FrozenInstanceError):
+            decl.plans = ()
+
+    def test_defaults(self):
+        plan = PlanSpec(labels=("main",), spans={"main": (2, 3)})
+        assert plan.is_causal and plan.write_store and plan.commit
+        assert not plan.dense_gen and not plan.rope
+        assert plan.pos_advance is None
+        decl = StepDeclaration(plans=(plan,))
+        assert decl.pre_forks == () and decl.post_forks == ()
+
+
+class TestDrive:
+    def test_plain_plan_with_rope(self):
+        cm = _RecordingManager(["r0", "r1"])
+        decl = StepDeclaration(plans=(
+            PlanSpec(
+                labels=("main",), spans={"main": (3, 1)},
+                is_causal=True, write_store=True, rope=True,
+            ),
+        ))
+        StepRunner().drive(decl, cm)
+        assert [c[0] for c in cm.calls] == ["plan_attention", "plan_rope"]
+        attn = cm.calls[0][1]
+        assert attn["seq_lens"] == [3, 1]
+        assert attn["label"] == "main"
+        assert attn["is_causal"] and attn["write_store"]
+        assert attn["dense_gen"] is False
+        rope = cm.calls[1][1]
+        assert rope["seq_lens"] == [3, 1]
+        assert rope["label"] == "main" and rope["pos_ids"] is None
+
+    def test_combined_plan_orders_forks_first(self):
+        cm = _RecordingManager(["r0"])
+        decl = StepDeclaration(
+            plans=(
+                PlanSpec(
+                    labels=("main", "uncond"),
+                    spans={"main": (5,), "uncond": (7,)},
+                    is_causal=False, write_store=False,
+                    dense_gen=True, rope=True,
+                    rope_pos_ids={"main": []},
+                ),
+            ),
+            pre_forks=(("main", "cfg_text"),),
+        )
+        StepRunner().drive(decl, cm)
+        assert [c[0] for c in cm.calls] == [
+            "snapshot_all", "plan_attention_batched_cfg", "plan_rope_batched_cfg",
+        ]
+        assert cm.calls[0][1:] == ("main", "cfg_text")
+        attn = cm.calls[1][1]
+        assert attn["labels"] == ["main", "uncond"]
+        assert attn["seq_lens"] == {"main": [5], "uncond": [7]}
+        assert not attn["is_causal"] and not attn["write_store"]
+        assert attn["dense_gen"] is True
+        assert attn["combined_label"] == "_cfg_batched"
+        rope = cm.calls[2][1]
+        assert rope["per_label_pos_ids"] == {"main": []}
+
+    def test_plans_run_in_declaration_order(self):
+        cm = _RecordingManager(["r0"])
+        decl = StepDeclaration(plans=(
+            PlanSpec(labels=("main",), spans={"main": (1,)}, rope=True),
+            PlanSpec(labels=("cfg_img",), spans={"cfg_img": (1,)}, rope=True),
+        ))
+        StepRunner().drive(decl, cm)
+        labels = [
+            c[1]["label"] for c in cm.calls
+            if c[0] in ("plan_attention", "plan_rope")
+        ]
+        assert labels == ["main", "main", "cfg_img", "cfg_img"]
+
+    def test_no_rope_plans_attention_only(self):
+        cm = _RecordingManager(["r0"])
+        decl = StepDeclaration(plans=(
+            PlanSpec(labels=("main",), spans={"main": (2,)}),
+        ))
+        StepRunner().drive(decl, cm)
+        assert [c[0] for c in cm.calls] == ["plan_attention"]
+
+
+class TestCommit:
+    def _pool_manager(self, request_ids, labels):
+        pool, manager = _make_pool()
+        for rid in request_ids:
+            pool.add_request(rid, labels)
+        return _PoolManager(pool, request_ids), manager
+
+    def test_commits_declared_spans(self):
+        cm, manager = self._pool_manager(["r0", "r1"], ["main"])
+        decl = StepDeclaration(plans=(
+            PlanSpec(labels=("main",), spans={"main": (3, 1)}),
+        ))
+        StepRunner().commit(decl, cm)
+        assert manager.get_state("r0", "main").seq_len == 3
+        assert manager.get_state("r0", "main").position_id_start == 3
+        assert manager.get_state("r1", "main").seq_len == 1
+        assert manager.get_state("r1", "main").position_id_start == 1
+
+    def test_pos_advance_overrides_span(self):
+        cm, manager = self._pool_manager(["r0"], ["main"])
+        decl = StepDeclaration(plans=(
+            PlanSpec(
+                labels=("main",), spans={"main": (10,)}, pos_advance=(14,),
+            ),
+        ))
+        StepRunner().commit(decl, cm)
+        state = manager.get_state("r0", "main")
+        assert state.seq_len == 10
+        assert state.position_id_start == 14
+
+    def test_non_committing_plan_is_skipped(self):
+        cm, manager = self._pool_manager(["r0"], ["main", "uncond"])
+        decl = StepDeclaration(plans=(
+            PlanSpec(
+                labels=("main", "uncond"),
+                spans={"main": (6,), "uncond": (6,)},
+                commit=False,
+            ),
+        ))
+        StepRunner().commit(decl, cm)
+        assert manager.get_state("r0", "main").seq_len == 0
+        assert manager.get_state("r0", "uncond").seq_len == 0
+
+    def test_combined_plan_commits_every_label(self):
+        cm, manager = self._pool_manager(["r0", "r1"], ["main", "uncond"])
+        decl = StepDeclaration(plans=(
+            PlanSpec(
+                labels=("main", "uncond"),
+                spans={"main": (4, 2), "uncond": (5, 3)},
+            ),
+        ))
+        StepRunner().commit(decl, cm)
+        assert manager.get_state("r0", "main").seq_len == 4
+        assert manager.get_state("r1", "main").seq_len == 2
+        assert manager.get_state("r0", "uncond").seq_len == 5
+        assert manager.get_state("r1", "uncond").seq_len == 3
+
+    def test_zero_span_is_a_no_op(self):
+        cm, manager = self._pool_manager(["r0"], ["main"])
+        decl = StepDeclaration(plans=(
+            PlanSpec(labels=("main",), spans={"main": (0,)}),
+        ))
+        StepRunner().commit(decl, cm)
+        state = manager.get_state("r0", "main")
+        assert state.seq_len == 0 and state.position_id_start == 0
+
+    def test_post_forks_apply_after_commits(self):
+        cm, manager = self._pool_manager(["r0"], ["main", "cfg_text"])
+        committed_at_fork = []
+        original = cm.snapshot_all
+
+        def snapshot_all(from_label, to_label):
+            committed_at_fork.append(manager.get_state("r0", "main").seq_len)
+            original(from_label, to_label)
+
+        cm.snapshot_all = snapshot_all
+        decl = StepDeclaration(
+            plans=(PlanSpec(labels=("main",), spans={"main": (8,)}),),
+            post_forks=(("main", "cfg_text"),),
+        )
+        StepRunner().commit(decl, cm)
+        assert cm.forks == [("main", "cfg_text")]
+        assert committed_at_fork == [8]
+
+
+class TestDeclareStepDefault:
+    def test_base_submodule_declares_nothing(self):
+        from mstar.model.submodule_base import NodeSubmodule
+
+        class _Stub(NodeSubmodule):
+            __abstractmethods__ = frozenset()
+
+        assert _Stub().declare_step("walk", None, []) is None

From 989051402a6d108e554956329c2e424bcf159f1f Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Mon, 10 Aug 2026 00:13:31 +0000
Subject: [PATCH 14/20] Adopt step declarations in cosmos3

---
 mstar/engine/resources/declare.py             |  21 +-
 mstar/engine/resources/step.py                |   2 +-
 mstar/model/cosmos3/components/transformer.py |  87 +++--
 mstar/model/cosmos3/submodules.py             | 305 ++++++++----------
 mstar/model/cosmos3/tests/test_action.py      | 107 +++---
 .../model/cosmos3/tests/test_engine_cache.py  |  80 +++--
 mstar/model/cosmos3/tests/test_sound.py       |   5 +-
 test/modular/test_step_declaration.py         |  15 +-
 8 files changed, 317 insertions(+), 305 deletions(-)

diff --git a/mstar/engine/resources/declare.py b/mstar/engine/resources/declare.py
index ef1a5defa..29ef6568c 100644
--- a/mstar/engine/resources/declare.py
+++ b/mstar/engine/resources/declare.py
@@ -16,15 +16,17 @@
 class PlanSpec:
     """One attention (and optional rope) plan of a declared step.
 
-    ``labels`` with a single entry plans that label alone; more than one
-    entry plans a single combined batch across them, label-major, under
-    ``combined_key``. ``spans`` gives each label's per-request token
-    counts in batch order. ``rope_pos_ids`` optionally supplies explicit
-    position ids (a tensor for a single label, a per-label list dict for
-    a combined plan); None lets the embedder build them from the pool's
-    counters. ``commit`` says whether the step's spans become part of the
-    streams' history; ``pos_advance`` overrides the per-request position
-    advance when it differs from the span.
+    A plain plan (``combined`` False) covers exactly one label. A combined
+    plan batches every (label, request) pair label-major into one plan
+    under ``combined_key``; it may carry a single label (a batched forward
+    whose guidance branch is off still runs against the combined key).
+    ``spans`` gives each label's per-request token counts in batch order.
+    ``rope_pos_ids`` optionally supplies explicit position ids (a tensor
+    for a plain plan, a per-label list dict for a combined one); None lets
+    the embedder build them from the pool's counters. ``commit`` says
+    whether the step's spans become part of the streams' history;
+    ``pos_advance`` overrides the per-request position advance when it
+    differs from the span.
     """
     labels: tuple[str, ...]
     spans: dict[str, tuple[int, ...]]
@@ -33,6 +35,7 @@ class PlanSpec:
     dense_gen: bool = False
     rope: bool = False
     rope_pos_ids: "dict[str, list[torch.Tensor]] | torch.Tensor | None" = None
+    combined: bool = False
     combined_key: str = "_cfg_batched"
     commit: bool = True
     pos_advance: tuple[int, ...] | None = None
diff --git a/mstar/engine/resources/step.py b/mstar/engine/resources/step.py
index cf7a23f25..97c151608 100644
--- a/mstar/engine/resources/step.py
+++ b/mstar/engine/resources/step.py
@@ -91,7 +91,7 @@ def drive(
         for from_label, to_label in declaration.pre_forks:
             cache_manager.snapshot_all(from_label, to_label)
         for plan in declaration.plans:
-            if len(plan.labels) > 1:
+            if plan.combined:
                 seq_lens = {
                     label: list(plan.spans[label]) for label in plan.labels
                 }
diff --git a/mstar/model/cosmos3/components/transformer.py b/mstar/model/cosmos3/components/transformer.py
index 3acded7bf..a75edaf03 100644
--- a/mstar/model/cosmos3/components/transformer.py
+++ b/mstar/model/cosmos3/components/transformer.py
@@ -23,6 +23,7 @@
 
 from __future__ import annotations
 
+import functools
 import math
 
 import torch
@@ -264,7 +265,10 @@ def forward(
     # handle's attention plan, not here.
     # ------------------------------------------------------------------
 
-    def forward_und(self, und_seq: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, cache_handle) -> torch.Tensor:
+    def forward_und(
+        self, und_seq: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+        cache_handle, layer_idx: int, label: str,
+    ) -> torch.Tensor:
         H, Hkv, D = self.num_attention_heads, self.num_key_value_heads, self.head_dim
         q = self.norm_q(self.to_q(und_seq).view(-1, H, D))
         k = self.norm_k(self.to_k(und_seq).view(-1, Hkv, D))
@@ -279,15 +283,20 @@ def forward_und(self, und_seq: torch.Tensor, cos: torch.Tensor, sin: torch.Tenso
             q = sp_head_slice(self.sp_group, q)
             k = sp_head_slice(self.sp_group, k)
             v = sp_head_slice(self.sp_group, v)
-            out = cache_handle.run_attention(q=q, k=k, v=v)
+            out = cache_handle.run_attention(
+                q=q, k=k, v=v, layer_idx=layer_idx, label=label,
+            )
             out = sp_head_gather(self.sp_group, out).reshape(-1, H * D)
         else:
-            out = cache_handle.run_attention(q=q, k=k, v=v).reshape(-1, H * D)
+            out = cache_handle.run_attention(
+                q=q, k=k, v=v, layer_idx=layer_idx, label=label,
+            ).reshape(-1, H * D)
         return self.to_out(out)
 
     def forward_gen(
         self, gen_seq: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
-        cache_handle, seq_sizes: list[int] | None = None,
+        cache_handle, layer_idx: int, label: str,
+        seq_sizes: list[int] | None = None,
         prefer_all_gather: bool = False,
     ) -> torch.Tensor:
         H, Hkv, D = self.num_attention_heads, self.num_key_value_heads, self.head_dim
@@ -303,7 +312,11 @@ def forward_gen(
         # denoise forward sets prefer_all_gather (the all-to-all does not replay
         # from a CUDA graph; all-gather does).
         out = ulysses_attention(
-            self.sp_group, q, k, v, cache_handle.run_attention, seq_sizes,
+            self.sp_group, q, k, v,
+            functools.partial(
+                cache_handle.run_attention, layer_idx=layer_idx, label=label,
+            ),
+            seq_sizes,
             prefer_all_gather=prefer_all_gather,
         ).reshape(-1, H * D)
         return self.to_add_out(out)
@@ -363,20 +376,27 @@ def forward(
 
         return residual_und + mlp_out_und, residual_gen + mlp_out_gen
 
-    def forward_und(self, und_seq: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, cache_handle) -> torch.Tensor:
+    def forward_und(
+        self, und_seq: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+        cache_handle, layer_idx: int, label: str,
+    ) -> torch.Tensor:
         und_norm = self.input_layernorm(und_seq)
-        attn_out = self.self_attn.forward_und(und_norm, cos, sin, cache_handle)
+        attn_out = self.self_attn.forward_und(
+            und_norm, cos, sin, cache_handle, layer_idx, label
+        )
         residual = und_seq + attn_out
         return residual + self.mlp(self.post_attention_layernorm(residual))
 
     def forward_gen(
         self, gen_seq: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
-        cache_handle, seq_sizes: list[int] | None = None,
+        cache_handle, layer_idx: int, label: str,
+        seq_sizes: list[int] | None = None,
         prefer_all_gather: bool = False,
     ) -> torch.Tensor:
         gen_norm = self.input_layernorm_moe_gen(gen_seq)
         attn_out = self.self_attn.forward_gen(
-            gen_norm, cos, sin, cache_handle, seq_sizes, prefer_all_gather
+            gen_norm, cos, sin, cache_handle, layer_idx, label,
+            seq_sizes, prefer_all_gather
         )
         residual = gen_seq + attn_out
         return residual + self.mlp_moe_gen(self.post_attention_layernorm_moe_gen(residual))
@@ -803,10 +823,11 @@ def forward(
     # K/V; the generation tower then runs per denoising step, re-reading that
     # frozen K/V. Because the text tokens never receive a timestep embedding,
     # their K/V is step-independent, so caching it once is exact. ``cache_handle``
-    # is a paged attention handle (set_layer_idx / run_attention / advance_seq_lens);
-    # the attention plan (causal vs not, which label) is configured by the caller.
-    # These entry points pack text + vision, plus an optional action or sound
-    # band appended to the generation block (matching the reference ``forward``'s
+    # is the step's attention surface; every call names its layer and plan key
+    # (``label``), and the plan itself (causal vs not, which pages) was driven
+    # from the step declaration before the forward. These entry points pack
+    # text + vision, plus an optional action or sound band appended to the
+    # generation block (matching the reference ``forward``'s
     # ``[vision | action | sound]`` order).
     # ------------------------------------------------------------------
 
@@ -816,19 +837,18 @@ def _rotary(self, position_ids: torch.Tensor, device, dtype):
         return cos.squeeze(0), sin.squeeze(0)
 
     def prefill_und(
-        self, input_ids: torch.Tensor, position_ids: torch.Tensor, cache_handle
+        self, input_ids: torch.Tensor, position_ids: torch.Tensor, cache_handle,
+        label: str,
     ) -> None:
-        """Run the understanding tower over the text prefix, writing per-layer K/V
-        to the cache under the active label and committing the prefix length.
-        ``position_ids`` are the text segment's 3D mRoPE ids ([3, und_len])."""
+        """Run the understanding tower over the text prefix, writing per-layer
+        K/V under ``label`` (the plan key the step declared). ``position_ids``
+        are the text segment's 3D mRoPE ids ([3, und_len])."""
         und_seq = self.embed_tokens(input_ids)
         cos, sin = self._rotary(position_ids, und_seq.device, und_seq.dtype)
         for i, layer in enumerate(self.layers):
-            cache_handle.set_layer_idx(i)
-            und_seq = layer.forward_und(und_seq, cos, sin, cache_handle)
-        cache_handle.advance_seq_lens()
+            und_seq = layer.forward_und(und_seq, cos, sin, cache_handle, i, label)
 
-    def _sp_run_gen_layers(self, gen_seq, cos, sin, cache_handle, prefer_all_gather=False):
+    def _sp_run_gen_layers(self, gen_seq, cos, sin, cache_handle, label, prefer_all_gather=False):
         """Run the generation layer stack, sequence-parallel-sharded across the
         SP group when active. ``gen_seq``/``cos``/``sin`` are the FULL sequence
         (identical on every SP rank); returns the FULL post-layer sequence.
@@ -844,17 +864,15 @@ def _sp_run_gen_layers(self, gen_seq, cos, sin, cache_handle, prefer_all_gather=
         sp = self.sp_group
         if sp.world_size == 1:
             for i, layer in enumerate(self.layers):
-                cache_handle.set_layer_idx(i)
-                gen_seq = layer.forward_gen(gen_seq, cos, sin, cache_handle)
+                gen_seq = layer.forward_gen(gen_seq, cos, sin, cache_handle, i, label)
             return gen_seq
         seq_sizes = sp_seq_split(gen_seq.shape[0], sp.world_size)
         gen_seq = scatter_sequence(sp, gen_seq, seq_sizes)
         cos = scatter_sequence(sp, cos, seq_sizes)
         sin = scatter_sequence(sp, sin, seq_sizes)
         for i, layer in enumerate(self.layers):
-            cache_handle.set_layer_idx(i)
             gen_seq = layer.forward_gen(
-                gen_seq, cos, sin, cache_handle, seq_sizes, prefer_all_gather
+                gen_seq, cos, sin, cache_handle, i, label, seq_sizes, prefer_all_gather
             )
         return gather_sequence(sp, gen_seq, seq_sizes)
 
@@ -867,6 +885,7 @@ def denoise_step(
         vision_noisy_frame_indexes: list[torch.Tensor],
         vision_mse_loss_indexes: torch.Tensor,
         cache_handle,
+        label: str,
         action_latents: torch.Tensor | None = None,
         action_token_shapes: list[tuple[int, int, int]] | None = None,
         action_noisy_frame_indexes: list[torch.Tensor] | None = None,
@@ -883,7 +902,7 @@ def denoise_step(
 
         Patchifies ``latents`` ([1, C, T, H, W]), scatter-adds the timestep
         embedding to the noisy tokens, runs the generation layers (each reading
-        the active label's cached understanding K/V plus its own freshly written
+        ``label``'s cached understanding K/V plus its own freshly written
         K/V), and decodes the flow velocity. ``position_ids`` are the generation
         segment's 3D mRoPE ids ([3, num_gen]) — the vision band, then the action
         or sound band when present. ``vision_mse_loss_indexes`` /
@@ -919,7 +938,7 @@ def denoise_step(
             gen_seq = torch.cat([gen_seq, sound_seq], dim=0)
 
         cos, sin = self._rotary(position_ids, gen_seq.device, gen_seq.dtype)
-        gen_seq = self._sp_run_gen_layers(gen_seq, cos, sin, cache_handle)
+        gen_seq = self._sp_run_gen_layers(gen_seq, cos, sin, cache_handle, label)
         gen_out = self.norm_moe_gen(gen_seq)
         preds_packed = self.proj_out(gen_out[vision_mse_loss_indexes])
         preds = self._unpatchify_and_unpack_latents(
@@ -950,6 +969,7 @@ def denoise_step_batched_cfg(
         vision_noisy_frame_indexes: list[torch.Tensor],
         vision_mse_loss_indexes: torch.Tensor,
         cache_handle,
+        label: str,
         action_latents: torch.Tensor | None = None,
         action_token_shapes: list[tuple[int, int, int]] | None = None,
         action_noisy_frame_indexes: list[torch.Tensor] | None = None,
@@ -1010,7 +1030,8 @@ def denoise_step_batched_cfg(
         sin = torch.cat([sin_c, sin_u], dim=0)
 
         gen_seq = self._sp_run_gen_layers(
-            gen_seq, cos, sin, cache_handle, prefer_all_gather=prefer_all_gather
+            gen_seq, cos, sin, cache_handle, label,
+            prefer_all_gather=prefer_all_gather,
         )
         gen_out = self.norm_moe_gen(gen_seq)
 
@@ -1036,7 +1057,7 @@ def _decode(out):
 
         return _decode(gen_out[:n]), _decode(gen_out[n:])
 
-    def denoise_step_batched(self, requests: list[dict], cache_handle):
+    def denoise_step_batched(self, requests: list[dict], cache_handle, label: str):
         """Denoise one step for several requests at once (image / video).
 
         Each request carries its own latents, timestep, rotary positions (which
@@ -1087,7 +1108,7 @@ def denoise_step_batched(self, requests: list[dict], cache_handle):
         all_gen = torch.cat(gen_seqs + gen_seqs, dim=0)
         cos = torch.cat(cos_cond + cos_uncond, dim=0)
         sin = torch.cat(sin_cond + sin_uncond, dim=0)
-        all_gen = self._sp_run_gen_layers(all_gen, cos, sin, cache_handle)
+        all_gen = self._sp_run_gen_layers(all_gen, cos, sin, cache_handle, label)
         gen_out = self.norm_moe_gen(all_gen)
 
         sizes = [g.shape[0] for g in gen_seqs]
@@ -1119,7 +1140,9 @@ def _decode(out, req, original_latent_shapes):
             results.append((cond_v, uncond_v))
         return results
 
-    def denoise_step_action_batched(self, requests: list[dict], cache_handle, with_cfg: bool):
+    def denoise_step_action_batched(
+        self, requests: list[dict], cache_handle, label: str, with_cfg: bool,
+    ):
         """Joint ``[video | action]`` denoise for several action requests at once.
 
         The action analogue of ``denoise_step_batched``. Each request carries its
@@ -1183,7 +1206,7 @@ def denoise_step_action_batched(self, requests: list[dict], cache_handle, with_c
             cos = torch.cat(cos_cond, dim=0)
             sin = torch.cat(sin_cond, dim=0)
 
-        all_gen = self._sp_run_gen_layers(all_gen, cos, sin, cache_handle)
+        all_gen = self._sp_run_gen_layers(all_gen, cos, sin, cache_handle, label)
         gen_out = self.norm_moe_gen(all_gen)
 
         sizes = [g.shape[0] for g in gen_seqs]
diff --git a/mstar/model/cosmos3/submodules.py b/mstar/model/cosmos3/submodules.py
index 4a2b61257..f07d8edd9 100644
--- a/mstar/model/cosmos3/submodules.py
+++ b/mstar/model/cosmos3/submodules.py
@@ -43,6 +43,7 @@
     build_static_inputs,
     vision_condition_frame_indexes,
 )
+from mstar.engine.resources import PlanSpec, StepDeclaration
 from mstar.model.cosmos3.constants import (
     ACTION_GEN_WALK,
     ACTION_VIDEO_GEN_WALK,
@@ -99,8 +100,8 @@
 COND_LABEL = "main"
 UNCOND_LABEL = "uncond"
 
-# Combined label for the single FlashInfer plan that runs both guidance branches
-# in one forward (see cache_manager.plan_attention_batched_cfg).
+# Combined plan key for the single FlashInfer plan that runs both guidance
+# branches in one forward (a combined PlanSpec in the step declaration).
 CFG_BATCHED_LABEL = "_cfg_batched"
 
 
@@ -600,59 +601,110 @@ def _prepare_action_gen(self, fwd_info, inputs, device) -> ARNodeInputs:
         )
 
     # ------------------------------------------------------------------
-    # preprocess: plan paged attention for the labels this walk touches.
+    # declare_step: what each walk's step is. The runner drives the plans
+    # and commits; preprocess below only marshals data.
     # ------------------------------------------------------------------
 
-    def _plan_gen(self, cm, st, num_gen: int, cfg_active: bool = True) -> None:
-        """Plan a denoise step's non-causal attention: one batched plan covering
-        both guidance branches when they run together, else a plan per label.
-        ``cfg_active`` False (a guidance_interval out-of-interval step, or
-        gs==1) plans the conditional branch alone — matching the cond-only
-        forward — so an interval step costs no wasted uncond/batched plan."""
-        if st["uncond"] is None or not cfg_active:
-            cm.plan_attention(
-                seq_lens=[num_gen], is_causal=False, label=COND_LABEL,
-                write_store=False, dense_gen=True,
-            )
-        elif self.batched_cfg:
-            cm.plan_attention_batched_cfg(
-                labels=[COND_LABEL, UNCOND_LABEL], seq_lens=[num_gen],
-                is_causal=False, write_store=False, dense_gen=True,
+    def declare_step(
+        self, graph_walk, engine_inputs: ModelInputsFromEngine, inputs,
+    ) -> StepDeclaration:
+        cm = engine_inputs.cache_manager
+
+        if graph_walk in PREFILL_WALKS:
+            # The text prefix is written once and committed; both guidance
+            # branches prefill together as one combined plan when present.
+            if inputs[0].kwargs.get("cfg"):
+                labels = (COND_LABEL, UNCOND_LABEL)
+                return StepDeclaration(plans=(PlanSpec(
+                    labels=labels,
+                    spans={
+                        label: tuple(inp.kwargs["seq_lens"][label] for inp in inputs)
+                        for label in labels
+                    },
+                    is_causal=True,
+                    write_store=False,
+                    combined=True,
+                ),))
+            return StepDeclaration(plans=(PlanSpec(
+                labels=(COND_LABEL,),
+                spans={COND_LABEL: tuple(inp.input_seq_len for inp in inputs)},
+                is_causal=True,
+                write_store=False,
+            ),))
+
+        if (
+            graph_walk not in GEN_WALKS
+            and graph_walk not in SOUND_WALKS
+            and graph_walk not in ACTION_WALKS
+        ):
+            raise ValueError(f"Unknown Cosmos3 DiT graph walk: {graph_walk!r}")
+
+        # Denoise steps read the frozen prefix and overwrite the same
+        # generation span every step, so nothing commits.
+        spans = tuple(inp.input_seq_len for inp in inputs)
+
+        def gen_plan(labels, combined):
+            return PlanSpec(
+                labels=labels,
+                spans={label: spans for label in labels},
+                is_causal=False,
+                write_store=False,
+                dense_gen=True,
+                combined=combined,
+                commit=False,
             )
-        else:
-            cm.plan_attention(
-                seq_lens=[num_gen], is_causal=False, label=COND_LABEL,
-                write_store=False, dense_gen=True,
+
+        if cm is not None and cm.is_captured:
+            # The captured denoise graph has a fixed shape and always runs
+            # both guidance branches; there is no per-request state here.
+            return StepDeclaration(plans=(
+                gen_plan((COND_LABEL, UNCOND_LABEL), combined=True),
+            ))
+
+        states = self._states(engine_inputs)
+        sts = [states[rid] for rid in engine_inputs.request_ids]
+        if len(sts) > 1:
+            # Cross-request batch: one combined plan over every request's
+            # guidance branches (a single branch when guidance is off).
+            labels = (
+                (COND_LABEL, UNCOND_LABEL)
+                if sts[0]["uncond"] is not None else (COND_LABEL,)
             )
-            cm.plan_attention(
-                seq_lens=[num_gen], is_causal=False, label=UNCOND_LABEL,
-                write_store=False, dense_gen=True,
+            return StepDeclaration(plans=(gen_plan(labels, combined=True),))
+
+        st = sts[0]
+        cfg_active = True
+        if graph_walk in GEN_WALKS or graph_walk in SOUND_WALKS:
+            step_index = int(
+                inputs[0].tensor_inputs["time_index"].reshape(-1)[0].item()
             )
+            cfg_active = self._cfg_active(st, step_index)
+        if st["uncond"] is None or not cfg_active:
+            # Guidance off, or a guidance_interval out-of-interval step: the
+            # conditional branch runs alone, so nothing else plans.
+            return StepDeclaration(plans=(gen_plan((COND_LABEL,), combined=False),))
+        if self.batched_cfg:
+            return StepDeclaration(plans=(
+                gen_plan((COND_LABEL, UNCOND_LABEL), combined=True),
+            ))
+        return StepDeclaration(plans=(
+            gen_plan((COND_LABEL,), combined=False),
+            gen_plan((UNCOND_LABEL,), combined=False),
+        ))
 
-    def _preprocess_image_gen_captured(self, cm, inputs) -> dict:
-        """Plan a denoise step for the CUDA-graph path.
+    def _preprocess_image_gen_captured(self, inputs) -> dict:
+        """Pack a denoise step's inputs for the CUDA-graph path.
 
-        Runs with synthetic request ids (no per-request state), so it derives the
-        token count from ``input_seq_len``. Both guidance branches are planned as
-        one combined attention (``plan_attention_batched_cfg``) so the captured
-        forward runs a single transformer pass over both — one weight load instead
-        of two. The static-input tensors (latents, timestep, rotary positions) are
+        Runs with synthetic request ids (no per-request state). The
+        static-input tensors (latents, timestep, rotary positions) are
         stacked on a leading batch dim, so one captured graph spans a whole
-        concurrent batch (a batch of one for the single-request latency path); the
-        replay side copies each request's tensors into these fixed buffers.
-
-        Separate from the eager ``preprocess`` by contract, not just output
-        formation: this path takes pre-derived timestep/rotary buffers with no
-        per-request scheduler state and must always plan both guidance branches
-        (the graph shape is fixed), while the eager path derives its step inputs
-        from ``time_index`` + per-request state and skips the uncond plan on
+        concurrent batch (a batch of one for the single-request latency
+        path); the replay side copies each request's tensors into these
+        fixed buffers. The attention plan comes from the step declaration,
+        which always covers both guidance branches here (the graph shape is
+        fixed), while the eager declaration skips the uncond plan on
         guidance-interval steps.
         """
-        seq_lens = [inp.input_seq_len for inp in inputs]
-        cm.plan_attention_batched_cfg(
-            labels=[COND_LABEL, UNCOND_LABEL], seq_lens=seq_lens,
-            is_causal=False, write_store=False, dense_gen=True,
-        )
         return {
             "latents": torch.stack([inp.tensor_inputs["latents"] for inp in inputs]),
             "vision_timesteps": torch.stack([inp.tensor_inputs["vision_timesteps"] for inp in inputs]),
@@ -664,30 +716,17 @@ def preprocess(
         self, graph_walk, engine_inputs: ModelInputsFromEngine,
         inputs: list[ARNodeInputs]
     ) -> dict:
-        cm = engine_inputs.cache_manager
-
-        if graph_walk == IMAGE_GEN_WALK and getattr(cm, "_cuda_graph_mode", False):
-            return self._preprocess_image_gen_captured(cm, inputs)
+        if (
+            graph_walk == IMAGE_GEN_WALK
+            and engine_inputs.cache_manager.is_captured
+        ):
+            return self._preprocess_image_gen_captured(inputs)
 
         if graph_walk in PREFILL_WALKS:
             has_cfg = inputs[0].kwargs.get("cfg")
             labels = [COND_LABEL, UNCOND_LABEL] if has_cfg else [COND_LABEL]
-            if has_cfg:
-                cm.plan_attention_batched_cfg(
-                    labels=labels,
-                    seq_lens={
-                        key: [inp.kwargs["seq_lens"][key] for inp in inputs] \
-                            for key in labels
-                    }, is_causal=True, write_store=False
-                )
-            else:
-                cm.plan_attention(
-                    seq_lens=[inp.input_seq_len for inp in inputs],
-                    is_causal=True, label=COND_LABEL,
-                    write_store=False
-                )
-            # Pack label-major (all cond segments, then all uncond) to match the
-            # (label, request) batch order of the plan above.
+            # Pack label-major (all cond segments, then all uncond) to match
+            # the (label, request) batch order of the declared plan.
             return {
                 "cfg": has_cfg,
                 "input_ids": torch.concat([
@@ -700,47 +739,20 @@ def preprocess(
                 ], dim=1),
             }
 
-        states = self._states(engine_inputs)
-        st = states[engine_inputs.request_ids[0]]
-
+        rids = engine_inputs.request_ids
         if graph_walk in GEN_WALKS:
-            rids = engine_inputs.request_ids
             if len(rids) > 1:
-                # Cross-request batch: one batched plan over every request's two
-                # guidance branches, each with its own page set and token count.
-                cm.plan_attention_batched_cfg(
-                    labels=[COND_LABEL, UNCOND_LABEL],
-                    seq_lens=[states[r]["cond"]["num_vision_tokens"] for r in rids],
-                    is_causal=False, write_store=False, dense_gen=True,
-                )
                 return {
                     "latents": {r: inp.tensor_inputs["latents"] for r, inp in zip(rids, inputs, strict=True)},
                     "time_index": {r: inp.tensor_inputs["time_index"] for r, inp in zip(rids, inputs, strict=True)},
                 }
-            ti = inputs[0].tensor_inputs["time_index"]
-            step_index = int(ti.reshape(-1)[0].item())
-            self._plan_gen(
-                cm, st, st["cond"]["num_vision_tokens"], cfg_active=self._cfg_active(st, step_index)
-            )
             return {
                 "latents": inputs[0].tensor_inputs["latents"],
-                "time_index": ti,
+                "time_index": inputs[0].tensor_inputs["time_index"],
             }
 
         if graph_walk in SOUND_WALKS:
-            rids = engine_inputs.request_ids
             if len(rids) > 1:
-                # Cross-request batch: one batched plan over every request's
-                # [vision | sound] block (sound requests batch only with sound
-                # requests — batches are per graph walk).
-                cm.plan_attention_batched_cfg(
-                    labels=[COND_LABEL, UNCOND_LABEL],
-                    seq_lens=[
-                        states[r]["cond"]["num_vision_tokens"] + states[r]["num_sound"]
-                        for r in rids
-                    ],
-                    is_causal=False, write_store=False, dense_gen=True,
-                )
                 return {
                     "latents": {r: inp.tensor_inputs["latents"] for r, inp in zip(rids, inputs, strict=True)},
                     "sound_latents": {
@@ -748,38 +760,14 @@ def preprocess(
                     },
                     "time_index": {r: inp.tensor_inputs["time_index"] for r, inp in zip(rids, inputs, strict=True)},
                 }
-            ti = inputs[0].tensor_inputs["time_index"]
-            step_index = int(ti.reshape(-1)[0].item())
-            self._plan_gen(
-                cm, st, st["cond"]["num_vision_tokens"] + st["num_sound"],
-                cfg_active=self._cfg_active(st, step_index),
-            )
             return {
                 "latents": inputs[0].tensor_inputs["latents"],
                 "sound_latents": inputs[0].tensor_inputs["sound_latents"],
-                "time_index": ti,
+                "time_index": inputs[0].tensor_inputs["time_index"],
             }
 
         if graph_walk in ACTION_WALKS:
-            rids = engine_inputs.request_ids
             if len(rids) > 1:
-                # Cross-request batch: one batched plan over every request's joint
-                # [video | action] block, each with its own page set and token
-                # count. A single label when guidance is off (the common
-                # guidance-scale-1 case), both labels with classifier-free
-                # guidance.
-                sts = [states[r] for r in rids]
-                labels = (
-                    [COND_LABEL, UNCOND_LABEL] if sts[0]["uncond"] is not None else [COND_LABEL]
-                )
-                cm.plan_attention_batched_cfg(
-                    labels=labels,
-                    seq_lens=[
-                        s["cond"]["num_vision_tokens"] + s["cond"]["num_action_tokens"]
-                        for s in sts
-                    ],
-                    is_causal=False, write_store=False, dense_gen=True,
-                )
                 return {
                     "latents": {r: inp.tensor_inputs["latents"] for r, inp in zip(rids, inputs, strict=True)},
                     "action_latents": {
@@ -787,7 +775,6 @@ def preprocess(
                     },
                     "time_index": {r: inp.tensor_inputs["time_index"] for r, inp in zip(rids, inputs, strict=True)},
                 }
-            self._plan_gen(cm, st, st["cond"]["num_vision_tokens"] + st["cond"]["num_action_tokens"])
             return {
                 "latents": inputs[0].tensor_inputs["latents"],
                 "action_latents": inputs[0].tensor_inputs["action_latents"],
@@ -832,15 +819,11 @@ def _forward_prefill(
         cfg=False,
         **kwargs
     ) -> dict:
-        if cfg:
-            cm.set_active_label(CFG_BATCHED_LABEL)
-        else:
-            cm.set_active_label(COND_LABEL)
-
-        self.transformer.prefill_und(input_ids, text_mrope_ids, cm)
+        label = CFG_BATCHED_LABEL if cfg else COND_LABEL
+        self.transformer.prefill_und(input_ids, text_mrope_ids, cm, label)
         return {}
 
-    def _denoise(self, cm, static, latents, vision_timesteps):
+    def _denoise(self, cm, static, latents, vision_timesteps, label):
         return self.transformer.denoise_step(
             latents,
             vision_timesteps,
@@ -849,6 +832,7 @@ def _denoise(self, cm, static, latents, vision_timesteps):
             static["vision_noisy_frame_indexes"],
             static["mse_gen_indexes"],
             cm,
+            label,
         )
 
     def _cfg_active(self, st, step_index: int) -> bool:
@@ -886,10 +870,8 @@ def _forward_image_gen(self, cm, st, latents, time_index, **kwargs) -> dict:
         cfg_active = self._cfg_active(st, step_index)
 
         if not cfg_active:
-            cm.set_active_label(COND_LABEL)
-            velocity = self._denoise(cm, st["cond"], latents, vision_timesteps)
+            velocity = self._denoise(cm, st["cond"], latents, vision_timesteps, COND_LABEL)
         elif self.batched_cfg:
-            cm.set_active_label(CFG_BATCHED_LABEL)
             cond_v, uncond_v = self.transformer.denoise_step_batched_cfg(
                 latents,
                 vision_timesteps,
@@ -899,13 +881,12 @@ def _forward_image_gen(self, cm, st, latents, time_index, **kwargs) -> dict:
                 st["cond"]["vision_noisy_frame_indexes"],
                 st["cond"]["mse_gen_indexes"],
                 cm,
+                CFG_BATCHED_LABEL,
             )
             velocity = uncond_v + st["gs"] * (cond_v - uncond_v)
         else:
-            cm.set_active_label(COND_LABEL)
-            cond_v = self._denoise(cm, st["cond"], latents, vision_timesteps)
-            cm.set_active_label(UNCOND_LABEL)
-            uncond_v = self._denoise(cm, st["uncond"], latents, vision_timesteps)
+            cond_v = self._denoise(cm, st["cond"], latents, vision_timesteps, COND_LABEL)
+            uncond_v = self._denoise(cm, st["uncond"], latents, vision_timesteps, UNCOND_LABEL)
             velocity = uncond_v + st["gs"] * (cond_v - uncond_v)
 
         new_latents = scheduler.step(
@@ -958,39 +939,35 @@ def _forward_video_sound_gen(self, cm, st, latents, sound_latents, time_index, *
         cfg_active = self._cfg_active(st, step_index)
 
         if not cfg_active:
-            cm.set_active_label(COND_LABEL)
             velocity, sound_v = self.transformer.denoise_step(
                 latents, vts, st["cond"]["position_ids"][:, st["cond"]["und_len"]:],
                 st["cond"]["vision_token_shapes"], st["cond"]["vision_noisy_frame_indexes"],
-                st["cond"]["mse_gen_indexes"], cm,
+                st["cond"]["mse_gen_indexes"], cm, COND_LABEL,
                 **self._sound_kwargs(st["cond"], sound_latents, sts),
             )
         elif self.batched_cfg:
-            cm.set_active_label(CFG_BATCHED_LABEL)
             (cond_v, s_c), (uncond_v, s_u) = self.transformer.denoise_step_batched_cfg(
                 latents, vts,
                 st["cond"]["position_ids"][:, st["cond"]["und_len"]:],
                 st["uncond"]["position_ids"][:, st["uncond"]["und_len"]:],
                 st["cond"]["vision_token_shapes"],
                 st["cond"]["vision_noisy_frame_indexes"],
-                st["cond"]["mse_gen_indexes"], cm,
+                st["cond"]["mse_gen_indexes"], cm, CFG_BATCHED_LABEL,
                 **self._sound_kwargs(st["cond"], sound_latents, sts),
             )
             velocity = uncond_v + st["gs"] * (cond_v - uncond_v)
             sound_v = s_u + st["gs"] * (s_c - s_u)
         else:
-            cm.set_active_label(COND_LABEL)
             cond_v, s_c = self.transformer.denoise_step(
                 latents, vts, st["cond"]["position_ids"][:, st["cond"]["und_len"]:],
                 st["cond"]["vision_token_shapes"], st["cond"]["vision_noisy_frame_indexes"],
-                st["cond"]["mse_gen_indexes"], cm,
+                st["cond"]["mse_gen_indexes"], cm, COND_LABEL,
                 **self._sound_kwargs(st["cond"], sound_latents, sts),
             )
-            cm.set_active_label(UNCOND_LABEL)
             uncond_v, s_u = self.transformer.denoise_step(
                 latents, vts, st["uncond"]["position_ids"][:, st["uncond"]["und_len"]:],
                 st["uncond"]["vision_token_shapes"], st["uncond"]["vision_noisy_frame_indexes"],
-                st["uncond"]["mse_gen_indexes"], cm,
+                st["uncond"]["mse_gen_indexes"], cm, UNCOND_LABEL,
                 **self._sound_kwargs(st["uncond"], sound_latents, sts),
             )
             velocity = uncond_v + st["gs"] * (cond_v - uncond_v)
@@ -1010,7 +987,7 @@ def _forward_video_sound_gen(self, cm, st, latents, sound_latents, time_index, *
             "time_index": [time_index + 1],
         }
 
-    def _denoise_action(self, cm, static, latents, action_latents, vts, ats, domain):
+    def _denoise_action(self, cm, static, latents, action_latents, vts, ats, domain, label):
         und_len = static["und_len"]
         return self.transformer.denoise_step(
             latents,
@@ -1020,6 +997,7 @@ def _denoise_action(self, cm, static, latents, action_latents, vts, ats, domain)
             static["vision_noisy_frame_indexes"],
             static["mse_gen_indexes"],
             cm,
+            label,
             action_latents=action_latents,
             action_token_shapes=static["action_token_shapes"],
             action_noisy_frame_indexes=static["action_noisy_frame_indexes"],
@@ -1070,10 +1048,10 @@ def _forward_action_gen(self, cm, st, latents, action_latents, time_index, **kwa
         domain = st["domain_t"]
 
         if st["uncond"] is None:
-            cm.set_active_label(COND_LABEL)
-            video_v, action_v = self._denoise_action(cm, st["cond"], latents, action_latents, vts, ats, domain)
+            video_v, action_v = self._denoise_action(
+                cm, st["cond"], latents, action_latents, vts, ats, domain, COND_LABEL
+            )
         elif self.batched_cfg:
-            cm.set_active_label(CFG_BATCHED_LABEL)
             (video_v, action_v), (v_u, a_u) = self.transformer.denoise_step_batched_cfg(
                 latents,
                 vts,
@@ -1083,6 +1061,7 @@ def _forward_action_gen(self, cm, st, latents, action_latents, time_index, **kwa
                 st["cond"]["vision_noisy_frame_indexes"],
                 st["cond"]["mse_gen_indexes"],
                 cm,
+                CFG_BATCHED_LABEL,
                 action_latents=action_latents,
                 action_token_shapes=st["cond"]["action_token_shapes"],
                 action_noisy_frame_indexes=st["cond"]["action_noisy_frame_indexes"],
@@ -1093,10 +1072,12 @@ def _forward_action_gen(self, cm, st, latents, action_latents, time_index, **kwa
             video_v = v_u + st["gs"] * (video_v - v_u)
             action_v = a_u + st["gs"] * (action_v - a_u)
         else:
-            cm.set_active_label(COND_LABEL)
-            video_v, action_v = self._denoise_action(cm, st["cond"], latents, action_latents, vts, ats, domain)
-            cm.set_active_label(UNCOND_LABEL)
-            v_u, a_u = self._denoise_action(cm, st["uncond"], latents, action_latents, vts, ats, domain)
+            video_v, action_v = self._denoise_action(
+                cm, st["cond"], latents, action_latents, vts, ats, domain, COND_LABEL
+            )
+            v_u, a_u = self._denoise_action(
+                cm, st["uncond"], latents, action_latents, vts, ats, domain, UNCOND_LABEL
+            )
             video_v = v_u + st["gs"] * (video_v - v_u)
             action_v = a_u + st["gs"] * (action_v - a_u)
 
@@ -1157,12 +1138,8 @@ def forward_batched(
     ):
         cm = engine_inputs.cache_manager
         if graph_walk in PREFILL_WALKS:
-            if kwargs.get("cfg"):
-                cm.set_active_label(CFG_BATCHED_LABEL)
-            else:
-                cm.set_active_label(COND_LABEL)
-
-            self.transformer.prefill_und(input_ids, text_mrope_ids, cm)
+            label = CFG_BATCHED_LABEL if kwargs.get("cfg") else COND_LABEL
+            self.transformer.prefill_und(input_ids, text_mrope_ids, cm, label)
             return {}
         if graph_walk in ACTION_WALKS:
             return self._forward_batched_action(engine_inputs, latents, action_latents, time_index)
@@ -1170,7 +1147,6 @@ def forward_batched(
             return self._forward_batched_sound(engine_inputs, latents, sound_latents, time_index)
         if graph_walk not in GEN_WALKS:
             raise ValueError(f"Cosmos3 batched forward only supports generation walks, got {graph_walk!r}")
-        cm.set_active_label(CFG_BATCHED_LABEL)
         states = self._states(engine_inputs)
         reqs, meta = [], []
         for rid in engine_inputs.request_ids:
@@ -1194,7 +1170,7 @@ def forward_batched(
             })
             meta.append((rid, st, lat, ti, t))
 
-        results = self.transformer.denoise_step_batched(reqs, cm)
+        results = self.transformer.denoise_step_batched(reqs, cm, CFG_BATCHED_LABEL)
 
         out = {}
         for (rid, st, lat, ti, t), (cond_v, uncond_v) in zip(meta, results, strict=True):
@@ -1215,7 +1191,6 @@ def _forward_batched_sound(self, engine_inputs, latents, sound_latents, time_ind
         one batched transformer pass, then per request the guidance combine and
         its own joint [video | sound] scheduler step."""
         cm = engine_inputs.cache_manager
-        cm.set_active_label(CFG_BATCHED_LABEL)
         states = self._states(engine_inputs)
         reqs, meta = [], []
         for rid in engine_inputs.request_ids:
@@ -1245,7 +1220,7 @@ def _forward_batched_sound(self, engine_inputs, latents, sound_latents, time_ind
             })
             meta.append((rid, st, lat, snd, ti, t))
 
-        results = self.transformer.denoise_step_batched(reqs, cm)
+        results = self.transformer.denoise_step_batched(reqs, cm, CFG_BATCHED_LABEL)
 
         out = {}
         for (rid, st, lat, snd, ti, t), ((cond_v, s_c), (uncond_v, s_u)) in zip(meta, results, strict=True):
@@ -1270,7 +1245,6 @@ def _forward_batched_action(self, engine_inputs, latents, action_latents, time_i
         run one batched transformer pass, then per request combine the guidance
         branches (when present) and apply its own joint scheduler step."""
         cm = engine_inputs.cache_manager
-        cm.set_active_label(CFG_BATCHED_LABEL)
         states = self._states(engine_inputs)
         rids = engine_inputs.request_ids
         with_cfg = states[rids[0]]["uncond"] is not None
@@ -1307,7 +1281,9 @@ def _forward_batched_action(self, engine_inputs, latents, action_latents, time_i
             reqs.append(req)
             meta.append((rid, st, lat, act, ti, t))
 
-        results = self.transformer.denoise_step_action_batched(reqs, cm, with_cfg)
+        results = self.transformer.denoise_step_action_batched(
+            reqs, cm, CFG_BATCHED_LABEL, with_cfg
+        )
 
         out = {}
         for (rid, st, lat, act, ti, t), branches in zip(meta, results, strict=True):
@@ -1516,7 +1492,6 @@ def forward_captured(
         (the same compute as the eager cross-request forward), one transformer pass
         over the whole batch."""
         cm = engine_inputs.cache_manager
-        cm.set_active_label(CFG_BATCHED_LABEL)
         layout = self._capture_layout[tuple(latents.shape[1:])]
         rids = engine_inputs.request_ids
         if latents.shape[0] == 1:
@@ -1527,7 +1502,7 @@ def forward_captured(
             cond_v, uncond_v = self.transformer.denoise_step_batched_cfg(
                 latents[0], vision_timesteps[0], position_ids_cond[0], position_ids_uncond[0],
                 layout["vision_token_shapes"], layout["vision_noisy_frame_indexes"],
-                layout["mse_gen_indexes"], cm, prefer_all_gather=True,
+                layout["mse_gen_indexes"], cm, CFG_BATCHED_LABEL, prefer_all_gather=True,
             )
             return {rids[0]: {"cond_v": [cond_v], "uncond_v": [uncond_v]}}
         reqs = [
@@ -1542,7 +1517,7 @@ def forward_captured(
             }
             for i in range(latents.shape[0])
         ]
-        results = self.transformer.denoise_step_batched(reqs, cm)
+        results = self.transformer.denoise_step_batched(reqs, cm, CFG_BATCHED_LABEL)
         return {
             rid: {"cond_v": [cond_v], "uncond_v": [uncond_v]}
             for rid, (cond_v, uncond_v) in zip(rids, results, strict=True)
diff --git a/mstar/model/cosmos3/tests/test_action.py b/mstar/model/cosmos3/tests/test_action.py
index 04a422d32..3d5878789 100644
--- a/mstar/model/cosmos3/tests/test_action.py
+++ b/mstar/model/cosmos3/tests/test_action.py
@@ -29,6 +29,7 @@
 )
 from mstar.model.cosmos3.components.transformer import Cosmos3OmniTransformer
 from mstar.model.cosmos3.config import Cosmos3Config
+from mstar.model.cosmos3.tests.test_engine_cache import _forward_step
 
 
 # --- verbatim vllm-omni references (transformer_cosmos3.py / action.py) ------
@@ -73,37 +74,34 @@ def _cfg() -> Cosmos3Config:
 
 
 class _SdpaCache:
-    """In-process cache-once handle (stored K/V + sdpa), the BatchedCacheManager
-    surface the DiT uses. Prefill stashes the understanding K/V; the denoise step
-    re-reads it. Also models the batched-CFG plan: under the combined label the
-    packed sequence is split into one block per batched label, each routed to its
-    own committed prefix (so a single-label batch of one request equals the plain
-    single-request path)."""
-
-    def __init__(self):
-        self.active, self.layer = "main", 0
+    """In-process cache-once handle (stored K/V + sdpa) with the step-surface
+    calls the DiT uses. Prefill stashes the understanding K/V; the denoise
+    step re-reads it. Declared plans land on ``plan_attention*``; the
+    runner's commit lands on the pool surface (``kv_pool.commit``), which
+    promotes the pending K/V to committed. Also models the batched-CFG
+    plan: under the combined key the packed sequence is split into one
+    block per batched label, each routed to its own committed prefix (so a
+    single-label batch of one request equals the plain single-request
+    path)."""
+
+    is_captured = False
+
+    def __init__(self, request_ids=("r0",)):
+        self.request_ids = list(request_ids)
+        self.kv_pool = self
         self.committed, self.pending, self.is_causal = {}, {}, {}
         self.batched_labels = None
 
-    def set_active_label(self, label):
-        self.active = label
-
-    def set_layer_idx(self, i):
-        self.layer = i
-
-    def plan(self, is_causal):
-        self.is_causal[self.active] = is_causal
-
-    # Engine-facing surface (used when the DiT submodule drives the cache).
     def plan_attention(self, seq_lens=None, dtype=None, is_causal=True, write_store=True, label=None, **kwargs):
-        self.is_causal[label or self.active] = is_causal
+        self.is_causal[label] = is_causal
 
     def plan_attention_batched_cfg(self, labels, seq_lens, is_causal=False, write_store=False, **kwargs):
         self.batched_labels = list(labels)
         self.is_causal["_cfg_batched"] = is_causal
 
-    def plan_rope(self, *a, **k):
-        pass
+    def commit(self, segment=None, pos_advance=None):
+        self.committed.update(self.pending)
+        self.pending = {}
 
     @staticmethod
     def _sdpa(q, k, v, c):
@@ -120,21 +118,16 @@ def _attend_label(self, label, layer, q, k, v, causal):
         self.pending[key] = (k, v)
         return self._sdpa(q, k, v, causal)
 
-    def run_attention(self, q, k, v, layer_idx=None):
-        layer = self.layer if layer_idx is None else layer_idx
-        if self.active == "_cfg_batched":
+    def run_attention(self, q, k, v, layer_idx=None, label=None):
+        if label == "_cfg_batched":
             causal = self.is_causal["_cfg_batched"]
             n = q.shape[0] // len(self.batched_labels)
             outs = []
-            for bi, label in enumerate(self.batched_labels):
+            for bi, blabel in enumerate(self.batched_labels):
                 sl = slice(bi * n, (bi + 1) * n)
-                outs.append(self._attend_label(label, layer, q[sl], k[sl], v[sl], causal))
+                outs.append(self._attend_label(blabel, layer_idx, q[sl], k[sl], v[sl], causal))
             return torch.cat(outs, 0)
-        return self._attend_label(self.active, layer, q, k, v, self.is_causal[self.active])
-
-    def advance_seq_lens(self, pos_id_ns=None):
-        self.committed.update(self.pending)
-        self.pending = {}
+        return self._attend_label(label, layer_idx, q, k, v, self.is_causal[label])
 
 
 _MODES = ("inverse_dynamics", "forward_dynamics", "policy")
@@ -234,18 +227,15 @@ def test_action_denoise_step_matches_fused() -> None:
         )
         cache = _SdpaCache()
         und_len = s["und_len"]
-        cache.set_active_label("main")
-        cache.plan(is_causal=True)
-        _cos, _sin = model._rotary(
-            s["text_mrope_ids"], s["input_ids"].device, model.embed_tokens.weight.dtype
-        )
-        model.prefill_und(s["input_ids"], _cos, _sin, cache)
-        cache.plan(is_causal=False)
+        cache.plan_attention(label="main", is_causal=True)
+        model.prefill_und(s["input_ids"], s["text_mrope_ids"], cache, "main")
+        cache.commit()
+        cache.plan_attention(label="main", is_causal=False)
         with torch.no_grad():
             dv, da = model.denoise_step(
                 latents, vts, s["position_ids"][:, und_len:],
                 s["vision_token_shapes"], s["vision_noisy_frame_indexes"],
-                s["vision_mse_loss_indexes"] - und_len, cache,
+                s["vision_mse_loss_indexes"] - und_len, cache, "main",
                 action_latents=action_lat, action_token_shapes=s["action_token_shapes"],
                 action_noisy_frame_indexes=s["action_noisy_frame_indexes"],
                 action_mse_gen_indexes=s["action_mse_loss_indexes"] - und_len,
@@ -273,18 +263,15 @@ def test_action_batched_one_matches_single() -> None:
         und_len = s["und_len"]
         # Reference: the single-request joint denoise step.
         cache = _SdpaCache()
-        cache.set_active_label("main")
-        cache.plan(is_causal=True)
-        _cos, _sin = model._rotary(
-            s["text_mrope_ids"], s["input_ids"].device, model.embed_tokens.weight.dtype
-        )
-        model.prefill_und(s["input_ids"], _cos, _sin, cache)
-        cache.plan(is_causal=False)
+        cache.plan_attention(label="main", is_causal=True)
+        model.prefill_und(s["input_ids"], s["text_mrope_ids"], cache, "main")
+        cache.commit()
+        cache.plan_attention(label="main", is_causal=False)
         with torch.no_grad():
             dv, da = model.denoise_step(
                 latents, vts, s["position_ids"][:, und_len:],
                 s["vision_token_shapes"], s["vision_noisy_frame_indexes"],
-                s["vision_mse_loss_indexes"] - und_len, cache,
+                s["vision_mse_loss_indexes"] - und_len, cache, "main",
                 action_latents=action_lat, action_token_shapes=s["action_token_shapes"],
                 action_noisy_frame_indexes=s["action_noisy_frame_indexes"],
                 action_mse_gen_indexes=s["action_mse_loss_indexes"] - und_len,
@@ -292,18 +279,14 @@ def test_action_batched_one_matches_single() -> None:
             )
         # Batched path with one request and no guidance (single-label batch).
         cache2 = _SdpaCache()
-        cache2.set_active_label("main")
-        cache2.plan(is_causal=True)
-        _cos, _sin = model._rotary(
-            s["text_mrope_ids"], s["input_ids"].device, model.embed_tokens.weight.dtype
-        )
-        model.prefill_und(s["input_ids"], _cos, _sin, cache2)
+        cache2.plan_attention(label="main", is_causal=True)
+        model.prefill_und(s["input_ids"], s["text_mrope_ids"], cache2, "main")
+        cache2.commit()
         cache2.plan_attention_batched_cfg(
             labels=["main"],
             seq_lens=[s["num_vision_tokens"] + s["num_action_tokens"]],
             is_causal=False,
         )
-        cache2.set_active_label("_cfg_batched")
         req = {
             "latents": latents, "action_latents": action_lat,
             "vision_timesteps": vts, "action_timesteps": ats,
@@ -317,7 +300,9 @@ def test_action_batched_one_matches_single() -> None:
             "action_domain_id": domain,
         }
         with torch.no_grad():
-            ((bv, ba),), = model.denoise_step_action_batched([req], cache2, with_cfg=False)
+            ((bv, ba),), = model.denoise_step_action_batched(
+                [req], cache2, "_cfg_batched", with_cfg=False
+            )
         assert (dv - bv).abs().max().item() < 1e-5, mode
         assert (da - ba).abs().max().item() < 1e-5, mode
 
@@ -430,14 +415,14 @@ def test_action_engine_matches_fused() -> None:
     cm = _SdpaCache()
     ei = ModelInputsFromEngine(request_ids=[rid], per_request_info={rid: fwd}, cache_manager=cm)
     ni = dit.prepare_inputs("prefill", fwd, {"text_inputs": [torch.tensor(cond_ids, dtype=torch.long, device=device)]})
-    dit.forward("prefill", ei, **dit.preprocess("prefill", ei, [ni]))
+    _forward_step(dit, "prefill", ei, [ni])
     fwd.graph_walk = "action_gen"
     latents, action_latents = cond_latent.clone(), a_noise.clone()
     time_index = torch.zeros(1, dtype=torch.long, device=device)
     for _ in range(steps):
         ni = dit.prepare_inputs("action_gen", fwd, {
             "latents": [latents], "action_latents": [action_latents], "time_index": [time_index]})
-        out = dit.forward("action_gen", ei, **dit.preprocess("action_gen", ei, [ni]))
+        out = _forward_step(dit, "action_gen", ei, [ni])
         latents, action_latents, time_index = out["latents"][0], out["action_latents"][0], out["time_index"][0]
     dit.cleanup_request(rid)
     # The loop emits the full action latents (self-edge); trim to the raw action
@@ -616,7 +601,7 @@ def _prefill(rid, idx, cm):
         ni = dit.prepare_inputs("prefill", fwd, {
             "text_inputs": [torch.tensor(conds[idx], dtype=torch.long, device=device)],
         })
-        dit.forward("prefill", ei, **dit.preprocess("prefill", ei, [ni]))
+        _forward_step(dit, "prefill", ei, [ni])
         fwd.graph_walk = "action_gen"
         if rid not in cond_lat:
             cond_lat[rid] = _encode(rid)
@@ -632,7 +617,7 @@ def _run_one(rid, idx):
             inp = ({"cond_latents": [cond_lat[rid]]} if lat is None
                    else {"latents": [lat], "action_latents": [act], "time_index": [ti]})
             ni = dit.prepare_inputs("action_gen", fwd, inp)
-            out = dit.forward("action_gen", ei, **dit.preprocess("action_gen", ei, [ni]))
+            out = _forward_step(dit, "action_gen", ei, [ni])
             lat, act, ti = out["latents"][0], out["action_latents"][0], out["time_index"][0]
         dit.cleanup_request(rid)
         return act[:, :, :raw].float().cpu()
@@ -653,7 +638,7 @@ def _run_batched():
                 inp = {"cond_latents": [cond_lat[rid]]} if lat[rid] is None else {
                     "latents": [lat[rid]], "action_latents": [act[rid]], "time_index": [ti[rid]]}
                 inputs.append(dit.prepare_inputs("action_gen", fwds[rid], inp))
-            out = dit.forward_batched("action_gen", eiN, **dit.preprocess("action_gen", eiN, inputs))
+            out = _forward_step(dit, "action_gen", eiN, inputs, batched=True)
             for rid in rids:
                 o = out[rid]
                 lat[rid], act[rid], ti[rid] = o["latents"][0], o["action_latents"][0], o["time_index"][0]
diff --git a/mstar/model/cosmos3/tests/test_engine_cache.py b/mstar/model/cosmos3/tests/test_engine_cache.py
index bc2c67739..1584bca04 100644
--- a/mstar/model/cosmos3/tests/test_engine_cache.py
+++ b/mstar/model/cosmos3/tests/test_engine_cache.py
@@ -41,33 +41,32 @@
 
 
 class _SdpaCacheHandle:
-    """In-process reference cache with the ``BatchedCacheManager`` surface the
-    DiT uses, backed by stored tensors + sdpa (same kernel as the fused pipeline).
-    Prefill stashes each layer's understanding K/V; every denoise step re-reads it.
-
-    Also models the batched classifier-free-guidance plan: when both guidance
-    branches run in one forward, ``run_attention`` receives the two branches
-    concatenated and routes each half to its own label's cached prefix, so the
-    batched result equals running the branches sequentially.
+    """In-process reference cache with the step-surface calls the DiT uses,
+    backed by stored tensors + sdpa (same kernel as the fused pipeline).
+    Prefill stashes each layer's understanding K/V; every denoise step
+    re-reads it. Declared plans land on ``plan_attention*``; the runner's
+    commit lands on the pool surface (``kv_pool.commit``), which promotes
+    the prefill's pending K/V to committed.
+
+    Also models the batched classifier-free-guidance plan: under the
+    combined key ``run_attention`` receives the branches concatenated and
+    routes each block to its own label's cached prefix, so the batched
+    result equals running the branches sequentially.
     """
 
-    def __init__(self):
-        self.active = "main"
-        self.layer = 0
+    is_captured = False
+
+    def __init__(self, request_ids=("r0",)):
+        self.request_ids = list(request_ids)
+        self.kv_pool = self
         self.committed: dict[tuple[str, int], tuple[torch.Tensor, torch.Tensor]] = {}
         self.pending: dict[tuple[str, int], tuple[torch.Tensor, torch.Tensor]] = {}
         self.is_causal: dict[str, bool] = {}
         self.batched_labels: list[str] | None = None
         self.batched_splits: list[int] | None = None
 
-    def set_active_label(self, label):
-        self.active = label
-
-    def set_layer_idx(self, i):
-        self.layer = i
-
     def plan_attention(self, seq_lens=None, dtype=None, is_causal=True, write_store=True, label=None, **kwargs):
-        self.is_causal[label or self.active] = is_causal
+        self.is_causal[label] = is_causal
 
     def plan_attention_batched_cfg(self, labels, seq_lens, is_causal=False, write_store=False, **kwargs):
         self.batched_labels = list(labels)
@@ -81,8 +80,9 @@ def plan_attention_batched_cfg(self, labels, seq_lens, is_causal=False, write_st
         else:
             self.batched_splits = [int(sum(seq_lens))] * len(labels)
 
-    def plan_rope(self, *args, **kwargs):
-        pass
+    def commit(self, segment=None, pos_advance=None):
+        self.committed.update(self.pending)
+        self.pending = {}
 
     @staticmethod
     def _sdpa(q, k, v, is_causal):
@@ -100,21 +100,33 @@ def _attend_label(self, label, layer, q, k, v, causal):
         self.pending[key] = (k, v)
         return self._sdpa(q, k, v, causal)
 
-    def run_attention(self, q, k, v, layer_idx=None):
-        layer = self.layer if layer_idx is None else layer_idx
-        if self.active == "_cfg_batched":
+    def run_attention(self, q, k, v, layer_idx=None, label=None):
+        if label == "_cfg_batched":
             causal = self.is_causal["_cfg_batched"]
             outs, off = [], 0
-            for bi, label in enumerate(self.batched_labels):
+            for bi, blabel in enumerate(self.batched_labels):
                 sl = slice(off, off + self.batched_splits[bi])
                 off += self.batched_splits[bi]
-                outs.append(self._attend_label(label, layer, q[sl], k[sl], v[sl], causal))
+                outs.append(self._attend_label(blabel, layer_idx, q[sl], k[sl], v[sl], causal))
             return torch.cat(outs, 0)
-        return self._attend_label(self.active, layer, q, k, v, self.is_causal[self.active])
+        return self._attend_label(label, layer_idx, q, k, v, self.is_causal[label])
 
-    def advance_seq_lens(self, pos_id_ns=None):
-        self.committed.update(self.pending)
-        self.pending = {}
+
+def _forward_step(dit, walk, ei, inputs, batched=False):
+    """Drive one step the way the engine does: declaration, plans, data
+    preprocess, forward, commit."""
+    from mstar.engine.resources.step import StepRunner
+
+    runner = StepRunner()
+    declaration = dit.declare_step(walk, ei, inputs)
+    if declaration is not None:
+        runner.drive(declaration, ei.cache_manager)
+    pre = dit.preprocess(walk, ei, inputs)
+    forward = dit.forward_batched if batched else dit.forward
+    out = forward(walk, ei, **pre)
+    if declaration is not None:
+        runner.commit(declaration, ei.cache_manager)
+    return out
 
 
 def _flashinfer_cache(model, rid, device, dtype, backend=None):
@@ -159,14 +171,14 @@ def _run_cache_once(model, dit, cm, init, cond_ids, uncond_ids, device, num_fram
         torch.tensor(uncond_ids, dtype=torch.long, device=device),
     ]
     ni = dit.prepare_inputs("prefill", fwd, {"text_inputs": text_inputs})
-    dit.forward("prefill", ei, **dit.preprocess("prefill", ei, [ni]))
+    _forward_step(dit, "prefill", ei, [ni])
 
     latents = init.clone()
     time_index = torch.zeros(1, dtype=torch.long, device=device)
     fwd.graph_walk = "image_gen"
     for _ in range(STEPS):
         ni = dit.prepare_inputs("image_gen", fwd, {"latents": [latents], "time_index": [time_index]})
-        out = dit.forward("image_gen", ei, **dit.preprocess("image_gen", ei, [ni]))
+        out = _forward_step(dit, "image_gen", ei, [ni])
         latents, time_index = out["latents"][0], out["time_index"][0]
     dit.cleanup_request(rid)
     return latents
@@ -226,7 +238,7 @@ def _run_batched(model, dit, shared, init, conds, unconds, device, rids):
         ti = [torch.tensor(conds[i], dtype=torch.long, device=device),
               torch.tensor(unconds[i], dtype=torch.long, device=device)]
         ni = dit.prepare_inputs("prefill", fwd, {"text_inputs": ti})
-        dit.forward("prefill", ei1, **dit.preprocess("prefill", ei1, [ni]))
+        _forward_step(dit, "prefill", ei1, [ni])
 
     cmN = _mk_cm(shared, rids)
     eiN = ModelInputsFromEngine(request_ids=rids, per_request_info=fwds, cache_manager=cmN)
@@ -240,7 +252,7 @@ def _run_batched(model, dit, shared, init, conds, unconds, device, rids):
                                {"latents": [latents[rid]], "time_index": [time_index[rid]]})
             for rid in rids
         ]
-        out = dit.forward_batched("image_gen", eiN, **dit.preprocess("image_gen", eiN, inputs))
+        out = _forward_step(dit, "image_gen", eiN, inputs, batched=True)
         for rid in rids:
             latents[rid], time_index[rid] = out[rid]["latents"][0], out[rid]["time_index"][0]
     for rid in rids:
@@ -645,7 +657,7 @@ def _run_cuda_graph_denoise(ctx):
     ti = [torch.tensor(ctx["cond"], dtype=torch.long, device=device),
           torch.tensor(ctx["uncond"], dtype=torch.long, device=device)]
     ni = dit.prepare_inputs("prefill", fwd, {"text_inputs": ti})
-    dit.forward("prefill", ei, **dit.preprocess("prefill", ei, [ni]))
+    _forward_step(dit, "prefill", ei, [ni])
 
     runner = CudaGraphRunner(
         submodule_name="dit", submodule=dit, kv_cache_config=shared["cfg"],
diff --git a/mstar/model/cosmos3/tests/test_sound.py b/mstar/model/cosmos3/tests/test_sound.py
index ffc38b8d8..cd46ce4bf 100644
--- a/mstar/model/cosmos3/tests/test_sound.py
+++ b/mstar/model/cosmos3/tests/test_sound.py
@@ -26,6 +26,7 @@
 from mstar.model.cosmos3.submodules import Cosmos3DiTSubmodule
 from mstar.model.cosmos3.tests.test_engine_cache import (
     _flashinfer_cache,
+    _forward_step,
     _load,
     _SdpaCacheHandle,
 )
@@ -252,7 +253,7 @@ def _run_cache_once_sound(model, dit, cm, init, sound_init, cond_ids, uncond_ids
         torch.tensor(uncond_ids, dtype=torch.long, device=device),
     ]
     ni = dit.prepare_inputs("prefill", fwd, {"text_inputs": text_inputs})
-    dit.forward("prefill", ei, **dit.preprocess("prefill", ei, [ni]))
+    _forward_step(dit, "prefill", ei, [ni])
 
     latents, sound_latents = init.clone(), sound_init.clone()
     time_index = torch.zeros(1, dtype=torch.long, device=device)
@@ -261,7 +262,7 @@ def _run_cache_once_sound(model, dit, cm, init, sound_init, cond_ids, uncond_ids
         ni = dit.prepare_inputs(VIDEO_SOUND_GEN_WALK, fwd, {
             "latents": [latents], "sound_latents": [sound_latents], "time_index": [time_index],
         })
-        out = dit.forward(VIDEO_SOUND_GEN_WALK, ei, **dit.preprocess(VIDEO_SOUND_GEN_WALK, ei, [ni]))
+        out = _forward_step(dit, VIDEO_SOUND_GEN_WALK, ei, [ni])
         latents = out["latents"][0]
         sound_latents = out["sound_latents"][0]
         time_index = out["time_index"][0]
diff --git a/test/modular/test_step_declaration.py b/test/modular/test_step_declaration.py
index 9e2903197..a16c1fdb7 100644
--- a/test/modular/test_step_declaration.py
+++ b/test/modular/test_step_declaration.py
@@ -99,7 +99,7 @@ def test_declaration_is_frozen(self):
     def test_defaults(self):
         plan = PlanSpec(labels=("main",), spans={"main": (2, 3)})
         assert plan.is_causal and plan.write_store and plan.commit
-        assert not plan.dense_gen and not plan.rope
+        assert not plan.dense_gen and not plan.rope and not plan.combined
         assert plan.pos_advance is None
         decl = StepDeclaration(plans=(plan,))
         assert decl.pre_forks == () and decl.post_forks == ()
@@ -135,6 +135,7 @@ def test_combined_plan_orders_forks_first(self):
                     is_causal=False, write_store=False,
                     dense_gen=True, rope=True,
                     rope_pos_ids={"main": []},
+                    combined=True,
                 ),
             ),
             pre_forks=(("main", "cfg_text"),),
@@ -174,6 +175,18 @@ def test_no_rope_plans_attention_only(self):
         StepRunner().drive(decl, cm)
         assert [c[0] for c in cm.calls] == ["plan_attention"]
 
+    def test_single_label_combined_plan_batches(self):
+        cm = _RecordingManager(["r0"])
+        decl = StepDeclaration(plans=(
+            PlanSpec(
+                labels=("main",), spans={"main": (6,)}, combined=True,
+                is_causal=False, write_store=False,
+            ),
+        ))
+        StepRunner().drive(decl, cm)
+        assert [c[0] for c in cm.calls] == ["plan_attention_batched_cfg"]
+        assert cm.calls[0][1]["labels"] == ["main"]
+
 
 class TestCommit:
     def _pool_manager(self, request_ids, labels):

From aba6041fb6d6a1699a0a265437eb482df685c242 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Mon, 10 Aug 2026 00:22:30 +0000
Subject: [PATCH 15/20] Adopt step declarations in qwen3_omni and drop the
 pos-advance side channel

---
 mstar/engine/cache_manager.py                 |  65 +--------
 mstar/engine/resources/attention.py           |  24 ++--
 .../model/qwen3_omni/components/attention.py  |  21 ++-
 mstar/model/qwen3_omni/components/talker.py   |  21 ++-
 mstar/model/qwen3_omni/components/thinker.py  |  37 ++---
 mstar/model/qwen3_omni/submodules.py          | 136 +++++++++---------
 test/integration/test_prefill_cuda_graph.py   |   7 +
 test/modular/test_position_embedder.py        |  15 +-
 8 files changed, 145 insertions(+), 181 deletions(-)

diff --git a/mstar/engine/cache_manager.py b/mstar/engine/cache_manager.py
index 3eca3d928..37a27469b 100644
--- a/mstar/engine/cache_manager.py
+++ b/mstar/engine/cache_manager.py
@@ -217,17 +217,9 @@ def set_layer_idx(self, layer_idx: int):
 
     @torch.compiler.disable
     def get_qo_indptr_buf(self, label: str = "main") -> torch.Tensor | None:
-        """Return the persistent qo_indptr static buffer for a CUDA-graph
-        prefill wrapper, or None if not in CUDA-graph mode / wrong wrapper.
-
-        Captured prefill paths read this to recover per-request token boundaries
-        from inside the captured region — plan_attention updates the buffer via
-        .copy_() outside the graph, so the address stays stable across replay.
-        """
-        ps = self._plan_states.get(label)
-        if ps is None or ps.wrapper is None:
-            return None
-        return getattr(ps.wrapper, "_qo_indptr_buf", None)
+        """The attention manager's ``qo_indptr_buf`` for ``label`` (see
+        ``FlashInferAttentionManager.qo_indptr_buf``)."""
+        return self.attention.qo_indptr_buf(label)
 
     @abstractmethod
     def plan_attention(
@@ -483,49 +475,13 @@ def advance_seq_len(self, n: int | None = None, pos_id_n: int | None = None) ->
                 pos_advance=pos_id_n if pos_id_n is not None else n,
             )
 
-    @torch.compiler.disable
-    def set_custom_pos_advance(
-        self, pos_advance: list[int] | None, label: str | None = None,
-    ) -> None:
-        """Stash a per-request position-id advance for the next
-        ``advance_seq_lens()`` call to consume.
-
-        Resolves ``label`` the same way ``plan_attention`` does: explicit
-        label wins; otherwise the (single) currently-active label is used.
-
-        Used by submodules whose forward advances ``position_id_start`` by
-        something other than ``seq_len`` (e.g. Qwen3-Omni's prefill_vision
-        passes the MRoPE 3D-grid span here). Auto-cleared by
-        ``advance_seq_lens`` after use, so it does not leak across calls.
-
-        Pass ``pos_advance=None`` to clear an earlier set explicitly.
-        """
-        effective_label = label
-        if effective_label is None:
-            labels = list(self.active_labels.values())
-            if not labels:
-                return
-            assert len(set(labels)) == 1, (
-                f"All active labels must be the same to omit ``label``, got {labels}"
-            )
-            effective_label = labels[0]
-        ps = self._plan_states.get(effective_label)
-        if ps is None:
-            return
-        ps.custom_pos_advance = (
-            list(pos_advance) if pos_advance is not None else None
-        )
-
     @torch.compiler.disable
     def advance_seq_lens(self, pos_id_ns: list[int] | int | None = None) -> None:
         """Advance seq_len for each request by different amounts.
 
-        When ``pos_id_ns`` is None, falls back to a per-label side-channel
-        (``_PlanState.custom_pos_advance``, set via
-        ``set_custom_pos_advance``) for walks whose position-id span
-        differs from seq_len (e.g. Qwen3-Omni prefill_vision). The
-        side-channel is auto-cleared after use so it doesn't leak across
-        calls.
+        ``pos_id_ns`` overrides the position-id advance for walks whose
+        position span differs from seq_len; None advances positions by
+        each request's planned seq_len.
         """
 
         if self._batched_cfg_info:
@@ -549,10 +505,7 @@ def advance_seq_lens(self, pos_id_ns: list[int] | int | None = None) -> None:
                     continue
                 n = ps.seq_lens[i]
                 if pos_id_ns is None:
-                    if ps.custom_pos_advance is not None:
-                        pos_advance = ps.custom_pos_advance[i]
-                    else:
-                        pos_advance = n
+                    pos_advance = n
                 elif isinstance(pos_id_ns, int):
                     pos_advance = pos_id_ns
                 else:
@@ -560,10 +513,6 @@ def advance_seq_lens(self, pos_id_ns: list[int] | int | None = None) -> None:
                 self.kv_pool.commit(
                     Segment(rid, label, n), pos_advance=pos_advance
                 )
-        # Clear the side-channel on every consumer so a stale value can't
-        # bleed into a subsequent walk.
-        for ps in self._plan_states.values():
-            ps.custom_pos_advance = None
 
     @torch.compiler.disable
     def snapshot_all(
diff --git a/mstar/engine/resources/attention.py b/mstar/engine/resources/attention.py
index b71b907e9..3c4731c71 100644
--- a/mstar/engine/resources/attention.py
+++ b/mstar/engine/resources/attention.py
@@ -91,24 +91,11 @@ class _PlanState:
     In CUDA graph mode, wrapper is a persistent FlashInferPrefillWrapper or
     FlashInferDecodeWrapper created once during capture. plan_attention()
     calls wrapper.plan() which updates static buffers via .copy_().
-
-    ``custom_pos_advance`` is a generic out-of-band channel for prefill
-    walks whose position-id span differs from the seq_len being prefilled
-    (e.g. Qwen3-Omni's ``prefill_vision``, where the 3D-grid MRoPE span is
-    larger than the number of tokens). The submodule writes a per-request
-    list here via ``BatchedCacheManager.set_custom_pos_advance``;
-    ``advance_seq_lens`` reads it when ``pos_id_ns`` is None and advances
-    ``position_id_start`` by these values instead of by ``seq_len``.
-    Auto-cleared by ``advance_seq_lens`` so it doesn't leak across calls.
-    The CUDA-graph runner's post-replay ``advance_seq_lens()`` call is what
-    actually consumes this — the model's inner ``advance_seq_lens(pos_id_ns=...)``
-    runs at capture time only and is not replayed.
     """
     wrapper: FlashInferPrefillWrapper | FlashInferDecodeWrapper | None = None
     pos_ids: torch.Tensor | None = None
     seq_lens: list[int] | None = None
     write_store: bool = True
-    custom_pos_advance: list[int] | None = None
     # Plan memo: fingerprint of the last wrapper.plan() inputs for this label;
     # when it matches, the re-plan is skipped. Only the cross-attention path
     # sets it today (its context pages are immutable after add_cross_attn_kv),
@@ -413,6 +400,17 @@ def plan_batched_cfg(
         ps.write_store = write_store
         ps.dense_gen = None
 
+    def qo_indptr_buf(self, label: str) -> torch.Tensor | None:
+        """The persistent qo_indptr static buffer of ``label``'s CUDA-graph
+        prefill wrapper, or None outside that mode. Captured prefill paths
+        read it to recover per-request token boundaries from inside the
+        captured region: ``plan`` updates the buffer via ``.copy_()``
+        outside the graph, so the address stays stable across replay."""
+        ps = self.states.get(label)
+        if ps is None or ps.wrapper is None:
+            return None
+        return getattr(ps.wrapper, "_qo_indptr_buf", None)
+
     def write_kv(self, k: torch.Tensor, v: torch.Tensor, layer_idx: int, label: str) -> None:
         """Write this step's K/V into the paged cache at the label's planned
         positions. Separate from ``run`` so strategies that attend without
diff --git a/mstar/model/qwen3_omni/components/attention.py b/mstar/model/qwen3_omni/components/attention.py
index e79bfd152..fa151277b 100644
--- a/mstar/model/qwen3_omni/components/attention.py
+++ b/mstar/model/qwen3_omni/components/attention.py
@@ -5,8 +5,9 @@
 is the 3D MRoPE path used by the Thinker — for ``use_mrope=True`` the
 RoPE call goes through ``apply_interleaved_mrope`` with externally
 provided ``cos_sin_3d`` instead of the cache handle. Talker uses
-standard 1D RoPE (``use_mrope=False``) and inherits the parent's
-``_apply_rope`` as-is.
+standard 1D RoPE (``use_mrope=False``) through the step surface's
+``apply_rope``. Every rope and attention call names its layer and plan
+key explicitly.
 
 Follows the same shape conventions as the shared attention:
   q: [tokens, num_heads, head_dim]
@@ -63,6 +64,8 @@ def forward(
         cache_handle: BatchedCacheManager,
         cos_sin_3d: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         mrope_section: Optional[list[int]] = None,
+        layer_idx: int | None = None,
+        label: str | None = None,
     ) -> torch.Tensor:
         num_tokens = hidden_states.shape[0]
         q, k, v = self._project_qkv(hidden_states)
@@ -73,8 +76,18 @@ def forward(
             cos, sin = cos_sin_3d
             q, k = apply_interleaved_mrope(q, k, cos, sin)
         else:
-            q, k = self._apply_rope(q, k, cache_handle)
+            q, k = cache_handle.apply_rope(
+                q, k,
+                rope_theta=self.rope_theta,
+                rope_scale=self.rope_scale,
+                low_freq_factor=self.rope_low_freq_factor,
+                high_freq_factor=self.rope_high_freq_factor,
+                old_context_len=self.rope_old_context_len,
+                label=label,
+            )
 
-        attn_output = cache_handle.run_attention(q=q, k=k, v=v)
+        attn_output = cache_handle.run_attention(
+            q=q, k=k, v=v, layer_idx=layer_idx, label=label,
+        )
         attn_output = attn_output.reshape(num_tokens, self.num_heads * self.head_dim)
         return self.o_proj(attn_output)
diff --git a/mstar/model/qwen3_omni/components/talker.py b/mstar/model/qwen3_omni/components/talker.py
index 5d8713242..ac68b716f 100644
--- a/mstar/model/qwen3_omni/components/talker.py
+++ b/mstar/model/qwen3_omni/components/talker.py
@@ -103,6 +103,8 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         cache_handle: BatchedCacheManager,
+        layer_idx: int | None = None,
+        label: str | None = None,
     ) -> torch.Tensor:
         # ---------- self-attention with pre-norm ----------
         residual = hidden_states
@@ -110,6 +112,8 @@ def forward(
         hidden_states = self.self_attn(
             hidden_states=hidden_states,
             cache_handle=cache_handle,
+            layer_idx=layer_idx,
+            label=label,
         )
         hidden_states = residual + hidden_states
 
@@ -156,16 +160,17 @@ def forward(
         self,
         input_embeds: torch.Tensor,
         cache_handle: BatchedCacheManager,
+        label: str = "main",
     ) -> torch.Tensor:
         hidden_states = input_embeds
         for layer_idx, decoder_layer in enumerate(self.layers):
-            cache_handle.set_layer_idx(layer_idx)
             hidden_states = decoder_layer(
-                hidden_states=hidden_states, cache_handle=cache_handle
+                hidden_states=hidden_states,
+                cache_handle=cache_handle,
+                layer_idx=layer_idx,
+                label=label,
             )
 
-        cache_handle.advance_seq_lens()
-
         hidden_states = self.norm(hidden_states)
         return hidden_states
 
@@ -230,6 +235,7 @@ def forward(
         self,
         input_embeds: torch.Tensor,
         cache_handle: BatchedCacheManager,
+        label: str = "main",
     ) -> torch.Tensor:
         """Run the Talker backbone and return the final hidden states.
 
@@ -239,12 +245,15 @@ def forward(
         Args:
             input_embeds: [total_tokens, hidden_size] -- pre-embedded input
                 (may combine codec embeddings and projected Thinker states).
-            cache_handle: ``BatchedCacheManager`` for paged KV attention.
+            cache_handle: the step surface for paged KV attention.
+            label: the plan key every layer runs against.
 
         Returns:
             hidden_states: [total_tokens, hidden_size] after final RMS norm.
         """
-        return self.model(input_embeds=input_embeds, cache_handle=cache_handle)
+        return self.model(
+            input_embeds=input_embeds, cache_handle=cache_handle, label=label,
+        )
 
 
 # ---------------------------------------------------------------------------
diff --git a/mstar/model/qwen3_omni/components/thinker.py b/mstar/model/qwen3_omni/components/thinker.py
index ca2d2ba31..e555699f8 100644
--- a/mstar/model/qwen3_omni/components/thinker.py
+++ b/mstar/model/qwen3_omni/components/thinker.py
@@ -104,13 +104,17 @@ def forward(
         cache_handle: BatchedCacheManager,
         cos_sin_3d: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         mrope_section: Optional[list[int]] = None,
+        layer_idx: int | None = None,
+        label: str | None = None,
     ) -> torch.Tensor:
         """
         Args:
             hidden_states: [tokens, hidden_size]
-            cache_handle: BatchedCacheManager with pre-planned attention.
+            cache_handle: step surface with pre-planned attention.
             cos_sin_3d: (cos, sin) for 3D MRoPE, each [tokens, head_dim].
             mrope_section: section sizes for interleaved 3D MRoPE.
+            layer_idx: this layer's index in the stack.
+            label: the plan key this layer's attention runs against.
 
         Returns:
             hidden_states: [tokens, hidden_size]
@@ -123,6 +127,8 @@ def forward(
             cache_handle=cache_handle,
             cos_sin_3d=cos_sin_3d,
             mrope_section=mrope_section,
+            layer_idx=layer_idx,
+            label=label,
         )
         hidden_states = residual + hidden_states
 
@@ -191,23 +197,21 @@ def forward(
         cache_handle: BatchedCacheManager,
         cos_sin_3d: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         mrope_section: Optional[list[int]] = None,
-        mrope_pos_advance: Optional[list[int]] = None,
         deepstack_visual_embeds: list[torch.Tensor] | None = None,
+        label: str = "main",
     ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
         """
         Args:
             input_embeds: [tokens, hidden_size] -- pre-embedded input
                 (token embeddings possibly merged with multimodal features).
-            cache_handle: BatchedCacheManager with pre-planned attention
-                and RoPE.
+            cache_handle: step surface with pre-planned attention and RoPE.
             cos_sin_3d: (cos, sin) for 3D MRoPE, each [tokens, head_dim].
             mrope_section: section sizes for interleaved 3D MRoPE,
                 e.g. [24, 20, 20].
-            mrope_pos_advance: optional per-request MRoPE position advance
-                for ``advance_seq_lens``.  Vision prefill passes an explicit
-                value because the 3D-grid position span is larger than the
-                number of tokens; text / audio / decode leave it None and
-                ``position_id_start`` advances by ``seq_len``.
+            label: the plan key every layer runs against. The stream's
+                stored length and position counter advance when the runner
+                commits the declared step (vision prefill declares its
+                MRoPE 3D-grid span there).
 
         Returns:
             hidden_states: [tokens, hidden_size] -- final normed hidden states
@@ -222,12 +226,13 @@ def forward(
         layer_n_hidden = None
 
         for layer_idx, decoder_layer in enumerate(self.model.layers):
-            cache_handle.set_layer_idx(layer_idx)
             hidden_states = decoder_layer(
                 hidden_states,
                 cache_handle=cache_handle,
                 cos_sin_3d=cos_sin_3d,
                 mrope_section=mrope_section,
+                layer_idx=layer_idx,
+                label=label,
             )
 
             # add visual features to the hidden states of first several layers
@@ -241,18 +246,6 @@ def forward(
             if layer_idx == self.accept_hidden_layer:
                 layer_n_hidden = hidden_states.clone()
 
-        # Advance sequence lengths after all layers.  ``pos_id_ns`` decouples
-        # the position-id advance from the seq-len advance (needed for vision
-        # prefill where the 3D-grid span != number of tokens).
-        #
-        # NOTE: correct for eager + decode-only capture.  CudaGraphRunner
-        # does its own post-replay ``advance_seq_lens()`` at
-        # cuda_graph_runner.py:552 with no args, so this ``pos_id_ns`` is
-        # NOT honored on the replay path.  If we ever capture vision
-        # prefill, that runner call would need to accept a submodule-
-        # supplied ``pos_id_ns``.
-        cache_handle.advance_seq_lens(pos_id_ns=mrope_pos_advance)
-
         # Final layer norm
         hidden_states = self.model.norm(hidden_states)
 
diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py
index 84e70bf37..6d2021305 100644
--- a/mstar/model/qwen3_omni/submodules.py
+++ b/mstar/model/qwen3_omni/submodules.py
@@ -25,6 +25,7 @@
 from mstar.engine.cuda_graph_config import FlashInferPackedCudaGraphConfig
 from mstar.engine.cuda_graph_runner import BasicBatchedCudaGraphConfig
 from mstar.engine.kv_store import PositionInfo
+from mstar.engine.resources import PlanSpec, StepDeclaration
 from mstar.model.qwen3_omni.components.code2wav import Qwen3OmniMoeCode2Wav
 from mstar.model.qwen3_omni.components.rope import (
     compute_3d_cos_sin,
@@ -361,8 +362,8 @@ def prepare_inputs(
             embeds = self.model.model.embed_tokens(token_id)
 
             # Next MRoPE position for all 3 components: read from the
-            # per-request cache-manager state (kept in sync by the
-            # post-forward ``advance_seq_lens`` call in ``thinker.py``).
+            # stream's position counter (advanced by the runner's commit
+            # of each declared step).
             pos_ids = torch.tensor(
                 [[start_pos], [start_pos], [start_pos]],
                 dtype=torch.float,
@@ -391,8 +392,8 @@ def prepare_inputs(
             # per-modality helper instead of the full HF parser.
             #
             # ``start_pos`` is the next MRoPE position for this request,
-            # carried forward across walks by ``state.position_id_start``
-            # (advanced post-forward by ``advance_seq_lens``).
+            # carried forward across walks by the stream's position counter
+            # (advanced by the runner's commit of each declared step).
             pos_ids = get_rope_index_text(seq_len, start_pos, device)
             masks_for_talker = torch.stack([
                 torch.zeros(text_ids.shape, dtype=torch.bool, device=device), # multimodal
@@ -487,11 +488,10 @@ def prepare_inputs(
             )
 
             # Next MRoPE position after this vision block is ``end_pos_base
-            # + 1`` (one past the EOS token).  ``advance_seq_lens`` by
-            # default advances ``position_id_start`` by ``seq_len``, which
-            # for vision (= vision_len + 2) is typically smaller than the
-            # 3D-grid span.  Emit the correct per-request advance so the
-            # Thinker forward can pass ``pos_id_ns`` through.
+            # + 1`` (one past the EOS token). A commit advances positions by
+            # the span by default, which for vision (= vision_len + 2) is
+            # typically smaller than the 3D-grid span. Emit the correct
+            # per-request advance so ``declare_step`` can declare it.
             mrope_pos_advance = int(end_pos_base + 1 - start_pos)
             deepstack = []
             for deepstack_inp in inputs["deepstack"]:
@@ -510,6 +510,28 @@ def prepare_inputs(
                 }
             )
 
+    def declare_step(
+        self,
+        graph_walk: str,
+        engine_inputs: ModelInputsFromEngine,
+        inputs: list[ARNodeInputs],
+    ) -> StepDeclaration:
+        pos_advance = None
+        if graph_walk == "prefill_vision":
+            # The 3D-grid MRoPE span is larger than the token count; the
+            # commit advances the position counter by the declared span.
+            pos_advance = tuple(
+                int(inp.tensor_inputs.get("mrope_pos_advance", 0))
+                for inp in inputs
+            )
+        return StepDeclaration(plans=(PlanSpec(
+            labels=("main",),
+            spans={"main": tuple(inp.input_seq_len for inp in inputs)},
+            is_causal=True,
+            rope=True,
+            pos_advance=pos_advance,
+        ),))
+
     def preprocess(
         self,
         graph_walk: str,
@@ -538,15 +560,6 @@ def preprocess(
             target_dtype=input_embeds.dtype,
         )
 
-        # Plan FlashInfer attention and rope for the main cache label
-        cache_manager = engine_inputs.cache_manager
-        cache_manager.set_active_label("main")
-        assert cache_manager is not None
-        cache_manager.plan_attention(
-            seq_lens=seq_lens, is_causal=True, label="main"
-        )
-        cache_manager.plan_rope(seq_lens=seq_lens, pos_ids=None, label="main")
-
         extra_inputs = {}
         if graph_walk == "prefill_vision":
             assert len(inputs) == 1, \
@@ -578,19 +591,6 @@ def preprocess(
                 )
                 for i, t in enumerate(deepstack_list):
                     extra_inputs[f"deepstack_{i}"] = t
-            mrope_pos_advance = [
-                inp.tensor_inputs.get("mrope_pos_advance", 0)
-            ]
-            extra_inputs["mrope_pos_advance"] = mrope_pos_advance
-            # Side-channel: stash on the cache_manager's plan state via the
-            # public setter so the CUDA-graph runner's post-replay
-            # ``advance_seq_lens()`` (which is called with no args) advances
-            # ``position_id_start`` by the MRoPE 3D-grid span instead of by
-            # ``seq_len``. The eager path consumes ``mrope_pos_advance`` from
-            # the dict via model.forward → cache_handle.advance_seq_lens(
-            # pos_id_ns=...); both paths converge on the same per-request
-            # advance.
-            cache_manager.set_custom_pos_advance(mrope_pos_advance, label="main")
 
         return {
             "input_embeds": input_embeds,
@@ -642,7 +642,6 @@ def forward(
         cos_3d: torch.Tensor | None = None,
         sin_3d: torch.Tensor | None = None,
         mrope_section: list[int] | None = None,
-        mrope_pos_advance: list[int] | None = None,
         masks_for_talker: dict[str, torch.Tensor] | None = None,
         **kwargs,
     ) -> NameToTensorList:
@@ -670,8 +669,8 @@ def forward(
             cache_handle=engine_inputs.cache_manager,
             cos_sin_3d=cos_sin_3d,
             mrope_section=mrope_section,
-            mrope_pos_advance=mrope_pos_advance,
             deepstack_visual_embeds=deepstack,
+            label="main",
         )
 
         result: NameToTensorList = {}
@@ -731,8 +730,8 @@ def _build_prefill_text_packed(
         runner's static-buffer interning is tensor-only by design (non-tensor
         entries are model-static and don't need a per-bucket buffer), so
         ``forward_batched`` recovers ``mrope_section`` from a class constant
-        and reads token boundaries from ``cache_manager.get_qo_indptr_buf``
-        instead. Per-token cos/sin values come from running the real RoPE
+        and reads token boundaries from the attention manager's
+        ``qo_indptr_buf`` instead. Per-token cos/sin values come from running the real RoPE
         math on a sequential dummy position (3 components × num_tokens) so
         the captured kernels see non-degenerate inputs at trace time.
         """
@@ -765,10 +764,10 @@ def _build_prefill_vision_packed(
         Mirrors ``_build_prefill_text_packed`` and additionally provides
         ``deepstack_<i>`` (one tensor per ``vision.deepstack_visual_indexes``
         entry). Non-tensor extras (``mrope_section``, ``seq_lens``,
-        ``mrope_pos_advance``, ``masks_for_talker``) are intentionally absent:
-        the runner's static-buffer interning is tensor-only, so non-tensors
-        come back from ``submodule.preprocess`` at replay time. ``mrope_pos_advance``
-        flows through the ``_PlanState`` side-channel (see ``cache_manager._PlanState``).
+        ``masks_for_talker``) are intentionally absent: the runner's
+        static-buffer interning is tensor-only, so non-tensors come back
+        from ``submodule.preprocess`` at replay time. ``mrope_pos_advance``
+        rides the step declaration, which the runner commits post-replay.
 
         Visual_pos_masks is a length-``num_tokens`` bool tensor; at capture
         time we set it to all-False so the inner ``_deepstack_process``
@@ -909,10 +908,9 @@ def get_cuda_graph_configs(self, device: torch.device, tp_world_size: int = 1):
                 ),
             ),
             # prefill_vision: separate capture because its post-preprocess
-            # tensor signature has extras (deepstack_<i>)
-            # that prefill_text/audio don't. mrope_pos_advance flows
-            # out-of-band via ``BatchedCacheManager.set_custom_pos_advance``
-            # — see ``cache_manager._PlanState.custom_pos_advance``.
+            # tensor signature has extras (deepstack_<i>) that
+            # prefill_text/audio don't. mrope_pos_advance rides the step
+            # declaration, whose commit the runner applies post-replay.
             FlashInferPackedCudaGraphConfig(
                 capture_graph_walk="prefill_vision",
                 replay_graph_walks=["prefill_vision"],
@@ -946,7 +944,6 @@ def forward_batched(
         cos_3d: torch.Tensor | None = None,
         sin_3d: torch.Tensor | None = None,
         mrope_section: list[int] | None = None,
-        mrope_pos_advance: list[int] | None = None,
         masks_for_talker: dict[str, torch.Tensor] | None = None,
         **kwargs,
     ) -> dict[str, NameToTensorList]:
@@ -978,16 +975,10 @@ def forward_batched(
           ``prefill_text`` / ``prefill_audio`` share one capture (their
           post-preprocess tensor signature is identical: ``input_embeds`` +
           ``cos_3d`` + ``sin_3d``). ``prefill_vision`` has its own capture
-          because it adds  per-layer ``deepstack_<i>``
-          tensors. ``mrope_pos_advance`` flows out-of-band via
-          ``BatchedCacheManager.set_custom_pos_advance``, which
-          ``preprocess`` populates — see
-          ``cache_manager._PlanState.custom_pos_advance``. The model's inner
-          ``cache_handle.advance_seq_lens(pos_id_ns=mrope_pos_advance)`` call
-          executes only at capture time (it's a ``@torch.compiler.disable``'d
-          Python op so it's not replayed); the runner's post-replay
-          ``advance_seq_lens()`` is what advances the real state, and that
-          path reads the side-channel.
+          because it adds per-layer ``deepstack_<i>`` tensors.
+          ``mrope_pos_advance`` rides the step declaration; the runner
+          commits it post-replay (nothing in the captured region advances
+          state).
         """
 
         # Packed dict from FlashInferPackedCudaGraphConfig is tensor-only by
@@ -1013,12 +1004,12 @@ def forward_batched(
             cache_handle=cache_manager,
             cos_sin_3d=cos_sin_3d,
             mrope_section=mrope_section,
-            mrope_pos_advance=mrope_pos_advance,
             deepstack_visual_embeds=deepstack,
+            label="main",
         )
 
         if is_prefill:
-            qo_indptr_buf = cache_manager.get_qo_indptr_buf("main")
+            qo_indptr_buf = cache_manager.attention.qo_indptr_buf("main")
             assert qo_indptr_buf is not None, (
                 f"{graph_walk} forward_batched requires a properly initialized "
                 "FlashInferPrefillWrapper (qo_indptr static buffer); got None."
@@ -1416,26 +1407,28 @@ def prepare_inputs(
             input_seq_len=seq_len,
         )
 
+    def declare_step(
+        self,
+        graph_walk: str,
+        engine_inputs: ModelInputsFromEngine,
+        inputs: list[ARNodeInputs],
+    ) -> StepDeclaration:
+        return StepDeclaration(plans=(PlanSpec(
+            labels=("main",),
+            spans={"main": tuple(inp.input_seq_len for inp in inputs)},
+            is_causal=True,
+            rope=True,
+        ),))
+
     def preprocess(
         self,
         graph_walk: str,
         engine_inputs: ModelInputsFromEngine,
         inputs: list[ARNodeInputs],
     ) -> dict[str, torch.Tensor | Any]:
-        cache_manager = engine_inputs.cache_manager
-        assert cache_manager is not None
-        cache_manager.set_active_label("main")
-
         seq_lens = [
             inp.input_seq_len for inp in inputs
         ]
-        cache_manager.plan_attention(
-            seq_lens=seq_lens, is_causal=True, label="main"
-        )
-        cache_manager.plan_rope(
-            seq_lens=seq_lens, pos_ids=None, label="main"
-        )
-
         input_embeds = torch.cat([
             inp.input_embeds for inp in inputs
         ], dim=0)
@@ -1465,7 +1458,9 @@ def _forward_prefill(
         self, cache_handle: BatchedCacheManager,
         input_embeds: torch.Tensor,
     ):
-        self.model(input_embeds=input_embeds, cache_handle=cache_handle)
+        self.model(
+            input_embeds=input_embeds, cache_handle=cache_handle, label="main",
+        )
         return {}
 
     def _forward_decode_like(
@@ -1494,7 +1489,7 @@ def _forward_decode_like(
         non-batched (``hidden[-1:, :]``) branches.
         """
         hidden = self.model(
-            input_embeds=input_embeds, cache_handle=cache_handle
+            input_embeds=input_embeds, cache_handle=cache_handle, label="main",
         )
         if last_token_indices is not None:
             last_hidden = hidden.index_select(0, last_token_indices)
@@ -1620,6 +1615,7 @@ def forward_batched(
             hidden = self.model(
                 input_embeds=input_embeds,
                 cache_handle=cache_handle,
+                label="main",
             )
             return {
                 "__batched_talker_prefill_hidden__": hidden,
@@ -1735,7 +1731,7 @@ def get_cuda_graph_configs(self, device: torch.device, tp_world_size: int = 1):
         bs; single_request_inputs has input_seq_len=9 so total_tokens = bs * 9).
         ``total_tokens != bs`` forces ``_create_persistent_wrappers`` to use a
         ``FlashInferPrefillWrapper`` instead of the decode wrapper, which means
-        ``cache_handle.get_qo_indptr_buf("main")`` is non-None at replay so
+        the attention manager's ``qo_indptr_buf("main")`` is non-None at replay so
         ``forward_batched`` can ``index_select`` per-request last hidden out of
         the packed ``(bs * 9, talker_hidden)`` LLM output before codec_head.
         """
diff --git a/test/integration/test_prefill_cuda_graph.py b/test/integration/test_prefill_cuda_graph.py
index f72c4a93d..dc2e8f2c7 100644
--- a/test/integration/test_prefill_cuda_graph.py
+++ b/test/integration/test_prefill_cuda_graph.py
@@ -41,6 +41,7 @@
 from mstar.engine.cuda_graph_runner import CudaGraphKey, CudaGraphRunner  # noqa: E402
 from mstar.engine.kv_cache_engine import KVCacheEngine  # noqa: E402
 from mstar.engine.kv_store import TransferEngineInfo  # noqa: E402
+from mstar.engine.resources.step import StepRunner  # noqa: E402
 from mstar.model.submodule_base import ARNodeInputs, ModelInputsFromEngine  # noqa: E402
 
 QWEN3_OMNI_REPO = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
@@ -279,6 +280,11 @@ def _run_eager_per_rid(
         )
         with torch.no_grad():
             with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
+                runner = StepRunner()
+                declaration = submodule.declare_step(
+                    "prefill_text", engine_inputs, [inp],
+                )
+                runner.drive(declaration, cache_mgr)
                 preprocessed = submodule.preprocess(
                     graph_walk="prefill_text",
                     engine_inputs=engine_inputs,
@@ -289,6 +295,7 @@ def _run_eager_per_rid(
                     engine_inputs=engine_inputs,
                     **preprocessed,
                 )
+                runner.commit(declaration, cache_mgr)
         logits_chunks.append(out["logits"][0])           # (1, V)
         states_chunks.append(out["thinker_states"][0])   # (seq_len, 2*hidden)
     return (
diff --git a/test/modular/test_position_embedder.py b/test/modular/test_position_embedder.py
index a7fb56220..0b1f6d599 100644
--- a/test/modular/test_position_embedder.py
+++ b/test/modular/test_position_embedder.py
@@ -2,9 +2,9 @@
 
 ``RopeEmbedder.plan`` turns a segment list plus the pool's counters into a
 ``PositionPlan``; the cache manager's advance paths resolve each step's
-position delta (default span, explicit ``pos_id_ns``, or the
-``custom_pos_advance`` side channel) and record it only through
-``KVCachePool.commit``.
+position delta (default span, or explicit ``pos_id_ns``) and record it
+only through ``KVCachePool.commit``. A declared step's position advance
+rides ``PlanSpec.pos_advance`` instead (see test_step_declaration).
 """
 
 from __future__ import annotations
@@ -109,17 +109,16 @@ def test_plain_decode_advances_by_span(self):
         assert cm.kv_pool.positions("b", "main") == 1
         assert cm.kv_pool.view(Segment("a", "main", 0)).length == 1
 
-    def test_custom_pos_advance_overrides_span(self):
+    def test_pos_id_ns_list_override(self):
         """The vision-prefill shape: position span larger than the token
-        count, delivered through the side channel."""
+        count, passed per request."""
         cm = _make_cache_manager(["a"])
         _seed_planned_seq_lens(cm, "main", [4])
-        cm.set_custom_pos_advance([90], label="main")
 
-        cm.advance_seq_lens()
+        cm.advance_seq_lens(pos_id_ns=[90])
         assert cm.kv_pool.view(Segment("a", "main", 0)).length == 4
         assert cm.kv_pool.positions("a", "main") == 90
-        # The side channel is consumed, not persistent.
+        # The override is per call, not persistent.
         _seed_planned_seq_lens(cm, "main", [1])
         cm.advance_seq_lens()
         assert cm.kv_pool.positions("a", "main") == 91

From 93e343913524838f372ba09a36c6487329aaddad Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Mon, 10 Aug 2026 00:27:27 +0000
Subject: [PATCH 16/20] Adopt step declarations in bagel with declared cfg
 forks

---
 .../model/bagel/components/language_model.py  |  33 +--
 mstar/model/bagel/submodules.py               | 250 ++++++++----------
 2 files changed, 121 insertions(+), 162 deletions(-)

diff --git a/mstar/model/bagel/components/language_model.py b/mstar/model/bagel/components/language_model.py
index a57b5e428..7fe28f94d 100644
--- a/mstar/model/bagel/components/language_model.py
+++ b/mstar/model/bagel/components/language_model.py
@@ -144,6 +144,8 @@ def forward(
         self,
         query_sequence: torch.Tensor,
         cache_handle,
+        layer_idx: int,
+        label: str,
         mode="und",
         static_vae_idxs=True,
         vae_token_indexes=None,
@@ -204,17 +206,19 @@ def forward(
                 **mot_kwargs,
             )
 
-        # RoPE: pos_ids pre-computed by plan_rope before the LLM forward
+        # RoPE: pos_ids planned from the step declaration before the LLM forward
         query_states, key_states = cache_handle.apply_rope(
-            query_states, key_states, rope_theta=self.rope_theta
+            query_states, key_states, rope_theta=self.rope_theta, label=label,
         )
 
-        # Paged attention: plan (page alloc, FlashInfer index tensors) was
-        # done by plan_attention before the LLM forward
+        # Paged attention: the plan (page alloc, FlashInfer index tensors)
+        # was driven from the step declaration before the LLM forward
         attn_output = cache_handle.run_attention(
             q=query_states,
             k=key_states,
             v=value_states,
+            layer_idx=layer_idx,
+            label=label,
         )
 
         attn_output = attn_output.reshape(-1, self.hidden_size)
@@ -260,6 +264,8 @@ def forward(
         self,
         query_sequence: torch.Tensor,
         cache_handle,
+        layer_idx: int,
+        label: str,
         mode="und",
         static_vae_idxs=True,
         vae_token_indexes=None,
@@ -292,6 +298,8 @@ def forward(
         query_sequence = self.self_attn(
             query_sequence=query_sequence,
             cache_handle=cache_handle,
+            layer_idx=layer_idx,
+            label=label,
             mode=mode,
             static_vae_idxs=static_vae_idxs,
             vae_token_indexes=vae_token_indexes,
@@ -355,13 +363,12 @@ def forward(
         self,
         query_sequence: torch.Tensor,
         cache_handle: BatchedCacheManager,
-        write_cache=True,
+        label: str,
         mode="und",
         static_vae_idxs=None,
         vae_token_indexes=None,
         text_indexes=None,
         text_mask=None,
-        custom_advance_pos_id=None,
     ):
         extra_inputs = {}
         if self.use_moe:
@@ -377,18 +384,14 @@ def forward(
                 )
 
         for _layer_idx, decoder_layer in enumerate(self.layers):
-            # torch.compile complains about the attetion module having an integer layer_idx
-            # field that varies across the layers (forcing recompiles), so this is a workaround
-            cache_handle.set_layer_idx(_layer_idx)
             query_sequence = decoder_layer(
                 query_sequence=query_sequence,
                 cache_handle=cache_handle,
+                layer_idx=_layer_idx,
+                label=label,
                 **extra_inputs,
             )
 
-        if write_cache:
-            cache_handle.advance_seq_lens(pos_id_ns=custom_advance_pos_id)
-
         if self.use_moe:
             if mode == "und":
                 query_sequence = run_rms_norm(
@@ -437,26 +440,24 @@ def forward(
         self,
         query_sequence: torch.Tensor,
         cache_handle,
-        write_cache=True,
+        label: str,
         mode="und",
         static_vae_idxs=True,
         vae_token_indexes=None,
         text_indexes=None,
         text_mask=None,
-        custom_advance_pos_id=None,
         **kwargs
     ):
         assert mode in ["und", "gen"]
         outputs = self.model(
             query_sequence=query_sequence,
             cache_handle=cache_handle,
-            write_cache=write_cache,
+            label=label,
             mode=mode,
             vae_token_indexes=vae_token_indexes,
             text_indexes=text_indexes,
             text_mask=text_mask,
             static_vae_idxs=static_vae_idxs,
-            custom_advance_pos_id=custom_advance_pos_id,
         )
 
         return outputs
diff --git a/mstar/model/bagel/submodules.py b/mstar/model/bagel/submodules.py
index 4ca8f3b2b..e1771e335 100644
--- a/mstar/model/bagel/submodules.py
+++ b/mstar/model/bagel/submodules.py
@@ -17,6 +17,7 @@
 from mstar.engine.cuda_graph_config import FlashInferPackedCudaGraphConfig
 from mstar.engine.cuda_graph_runner import BasicBatchedCudaGraphConfig
 from mstar.engine.kv_store import PositionInfo
+from mstar.engine.resources import PlanSpec, StepDeclaration
 from mstar.model.bagel.components.language_model import BagelForCausalLM
 from mstar.model.bagel.components.modeling_utils import (
     ImageTransform,
@@ -396,13 +397,13 @@ class LLMSubmodule(ARNodeSubmodule):
     followed by an Euler step: x_{t+1} = x_t + v_final * dt.
 
     Multi-cache orchestration is driven by the requires_cfg flag in
-    per-request metadata. When True, graph walk methods manage 3 caches:
-      - prefill_text: snapshot main->cfg_text, forward [main, cfg_img]
-      - prefill_vit/vae: forward [main], snapshot main->cfg_text
+    per-request metadata. When True, the declared steps cover 3 caches:
+      - prefill_text: fork main->cfg_text before, forward [main, cfg_img]
+      - prefill_vit/vae: forward [main], fork main->cfg_text at commit
       - decode: forward [main, cfg_img]
       - image_gen: 3-pass CFG with conditional skip and renormalization
-    The CacheHandle (provided by KVCacheEngine) manages label switching, page
-    allocation, and KV data copying.
+    The runner drives the declaration (forks, plans, commits); every
+    forward call names its label.
     """
 
     # Node name → cache label mapping for image_gen_cfg
@@ -498,13 +499,12 @@ def get_cuda_graph_configs(
     ) -> list[BasicBatchedCudaGraphConfig | FlashInferPackedCudaGraphConfig]:
         """Declare CUDA graph captures for ``decode`` (cfg-off + cfg-on) and ``prefill_text`` (cfg-off only).
 
-        cfg-on prefill_text is intentionally NOT captured. BAGEL's
-        ``preprocess`` for prefill_text+cfg calls
-        ``cache_handle.snapshot_all("main", "cfg_text")``, an allocating
-        fork of the main stream. Replay addresses the cache manager at the
-        real request ids, so the snapshot would land on real state now,
-        but a captured cfg-on prefill has never been exercised and stays
-        off until verified on its own. cfg-on prefill_text continues to
+        cfg-on prefill_text is intentionally NOT captured. BAGEL's cfg-on
+        prefill_text declares a pre-forward fork of the main stream
+        (main to cfg_text), an allocating operation. Replay addresses the
+        cache manager at the real request ids, so the fork would land on
+        real state now, but a captured cfg-on prefill has never been
+        exercised and stays off until verified on its own. cfg-on prefill_text continues to
         use the eager path; downstream image_gen / decode+cfg captures are
         unaffected (they don't depend on this capture's snapshot
         semantics).
@@ -633,107 +633,97 @@ def prepare_inputs(
 
         return node_inputs
 
-    def preprocess(
+    def declare_step(
         self,
         graph_walk: str,
         engine_inputs: ModelInputsFromEngine,
         inputs: list[ARNodeInputs],
-    ) -> dict[str, torch.Tensor | Any]:
-
-        """Data transform + plan attention/rope for all relevant labels.
-
-        When cache_handle is provided (sequential execution), calls
-        plan_attention/plan_rope for every cache label needed by this
-        graph walk. This must happen outside forward() because plan
-        operations are CUDA graph incompatible.
-
-        When cache_handle is None (batched execution preprocesses per-request
-        without planning; planning is done separately via preprocess_batched).
-        """
-        cache_manager = engine_inputs.cache_manager
-
+    ) -> StepDeclaration:
         requires_cfg = self._batch_get_requires_cfg(
             engine_inputs.per_request_info
         )
         labels = self._get_active_labels(graph_walk, requires_cfg)
-        seq_lens = [inp.input_seq_len for inp in inputs]
-        per_label_custom_pos_ids = {
+        spans = tuple(inp.input_seq_len for inp in inputs)
+        per_label_pos_ids = {
             label: [
                 inp.custom_pos_ids[label] for inp in inputs \
                     if isinstance(inp.custom_pos_ids, dict) and label in inp.custom_pos_ids
             ] for label in labels
         }
 
-        result = {}
-
-
-        if graph_walk in ("image_gen", "image_gen_cfg"):
-            assert len(inputs) == 1 , "Batching not supported for image gen"
-
         if graph_walk == "image_gen" and requires_cfg:
-            # Batched CFG: plan a single FlashInfer batch across all 3 labels
-            # so that image_gen can run one forward pass instead of 3.
-            cache_manager.plan_attention_batched_cfg(
-                labels=labels,
-                seq_lens=seq_lens,
+            # Batched CFG: one combined plan across all 3 labels so image_gen
+            # runs one forward pass instead of 3. The frozen caches never
+            # grow during flow matching, so nothing commits.
+            return StepDeclaration(plans=(PlanSpec(
+                labels=tuple(labels),
+                spans={label: spans for label in labels},
                 is_causal=False,
                 write_store=False,
-            )
-            cache_manager.plan_rope_batched_cfg(
-                labels=labels,
-                seq_lens=seq_lens,
-                per_label_pos_ids=per_label_custom_pos_ids,
-            )
-        else:
-            self._plan_for_graph_walk(
-                cache_handle=cache_manager,
-                seq_lens=seq_lens,
-                per_label_custom_pos_ids=per_label_custom_pos_ids,
-                is_causal=graph_walk in [
-                    "prefill_text", "decode"
-                ],
-                labels=labels,
-                snapshots=[("main", "cfg_text")] if graph_walk == "prefill_text" and requires_cfg else [],
-                write_cache=graph_walk not in ("image_gen", "image_gen_cfg")
-            )
+                rope=True,
+                rope_pos_ids=per_label_pos_ids,
+                combined=True,
+                commit=False,
+            ),))
+
+        writes = graph_walk not in ("image_gen", "image_gen_cfg")
+        # Image blocks occupy one position regardless of their token count.
+        pos_advance = (
+            (1,) * len(inputs)
+            if graph_walk in ("prefill_vit", "prefill_vae") else None
+        )
+        plans = []
+        for label in labels:
+            pos_list = per_label_pos_ids.get(label)
+            plans.append(PlanSpec(
+                labels=(label,),
+                spans={label: spans},
+                is_causal=graph_walk in ("prefill_text", "decode"),
+                write_store=writes,
+                rope=True,
+                rope_pos_ids=torch.cat(pos_list) if pos_list else None,
+                commit=writes,
+                pos_advance=pos_advance,
+            ))
+
+        pre_forks: tuple = ()
+        post_forks: tuple = ()
+        if requires_cfg:
+            if graph_walk == "prefill_text":
+                # cfg_text keeps the pre-text context, so it forks before
+                # anything plans or writes.
+                pre_forks = (("main", "cfg_text"),)
+            elif graph_walk in ("prefill_vit", "prefill_vae"):
+                # cfg_text tracks the context including this image, so it
+                # forks at commit, after the step's writes have landed.
+                post_forks = (("main", "cfg_text"),)
+        return StepDeclaration(
+            plans=tuple(plans), pre_forks=pre_forks, post_forks=post_forks,
+        )
+
+    def preprocess(
+        self,
+        graph_walk: str,
+        engine_inputs: ModelInputsFromEngine,
+        inputs: list[ARNodeInputs],
+    ) -> dict[str, torch.Tensor | Any]:
+        """Data transform for the forward; the plans come from the step
+        declaration the runner drives before this runs."""
+        if graph_walk in ("image_gen", "image_gen_cfg"):
+            assert len(inputs) == 1 , "Batching not supported for image gen"
 
         # Concatenate lists of tensors into single tensors for each input name
         result = ARNodeInputs.collate(inputs, stacking_method=StackingMethod.CAT)
-        result["seq_lens"] = seq_lens
-        result["requires_cfg"] =  requires_cfg
+        result["seq_lens"] = [inp.input_seq_len for inp in inputs]
+        result["requires_cfg"] = self._batch_get_requires_cfg(
+            engine_inputs.per_request_info
+        )
         result["sample_token"] = any([
             info.step_metadata.get("sample_prefill_token", True) \
                 for info in engine_inputs.per_request_info.values()
         ])
         return result
 
-    def _plan_for_graph_walk(
-        self, cache_handle: BatchedCacheManager,
-        seq_lens: list[int],
-        per_label_custom_pos_ids: dict[str, list[torch.Tensor]] = {},
-        is_causal=True,
-        labels=["main"],
-        snapshots=[],
-        write_cache=True
-    ) -> None:
-        """Plan attention and rope for all cache labels needed by this graph walk."""
-        for snap in snapshots:
-            cache_handle.snapshot_all(*snap)
-
-        for label in labels:
-            pos_ids = per_label_custom_pos_ids.get(label)
-            if pos_ids is not None and len(pos_ids) > 0:
-                pos_ids = torch.cat(pos_ids)
-            else:
-                pos_ids = None
-            cache_handle.plan_attention(
-                seq_lens=seq_lens, is_causal=is_causal, label=label,
-                write_store=write_cache
-            )
-            cache_handle.plan_rope(
-                seq_lens=seq_lens, pos_ids=pos_ids, label=label
-            )
-
     def forward(
         self,
         graph_walk: str,
@@ -769,11 +759,9 @@ def _forward_prefill_text(
     ) -> NameToTensorList:
         """embed_tokens -> LLM forward (causal, mode='und') -> KV cache update.
 
-        When requires_cfg is True (image generation mode):
-        1. Snapshot main -> cfg_text BEFORE forward (done in preprocess)
-        2. Forward for main and cfg_img (both see the text tokens)
-
-        plan_attention/plan_rope are called in preprocess for all needed labels.
+        When requires_cfg is True (image generation mode), the declared
+        step forks main -> cfg_text before anything plans, and the forward
+        runs for main and cfg_img (both see the text tokens).
         """
         emb = self.embed_tokens(input_ids)
         requires_cfg = kwargs.pop("requires_cfg", False)
@@ -784,22 +772,19 @@ def _forward_prefill_text(
 
         if requires_cfg and cache_handle is not None:
             for label in ["main", "cfg_img"]:
-                cache_handle.set_active_label(label)
                 out = self.language_model(
                     emb, mode="und",
-                    cache_handle=cache_handle, **kwargs
+                    cache_handle=cache_handle, label=label, **kwargs
                 )
                 if label == "main":
                     hidden = out
         else:
-            if cache_handle is not None:
-                cache_handle.set_active_label("main")
             hidden = self.language_model(
                 emb, mode="und",
-                cache_handle=cache_handle, **kwargs
+                cache_handle=cache_handle, label="main", **kwargs
             )
         if sample_token:
-            qo_indptr_buf = cache_handle.get_qo_indptr_buf("main")
+            qo_indptr_buf = cache_handle.attention.qo_indptr_buf("main")
             assert qo_indptr_buf is not None
             last_token_indices = (qo_indptr_buf[1:] - 1).long()  # (padded_bs,)
             last_hidden = hidden.index_select(0, last_token_indices)
@@ -816,28 +801,22 @@ def _forward_prefill_vit(
     ) -> NameToTensorList:
         """Wrap img_emb with BOI/EOI tokens -> LLM forward (bidirectional).
 
-        When requires_cfg is True: forward for main only, then snapshot
-        main -> cfg_text (cfg_text = context including this image).
-
-        plan_attention/plan_rope are called in preprocess.
+        When requires_cfg is True the declared step forks main ->
+        cfg_text at commit (cfg_text = context including this image).
         """
-        requires_cfg = kwargs.pop("requires_cfg", False)
+        kwargs.pop("requires_cfg", None)
         kwargs.pop("cache_labels", None)
         kwargs.pop("snapshot_after", None)
         kwargs.pop("is_prefill", None)
         sample_token = kwargs.pop("sample_prefill_token", True)
 
-        cache_handle.set_active_label("main")
         hidden = self.language_model(
             input_embeds, mode="und",
-            custom_advance_pos_id=1,
-            cache_handle=cache_handle, **kwargs
+            cache_handle=cache_handle, label="main", **kwargs
         )
 
-        if requires_cfg:
-            cache_handle.snapshot_all("main", "cfg_text")
         if sample_token:
-            qo_indptr_buf = cache_handle.get_qo_indptr_buf("main")
+            qo_indptr_buf = cache_handle.attention.qo_indptr_buf("main")
             assert qo_indptr_buf is not None
             last_token_indices = (qo_indptr_buf[1:] - 1).long()  # (padded_bs,)
             last_hidden = hidden.index_select(0, last_token_indices)
@@ -857,33 +836,26 @@ def _forward_prefill_vae(
     ) -> NameToTensorList:
         """VAE image emb -> LLM forward (bidirectional, gen mode).
 
-        When requires_cfg is True: forward for main only, then snapshot
-        main -> cfg_text (cfg_text = context including this image).
-
-        plan_attention/plan_rope are called in preprocess.
+        When requires_cfg is True the declared step forks main ->
+        cfg_text at commit (cfg_text = context including this image).
         """
-        requires_cfg = kwargs.pop("requires_cfg", False)
+        kwargs.pop("requires_cfg", None)
         kwargs.pop("cache_labels", None)
         kwargs.pop("snapshot_after", None)
         kwargs.pop("is_prefill", None)
 
-        if cache_handle is not None:
-            cache_handle.set_active_label("main")
         self.language_model(
             input_embeds, mode="gen",
             cache_handle=cache_handle,
+            label="main",
             vae_token_indexes=vae_token_indexes,
             text_indexes=text_indexes,
             text_mask=text_mask,
-            custom_advance_pos_id=1,
             **kwargs
         )
 
         # NOTE: we will never sample a token from prefill_vae because it is alwaus
         # followed by prefill_vit
-
-        if requires_cfg and cache_handle is not None:
-            cache_handle.snapshot_all("main", "cfg_text")
         return {}
 
     def _forward_decode(
@@ -898,8 +870,6 @@ def _forward_decode(
 
         When requires_cfg is True: also forward for cfg_img to keep its
         KV cache in sync (cfg_img tracks all text, no images).
-
-        plan_attention/plan_rope are called in preprocess for all needed labels.
         """
         requires_cfg = kwargs.pop("requires_cfg", False)
         kwargs.pop("cache_labels", None)
@@ -910,18 +880,15 @@ def _forward_decode(
         kwargs.pop("top_p", None)
         emb = self.embed_tokens(input_ids)
 
-        if cache_handle is not None:
-            cache_handle.set_active_label("main")
         hidden = self.language_model(
             emb, mode="und",
-            cache_handle=cache_handle, **kwargs
+            cache_handle=cache_handle, label="main", **kwargs
         )
 
         if requires_cfg and cache_handle is not None:
-            cache_handle.set_active_label("cfg_img")
             self.language_model(
                 emb, mode="und",
-                cache_handle=cache_handle, **kwargs
+                cache_handle=cache_handle, label="cfg_img", **kwargs
             )
 
         logits = self.lm_head(hidden[-1:])
@@ -953,9 +920,10 @@ def _forward_image_gen(
     ) -> NameToTensorList:
         """Flow matching Euler step with optional 3-pass CFG.
 
-        Uses cache_handle to switch between the 3 frozen KV caches
-        (main, cfg_text, cfg_img). write_cache=False since caches are
-        frozen during flow matching.
+        Runs against the 3 frozen KV caches (main, cfg_text, cfg_img)
+        through one combined plan; the declared step writes no store and
+        commits nothing since the caches are frozen during flow
+        matching.
 
         When requires_cfg is False, runs a single forward pass (main only)
         without CFG, saving 2/3 of the compute.
@@ -1029,10 +997,9 @@ def _forward_image_gen(
             ])
             batched_text_mask = text_mask.repeat(3)
 
-            cache_handle.set_active_label("_cfg_batched")
             hidden = self.language_model(
                 batched_emb, mode="gen",
-                cache_handle=cache_handle, write_cache=False,
+                cache_handle=cache_handle, label="_cfg_batched",
                 vae_token_indexes=batched_vae_indexes,
                 text_indexes=batched_text_indexes,
                 text_mask=batched_text_mask,
@@ -1077,12 +1044,10 @@ def _forward_image_gen(
                     ).clamp(min=cfg_renorm_min, max=1.0)
                 v_final = v_combined * renorm_scale
         else:
-            # No CFG: single forward pass (plan done in preprocess)
-            if cache_handle is not None:
-                cache_handle.set_active_label("main")
+            # No CFG: single forward pass over the declared plan
             hidden = self.language_model(
                 empty_combined_emb, mode="gen",
-                cache_handle=cache_handle, write_cache=False,
+                cache_handle=cache_handle, label="main",
                 vae_token_indexes=vae_token_indexes,
                 text_indexes=text_indexes,
                 text_mask=text_mask,
@@ -1148,12 +1113,9 @@ def _forward_image_gen_single_branch(
         empty_combined_emb[1:-1] = latents_
 
         label = self._NODE_TO_CFG_LABEL.get(self.node_name, "main")
-        if cache_handle is not None:
-            cache_handle.set_active_label(label)
-
         hidden = self.language_model(
             empty_combined_emb, mode="gen",
-            cache_handle=cache_handle, write_cache=False,
+            cache_handle=cache_handle, label=label,
             vae_token_indexes=vae_token_indexes,
             text_indexes=text_indexes,
             text_mask=text_mask,
@@ -1252,25 +1214,21 @@ def _forward_decode_batched(
 
         Returns logits per request. Token sampling is done by the engine
         post-forward (outside CUDA graph capture).
-
-        plan_attention/plan_rope are called in preprocess_batched.
         """
         # 1. Embed and concatenate
         embs = self.embed_tokens(input_ids)
 
         # 2. Single LLM forward (main cache, already planned)
-        cache_manager.set_active_label("main")
         hidden = self.language_model(
             embs, mode="und",
-            cache_handle=cache_manager,
+            cache_handle=cache_manager, label="main",
         )
 
         # 3. CFG sync pass for cfg_img if needed (already planned)
         if requires_cfg:
-            cache_manager.set_active_label("cfg_img")
             self.language_model(
                 embs, mode="und",
-                cache_handle=cache_manager,
+                cache_handle=cache_manager, label="cfg_img",
             )
 
         # 4. Per-request lm_head -> logits (no sampling — done post-forward)

From 9d5a491c4c2ae4a85520252bfebbb43585bd57a4 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Mon, 10 Aug 2026 00:32:35 +0000
Subject: [PATCH 17/20] Add windowed retention and a block positional scheme;
 mark the facade internal

---
 mstar/engine/cache_manager.py          | 19 ++++---
 mstar/engine/resources/__init__.py     | 11 +++-
 mstar/engine/resources/kv_pool.py      | 76 +++++++++++++++++++++++---
 mstar/engine/resources/positions.py    | 24 ++++++++
 test/modular/test_kv_pool.py           | 75 +++++++++++++++++++++++++
 test/modular/test_position_embedder.py | 55 ++++++++++++++++++-
 6 files changed, 241 insertions(+), 19 deletions(-)

diff --git a/mstar/engine/cache_manager.py b/mstar/engine/cache_manager.py
index 37a27469b..782ba05be 100644
--- a/mstar/engine/cache_manager.py
+++ b/mstar/engine/cache_manager.py
@@ -48,16 +48,23 @@ class BatchedCfgInfo:
 
 
 class BatchedCacheManager(ABC):
-    """Model-facing facade over the engine's per-step resources.
+    """Internal: the runner's per-step surface over the engine's resources.
 
     Holds the step's addressing (which requests, which label each is on)
     and dispatches every call into the resources behind it: the KV cache
     pool (admission, views, commits, forks), the attention manager (plans,
     wrappers, workspaces), the positional embedder, and one cross-attention
     manager per declared source. The domain state lives on those resources;
-    what stays here is per-step bookkeeping the model or the graph runner
-    drives directly (active labels, the batched-CFG advance info, the
-    pre-plan short-circuit set).
+    what stays here is per-step bookkeeping (active labels, the batched-CFG
+    advance info, the pre-plan short-circuit set).
+
+    Adopted models (cosmos3, qwen3_omni, bagel) never sequence through this
+    surface: they declare their step (``NodeSubmodule.declare_step``), the
+    runner drives the plans, forks, and commits, and their forwards call
+    only the device-side ops with explicit labels (``run_attention``,
+    ``apply_rope``). The plan and advance methods here are transitional
+    surface for the families that have not adopted declarations yet; new
+    model code declares instead of calling them.
 
     Concrete backends pick the attention manager kind via
     ``ATTENTION_MANAGER_CLS``; ``ATTENTION_BACKENDS`` maps
@@ -65,9 +72,7 @@ class BatchedCacheManager(ABC):
     ``create_cache_manager`` instantiates the configured one.
 
     Constructed per batch: one facade serves the whole batch with a single
-    attention call per layer instead of N per-request calls. Complex paths
-    like image_gen (3-pass CFG with label switching) continue using
-    per-request construction.
+    attention call per layer instead of N per-request calls.
     """
 
     ATTENTION_MANAGER_CLS: type[FlashInferAttentionManager] = FlashInferAttentionManager
diff --git a/mstar/engine/resources/__init__.py b/mstar/engine/resources/__init__.py
index 661875f5c..673920c5b 100644
--- a/mstar/engine/resources/__init__.py
+++ b/mstar/engine/resources/__init__.py
@@ -11,12 +11,18 @@
     SequenceView,
 )
 from mstar.engine.resources.declare import PlanSpec, StepDeclaration
-from mstar.engine.resources.kv_pool import KVCachePool, PageArena, ScratchKVPool
-from mstar.engine.resources.positions import RopeEmbedder
+from mstar.engine.resources.kv_pool import (
+    KVCachePool,
+    PageArena,
+    RetentionPolicy,
+    ScratchKVPool,
+)
+from mstar.engine.resources.positions import BlockRopeEmbedder, RopeEmbedder
 from mstar.engine.resources.spec import NodeResourceSpec, ScratchKVSpec
 from mstar.engine.resources.step import StepPlan, StepRunner
 
 __all__ = [
+    "BlockRopeEmbedder",
     "CrossAttentionManager",
     "DenseGenAttentionManager",
     "FlashInferAttentionManager",
@@ -26,6 +32,7 @@
     "PlanSpec",
     "PositionPlan",
     "Reservation",
+    "RetentionPolicy",
     "RopeEmbedder",
     "ScratchKVPool",
     "ScratchKVSpec",
diff --git a/mstar/engine/resources/kv_pool.py b/mstar/engine/resources/kv_pool.py
index 2adec40d3..0f4c9446b 100644
--- a/mstar/engine/resources/kv_pool.py
+++ b/mstar/engine/resources/kv_pool.py
@@ -3,10 +3,12 @@
 Storage divides in two. The arena is the physical pool: the backing tensor
 and the page allocator over it. The pool is per-request accounting against
 an arena: which pages a stream holds, its stored length, its position
-counter. Several pools may share one arena; nothing above a pool sees the
-arena.
+counter, and its retention. Several pools may share one arena; nothing
+above a pool sees the arena.
 """
 
+from dataclasses import dataclass
+
 import torch
 
 from mstar.conductor.request_info import SequenceInfo
@@ -52,6 +54,18 @@ def total_pages(self) -> int:
         return self.allocator.max_num_pages
 
 
+@dataclass(frozen=True)
+class RetentionPolicy:
+    """Sliding-window retention for one cache stream: keep at most
+    ``context_budget`` tokens of physical history. The pool applies the
+    policy at commit, releasing whole pages from the front of the stream
+    as it grows past the budget; the next view's extent starts past the
+    released tokens. Configured per stream at request admission. Models
+    never call release. The protected-prefix windowed mode maps onto the
+    same release mechanism."""
+    context_budget: int
+
+
 class ScratchKVPool:
     """Fixed-shape scratch KV storage with a trivial lifecycle: no admit,
     no publish, no per-request lifetime. The tensor is overwritten every
@@ -91,6 +105,24 @@ def __init__(
             allocator=manager.page_allocator,
             page_size=manager.config.page_size,
         )
+        # Streams with a retention policy, and how many tokens each has
+        # released from its front (always whole pages). Empty in the
+        # common case, and every hot-path read is gated on that emptiness
+        # so unconfigured pools pay nothing.
+        self._retention: dict[tuple[str, str], RetentionPolicy] = {}
+        self._released: dict[tuple[str, str], int] = {}
+
+    def set_retention_policy(
+        self, request_id: str, label: str, policy: RetentionPolicy,
+    ) -> None:
+        """Configure one stream's retention at request admission."""
+        self._retention[(request_id, label)] = policy
+        self._released[(request_id, label)] = 0
+
+    def _released_tokens(self, request_id: str, label: str) -> int:
+        if not self._released:
+            return 0
+        return self._released.get((request_id, label), 0)
 
     @property
     def page_size(self) -> int:
@@ -106,12 +138,16 @@ def total_pages(self) -> int:
 
     def admit(self, segment: Segment) -> Reservation:
         """Reserve pages so the segment's stream can hold its history plus
-        this segment's span. Raises ``AllocationFailedError`` when the arena
-        cannot supply the pages; a zero-span segment reserves nothing."""
+        this segment's span. A retained stream sizes the reservation from
+        its physical remainder (released pages are gone for good). Raises
+        ``AllocationFailedError`` when the arena cannot supply the pages;
+        a zero-span segment reserves nothing."""
         state = self._manager.get_state(segment.request_id, segment.label)
         resident = state.seq_len
+        released = self._released_tokens(segment.request_id, segment.label)
         self._manager.alloc(
-            segment.request_id, segment.label, resident + segment.span
+            segment.request_id, segment.label,
+            resident + segment.span - released,
         )
         return Reservation(
             resident=resident,
@@ -122,24 +158,42 @@ def admit(self, segment: Segment) -> Reservation:
     def view(self, segment: Segment) -> SequenceView:
         """The stream as this step's plans must see it: every page backing
         it and the extent those pages cover once the segment's span lands.
-        Call after ``admit`` for spans that need new pages."""
+        The extent starts past any tokens retention released, so it is not
+        necessarily a prefix from zero. Call after ``admit`` for spans that
+        need new pages."""
         state = self._manager.get_state(segment.request_id, segment.label)
+        released = self._released_tokens(segment.request_id, segment.label)
         return SequenceView(
             pool=self,
             page_indices=tuple(state.page_indices),
-            start=0,
-            length=state.seq_len + segment.span,
+            start=released,
+            length=state.seq_len + segment.span - released,
         )
 
     def commit(self, segment: Segment, pos_advance: int | None = None) -> None:
         """Record that the segment's span was computed: stored length grows
         by the span, the position counter by ``pos_advance`` (defaults to
-        the span)."""
+        the span). A stream configured with a retention policy releases
+        what aged out, whole pages at a time, and the next view's extent
+        reflects the release."""
         state = self._manager.get_state(segment.request_id, segment.label)
         state.seq_len += segment.span
         state.position_id_start += (
             segment.span if pos_advance is None else pos_advance
         )
+        if not self._retention:
+            return
+        key = (segment.request_id, segment.label)
+        policy = self._retention.get(key)
+        if policy is None:
+            return
+        page_size = self.page_size
+        while (
+            state.seq_len - self._released[key] - policy.context_budget
+            >= page_size
+        ):
+            self._arena.free([state.page_indices.pop(0)])
+            self._released[key] += page_size
 
     def positions(self, request_id: str, label: str) -> int:
         """Current position counter for one stream, read-only."""
@@ -161,6 +215,10 @@ def remove_request(self, request_id: str) -> None:
         """Drop a request's accounting on every tier and free its pages."""
         if self._cpu_pool is not None:
             self._cpu_pool.remove_request(request_id)
+        if self._retention:
+            for key in [k for k in self._retention if k[0] == request_id]:
+                del self._retention[key]
+                del self._released[key]
         self._manager.remove_request(request_id)
 
     def set_write_policy(self, policy: StoreWritePolicy) -> None:
diff --git a/mstar/engine/resources/positions.py b/mstar/engine/resources/positions.py
index 3b43ea9fb..adb4f7ba4 100644
--- a/mstar/engine/resources/positions.py
+++ b/mstar/engine/resources/positions.py
@@ -83,3 +83,27 @@ def apply(
                 **llama31_params,
             )
         return q.to(orig_dtype), k.to(orig_dtype)
+
+
+class BlockRopeEmbedder(RopeEmbedder):
+    """Block positional scheme: every token of a segment shares the
+    stream's current position, and the counter advances by ``step`` per
+    non-empty segment regardless of span. This is the image-block shape,
+    where a block of N tokens occupies one position: the positions a step
+    consumes differ from the tokens it carries. ``apply`` is inherited
+    unchanged; only the identifiers and the advance rule differ."""
+
+    def __init__(self, step: int = 1):
+        self.step = step
+
+    def plan(self, segments: list[Segment], pool: KVCachePool) -> PositionPlan:
+        pos_ids: list[int] = []
+        advance: list[int] = []
+        for segment in segments:
+            start = pool.positions(segment.request_id, segment.label)
+            pos_ids.extend([start] * segment.span)
+            advance.append(self.step if segment.span else 0)
+        return PositionPlan(
+            pos_ids=torch.tensor(pos_ids, dtype=torch.long),
+            advance=tuple(advance),
+        )
diff --git a/test/modular/test_kv_pool.py b/test/modular/test_kv_pool.py
index 84fa9a5cb..770844725 100644
--- a/test/modular/test_kv_pool.py
+++ b/test/modular/test_kv_pool.py
@@ -32,6 +32,7 @@
     PageArena,
     PositionPlan,
     Reservation,
+    RetentionPolicy,
     Segment,
     SequenceView,
 )
@@ -599,5 +600,79 @@ def test_position_plan_frozen(self):
             plan.advance = (4,)
 
 
+class TestRetentionPolicy:
+    """A windowed retention policy as pool configuration: the pool applies
+    release at commit, whole pages at a time, and the next view's extent
+    reflects it. No attention manager, embedder, runner, or model code is
+    involved; unconfigured streams behave exactly as before."""
+
+    def _stream(self, budget: int):
+        pool, manager = _make_pool()
+        pool.add_request("r", ["main"])
+        pool.set_retention_policy("r", "main", RetentionPolicy(context_budget=budget))
+        return pool, manager
+
+    @staticmethod
+    def _step(pool, span: int) -> None:
+        segment = Segment("r", "main", span)
+        pool.admit(segment)
+        pool.view(segment)
+        pool.commit(segment)
+
+    def test_releases_whole_pages_past_the_budget(self):
+        pool, manager = self._stream(budget=16)
+        free0 = pool.num_free_pages
+        self._step(pool, 8)
+        self._step(pool, 8)
+        view = pool.view(Segment("r", "main", 0))
+        assert view.start == 0 and view.length == 16
+        assert pool.num_free_pages == free0 - 2
+
+        self._step(pool, 8)
+        view = pool.view(Segment("r", "main", 0))
+        assert view.start == 8 and view.length == 16
+        assert pool.num_free_pages == free0 - 2
+        state = manager.get_state("r", "main")
+        assert state.seq_len == 24
+        assert pool.positions("r", "main") == 24
+
+    def test_release_is_page_granular(self):
+        pool, _ = self._stream(budget=12)
+        for _ in range(3):
+            self._step(pool, 5)
+        assert pool.view(Segment("r", "main", 0)).start == 0
+        self._step(pool, 5)
+        view = pool.view(Segment("r", "main", 0))
+        assert view.start == 8 and view.length == 12
+
+    def test_admit_sizes_from_the_physical_remainder(self):
+        pool, manager = self._stream(budget=16)
+        for _ in range(6):
+            self._step(pool, 8)
+        state = manager.get_state("r", "main")
+        assert len(state.page_indices) == 2
+        assert pool.view(Segment("r", "main", 0)).length == 16
+
+        pool.admit(Segment("r", "main", 8))
+        assert len(state.page_indices) == 3
+
+    def test_unconfigured_stream_is_untouched(self):
+        pool, _ = _make_pool()
+        pool.add_request("r", ["main"])
+        for _ in range(4):
+            self._step(pool, 8)
+        view = pool.view(Segment("r", "main", 0))
+        assert view.start == 0 and view.length == 32
+
+    def test_remove_request_clears_retention_state(self):
+        pool, _ = self._stream(budget=16)
+        total = pool.num_free_pages
+        for _ in range(4):
+            self._step(pool, 8)
+        pool.remove_request("r")
+        assert pool.num_free_pages == total
+        assert pool._retention == {} and pool._released == {}
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))
diff --git a/test/modular/test_position_embedder.py b/test/modular/test_position_embedder.py
index 0b1f6d599..d2aee436c 100644
--- a/test/modular/test_position_embedder.py
+++ b/test/modular/test_position_embedder.py
@@ -23,7 +23,12 @@
     PagedAllocationManager,
     StoreWritePolicy,
 )
-from mstar.engine.resources import KVCachePool, RopeEmbedder, Segment
+from mstar.engine.resources import (
+    BlockRopeEmbedder,
+    KVCachePool,
+    RopeEmbedder,
+    Segment,
+)
 
 
 def _make_manager(max_num_pages: int = 16, page_size: int = 8) -> PagedAllocationManager:
@@ -175,6 +180,54 @@ def test_explicit_pos_ids_pass_through(self):
         assert cm._plan_states["main"].pos_ids.tolist() == [7, 9]
 
 
+class TestBlockRopeEmbedder:
+    """A second positional scheme as one new embedder implementation:
+    block positions, where every token of a segment shares the stream's
+    current position and the counter advances by one step per segment.
+    Nothing else changes: the pool commits the plan's advance, the
+    attention managers and models are untouched."""
+
+    def _pool(self):
+        alloc = _make_manager()
+        alloc.add_request("a", ["main"])
+        alloc.add_request("b", ["main"])
+        return KVCachePool(alloc)
+
+    def test_tokens_share_the_stream_position(self):
+        pool = self._pool()
+        pool.commit(Segment("a", "main", 0), pos_advance=4)  # counter at 4
+        embedder = BlockRopeEmbedder()
+
+        plan = embedder.plan(
+            [Segment("a", "main", 3), Segment("b", "main", 2)], pool
+        )
+        assert plan.pos_ids.tolist() == [4, 4, 4, 0, 0]
+        assert plan.advance == (1, 1)
+
+    def test_positions_differ_from_tokens_across_commits(self):
+        pool = self._pool()
+        embedder = BlockRopeEmbedder()
+        for expected_pos in (0, 1, 2):
+            segment = Segment("a", "main", 16)
+            plan = embedder.plan([segment], pool)
+            assert plan.pos_ids.tolist() == [expected_pos] * 16
+            pool.commit(segment, pos_advance=plan.advance[0])
+        state_pos = pool.positions("a", "main")
+        length = pool.view(Segment("a", "main", 0)).length
+        assert state_pos == 3 and length == 48
+
+    def test_zero_span_advances_nothing(self):
+        pool = self._pool()
+        plan = BlockRopeEmbedder().plan([Segment("a", "main", 0)], pool)
+        assert plan.pos_ids.numel() == 0
+        assert plan.advance == (0,)
+
+    def test_step_is_configurable(self):
+        pool = self._pool()
+        plan = BlockRopeEmbedder(step=3).plan([Segment("a", "main", 5)], pool)
+        assert plan.advance == (3,)
+
+
 if __name__ == "__main__":
     import pytest
     sys.exit(pytest.main([__file__, "-v"]))

From 6e7b1e2df42334acf84bd07b2b7e2482b2657a42 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Mon, 10 Aug 2026 18:27:14 +0000
Subject: [PATCH 18/20] Build the batched-cfg plan through build_paged_indptrs

---
 mstar/engine/resources/attention.py | 41 ++++++++---------------------
 1 file changed, 11 insertions(+), 30 deletions(-)

diff --git a/mstar/engine/resources/attention.py b/mstar/engine/resources/attention.py
index 3c4731c71..010761bab 100644
--- a/mstar/engine/resources/attention.py
+++ b/mstar/engine/resources/attention.py
@@ -321,32 +321,13 @@ def plan_batched_cfg(
         cfg = self.kv_cache_config
         page_size = cfg.page_size
 
-        # CPU-side accumulation (see plan for the same pattern).
-        qo_indptr_list = [0]
-        kv_indptr_list = [0]
-        all_page_indices = []
-        kv_last_page_lens = []
-        combined_seq_lens = []
-
-        for segment, view in zip(segments, views, strict=True):
-            qo_indptr_list.append(qo_indptr_list[-1] + segment.span)
-            all_page_indices.extend(view.page_indices)
-            kv_indptr_list.append(
-                kv_indptr_list[-1] + len(view.page_indices)
-            )
-
-            last_page_len = view.length % page_size or page_size
-            kv_last_page_lens.append(last_page_len)
-            combined_seq_lens.append(segment.span)
-
-        # CPU tensors — see comment in ``plan`` above. FlashInfer
-        # async-H2Ds these inside ``plan()``; passing GPU tensors would
-        # trigger a synchronous default-stream sync via the internal
-        # ``.to("cpu")`` call.
-        qo_indptr = torch.tensor(qo_indptr_list, dtype=torch.int32)
-        paged_kv_indptr = torch.tensor(kv_indptr_list, dtype=torch.int32)
-        paged_kv_indices = torch.tensor(all_page_indices, dtype=torch.int32)
-        paged_kv_last_page_len = torch.tensor(kv_last_page_lens, dtype=torch.int32)
+        combined_seq_lens = [segment.span for segment in segments]
+        indptrs = build_paged_indptrs(
+            q_seq_lens=combined_seq_lens,
+            page_indices_per_request=[view.page_indices for view in views],
+            context_lens=[view.length for view in views],
+            page_size=page_size,
+        )
 
         ps = self.states.get(combined_label)
         if self.cuda_graph_mode and ps is not None and ps.wrapper is not None:
@@ -389,10 +370,10 @@ def plan_batched_cfg(
             self.states[combined_label] = ps
 
         wrapper.plan(
-            qo_indptr=qo_indptr,
-            paged_kv_indptr=paged_kv_indptr,
-            paged_kv_indices=paged_kv_indices,
-            paged_kv_last_page_len=paged_kv_last_page_len,
+            qo_indptr=indptrs.qo_indptr,
+            paged_kv_indptr=indptrs.paged_kv_indptr,
+            paged_kv_indices=indptrs.paged_kv_indices,
+            paged_kv_last_page_len=indptrs.paged_kv_last_page_len,
             causal=is_causal,
             dtype=dtype,
         )

From 27251561f3a81b96662fadd8dbf18b1e1990b3e1 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Mon, 10 Aug 2026 18:27:14 +0000
Subject: [PATCH 19/20] Route request lifecycle through the pool fronts

---
 mstar/engine/kv_cache_engine.py   | 47 ++++++++++++++++---------------
 mstar/engine/resources/kv_pool.py |  9 ++++++
 2 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/mstar/engine/kv_cache_engine.py b/mstar/engine/kv_cache_engine.py
index f917adfce..223970a8e 100644
--- a/mstar/engine/kv_cache_engine.py
+++ b/mstar/engine/kv_cache_engine.py
@@ -55,16 +55,17 @@ class KVManagement:
     buffer_manager: WorkspaceBufferManager
     # source name -> cross-attention context pool (see KVCacheConfig.cross_attn)
     cross_pools: dict[str, CrossAttnPool] = field(default_factory=dict)
-    # Distinct cross-attention alloc managers (pools may be shared between
-    # sources); precomputed at build time since it can't change after
-    # startup, so per-request add/remove doesn't re-walk cross_pools.
-    cross_alloc_managers: list[PagedAllocationManager] = field(default_factory=list)
     # Persistent resource fronts over the managers above: the pool is the
     # engine's surface for admission, retrieval, offload tiers, position
     # reads, and publishing; the embedder owns position semantics. Built
     # once at load_model and shared by every step's facade.
     kv_pool: KVCachePool | None = None
     cross_kv_pools: dict[str, KVCachePool] = field(default_factory=dict)
+    # The distinct cross fronts (sources may share one physical pool, and
+    # share its front); precomputed at build time since it can't change
+    # after startup, so per-request add/remove doesn't re-walk
+    # cross_kv_pools.
+    distinct_cross_kv_pools: list[KVCachePool] = field(default_factory=list)
     rope_embedder: RopeEmbedder | None = None
 
 
@@ -273,12 +274,17 @@ def load_model(
             cross_pools = _build_cross_pools(
                 cfg, num_layers, device, kv_cache_type, transfer_engine_info,
             )
-            # Distinct alloc managers, deduped by identity (shared pools),
-            # computed once here rather than per request.
-            cross_alloc_managers: list[PagedAllocationManager] = []
-            for pool in cross_pools.values():
-                if all(pool.alloc_manager is not m for m in cross_alloc_managers):
-                    cross_alloc_managers.append(pool.alloc_manager)
+            # One accounting front per physical cross pool; sources sharing
+            # a pool share its front. The distinct list is computed once
+            # here rather than per request.
+            cross_kv_pool_by_id: dict[int, KVCachePool] = {}
+            cross_kv_pools: dict[str, KVCachePool] = {}
+            for source, pool in cross_pools.items():
+                front = cross_kv_pool_by_id.get(id(pool))
+                if front is None:
+                    front = KVCachePool(pool.alloc_manager)
+                    cross_kv_pool_by_id[id(pool)] = front
+                cross_kv_pools[source] = front
 
             alloc_manager = PagedAllocationManager(
                 config=cfg,
@@ -294,12 +300,9 @@ def load_model(
                     device=device,
                 ),
                 cross_pools=cross_pools,
-                cross_alloc_managers=cross_alloc_managers,
+                distinct_cross_kv_pools=list(cross_kv_pool_by_id.values()),
                 kv_pool=KVCachePool(alloc_manager, cpu_pool=cpu_page_pool),
-                cross_kv_pools={
-                    source: KVCachePool(pool.alloc_manager)
-                    for source, pool in cross_pools.items()
-                },
+                cross_kv_pools=cross_kv_pools,
                 rope_embedder=RopeEmbedder(),
             )
             self.kv_management[cfg.get_node_str()] = kv_mgmt
@@ -1335,8 +1338,8 @@ def add_request(
     ) -> None:
         for submodule_mgmt in self.submodule_management.values():
             submodule_mgmt.kv_management.kv_pool.add_request(request_id, cache_labels or ["main"])
-            for cross_mgr in submodule_mgmt.kv_management.cross_alloc_managers:
-                cross_mgr.add_request(request_id)
+            for cross_pool in submodule_mgmt.kv_management.distinct_cross_kv_pools:
+                cross_pool.add_request(request_id, [])
             submodule_mgmt.sampler.add_request(request_id)
             # Mirror into the cuda-graph runner's master sampler buffers so
             # the per-step path can index_select instead of rebuilding from
@@ -1348,8 +1351,8 @@ def remove_request(self, request_id: str) -> None:
         for submodule_mgmt in self.submodule_management.values():
             cache_mgmt = submodule_mgmt.kv_management
             cache_mgmt.kv_pool.remove_request(request_id)
-            for cross_mgr in cache_mgmt.cross_alloc_managers:
-                cross_mgr.remove_request(request_id)
+            for cross_pool in cache_mgmt.distinct_cross_kv_pools:
+                cross_pool.remove_request(request_id)
             submodule_mgmt.sampler.remove_request(request_id)
             submodule_mgmt.submodule.cleanup_request(request_id)
             if submodule_mgmt.cuda_graph_runner is not None:
@@ -1360,16 +1363,14 @@ def pause_request(
     ) -> None:
         """For interleaved loop: mark as paused, keep KV pages allocated."""
         for submodule_mgmt in self.submodule_management.values():
-            cache_mgmt = submodule_mgmt.kv_management
-            cache_mgmt.alloc_manager.get_state(request_id, cache_label).is_paused = True
+            submodule_mgmt.kv_management.kv_pool.pause(request_id, cache_label)
 
     def resume_request(
         self, request_id: str, cache_label: str = "main",
     ) -> None:
         """Resume from paused state for next LLM step in loop."""
         for submodule_mgmt in self.submodule_management.values():
-            cache_mgmt = submodule_mgmt.kv_management
-            cache_mgmt.alloc_manager.get_state(request_id, cache_label).is_paused = False
+            submodule_mgmt.kv_management.kv_pool.resume(request_id, cache_label)
 
     def shutdown(self) -> None:
         for submodule_mgmt in self.submodule_management.values():
diff --git a/mstar/engine/resources/kv_pool.py b/mstar/engine/resources/kv_pool.py
index 0f4c9446b..7ccef382d 100644
--- a/mstar/engine/resources/kv_pool.py
+++ b/mstar/engine/resources/kv_pool.py
@@ -221,6 +221,15 @@ def remove_request(self, request_id: str) -> None:
                 del self._released[key]
         self._manager.remove_request(request_id)
 
+    def pause(self, request_id: str, label: str) -> None:
+        """Mark one stream paused between turns of an interleaved loop;
+        its pages stay resident."""
+        self._manager.get_state(request_id, label).is_paused = True
+
+    def resume(self, request_id: str, label: str) -> None:
+        """Clear a stream's paused mark."""
+        self._manager.get_state(request_id, label).is_paused = False
+
     def set_write_policy(self, policy: StoreWritePolicy) -> None:
         """Set whether committed pages are pushed to the distributed store."""
         self._manager.write_policy = policy

From 9ddc9cbe335ec48f8affc3992f068fafe9e97dfe Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Mon, 10 Aug 2026 18:27:14 +0000
Subject: [PATCH 20/20] Note the node-resource end state on spec and model
 surfaces

---
 mstar/engine/resources/spec.py | 4 ++++
 mstar/model/base.py            | 4 +++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/mstar/engine/resources/spec.py b/mstar/engine/resources/spec.py
index e7fa696ca..cf3102f32 100644
--- a/mstar/engine/resources/spec.py
+++ b/mstar/engine/resources/spec.py
@@ -31,6 +31,10 @@ class NodeResourceSpec:
     backend, the rope embedder, and the cross-attention pools, exactly as
     it always has. ``scratch`` adds keyed fixed-shape caches built
     alongside them (resource key to spec).
+
+    TODO: once get_kv_cache_config retires, allow several named KV cache
+    configs per node and split the cross-attention and rope settings out
+    of KVCacheConfig into their own spec entries.
     """
     kv_cache_config: KVCacheConfig
     scratch: dict[str, ScratchKVSpec] = field(default_factory=dict)
diff --git a/mstar/model/base.py b/mstar/model/base.py
index dc987d837..59f3bbe2b 100644
--- a/mstar/model/base.py
+++ b/mstar/model/base.py
@@ -371,7 +371,9 @@ def get_node_resources(
         unchanged, so models that only define ``get_kv_cache_config``
         need no change here. Models with resources those configs cannot
         describe (e.g. a fixed-shape scratch cache) override this and
-        extend the returned specs.
+        extend the returned specs. Once every in-tree model migrates,
+        this replaces ``get_kv_cache_config`` as the resource surface
+        and the per-config pairing gives way to per-node declarations.
         """
         return [NodeResourceSpec(kv_cache_config=cfg) for cfg in kv_cache_config]