Resident decode planning: decode rows on device tables

.claude/skills/engine-resources/SKILL.md

@@ -32,6 +32,10 @@ Prototype a new kernel standalone against a dense reference before writing the r
 
 **`supports_preplan` pays off even with no kernel-side plan to hoist.** Host work counts. On Whisper, building a page table inline cost 0.052 ms on the critical path, while promoting a pre-planned one cost 0.009 ms.
 
+## Resident planning
+
+A walk a submodule names in `resident_walks` runs its captured steps off per-request tables on the device (`resources/resident.py`): `plan` stages nothing per row, and kernels recorded around the forward fill the buffers the forward reads. A resource joins by overriding `resident_refusal` and keeping a table; the base refuses, so a new resource is host-planned on its node until then. Per-step state is keyed by resource key, never one per planner: each owner reads its own stream's label (`ResidentWrites.streams[key]`), each KV pool its own block table and gathered lengths, and a pool the walk declares at span 0 is read only (no write addressing, no advance). An attention reports whether it read its pages off the tables (`reads_tables`) or through a host plan (`host_planned`); a bucket skips the host plan only if none needed one. Anything that changes a row other than a resident commit has to reach `ResidentPlanner.mark_stale`, or the next step reads the old row. In `KVManager` that is what `_touch_plan_cache(rid)` does. Writes to a table stay on the node's stream, in step order: the plan stream only writes the step's own staging. A slot is a cache entry, not a finite resource: an offload gives it up and a step short of slots takes the least recently run row's, so a table can only hold what a whole-row rewrite can rebuild from the host.
+
 ## Stateful resources move together
 
 KV pages and recurrent (GDN/Mamba) state slots are two forms of the same per-request state. A lifecycle feature added to one usually needs its analogue in the other. Forking for a request that branches, offload/reload, what a partially-run (chunked) step leaves behind: when you change one, check the others. Recurrent forking was initially missed when KV forking learned chunked-prefill semantics.

docs/adding_models.rst

@@ -1643,6 +1643,18 @@ and evict other requests. Without a declaration, the same situation raises a
 Returning ``None`` means that the submodule manages its own resources. This is the legacy
 path, kept for submodules that have not been migrated.
 
+A walk named in ``resident_walks()`` runs its captured steps off per-request tables on the
+device instead of a plan rebuilt per step (``mstar/engine/resources/resident.py``). The
+walk must declare, for each resource, one stream that takes one token per row, or one
+it only reads at span 0 (a cross-attention context), with no forks, retention or sink,
+and its ``preprocess`` must read no planned position or KV location, because the
+captured graph fills those in first. Each KV pool of the node keeps its own block table
+and lengths; positions read the pool their ``PositionConfig.kv_cache`` names. Capture raises on a declaration that breaks this. The engine
+turns it on only if every resource of the node can (``Resource.resident_refusal``) and
+the deployment left ``knobs: resident_planning`` at ``auto``; the worker log says why
+otherwise. Steps of other walks, and eager steps, keep the host plan and leave the row
+stale, so the first resident step after them writes it.
+
 A ``SubmoduleStep`` contains a list of ``Segment`` objects and one ``ResourceStep`` per
 resource key:
 

docs/environment_variables.rst

@@ -631,6 +631,19 @@ Worker scheduling
        allocating. ``50000`` removed them on Qwen3.5 0.8B at c=128 (an 8%
        shorter decode step), at the cost of holding at most that many young
        objects between collections. Read by ``knobs.gc_gen0_threshold``.
+   * - ``MSTAR_RESIDENT_PLANNING``
+     - config's ``knobs: resident_planning``, else ``auto``
+     - ``auto`` runs a node's captured decode steps off per-request tables on
+       the device where the submodule names resident walks
+       (``NodeSubmodule.resident_walks``) and every resource of the node can
+       (``Resource.resident_refusal``); the worker log says why a node that
+       names walks stays on the host. ``off`` plans every step on the host. The
+       tables are sized at startup, not by a setting: twice the walk's largest
+       capture batch in rows, plus one, since a request that finds no row takes
+       the least recently run one's. A row has a cell per page of the KV pool's
+       ``max_seq_len`` (the model's context, or ``resources: <kv>: max_seq_len``;
+       not the config's top-level key), at most the pool's pages, plus a few, 8
+       bytes each. Read by ``knobs.resident_planning``.
    * - ``MSTAR_KV_PLAN_CACHE``
      - config's ``knobs: kv_plan_cache``, else ``0``
      - The KV manager plans a captured decode step off the previous step's

docs/serving.rst

@@ -389,10 +389,11 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key
      - *(optional)* Engine knobs for this deployment, by name: ``device_loopback``,
        ``sampler_ingraph_scatter``, ``sampler_slots``, ``ingraph_decode_rope``,
        ``kv_chain_lazy_steps``, ``kv_plan_cache``, ``kv_pinned_indptrs``,
-       ``gc_gen0_threshold``, ``pre_plan``, ``step_thread_core``, ``idle_hold`` and
-       ``compile_deterministic`` (``true``/``false`` or an integer; ``step_thread_core`` takes
-       ``auto``, the default, ``off``, ``always``, ``siblings`` or a CPU number, and ``idle_hold``
-       ``off``, the default, or ``qos``; see :ref:`host-idle`). ``compile_deterministic`` reaches
+       ``gc_gen0_threshold``, ``pre_plan``, ``step_thread_core``, ``idle_hold``,
+       ``compile_deterministic`` and ``resident_planning`` (``true``/``false`` or an integer;
+       ``step_thread_core`` takes ``auto``, the default, ``off``, ``always``, ``siblings`` or a CPU number,
+       ``idle_hold`` ``off``, the default, or ``qos``, see :ref:`host-idle`, and ``resident_planning``
+       ``auto``, the default, or ``off``). ``compile_deterministic`` reaches
        Qwen3.5's text walks and no other model (see :doc:`models`). Each one's environment variable
        (``MSTAR_DEVICE_LOOPBACK`` and so on; see
        :doc:`environment_variables`) overrides it when set, and an unknown name

mstar/engine/accelerator_graph_runner.py

@@ -28,6 +28,7 @@
 from mstar.engine.node_dtype import node_scope
 from mstar.engine.resources import BucketKey, CGSlotSpec, Resource, SlotLease, StepContext, StepRunner
 from mstar.engine.resources.lora.config import LoraBaseOnly
+from mstar.engine.resources.resident import ResidentPlanner
 from mstar.engine.resources.step import PIECEWISE_WALK
 from mstar.model.submodule_base import (
     BatchedModelOutput,
@@ -696,6 +697,7 @@ def __init__(
         self._submodule_name = submodule_name
         self._submodule = submodule
         self._comm_group = joint_comm_group
+        self._resident: ResidentPlanner | None = None
 
         self._capture_configs: list[AcceleratorGraphConfig] = (
             submodule.get_accelerator_graph_configs(
@@ -776,6 +778,17 @@ def any_graphs(self):
     def _batch_sizes(self, config: AcceleratorGraphConfig) -> list[int]:
         return config.capture_batch_sizes or self.CAPTURE_BATCH_SIZES
 
+    def max_capture_bs(self, walks: Collection[str]) -> int:
+        """The widest capture of any of ``walks``, 0 when none of them has one."""
+        return max(
+            (max(self._batch_sizes(config)) for config in self._capture_configs if config.capture_graph_walk in walks),
+            default=0,
+        )
+
+    def bind_resident(self, planner: ResidentPlanner | None) -> None:
+        """Before any capture: a capture of a walk ``planner`` covers records its kernels."""
+        self._resident = planner
+
     # ── Capture ─────────────────────────────────────────────────────────
 
     def _get_slot_specs(self) -> list[CGSlotSpec]:
@@ -1186,13 +1199,24 @@ def prepare() -> dict[str, Any]:
             )
 
             forward = self._forward_for(spec)
+            planner = self._resident
+            resident = (
+                planner.capture_step(spec.bucket, spec.slot, step)
+                if planner is not None and planner.covers(spec.bucket)
+                else None
+            )
 
             def run_forward():
-                return forward(
+                if resident is not None:
+                    planner.gather(resident)
+                out = forward(
                     graph_walk=walk,
                     engine_inputs=engine_inputs,
                     **static_inputs,
                 )
+                if resident is not None:
+                    planner.commit_in_graph(resident)
+                return out
 
             self._graph_backend.set_device()
             self._graph_backend.synchronize()

mstar/engine/engine.py

@@ -43,6 +43,7 @@
 from mstar.engine.resources.kv.keys import fingerprint
 from mstar.engine.resources.kv.transfer import TransferEngineInfo
 from mstar.engine.resources.lora.config import LoraBaseOnly, node_capture_key
+from mstar.engine.resources.resident import ResidentPlanner, build_planner
 from mstar.engine.resources.spec import resolve_spec_dependencies
 from mstar.engine.resources.step import (
     ADMIT_OK,
@@ -188,6 +189,8 @@ class SubmoduleManagement:
     capture_runner: AcceleratorGraphRunner | None = None
     # the submodule's declared float32 settings; None keeps the process default
     fp32_precision: Fp32Precision | None = None
+    # the node's request slots and device tables, where it runs on them
+    resident: ResidentPlanner | None = None
 
     # Rotated globally, not per runner: slot-keyed buffers are shared across
     # buckets and regions, so per-runner counters let consecutive steps collide
@@ -693,6 +696,28 @@ def load_model(
                 # reads them back on the device (see DeviceLoopback)
                 resources[loopback.sampler].enable_device_loopback_reader()
 
+    def _build_resident_planner(self, node_name: str, runner: AcceleratorGraphRunner) -> None:
+        """Put ``node_name`` on resident tables if it can and the deployment
+        lets it; the startup log says why a node that named walks stays on the
+        host. In `order`, so an owner runs after what it depends on. Built here
+        because the tables are sized by the node's captures, and ahead of them
+        because a capture holds the tables' addresses."""
+        if knobs.resident_planning() == "off":
+            return
+        submodule_mgmt = self._submodules[node_name]
+        walks = submodule_mgmt.submodule.resident_walks()
+        if not walks:
+            return
+        resources = submodule_mgmt.resources
+        submodule_mgmt.resident, refused = build_planner(
+            node_name, walks,
+            {key: resources[key] for key in self._runner.order if key in resources},
+            self._resource_nodes, self._device, runner.max_capture_bs(walks),
+        )
+        if submodule_mgmt.resident is None:
+            logger.warning("resident planning stays off on %s: %s", node_name, refused)
+        runner.bind_resident(submodule_mgmt.resident)
+
     def paged_kv_caches(self) -> list[PagedKVCache]:
         """Every paged KV cache this engine built, with the nodes that use it."""
         from mstar.engine.resources.kv.manager import KVManager
@@ -805,7 +830,7 @@ def warmup(self) -> None:
         piecewise: dict[str, dict[str, PiecewiseAcceleratorGraphRunner]] = {}
         for node_name, submodule_mgmt in self._submodules.items():
             submodule = submodule_mgmt.submodule
-            ag_runners[node_name] = AcceleratorGraphRunner(
+            runner = ag_runners[node_name] = AcceleratorGraphRunner(
                 submodule_name=node_name,
                 submodule=submodule,
                 resources=submodule_mgmt.resources,
@@ -817,6 +842,7 @@ def warmup(self) -> None:
                 enable_nvtx=self._enable_nvtx,
                 device_capability=self.device_capability,
             )
+            self._build_resident_planner(node_name, runner)
             piecewise[node_name] = self._build_piecewise_runners(
                 node_name, submodule_mgmt
             )
@@ -1951,6 +1977,13 @@ def _drive_step(
                     "cg_key_info() and declare_step disagree"
                 )
             step.set_ctx(ctx)
+            # promoted and inline plans per resource, apart
+            plan_phase = "res.plan" if batch.preplan_event is not None else "res.plan_fresh"
+            if ctx.resident is not None:
+                ctx.plan_phase = plan_phase
+                t0 = phase_start()
+                ctx.resident.plan(step, ctx)
+                phase_end(f"{plan_phase}.resident", t0)
 
             # Admit was already done, can move on straight to plan
             if nvtx:
@@ -1961,10 +1994,7 @@ def _drive_step(
                 )
             t0 = phase_start()
             try:
-                # promoted and inline plans per resource, apart
-                self._runner.plan(
-                    step, "res.plan" if batch.preplan_event is not None else "res.plan_fresh",
-                )
+                self._runner.plan(step, plan_phase)
             finally:
                 phase_end(
                     "engine.plan.promoted" if batch.preplan_event is not None
@@ -2033,6 +2063,8 @@ def _drive_step(
             try:
                 self._runner.commit(step)
                 batch.committed = step
+                if submodule_mgmt.resident is not None:
+                    submodule_mgmt.resident.commit(step, ctx)
             finally:
                 phase_end("engine.commit", t0)
                 if nvtx:
@@ -2057,6 +2089,9 @@ def _forward(
         the GIL — the replay, or the submodule forward on the eager path."""
         graph_walk = batch.step_context.graph_walk
         if lease is not None:
+            ctx = batch.step_context
+            if ctx.resident is not None:
+                ctx.resident.flush(ctx.resident_plan)
             return cg_runner.run_forward(
                 lease, preprocessed, plan_done_event=batch.preplan_event,
                 launch_started_event=release_event,
@@ -2813,9 +2848,13 @@ def reserve_replay_slot(self, batch: ExecutingBatch) -> SlotLease | None:
                 graph_walk=batch.step_context.graph_walk, bs=len(batch.request_ids),
                 slot=batch.slot, num_tokens=num_tokens, cg_key_info=key,
             )
+        planner = submodule_mgmt.resident
+        ctx = batch.step_context
+        ctx.resident = planner if lease is not None and planner is not None and planner.covers(lease.bucket) else None
+        ctx.resident_plan = None
         if lease is not None:
             batch.lease_slot(lease)
-            batch.step_context.force_prefill = cg_runner.is_packed(lease)
+            ctx.force_prefill = cg_runner.is_packed(lease)
         return lease
 
     def can_pre_plan(self, node_name: str) -> bool:
@@ -2926,7 +2965,16 @@ def _pre_plan_for_batch(
                 self.reset_pre_plan_for_batch(batch)
                 return False
             stream = cg_runner.plan_stream()
+            ctx = batch.step_context
             with stream if stream is not None else nullcontext():
+                # A raise here takes what a raise from `pre_plan` below takes:
+                # `Worker._preplan_spec` resets the pre-plan, whose reservation
+                # stays with its streams for exec's admit to find.
+                if ctx.resident is not None:
+                    ctx.plan_phase = "res.pre_plan"
+                    t0 = phase_start()
+                    ctx.resident.plan(step, ctx)
+                    phase_end("res.pre_plan.resident", t0)
                 self._runner.pre_plan(step)
             laps.skip()
             batch.preplan_event = torch.Event(device=self._device)
@@ -3011,12 +3059,20 @@ def offload_request(self, node_name: str, request_id: int) -> int:
             for resource in self._submodules[node_name].resources.values()
             if resource.supports_eviction
         )
-        if freed > 0 and node_name in self._resident_delta:
-            self._resident_delta[node_name].add_offloaded(
-                self._graph_runtime.get_rid_string(request_id)
-            )
+        if freed > 0:
+            self._release_resident_slot(node_name, request_id)
+            if node_name in self._resident_delta:
+                self._resident_delta[node_name].add_offloaded(
+                    self._graph_runtime.get_rid_string(request_id)
+                )
         return freed
 
+    def _release_resident_slot(self, node_name: str, request_id: int) -> None:
+        """An offloaded request needs no slot: a reload is written whole anyway."""
+        planner = self._submodules[node_name].resident
+        if planner is not None:
+            planner.release(request_id)
+
     def reload_request(self, node_name: str, request_id: int) -> bool:
         """Bring it back. False when any resource can't fit it yet, in which
         case the request stays offloaded and the caller retries later."""
@@ -3148,6 +3204,7 @@ def apply_resident_delta(
                             wire_rid, node_name, len(delta),
                         )
                         return False
+                    self._release_resident_slot(node_name, rid)
                     logger.info(
                         "Replayed rank 0's offload of request %s from %s "
                         "(%d reclaimed)", wire_rid, node_name, freed,
@@ -3247,6 +3304,8 @@ def remove_request(self, request_id: str) -> set[str]:
         self._chunking.forget_request(request_id)
         for submodule_mgmt in self._submodules.values():
             submodule_mgmt.submodule.cleanup_request(request_id)
+            if submodule_mgmt.resident is not None:
+                submodule_mgmt.resident.release(request_id)
         return freed
 
     def shutdown(self):

mstar/engine/resources/attn/cross.py

@@ -1,6 +1,6 @@
 """Cross-attention over a context written once and never extended."""
 
-from typing import Any, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 import numpy as np
 import torch
@@ -37,6 +37,9 @@
 )
 from mstar.engine.resources.step import SlotLease, StepContext
 
+if TYPE_CHECKING:
+    from mstar.engine.resources.resident import CapturedResidentStep
+
 
 class _LastPlan(NamedTuple):
     """What a wrapper was last planned with. The context's KV manager hands
@@ -136,6 +139,15 @@ def __init__(
 
         self._workspaces = WorkspacePool(device, owner=kv_cache)
 
+    def resident_refusal(self) -> str | None:
+        if self._query_kv_cache_name == self._kv_cache_name:
+            return "a context in the query's own pool would be a second stream on one resident row"
+        return None
+
+    def resident_gather(self, step: "CapturedResidentStep") -> None:
+        # the context pool's row is read only, and this plans its views on the host
+        step.host_planned = True
+
     def depends_on(self):
         # a query side with no cache contributes no dependency: its packing
         # comes off the step, not another resource's plan