From 6d92b2e6a53661bf9dc24667418a39d2fba40a10 Mon Sep 17 00:00:00 2001 From: merceod Date: Sun, 1 Nov 2026 14:32:54 +0000 Subject: [PATCH 01/15] resident planning: decode rows on device tables A steady decode step only moves lengths by one and crosses a page now and then, so keep each request's block table row, length, GDN slot and position offset on the device. The captured step reads them by a row of slots and commits its own lengths; the host sends what changed and stays the truth. Mixed and eager steps keep the host plan and mark their rows. Behind knobs: resident_planning (auto|off), Qwen3.5 decode opts in, fa2 keeps its per-step plan. --- .claude/skills/engine-resources/SKILL.md | 4 + docs/adding_models.rst | 10 + docs/environment_variables.rst | 14 + docs/serving.rst | 7 +- mstar/engine/accelerator_graph_runner.py | 16 +- mstar/engine/engine.py | 44 +- mstar/engine/resources/attn/flashinfer.py | 5 + mstar/engine/resources/base.py | 41 ++ mstar/engine/resources/kv/manager.py | 87 ++- mstar/engine/resources/kv/resident.py | 109 +++ mstar/engine/resources/linear_attn/gdn.py | 7 + mstar/engine/resources/position/manager.py | 35 + mstar/engine/resources/recurrent/pool.py | 34 + mstar/engine/resources/resident.py | 512 +++++++++++++ mstar/engine/resources/sampler/resource.py | 4 + mstar/engine/resources/step.py | 5 + mstar/model/qwen3_5/submodules.py | 4 + mstar/model/submodule_base.py | 8 + mstar/utils/h2d.py | 33 + mstar/utils/knobs.py | 30 + test/modular/test_capture_budget.py | 2 + test/modular/test_engine_request_added.py | 2 +- test/modular/test_h2d_to_device.py | 21 +- test/modular/test_resident_planning.py | 794 +++++++++++++++++++++ 24 files changed, 1815 insertions(+), 13 deletions(-) create mode 100644 mstar/engine/resources/kv/resident.py create mode 100644 mstar/engine/resources/resident.py create mode 100644 test/modular/test_resident_planning.py diff --git a/.claude/skills/engine-resources/SKILL.md b/.claude/skills/engine-resources/SKILL.md index 60db4ed535..0beb983570 100644 --- a/.claude/skills/engine-resources/SKILL.md +++ b/.claude/skills/engine-resources/SKILL.md @@ -31,6 +31,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. 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. + ## 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. diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 00f72102b2..dc7370c58e 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -1281,6 +1281,16 @@ 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 one token per row in one stream, 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. 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: diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 9749609845..9f2b734e65 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -619,6 +619,20 @@ 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. + Read by ``knobs.resident_planning``. + * - ``MSTAR_RESIDENT_SLOTS`` + - config's ``knobs: resident_slots``, else ``4096`` + - Request rows of those tables, allocated at startup because a captured + graph holds their addresses; a step that needs a row past the last + fails. Each row costs a cell per page of ``max_seq_len`` plus a few, 8 + bytes each. Read by ``knobs.resident_slots``. * - ``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 diff --git a/docs/serving.rst b/docs/serving.rst index be13f258a8..f8a1694fd4 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -354,9 +354,10 @@ 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`` and ``step_thread_core`` - (``true``/``false`` or an integer; ``step_thread_core`` takes - ``off``, the default, ``auto`` or a CPU number). Each one's environment variable + ``gc_gen0_threshold``, ``pre_plan``, ``step_thread_core``, ``resident_planning`` + and ``resident_slots`` (``true``/``false`` or an integer; ``step_thread_core`` takes + ``off``, the default, ``auto`` or a CPU number, ``resident_planning`` + ``auto``, the default, or ``off``). Each one's environment variable (``MSTAR_DEVICE_LOOPBACK`` and so on; see :doc:`environment_variables`) overrides it when set, and an unknown name fails at startup. For example ``knobs: {kv_plan_cache: 1, diff --git a/mstar/engine/accelerator_graph_runner.py b/mstar/engine/accelerator_graph_runner.py index f4866dc86e..003780e751 100644 --- a/mstar/engine/accelerator_graph_runner.py +++ b/mstar/engine/accelerator_graph_runner.py @@ -27,6 +27,7 @@ ) from mstar.engine.node_dtype import node_scope from mstar.engine.resources import BucketKey, CGSlotSpec, Resource, SlotLease, StepContext, StepRunner +from mstar.engine.resources.resident import ResidentPlanner from mstar.model.submodule_base import BatchedModelOutput, ModelInputsFromEngine, NodeInputs, NodeSubmodule from mstar.platforms import current_platform from mstar.platforms.compile import with_inductor_overrides @@ -683,10 +684,12 @@ def __init__( num_slots: int, enable_nvtx: bool=False, device_capability: DeviceCapability | None = None, + resident: ResidentPlanner | None = None, ): self._submodule_name = submodule_name self._submodule = submodule self._comm_group = joint_comm_group + self._resident = resident self._capture_configs: list[AcceleratorGraphConfig] = ( submodule.get_accelerator_graph_configs( @@ -1168,13 +1171,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() diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 82132a9718..2e9607c992 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -41,6 +41,7 @@ from mstar.engine.resources.kv.config import instance_size, opening_prefix_caches, unrestorable from mstar.engine.resources.kv.keys import fingerprint from mstar.engine.resources.kv.transfer import TransferEngineInfo +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, @@ -143,6 +144,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 @@ -619,6 +622,26 @@ def load_model( enable = getattr(resource, "enable_device_loopback_reader", None) if enable is not None: enable() + self._build_resident_planner(node_name) + + def _build_resident_planner(self, node_name: str) -> 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.""" + 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, knobs.resident_slots(), + ) + if submodule_mgmt.resident is None: + logger.warning("resident planning stays off on %s: %s", node_name, refused) def paged_kv_caches(self) -> list[PagedKVCache]: """Every paged KV cache this engine built, with the nodes that use it.""" @@ -695,6 +718,7 @@ def warmup(self) -> None: num_slots=submodule_mgmt.num_slots, enable_nvtx=self._enable_nvtx, device_capability=self.device_capability, + resident=submodule_mgmt.resident, ) piecewise[node_name] = self._build_piecewise_runners( node_name, submodule_mgmt @@ -1686,6 +1710,8 @@ def _drive_step( "cg_key_info() and declare_step disagree" ) step.set_ctx(ctx) + if ctx.resident is not None: + ctx.resident.plan(step, ctx) # Admit was already done, can move on straight to plan if nvtx: @@ -1764,6 +1790,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: @@ -1788,6 +1816,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, @@ -2451,9 +2482,13 @@ def reserve_replay_slot(self, batch: ExecutingBatch) -> SlotLease | None: ), cg_key_info=batch.cg_key_info, ) + 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: @@ -2555,7 +2590,12 @@ 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(): + if ctx.resident is not None: + 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) @@ -2871,6 +2911,8 @@ def remove_request(self, request_id: str) -> set[str]: self._runner.remove_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): diff --git a/mstar/engine/resources/attn/flashinfer.py b/mstar/engine/resources/attn/flashinfer.py index 517f19313d..cd26908ce2 100644 --- a/mstar/engine/resources/attn/flashinfer.py +++ b/mstar/engine/resources/attn/flashinfer.py @@ -87,6 +87,11 @@ def __init__( def depends_on(self): return {self._kv_cache_name} + def resident_refusal(self) -> str | None: + if self._window_left >= 0: + return "a sliding window reads pages a resident row no longer holds" + return None + def _check_layer(self, layer_idx: int) -> None: """Over a pool of more than one group, ``layer_idx`` must be a layer this attention runs, since its plans address only that group's pages. Checked in Python, so at eager and diff --git a/mstar/engine/resources/base.py b/mstar/engine/resources/base.py index 19363d4b9d..d947c28cc6 100644 --- a/mstar/engine/resources/base.py +++ b/mstar/engine/resources/base.py @@ -19,6 +19,12 @@ PiecewiseAcceleratorGraphConfig, ) from mstar.engine.resources.kv.transfer import TransferEngineInfo + from mstar.engine.resources.resident import ( + CapturedResidentStep, + ResidentPlan, + ResidentPlanner, + ResidentWrites, + ) @dataclass(frozen=True) @@ -127,6 +133,41 @@ def depends_on(self) -> set[str]: def can_capture(self) -> bool: return True + # Resident planning (`resources/resident.py`). A captured decode step reads + # per-request state off tables on the device; each owner of that state + # keeps its own table. A resource overrides these only once it has one. + + def resident_refusal(self) -> str | None: + """Why a node holding this resource can't run on resident tables, or None if it can.""" + return f"{type(self).__name__} keeps no resident table" + + def bind_resident(self, planner: "ResidentPlanner") -> None: + """Declare this resource's tables on the node's planner, once, at load.""" + return + + def resident_row(self, rid: int, slot: int, row: list[int], writes: "ResidentWrites") -> None: + """Set this resource's cells of ``rid``'s row: it is new or stale. ``row`` + has the planner's narrow columns; the owner of the wide table appends.""" + return + + def resident_stage(self, plan: "ResidentPlan") -> None: + """Queue this step's per-row changes on ``plan.writes``. Host work only: + this can run beside a live graph.""" + return + + def resident_gather(self, step: "CapturedResidentStep") -> None: + """Record, inside the capture, the reads that fill what the step's forward + addresses. Runs at warmup and capture only; a replay repeats the kernels.""" + return + + def resident_commit(self, step: "CapturedResidentStep") -> None: + """Record, inside the capture and after the forward, the step's table updates.""" + return + + def resident_advance(self, plan: "ResidentPlan") -> None: + """Advance the host's copy of the table by the rule ``resident_commit`` recorded.""" + return + # stays correct when a hit skips the front of a request; required of every resource on a keyed node prefix_skip_safe: bool = False diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index b52671248b..2ad655c730 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -6,7 +6,7 @@ from collections.abc import Callable, Collection from concurrent.futures import Future, wait from dataclasses import dataclass, field -from typing import Any, NamedTuple +from typing import TYPE_CHECKING, Any, NamedTuple import numpy as np import torch @@ -65,6 +65,15 @@ from mstar.utils.knobs import kv_chain_lazy_steps, kv_pinned_indptrs, kv_plan_cache from mstar.utils.profiler import PhaseLaps +if TYPE_CHECKING: + from mstar.engine.resources.kv.resident import KVResidentTables + from mstar.engine.resources.resident import ( + CapturedResidentStep, + ResidentPlan, + ResidentPlanner, + ResidentWrites, + ) + logger = logging.getLogger(__name__) # Off by default: `assert_pages_conserved` walks every stream of every live @@ -493,6 +502,8 @@ def rows(self) -> frozenset: class KVManager(AttentionResource): prefix_skip_safe = True + # this pool's resident tables, once a node's planner bound them + _resident: "KVResidentTables | None" = None def __init__( self, @@ -534,6 +545,7 @@ def __init__( cfg, device, dtype ) self.name = name + self._needs_remote_transfer = needs_remote_transfer self._arena = PageArena( kv_cache=self.kv_cache, @@ -1630,7 +1642,9 @@ def admit_retrieve( graph_walk: str, published: PublishedKVInfo | None ) -> AdmitOutcome: - self._touch_plan_cache(rid) + # asked of a waiting row every readiness pass: a pool that retrieves + # from another node refuses resident planning, so its row stands + self._touch_plan_cache(rid, row=False) if self._host_copies: with self._lock: self._land_host_copies() @@ -2402,15 +2416,21 @@ def _drop_readonly_plans(self) -> None: self._readonly_plans = {} self._readonly_rows = frozenset() - def _touch_plan_cache(self, rid: str | None = None) -> None: + def _touch_plan_cache(self, rid: str | None = None, row: bool = True) -> None: """A stream of ``rid`` (any stream, when None) changed other than by a commit: the plan caches no longer describe the streams if they hold - that request. A padding row's reset after a replay holds nothing.""" + that request. A padding row's reset after a replay holds nothing. + + ``row`` also marks its resident row stale. A touch of every stream + doesn't: those are refused admits and abandoned pre-plans, which leave + only pages past the row's committed columns.""" if rid is None or rid in self._readonly_rows: self._drop_readonly_plans() cache = self._decode_plan_cache if cache is not None and (rid is None or any(rid == seg[0] for seg in cache.real_segments)): self._plan_epoch += 1 + if row and rid is not None and self._resident is not None: + self._resident.planner.mark_stale(rid) def _split_padding(self, step: KVStep, ctx: StepContext): """``(n_real, (n_pad, span))`` of a step whose real segments come @@ -2812,8 +2832,13 @@ def plan(self, step: KVStep, ctx: StepContext) -> dict[str, KVPlanOutput]: if self._plan_cache_mode and not pending: self._remember_decode_plan(step, ctx, res) laps.lap("full") - for g, outputs in enumerate(res.groups): - self._setup_plan_states(outputs, ctx, ctx.slot_lease, g, eager_write=step.eager_write) + if ctx.resident is None: + for g, outputs in enumerate(res.groups): + self._setup_plan_states(outputs, ctx, ctx.slot_lease, g, eager_write=step.eager_write) + else: + # the replay fills the slot's addressing, so no reuse may skip re-staging it + for label in res: + self._staged_outputs.pop((ctx.slot_lease.slot, (label, 0)), None) laps.lap("setup_states") if readonly: self._readonly_plans[slot] = (key, epochs, res) @@ -3093,6 +3118,7 @@ def _apply_retention( return 0 if touch: self._touch_plan_cache() + self._touch_plan_cache(rid) return self._release_oldest_locked(rid, stream, excess, policy.protected_prefix, group) def _settle_window_allotment(self, segment: Segment, stream: CacheStream) -> None: @@ -3179,6 +3205,7 @@ def release_oldest( "pages are offloaded or in transfer" ) self._touch_plan_cache() + self._touch_plan_cache(request_id) return self._release_oldest_locked(request_id, stream, num_tokens, stream.protected_prefix) # Eviction @@ -4359,6 +4386,54 @@ def _report_overrun(self, rid: str, label: str, stream: CacheStream) -> None: # them. TODO: rename the `set_layer_idx` call sites and drop this alias. set_layer_idx = AttentionResource.set_default_layer_idx + # Resident planning + + def resident_refusal(self) -> str | None: + if len(self._groups) > 1: + return "a pool with sliding-window layers holds its pages per group, which a block table row can't" + if self.kv_cache.layout != KVLayout.NHD: + return f"the {self.kv_cache.layout.value} layout has no resident write path" + if self._needs_remote_transfer: + return "a pool read across nodes takes pages a resident row doesn't track" + return None + + def bind_resident(self, planner: "ResidentPlanner") -> None: + from mstar.engine.resources.kv.resident import KVResidentTables + + page_size = self._groups[0].page_size + self._resident = KVResidentTables(planner, page_size, -(-self.config.max_seq_len // page_size)) + + def resident_row(self, rid: int, slot: int, row: list[int], writes: "ResidentWrites") -> None: + stream = self._streams[rid][writes.label] + if stream.released: + raise ValueError( + f"KV {self.name}: {rid}/{writes.label} released {stream.released} tokens, " + "which a block table row can't address" + ) + writes.kv_len = stream.stored_len + self._resident.row(stream.stored_len, stream.page_indices, slot, row) + + def resident_stage(self, plan: "ResidentPlan") -> None: + tables, streams, label = self._resident, self._streams, plan.writes.label + planner = tables.planner + for rid in plan.moved: + tables.refresh_len(streams[rid][label].stored_len, planner.slot_of(rid), plan.writes) + tables.stage_step(plan, lambda rid: streams[rid][label].page_indices, planner.rewrite) + if _DEBUG_ASSERTS: + for rid, slot in zip(plan.rids, plan.slots.tolist(), strict=True): + held, stored = int(tables.host_len[slot]), streams[rid][label].stored_len + assert held == stored, f"KV {self.name}: resident row of {rid} at {held}, its stream at {stored}" + + def resident_gather(self, step: "CapturedResidentStep") -> None: + state = self._static_plan_state(step.slot, (step.label, 0)) + self._resident.gather(step, state.token_to_page, state.token_to_cache) + + def resident_commit(self, step: "CapturedResidentStep") -> None: + self._resident.commit(step) + + def resident_advance(self, plan: "ResidentPlan") -> None: + self._resident.advance(plan) + def reset_default_cursors(self) -> None: super().reset_default_cursors() # unlike attention, every read here needs a usable index diff --git a/mstar/engine/resources/kv/resident.py b/mstar/engine/resources/kv/resident.py new file mode 100644 index 0000000000..4533898f2c --- /dev/null +++ b/mstar/engine/resources/kv/resident.py @@ -0,0 +1,109 @@ +"""The KV pool's resident tables: each request's block table row and stored length.""" +from __future__ import annotations + +import numpy as np +import torch +import triton +import triton.language as tl + +from mstar.engine.resources.resident import ( + CapturedResidentStep, + ResidentPlan, + ResidentPlanner, + ResidentWrites, + advance_column, + gather_column, + scatter_cells, +) + + +@triton.jit +def _gather_write_slots_kernel( + slots_ptr, kv_len_ptr, store_ptr, page_ptr, offset_ptr, + STRIDE: tl.constexpr, BLOCK: tl.constexpr, PAGE: tl.constexpr, +): + i = tl.program_id(0) + n = tl.load(kv_len_ptr + i) + tl.store(page_ptr + i, tl.load(store_ptr + tl.load(slots_ptr + i) * STRIDE + BLOCK + n // PAGE)) + tl.store(offset_ptr + i, n % PAGE) + + +class KVResidentTables: + """``block_table[slot, col]`` is a stream's col-th page, ``kv_len[slot]`` its + stored length: what a step appending one token per row reads to find where + it writes and what it attends to. + + The stream stays the truth. ``host_len`` advances by the rule the graph + does, so nothing comes back from the device, and ``host_cols`` counts the + columns a row holds, which is how the host sees a crossing. + """ + + def __init__(self, planner: ResidentPlanner, page_size: int, max_pages: int): + self.planner = planner + self.page_size = page_size + self.block_table = planner.table(max_pages) + self.kv_len = planner.table() + self.host_len = np.zeros(planner.max_slots, dtype=np.int64) + self.host_cols = np.zeros(planner.max_slots, dtype=np.int64) + + def row(self, stored_len: int, pages: list[int], slot: int, row: list[int]) -> None: + """The cells of a stream about to append one token. Its host mirrors + move here, not at the flush: a plan that never runs leaves the row stale.""" + cols = stored_len // self.page_size + 1 + row[self.kv_len.col] = stored_len + row.extend(pages[:cols]) + self.host_len[slot] = stored_len + self.host_cols[slot] = cols + + def refresh_len(self, stored_len: int, slot: int, writes: ResidentWrites) -> None: + """A row a host-planned step advanced; its columns stay until `stage_step`.""" + writes.put(self.kv_len.cell(slot), stored_len) + self.host_len[slot] = stored_len + + def stage_step(self, plan: ResidentPlan, stream_pages, rewrite) -> None: + """Bring each row's columns up to this step. The usual case, one new + column, rides the graph's scatter, ahead of its gather: only columns + past the row's length change, so a step still replaying reads none of + them. A row further behind (a host-planned step took its pages) is + rewritten whole.""" + slots = plan.slots + cols = self.host_cols[slots] + need = self.host_len[slots] // self.page_size + 1 + rows = np.flatnonzero(need > cols) + if not rows.size: + return + for i in rows.tolist(): + rid, slot, last = plan.rids[i], int(slots[i]), int(need[i]) - 1 + if last > cols[i]: + plan.taken.add(rid) + rewrite(rid, slot, plan.writes) + else: + plan.writes.put(self.block_table.cell(slot, last), stream_pages(rid)[last]) + plan.crossed = (rows, need[rows]) + + def advance(self, plan: ResidentPlan) -> None: + self.host_len[plan.slots] += 1 + if plan.crossed is not None: + rows, need = plan.crossed + self.host_cols[plan.slots[rows]] = need + + def gather(self, step: CapturedResidentStep, token_to_page: torch.Tensor, token_to_cache: torch.Tensor) -> None: + store, table = self.planner.store, self.block_table + stride, page_size = table.stride, self.page_size + # before the reads: a row that crossed writes into its new page + scatter_cells(store, step.buffers.cells) + kv_len = step.buffers.kv_len + gather_column(store, step.slots, stride, self.kv_len.col, kv_len) + if store.is_cuda: + _gather_write_slots_kernel[(step.bucket.bs,)]( + step.slots, kv_len, store, token_to_page, token_to_cache, + STRIDE=stride, BLOCK=table.col, PAGE=page_size, num_warps=1, + ) + return + col = torch.div(kv_len, page_size, rounding_mode="floor") + bs = step.bucket.bs + torch.index_select(store, 0, step.slots * stride + col + table.col, out=token_to_page[:bs]) + token_to_cache[:bs].copy_(kv_len - col * page_size) + + def commit(self, step: CapturedResidentStep) -> None: + advance_column(self.planner.store, step.slots, self.kv_len.stride, self.kv_len.col) diff --git a/mstar/engine/resources/linear_attn/gdn.py b/mstar/engine/resources/linear_attn/gdn.py index 69c8df9763..4493e414db 100644 --- a/mstar/engine/resources/linear_attn/gdn.py +++ b/mstar/engine/resources/linear_attn/gdn.py @@ -141,6 +141,10 @@ def depends_on(self) -> set[str]: # the pool plans first; its addressing arrives via `ctx.plan_results` return {self._pool_key} + def resident_refusal(self) -> str | None: + # decode addresses the pool by slot alone, which the pool's table holds + return None + def bind_dependencies(self, resources: Mapping[str, Resource]) -> None: self._pool = resources.get(self._pool_key) @@ -155,6 +159,9 @@ def plan(self, step: LinearAttnStep, ctx: StepContext): "linear-attn preplan is already pending; clear_preplan before " "planning a different step ahead" ) + if ctx.resident is not None: + # a replayed decode wrapper reads only the pool's slots, which its graph fills + return self._current self.reset_default_cursors() if self._preplanned: self._current = self._cached_plan_output diff --git a/mstar/engine/resources/position/manager.py b/mstar/engine/resources/position/manager.py index 187af7f7cd..26d2ea6f34 100644 --- a/mstar/engine/resources/position/manager.py +++ b/mstar/engine/resources/position/manager.py @@ -3,6 +3,7 @@ has ownership of per-(request,label) position counter""" from dataclasses import dataclass, field, fields +from typing import TYPE_CHECKING import torch @@ -27,8 +28,16 @@ from mstar.engine.resources.position.rope import ( rope_apply_qk_inplace, # noqa: F401 (registers mstar::rope_apply_qk_inplace) ) +from mstar.engine.resources.resident import gather_column from mstar.engine.resources.step import ADMIT_OK, AdmitOutcome, StepContext +if TYPE_CHECKING: + from mstar.engine.resources.resident import ( + CapturedResidentStep, + ResidentPlanner, + ResidentWrites, + ) + @dataclass class PublishedPositionInfo(PublishedInfo): @@ -302,6 +311,9 @@ def plan(self, step: PositionStep, ctx: StepContext) -> dict[str, torch.Tensor]: "position preplan is already pending; clear_preplan before " "planning a different step ahead" ) + if ctx.resident is not None: + # the replay fills the slot's buffer from the tables itself + return self._current_pos_ids if self._preplanned: # staged a step early against this same step's KV plan self._current_pos_ids = self._preplan_pos_ids @@ -346,6 +358,29 @@ def plan(self, step: PositionStep, ctx: StepContext) -> dict[str, torch.Tensor]: self._preplanned = ctx.is_preplan return pos_ids_out + # Resident planning: a stream's position runs ahead of its stored length by a + # fixed amount between host-planned steps, so a row keeps that offset + + def resident_refusal(self) -> str | None: + if self._config.scheme != PosScheme.SEQUENTIAL: + return f"position scheme {self._config.scheme} doesn't advance with the stream" + return None + + def bind_resident(self, planner: "ResidentPlanner") -> None: + self._resident_store = planner + self._resident_delta = planner.table() + + def resident_row(self, rid: int, slot: int, row: list[int], writes: "ResidentWrites") -> None: + row[self._resident_delta.col] = self.position(rid, writes.label) - writes.kv_len + + def resident_gather(self, step: "CapturedResidentStep") -> None: + key = CGSlotKey(bucket=step.bucket, slot=step.slot, label=step.label) + table = self._resident_delta + gather_column( + self._resident_store.store, step.slots, table.stride, table.col, + self._static_buffer(key, step.bucket.num_tokens), plus=step.buffers.kv_len, + ) + def commit(self, step: PositionStep, ctx: StepContext): for segment, advance in zip( step.segments, self._advances(step), strict=True diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index 7787f8b5ae..0ea21fad25 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -30,6 +30,7 @@ RecurrentStateSpec, RecurrentStep, ) +from mstar.engine.resources.resident import gather_column from mstar.engine.resources.spec import ResourceReqConfig from mstar.engine.resources.step import ( ADMIT_OK, @@ -44,6 +45,12 @@ if TYPE_CHECKING: from mstar.engine.resources.kv.manager import KVManager + from mstar.engine.resources.resident import ( + CapturedResidentStep, + ResidentPlanner, + ResidentTable, + ResidentWrites, + ) logger = logging.getLogger(__name__) @@ -132,6 +139,10 @@ class RecurrentAddressing: class RecurrentStatePool(Resource): + # the node's planner and this pool's table of slots, once bound + _resident: "ResidentPlanner | None" = None + _resident_table: "ResidentTable | None" = None + @classmethod def build( cls, spec: RecurrentStateSpec, info: EngineResourceInfo, @@ -554,6 +565,9 @@ def plan(self, step: RecurrentStep, ctx: StepContext): "recurrent preplan is already pending; clear_preplan before " "planning a different step ahead" ) + if ctx.resident is not None: + # the replay's kernels read the buffers its graph fills from the table + return self._current if self._preplanned: # Promotion: apply the fork copies staging left undone; see # `_pending_fork_state`. @@ -1044,6 +1058,26 @@ def addressing(self, label: str = "main") -> RecurrentAddressing: ) return found + # Resident planning + + def resident_refusal(self) -> str | None: + if self.config.disable_sink_slot: + return "a pool with no sink slot has nowhere for a padding row" + return None + + def bind_resident(self, planner: "ResidentPlanner") -> None: + self._resident = planner + self._resident_table = planner.table() + + def resident_row(self, rid: int, slot: int, row: list[int], writes: "ResidentWrites") -> None: + row[self._resident_table.col] = self._slots[rid][writes.label].index + + def resident_gather(self, step: "CapturedResidentStep") -> None: + key = CGSlotKey(bucket=step.bucket, slot=step.slot, label=step.label) + table = self._resident_table + # a padding row reads the trash row's zero, which is SINK_SLOT + gather_column(self._resident.store, step.slots, table.stride, table.col, self._cg_addressing[key].slot_indices) + # Engine lifecycle def build_accelerator_graph_buffers( diff --git a/mstar/engine/resources/resident.py b/mstar/engine/resources/resident.py new file mode 100644 index 0000000000..94ed1488fd --- /dev/null +++ b/mstar/engine/resources/resident.py @@ -0,0 +1,512 @@ +"""Request state kept on the device across decode steps. + +Two steps of one batch differ by "every length +1" and "a row crossed a page +boundary", yet a host plan rebuilds each row's write location, position and +recurrent slot every step, and the plan thread shares a core with the other two +step threads. Here each request holds a slot, a row in fixed-address tables on +the device. The owner of a state keeps its own table (the KV cache its block +table and lengths, the recurrent pool its slot, positions their offset), and a +captured step reads them by a per-step row of slots, updates them itself, and +needs from the host only what changed. + +Every table write lands on the node's stream in step order: the graph's own +scatters, and the row copies `ResidentPlanner.flush` issues right before a +replay. The plan stream only ever writes a step's own staging buffers. So a row +is never rewritten under a graph that reads it, an abandoned pre-plan leaves the +tables alone, and nothing on the host path needs a kernel: a slot's fields sit +side by side in one row of one store, so rewriting it is one H2D. + +The host stays the source of truth. A row is written whole the first time a +step reads it and whenever something other than a resident commit changed it +(a host-planned step, an offload, a reload, a fork); see `mark_stale`. +""" +from __future__ import annotations + +import logging +import threading +from collections.abc import Collection, Iterable, Mapping +from dataclasses import dataclass +from itertools import chain + +import numpy as np +import torch +import triton +import triton.language as tl + +from mstar.engine.resources.base import Resource +from mstar.engine.resources.step import BucketKey, SubmoduleStep +from mstar.utils.h2d import PinnedStager +from mstar.utils.profiler import PhaseLaps + +logger = logging.getLogger(__name__) + +# Row 0 of every table, never a request's. Padding rows read and write it: its +# cells stay zero, which every owner's table takes to mean the sink. +TRASH_SLOT = 0 + + +# What a captured step records: a small kernel apiece. The torch ops for the same +# work take several times the GPU time, and the replay's launch on the GPU thread. + +@triton.jit +def _scatter_cells_kernel(cells_ptr, values_ptr, store_ptr): + i = tl.program_id(0) + tl.store(store_ptr + tl.load(cells_ptr + i), tl.load(values_ptr + i)) + + +@triton.jit +def _gather_column_kernel( + slots_ptr, store_ptr, out_ptr, plus_ptr, STRIDE: tl.constexpr, COL: tl.constexpr, PLUS: tl.constexpr, +): + i = tl.program_id(0) + value = tl.load(store_ptr + tl.load(slots_ptr + i) * STRIDE + COL) + if PLUS: + value += tl.load(plus_ptr + i) + tl.store(out_ptr + i, value.to(out_ptr.dtype.element_ty)) + + +@triton.jit +def _advance_column_kernel(slots_ptr, store_ptr, STRIDE: tl.constexpr, COL: tl.constexpr, TRASH: tl.constexpr): + i = tl.program_id(0) + slot = tl.load(slots_ptr + i) + cell = slot * STRIDE + COL + tl.store(store_ptr + cell, tl.load(store_ptr + cell) + tl.where(slot != TRASH, 1, 0)) + + +def scatter_cells(store: torch.Tensor, cells: torch.Tensor) -> None: + """``store[cells[0]] = cells[1]``.""" + if store.is_cuda: + _scatter_cells_kernel[(cells.shape[1],)](cells[0], cells[1], store, num_warps=1) + else: + store.index_copy_(0, cells[0], cells[1]) + + +def gather_column( + store: torch.Tensor, slots: torch.Tensor, stride: int, col: int, out: torch.Tensor, + plus: torch.Tensor | None = None, +) -> None: + """``out[i] = store[slots[i], col]``, plus ``plus[i]``.""" + n = slots.numel() + if store.is_cuda: + _gather_column_kernel[(n,)]( + slots, store, out, plus if plus is not None else slots, + STRIDE=stride, COL=col, PLUS=plus is not None, num_warps=1, + ) + return + value = store.index_select(0, slots * stride + col) + out[:n].copy_(value if plus is None else value + plus) + + +def advance_column(store: torch.Tensor, slots: torch.Tensor, stride: int, col: int) -> None: + """``store[slots[i], col] += 1`` for every row but the trash row's.""" + if store.is_cuda: + _advance_column_kernel[(slots.numel(),)](slots, store, STRIDE=stride, COL=col, TRASH=TRASH_SLOT, num_warps=1) + else: + cells = slots * stride + col + store.index_copy_(0, cells, store.index_select(0, cells) + (slots != TRASH_SLOT)) + + +class ResidentTable: + """``width`` columns of the planner's store, one cell per slot per column.""" + + __slots__ = ("col", "width", "stride", "view") + + def __init__(self, width: int): + self.width = width + # set when the planner seals the store + self.col = 0 + self.stride = 0 + self.view: torch.Tensor | None = None # `[slots, width]` + + def cell(self, slot: int, col: int = 0) -> int: + """A column's cell in the flat store.""" + return slot * self.stride + self.col + col + + +class ResidentWrites: + """What one step changes. ``rows`` are rewritten whole by `flush`, from + their first column; ``cells`` and ``values`` are single cells, which ride + the graph's own scatter.""" + + __slots__ = ("label", "rows", "cells", "values", "kv_len") + + def __init__(self, label: str): + self.label = label + self.rows: list[tuple[int, list[int]]] = [] + self.cells: list[int] = [] + self.values: list[int] = [] + # the stored length of the row being rewritten, for the owners after the KV cache + self.kv_len = 0 + + def put(self, cell: int, value: int) -> None: + self.cells.append(cell) + self.values.append(value) + + +@dataclass(slots=True) +class ResidentBuffers: + """What one (bucket, slot) capture reads per step. ``cells`` is `[2, 2 * bs]`, + flat cells then values, so a step's single-cell writes are one copy: a row + has at most a length and a new column to write. Unused entries write zero + to cell zero, a trash cell, and ``dirty`` says whether any entry is in use.""" + slots: torch.Tensor + cells: torch.Tensor + # the rows' stored lengths this step: the graph writes them, for what reads them after + kv_len: torch.Tensor + # the host row ``slots`` holds: a steady batch keeps one array, so no H2D and no compare + staged: np.ndarray | None = None + dirty: bool = False + + +@dataclass(slots=True) +class ResidentPlan: + """One step's row of slots and what it changes.""" + key: tuple + rids: tuple[int, ...] + slots: np.ndarray + seq: int + commits: int + writes: ResidentWrites + # rows written whole by this plan, and rows whose length alone it refreshes + taken: set[int] + moved: set[int] + # (rows, pages after the step) of the step's crossings, for `resident_advance` + crossed: tuple[np.ndarray, np.ndarray] | None = None + + +@dataclass(slots=True) +class CapturedResidentStep: + """What one capture's kernels share, at warmup and capture only.""" + bucket: BucketKey + slot: int + label: str + slots: torch.Tensor + buffers: ResidentBuffers + + +class ResidentPlanner: + """One node's request slots, tables and per-step rows.""" + + def __init__( + self, node: str, walks: Collection[str], resources: Mapping[str, Resource], + device: torch.device, max_slots: int, + ): + self.node = node + self.walks = frozenset(walks) + self.device = device + # request rows plus the trash row + self.max_slots = max_slots + 1 + self._owners = tuple(resources.values()) + self._tables: list[ResidentTable] = [] + self._store: torch.Tensor | None = None + # columns of the one-cell tables, which a rewritten row always covers + self.narrow_cols = 0 + self._stager = PinnedStager(torch.long, numel=1 << 16) + self._buffers: dict[tuple[BucketKey, int], ResidentBuffers] = {} + + self._lock = threading.Lock() + # rid -> slot. Handle-keyed: `release` purges it. + self._slot_of: dict[int, int] = {} + self._free: list[int] = list(range(max_slots, TRASH_SLOT, -1)) + # (slot, seq of the last step that read it), freed under a step not yet committed + self._cooling: list[tuple[int, int]] = [] + self._last_read = np.zeros(self.max_slots, dtype=np.int64) + # Handle-keyed, like `_slot_of`. Rows to write whole, and rows a + # host-planned step only advanced as a resident one would have. + self._stale: set[int] = set() + self._moved: set[int] = set() + self._plan_seq = 0 + self._committed_seq = 0 + self._commits = 0 + self._staged: ResidentPlan | None = None + # the last batch's rows, so a steady batch skips the per-row lookup + self._last_rids: tuple[int, ...] | None = None + self._last_slots: np.ndarray | None = None + + for resource in self._owners: + resource.bind_resident(self) + self._seal() + + # Tables + + def table(self, width: int = 1) -> ResidentTable: + """A table of ``width`` columns. Only from `bind_resident`.""" + if self._store is not None: + raise RuntimeError("resident tables are fixed once the planner is built") + table = ResidentTable(width) + self._tables.append(table) + return table + + def _seal(self) -> None: + # the one wide table last, so a row's used cells are contiguous from column 0 + wide = [t for t in self._tables if t.width > 1] + if len(wide) > 1: + raise ValueError("a resident row holds one wide table") + narrow = [t for t in self._tables if t.width == 1] + self.narrow_cols = len(narrow) + stride = sum(t.width for t in self._tables) + # int64: what the KV write addressing and the positions are, and index ops take + self._store = torch.zeros(self.max_slots * stride, dtype=torch.long, device=self.device) + rows = self._store.view(self.max_slots, stride) + col = 0 + for t in (*narrow, *wide): + t.col, t.stride = col, stride + t.view = rows[:, col: col + t.width] + col += t.width + logger.info( + "resident planning on %s: %d slots of %d cells, %.1f MiB, walks %s", + self.node, self.max_slots - 1, stride, self._store.nbytes / 2**20, sorted(self.walks), + ) + + @property + def store(self) -> torch.Tensor: + return self._store + + def covers(self, bucket: BucketKey) -> bool: + return bucket.graph_walk in self.walks + + # Slots + + def slot_of(self, rid: int) -> int | None: + return self._slot_of.get(rid) + + def _take_slot(self) -> int: + """Caller holds the lock. A freed slot comes back once the last step + that read it has committed: until then that step's commit still + advances the slot's host lengths.""" + if self._cooling: + done = self._committed_seq + ready = [slot for slot, seq in self._cooling if seq <= done] + if ready: + self._cooling = [(slot, seq) for slot, seq in self._cooling if seq > done] + self._free.extend(ready) + if not self._free: + raise RuntimeError( + f"resident planning on {self.node}: all {self.max_slots - 1} slots are held; " + "raise `knobs: resident_slots`" + ) + return self._free.pop() + + def release(self, rid: int) -> None: + """The request is gone. Its slot cools until the last step that read it + has committed; on the device the row is only read again behind that + step, by stream order.""" + with self._lock: + self._stale.discard(rid) + self._moved.discard(rid) + slot = self._slot_of.pop(rid, None) + if slot is None: + return + # a handle comes back as another request, under a batch tuple equal to the old one + self._last_rids = None + if self._staged is not None and rid in self._staged.rids: + self._staged.commits = -1 + seq = int(self._last_read[slot]) + if seq <= self._committed_seq: + self._free.append(slot) + else: + self._cooling.append((slot, seq)) + + def mark_stale(self, rid: int) -> None: + """Something other than a resident commit changed ``rid``'s row.""" + if rid in self._slot_of: + with self._lock: + self._stale.add(rid) + + # The step + + def buffers(self, bucket: BucketKey, slot: int) -> ResidentBuffers: + """The per-step rows of one capture, every slot at the trash row until staged.""" + key = (bucket, slot) + found = self._buffers.get(key) + if found is None: + found = self._buffers[key] = ResidentBuffers( + slots=torch.full((bucket.bs,), TRASH_SLOT, dtype=torch.long, device=self.device), + cells=torch.zeros((2, 2 * bucket.bs), dtype=torch.long, device=self.device), + kv_len=torch.zeros(bucket.bs, dtype=torch.long, device=self.device), + ) + return found + + def plan(self, step: SubmoduleStep, ctx) -> None: + """Stage this step's rows; on the plan thread for a pre-plan. Its admit + has run, so every page the step writes is on its stream.""" + lease = ctx.slot_lease + rids = tuple(ctx.request_ids) + key = (rids, lease.bucket, lease.slot) + staged = self._staged + if staged is not None: + self._staged = None + if not ctx.is_preplan and staged.key == key and staged.commits == self._commits: + ctx.resident_plan = staged + return + # the step it was for never ran, or something ran first + self._restore(staged) + plan = self._build(rids, key, step.segments[0].label, ctx) + ctx.resident_plan = plan + if ctx.is_preplan: + self._staged = plan + + def _restore(self, plan: ResidentPlan) -> None: + """Hand back the rows of a plan that won't run: it wrote nothing.""" + with self._lock: + self._stale |= {rid for rid in plan.taken if rid in self._slot_of} + self._moved |= {rid for rid in plan.moved if rid in self._slot_of} + + def _build(self, rids: tuple[int, ...], key: tuple, label: str, ctx) -> ResidentPlan: + lease = ctx.slot_lease + laps = PhaseLaps("res.pre_plan.resident." if ctx.is_preplan else "res.plan.resident.") + with self._lock: + if rids == self._last_rids: + slots = self._last_slots + else: + get = self._slot_of.get + row = [get(rid) for rid in rids] + if None in row: + for i, slot in enumerate(row): + if slot is None: + row[i] = self._slot_of[rids[i]] = self._take_slot() + self._stale.add(rids[i]) + slots = np.array(row, dtype=np.int64) + self._last_rids, self._last_slots = rids, slots + stale = self._stale + taken = stale.intersection(rids) if stale else set() + stale -= taken + moved = self._moved + refreshed = moved.intersection(rids) if moved else set() + moved -= refreshed + refreshed -= taken + self._plan_seq += 1 + seq = self._plan_seq + self._last_read[slots] = seq + rows = [(rid, self._slot_of[rid]) for rid in taken] + plan = ResidentPlan(key, rids, slots, seq, self._commits, ResidentWrites(label), taken, refreshed) + laps.lap("rows") + for rid, slot in rows: + self.rewrite(rid, slot, plan.writes) + laps.lap("rewrites") + buffers = self.buffers(lease.bucket, lease.slot) + if buffers.staged is not slots: + self._stager.copy_(buffers.slots, slots, pad_value=TRASH_SLOT) + buffers.staged = slots + laps.lap("slot_row") + for resource in self._owners: + resource.resident_stage(plan) + laps.lap("stage") + self._stage_cells(plan, buffers) + laps.lap("cells") + return plan + + def rewrite(self, rid: int, slot: int, writes: ResidentWrites) -> None: + """Queue ``rid``'s row whole. The narrow columns come first, then the + wide table's used cells, which its owner appends.""" + row = [0] * self.narrow_cols + for resource in self._owners: + resource.resident_row(rid, slot, row, writes) + writes.rows.append((slot, row)) + + def _stage_cells(self, plan: ResidentPlan, buffers: ResidentBuffers) -> None: + """One copy, on the plan stream: the graph scatters these first thing.""" + cells = plan.writes.cells + if not cells and not buffers.dirty: + return + n = buffers.cells.shape[1] + staged = np.zeros(2 * n, dtype=np.int64) + staged[: len(cells)] = cells + staged[n: n + len(cells)] = plan.writes.values + buffers.dirty = bool(cells) + self._stager.copy_(buffers.cells, staged) + + def flush(self, plan: ResidentPlan) -> None: + """The rows `plan` rewrites, on the step's stream: behind every earlier + step, ahead of this one. On the GPU thread, right before the replay.""" + rows = plan.writes.rows + if not rows: + return + stride = self._tables[0].stride + dsts = [self._store[slot * stride: slot * stride + len(row)] for slot, row in rows] + self._stager.copy_runs_(dsts, np.fromiter(chain.from_iterable(r for _, r in rows), dtype=np.int64)) + + def commit(self, step: SubmoduleStep, ctx) -> None: + """After the resources' commit, before the step's commit event: the + next pre-plan reads what this leaves.""" + plan = ctx.resident_plan if ctx.resident is not None else None + self._commits += 1 + if plan is None: + self._note_host_step(step, ctx) + return + for resource in self._owners: + resource.resident_advance(plan) + self._committed_seq = plan.seq + + def _note_host_step(self, step: SubmoduleStep, ctx) -> None: + """A host-planned step moved its rows without their tables. A + one-token row moved as a resident step would have, and needs only its + length and any new column; anything else is written whole.""" + held = self._slot_of + if not held: + return + segments = step.segments or [seg for s in step.steps.values() for seg in s.segments or ()] + one = [seg.request_id for seg in segments if seg.span == 1 and seg.request_id in held] + rest = [seg.request_id for seg in segments if seg.span != 1 and seg.request_id in held] + with self._lock: + self._moved.update(one) + self._stale.update(rest) + + # The capture + + def capture_step(self, bucket: BucketKey, slot: int, step: SubmoduleStep) -> CapturedResidentStep: + """What a capture of ``bucket`` records around its forward. Refuses a + walk whose declared step isn't one token per row in one stream, each + row advancing alone: a captured step can't hold anything else.""" + if step is None: + raise ValueError(f"resident walk {bucket.graph_walk!r} on {self.node} declares no step") + segments = step.segments or next(iter(step.steps.values())).segments + labels = {seg.label for seg in segments} + if len(labels) != 1 or any(seg.span != 1 for seg in segments): + raise ValueError( + f"resident walk {bucket.graph_walk!r} on {self.node} declares " + f"{len(segments)} segments over labels {sorted(labels)}; a resident step is " + "one token per row in one stream" + ) + for key, resource_step in step.steps.items(): + refusal = next( + ( + name for name in ("pre_forks", "post_forks", "retention", "sink", "advance", "pos_ids", "fork_rids") + if getattr(resource_step, name, None) + ), + None, + ) + if refusal is not None: + raise ValueError( + f"resident walk {bucket.graph_walk!r} on {self.node}: resource {key!r} declares " + f"{refusal}, which a resident step can't carry" + ) + buffers = self.buffers(bucket, slot) + return CapturedResidentStep(bucket, slot, labels.pop(), buffers.slots, buffers) + + def gather(self, step: CapturedResidentStep) -> None: + for resource in self._owners: + resource.resident_gather(step) + + def commit_in_graph(self, step: CapturedResidentStep) -> None: + for resource in self._owners: + resource.resident_commit(step) + + +def build_planner( + node: str, walks: Iterable[str], resources: Mapping[str, Resource], + resource_nodes: Mapping[str, Collection[str]], device: torch.device, max_slots: int, +) -> tuple[ResidentPlanner | None, str]: + """A planner for ``node``, or None and why not. Every resource of the node + must be able to, since a capture is one graph: one host-planned resource + would need the row a resident step no longer stages.""" + walks = frozenset(walks) + if not walks: + return None, "its submodule names no resident walk" + for key, resource in resources.items(): + if len(resource_nodes.get(key, ())) > 1: + return None, f"resource {key!r} is shared by nodes {sorted(resource_nodes[key])}" + refusal = resource.resident_refusal() + if refusal is not None: + return None, f"resource {key!r}: {refusal}" + return ResidentPlanner(node, walks, resources, device, max_slots), "" diff --git a/mstar/engine/resources/sampler/resource.py b/mstar/engine/resources/sampler/resource.py index 9f75e654cd..6db7d6d49d 100644 --- a/mstar/engine/resources/sampler/resource.py +++ b/mstar/engine/resources/sampler/resource.py @@ -387,6 +387,10 @@ def _assert_penalty_inert(self, ctx: StepContext) -> None: "with the sampler's configs (a mid-request config update?)" ) + def resident_refusal(self) -> str | None: + # its masters already live by slot on the device, on numbers of its own + return None + @property def supports_preplan(self): return True diff --git a/mstar/engine/resources/step.py b/mstar/engine/resources/step.py index ef4c0a4412..456472ca92 100644 --- a/mstar/engine/resources/step.py +++ b/mstar/engine/resources/step.py @@ -80,6 +80,11 @@ class StepContext: # plan prefill attention even when every row is one token: a packed # capture's forward reads per-row token offsets force_prefill: bool = False + # the node's `ResidentPlanner` when this step's capture reads the resident + # tables: a resource then stages nothing per row (see `resources/resident.py`) + resident: Any = None + # what that planner staged for this step + resident_plan: Any = None def walk_of(self, rid: int) -> str: return self.request_walks.get(rid, self.graph_walk) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index f9b302fcfb..3cf75eb3a6 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -349,6 +349,10 @@ def _device_loopback(self) -> bool: and self.node_resources[SAMPLER].has_slot_masters ) + def resident_walks(self) -> frozenset[str]: + # with the tables staged outside the graph, `preprocess` reads positions + return frozenset({"decode"}) if ingraph_decode_rope() else frozenset() + def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: if graph_walk == "decode" and self._device_loopback(): return frozenset({"text_inputs"}) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index df195c697d..3ab28b77a6 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -864,6 +864,14 @@ def uniform_row_inputs( del graph_walk, has_loopback return None + def resident_walks(self) -> frozenset[str]: + """Walks whose captured steps may read their rows off the resident + tables (`mstar.engine.resources.resident`). A walk here runs one token + per row in one stream, with no forks or retention, and its `preprocess` + reads no planned position or KV location: the forward's graph fills + those in. The engine still refuses if a resource can't.""" + return frozenset() + def step_is_input_free(self, graph_walk: str) -> bool: """Whether ``declare_step`` for this walk reads nothing from the inputs but their spans, and every real row has the span of the template row. diff --git a/mstar/utils/h2d.py b/mstar/utils/h2d.py index 3bd2857f75..60567275bc 100644 --- a/mstar/utils/h2d.py +++ b/mstar/utils/h2d.py @@ -138,6 +138,39 @@ def copy_( mirror.record(key, values) + def copy_runs_(self, dsts: Sequence[torch.Tensor], values: np.ndarray) -> None: + """Fill each of ``dsts`` from consecutive ``values``, in order: one + non-blocking copy per destination, all out of one pinned buffer. On a + stream behind a running graph the copies queue there, and a buffer per + copy would have ``copy_`` wait out that graph once the ring wraps.""" + if not dsts: + return + capability = DeviceCapability.of(dsts[0].device) + if not capability.is_accelerator or not self._pinned: + at = 0 + for dst in dsts: + dst.view(-1).copy_(torch.as_tensor(values[at: at + dst.numel()], dtype=self.dtype)) + at += dst.numel() + return + with self._lock: + if len(values) > self._numel: + self._grow(len(values)) + i = self._next + self._next = (i + 1) % self._depth + ev = self._events[i] + if ev is not None: + ev.synchronize() + self._views[i][: len(values)] = values + at = 0 + for dst in dsts: + n = dst.numel() + dst.view(-1).copy_(self._bufs[i][at: at + n], non_blocking=True) + at += n + if ev is None: + ev = self._events[i] = torch.Event(device=dsts[0].device) + ev.record(torch.accelerator.current_stream(dsts[0].device)) + + def to_device_nonblocking( data: torch.Tensor | Sequence, device: torch.device | str, diff --git a/mstar/utils/knobs.py b/mstar/utils/knobs.py index 6f6e42bc42..a799164c82 100644 --- a/mstar/utils/knobs.py +++ b/mstar/utils/knobs.py @@ -20,6 +20,8 @@ "gc_gen0_threshold": "MSTAR_GC_GEN0_THRESHOLD", "pre_plan": "MSTAR_PRE_PLAN_SPEC", "step_thread_core": "MSTAR_STEP_THREAD_CORE", + "resident_planning": "MSTAR_RESIDENT_PLANNING", + "resident_slots": "MSTAR_RESIDENT_SLOTS", } # this process's ``knobs:`` block, as strings in the env vars' format @@ -44,6 +46,8 @@ def check(settings: Mapping[str, object] | None) -> dict[str, str]: for key, value in settings.items(): if key == "step_thread_core": out[key] = _step_thread_core_value(value) + elif key == "resident_planning": + out[key] = _resident_planning_value(value) elif isinstance(value, bool): out[key] = "1" if value else "0" elif isinstance(value, int): @@ -55,6 +59,14 @@ def check(settings: Mapping[str, object] | None) -> dict[str, str]: return out +def _resident_planning_value(value: object) -> str: + # yaml reads a bare ``off`` as False + text = "off" if value is False else str(value).strip().lower() + if text in ("auto", "off"): + return text + raise ValueError(f"`knobs: resident_planning` must be auto or off, got {value!r}") + + def _step_thread_core_value(value: object) -> str: # yaml reads a bare ``off`` as False if isinstance(value, bool): @@ -194,3 +206,21 @@ def step_thread_core() -> str | int: raise ValueError( f"MSTAR_STEP_THREAD_CORE must be auto, off or a CPU number, got {value!r}" ) + + +def resident_planning() -> str: + """``resident_planning`` / ``MSTAR_RESIDENT_PLANNING`` (default ``auto``): + a captured decode step reads its rows off per-request tables held on the + device instead of having them rebuilt on the plan thread + (``mstar.engine.resources.resident``). ``auto`` turns it on for a node + whose submodule names resident walks and every resource of which is + resident-capable, and says why at startup when one is not; ``off`` keeps + the host plan everywhere.""" + return _resident_planning_value(_setting("resident_planning", "auto")) + + +def resident_slots() -> int: + """``resident_slots`` / ``MSTAR_RESIDENT_SLOTS`` (default 4096): request + rows of the resident tables, allocated once at startup since a captured + graph holds their addresses.""" + return int(_setting("resident_slots", "4096")) diff --git a/test/modular/test_capture_budget.py b/test/modular/test_capture_budget.py index 8a409f7e0c..6469100604 100644 --- a/test/modular/test_capture_budget.py +++ b/test/modular/test_capture_budget.py @@ -72,6 +72,7 @@ def __init__(self, device: torch.device, steps: dict[BucketKey, CaptureCost]): self._device = device self._autocast_dtype = None self._submodule_name = "Thinker" + self._resident = None self._num_slots = 2 self._capture_costs = {} self._lent = set() @@ -788,6 +789,7 @@ def __init__(self, buckets: list[BucketKey], peer: list[bool]): self._capture_costs = {bucket: _cost(100) for bucket in buckets} self._device = torch.device("cpu") self._submodule_name = "Thinker" + self._resident = None self._num_slots = 1 self._buckets = {} self._comm_group = JointGroups(tp_group=_PeerGroup(peer), sp_group=CommGroup.trivial()) diff --git a/test/modular/test_engine_request_added.py b/test/modular/test_engine_request_added.py index 622dbb4a6e..bc0480a3f3 100644 --- a/test/modular/test_engine_request_added.py +++ b/test/modular/test_engine_request_added.py @@ -41,7 +41,7 @@ def _engine(*submodules): engine._resources = {} engine._loopback_steps = {} engine._submodules = { - name: SimpleNamespace(submodule=_Submodule(log, name, raises)) for name, raises in submodules + name: SimpleNamespace(submodule=_Submodule(log, name, raises), resident=None) for name, raises in submodules } return engine, log diff --git a/test/modular/test_h2d_to_device.py b/test/modular/test_h2d_to_device.py index bd6f1c18dc..f3a648adb1 100644 --- a/test/modular/test_h2d_to_device.py +++ b/test/modular/test_h2d_to_device.py @@ -7,11 +7,12 @@ import types +import numpy as np import pytest import torch from mstar.utils import h2d -from mstar.utils.h2d import to_device_nonblocking +from mstar.utils.h2d import PinnedStager, to_device_nonblocking @pytest.fixture(params=["cuda", "xpu"]) @@ -71,3 +72,21 @@ def spy_empty(*args, pin_memory=False, **kwargs): # a non-contiguous one is packed into a pinned buffer: ``.to`` would pack it pageable and sync assert (pinned_allocs == []) == contiguous and (got.data_ptr() == source.data_ptr()) == contiguous assert torch.equal(got, source) + + +def test_runs_fill_their_destinations_in_order(): + stager = PinnedStager(torch.long, numel=4) + store = torch.zeros(10, dtype=torch.long) + stager.copy_runs_([store[1:3], store[6:9]], [1, 2, 3, 4, 5]) + assert store.tolist() == [0, 1, 2, 0, 0, 0, 3, 4, 5, 0] + + +def test_runs_past_the_staging_buffer_grow_it(accelerator_device): + stager = PinnedStager(torch.long, numel=4) + store = torch.zeros(64, dtype=torch.long, device=accelerator_device) + values = np.arange(1, 33, dtype=np.int64) + stager.copy_runs_([store[:16], store[32:48]], values) + torch.accelerator.synchronize(accelerator_device) + assert store[:16].tolist() == values[:16].tolist() + assert store[16:32].tolist() == [0] * 16 + assert store[32:48].tolist() == values[16:].tolist() diff --git a/test/modular/test_resident_planning.py b/test/modular/test_resident_planning.py new file mode 100644 index 0000000000..c98c73f4be --- /dev/null +++ b/test/modular/test_resident_planning.py @@ -0,0 +1,794 @@ +"""Resident planning: a captured decode step reads its rows off device tables and +updates them itself, and what its graph computes matches what a host plan would +have staged. + +The resources here are the real ones on the CPU. ``gather`` and ``commit_in_graph`` +run the ops a capture records, so each step below is what a replay does; a second +set of resources plans the same requests on the host as the reference. +""" +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import pytest +import torch + +sys.path.insert(0, ".") + +from mstar.engine.resources.attn.config import AttentionConfig, AttentionSpec # noqa: E402 +from mstar.engine.resources.base import CGSlotKey, Resource # noqa: E402 +from mstar.engine.resources.kv import manager as manager_mod # noqa: E402 +from mstar.engine.resources.kv.config import KVLayout, KVReqConfig, KVSpec, KVStep, PagedKVConfig # noqa: E402 +from mstar.engine.resources.kv.manager import KVManager # noqa: E402 +from mstar.engine.resources.kv.plan import SINK_PAGE # noqa: E402 +from mstar.engine.resources.position.config import PositionConfig, PositionStep, PosScheme # noqa: E402 +from mstar.engine.resources.position.manager import RopeManager # noqa: E402 +from mstar.engine.resources.recurrent.config import ( # noqa: E402 + RecurrentBlockConfig, + RecurrentStateConfig, + RecurrentStep, +) +from mstar.engine.resources.recurrent.pool import SINK_SLOT, RecurrentStatePool # noqa: E402 +from mstar.engine.resources.resident import TRASH_SLOT, ResidentPlanner, build_planner # noqa: E402 +from mstar.engine.resources.spec import resolve_spec_dependencies # noqa: E402 +from mstar.engine.resources.step import BucketKey, Segment, SlotLease, StepContext, SubmoduleStep # noqa: E402 +from mstar.utils import knobs # noqa: E402 + +PAGE = 4 +WALK = "decode" +CPU = torch.device("cpu") + + +class _StubTransfer: + def __init__(self, transfer_engine_info, kv_cache, **kwargs): + del transfer_engine_info, kv_cache, kwargs + + def get_kv_transfer_info(self, **kwargs): + del kwargs + + def remove_request(self, request_id): + del request_id + + def cleanup(self): + pass + + +@pytest.fixture(autouse=True) +def _stub_transfer(monkeypatch): + monkeypatch.setattr(manager_mod, "KVTransferManager", _StubTransfer) + + +def _resources(cpu_offload_pages: int = 0, max_num_pages: int = 64): + kv = KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=512, + max_num_pages=max_num_pages, page_size=PAGE, cpu_offload_pages=cpu_offload_pages, + ), + name="kv", joint_comm_group=None, transfer_engine_info=None, device=CPU, dtype=torch.float32, + ) + kv._cg_max_seq_len = 64 + pool = RecurrentStatePool(device=CPU, config=RecurrentStateConfig( + num_layers=1, max_slots=9, blocks={"state": RecurrentBlockConfig(shape=(2,), dtype=torch.float32)}, + )) + rope = RopeManager(PositionConfig(kv_cache="kv"), CPU, max_seq_len=512, head_dim=8) + # page-locked memory needs a GPU, and nothing here copies out of it asynchronously + rope._pinned_pos_ids_buffer = lambda key, n: torch.empty(n, dtype=torch.long) + return {"kv": kv, "pool": pool, "rope": rope} + + +def _ingest(res: dict[str, Resource], rid: int) -> None: + res["kv"].ingest_request(rid, KVReqConfig()) + res["pool"].ingest_request(rid, None) + res["rope"].ingest_request(rid, None) + + +def _host_step(res, rids, spans, walk="prefill", advance=None): + """A host-planned step, eager, as a prefill or mixed step runs.""" + segments = tuple(Segment(r, "main", n) for r, n in zip(rids, spans, strict=True)) + ctx = StepContext(request_ids=tuple(rids), graph_walk=walk, slot=0, capture=False) + kv, pool = KVStep(segments=segments), RecurrentStep(segments=segments) + rope = PositionStep(segments=segments, advance=advance) + assert res["kv"].admit(kv, ctx).ok + assert res["pool"].admit(pool, ctx).ok + ctx.plan_results["kv"] = res["kv"].plan(kv, ctx) + res["pool"].plan(pool, ctx) + res["rope"].plan(rope, ctx) + res["kv"].commit(kv, ctx) + res["pool"].commit(pool, ctx) + res["rope"].commit(rope, ctx) + return SubmoduleStep(steps={"kv": kv}, segments=list(segments)), ctx + + +def _captured_step(res, planner: ResidentPlanner | None, rids, slot, bs, preplan=False, run=True): + """One decode step on a bucket of ``bs`` rows. With a planner, the replay of a + resident capture; without one, the host plan of the same step.""" + padded = tuple(rids) + tuple(-(i + 1) for i in range(bs - len(rids))) + bucket = BucketKey(graph_walk=WALK, bs=bs, num_tokens=bs) + lease = SlotLease(slot=slot, bucket=bucket) + ctx = StepContext(request_ids=tuple(rids), graph_walk=WALK, slot=slot, capture=False, slot_lease=lease) + ctx.set_padded_rids(padded) + segments = tuple(Segment(r, "main", 1) for r in padded) + kv, pool, rope = KVStep(segments=segments), RecurrentStep(segments=segments), PositionStep(segments=segments) + step = SubmoduleStep(steps={"kv": kv}, segments=list(segments)) + assert res["kv"].admit(kv, ctx).ok + assert res["pool"].admit(pool, ctx).ok + # a capture's first plan builds the static buffers on the host + first = StepContext(request_ids=tuple(padded), graph_walk=WALK, slot=slot, capture=True, slot_lease=lease) + if (bucket, slot) not in res.setdefault("captured", set()): + res["captured"].add((bucket, slot)) + first.plan_results["kv"] = res["kv"].plan(kv, first) + res["pool"].plan(pool, first) + res["rope"].plan(rope, first) + captured = None + if planner is not None: + captured = planner.capture_step(bucket, slot, step) + ctx.resident = planner + if planner is None or preplan: + ctx.is_preplan = preplan + if planner is not None: + planner.plan(step, ctx) + ctx.plan_results["kv"] = res["kv"].plan(kv, ctx) + res["pool"].plan(pool, ctx) + res["rope"].plan(rope, ctx) + ctx.is_preplan = False + out = SimpleNamespace(step=step, ctx=ctx, kv=kv, pool=pool, rope=rope, captured=captured) + if run: + _replay(res, planner, out) + return out + + +def _promote(res, planner, s): + ctx = s.ctx + planner.plan(s.step, ctx) + ctx.plan_results["kv"] = res["kv"].plan(s.kv, ctx) + res["pool"].plan(s.pool, ctx) + res["rope"].plan(s.rope, ctx) + + +def _replay(res, planner, s, promote=True): + """What the GPU thread does from the promote on.""" + ctx, bucket, lease = s.ctx, s.ctx.slot_lease.bucket, s.ctx.slot_lease + if planner is not None: + if promote: + _promote(res, planner, s) + planner.flush(ctx.resident_plan) + planner.gather(s.captured) + state = res["kv"]._static_plan_state(lease.slot, ("main", 0)) + ck = CGSlotKey(bucket=bucket, slot=lease.slot, label="main") + s.pages = state.token_to_page[: bucket.bs].tolist() + s.offsets = state.token_to_cache[: bucket.bs].tolist() + s.positions = res["rope"]._static_pos_ids[ck][: bucket.bs].tolist() + s.pool_slots = res["pool"]._cg_addressing[ck].slot_indices[: bucket.bs].tolist() + if planner is not None: + planner.commit_in_graph(s.captured) + res["kv"].commit(s.kv, ctx) + res["pool"].commit(s.pool, ctx) + res["rope"].commit(s.rope, ctx) + if planner is not None: + planner.commit(s.step, ctx) + return s + + +def _setup(prompts: dict[int, int], cpu_offload_pages: int = 0, max_num_pages: int = 64): + """The same requests prefilled on resident resources and on a host-planned reference.""" + live, ref = _resources(cpu_offload_pages, max_num_pages), _resources(cpu_offload_pages, max_num_pages) + planner = ResidentPlanner("n", {WALK}, {k: live[k] for k in ("kv", "pool", "rope")}, CPU, max_slots=8) + for res in (live, ref): + for rid, n in prompts.items(): + _ingest(res, rid) + _host_step(res, [rid], [n]) + return live, ref, planner + + +def _same(a, b, real: int): + assert a.pages[:real] == b.pages[:real] + assert a.offsets[:real] == b.offsets[:real] + assert a.positions[:real] == b.positions[:real] + assert a.pool_slots[:real] == b.pool_slots[:real] + + +def _row(planner: ResidentPlanner, live, rid: int): + tables = live["kv"]._resident + slot = planner.slot_of(rid) + n = int(tables.kv_len.view[slot, 0]) + cols = n // PAGE + (1 if n % PAGE else 0) + return n, tables.block_table.view[slot, :cols].tolist() + + +def _assert_row(planner, live, rid): + stream = live["kv"]._streams[rid]["main"] + n, pages = _row(planner, live, rid) + assert n == stream.stored_len + assert pages == stream.page_indices[: -(-n // PAGE)] + + +# ---- the rows ---- + +def test_a_new_row_is_written_whole_by_its_first_step(): + live, ref, planner = _setup({1: 9}) + assert planner.slot_of(1) is None + a = _captured_step(live, planner, [1], slot=0, bs=1) + b = _captured_step(ref, None, [1], slot=0, bs=1) + _same(a, b, 1) + _assert_row(planner, live, 1) + assert planner.slot_of(1) != TRASH_SLOT + + +def test_steps_with_crossings_match_the_host_plan_and_the_host_lengths(): + live, ref, planner = _setup({1: 3, 2: 5, 3: 8}) + rids = [1, 2, 3] + for t in range(40): + a = _captured_step(live, planner, rids, slot=t % 2, bs=4) + b = _captured_step(ref, None, rids, slot=t % 2, bs=4) + _same(a, b, 3) + tables = live["kv"]._resident + for rid in rids: + _assert_row(planner, live, rid) + slot = planner.slot_of(rid) + assert tables.host_len[slot] == live["kv"]._streams[rid]["main"].stored_len + assert tables.host_cols[slot] == live["kv"]._streams[rid]["main"].stored_len // PAGE + ( + 1 if live["kv"]._streams[rid]["main"].stored_len % PAGE else 0 + ) + + +def test_a_steady_batch_keeps_its_row(): + live, ref, planner = _setup({1: 1, 2: 1}) + del ref + _captured_step(live, planner, [1, 2], slot=0, bs=2) + buffers = planner.buffers(BucketKey(WALK, 2, 2), 0) + before = buffers.slots.clone() + last = planner._last_slots + _captured_step(live, planner, [1, 2], slot=0, bs=2) + assert planner._last_slots is last + assert torch.equal(buffers.slots, before) + + +def test_padding_rows_read_the_trash_row_and_leave_it_alone(): + live, ref, planner = _setup({1: 6}) + for t in range(9): + a = _captured_step(live, planner, [1], slot=t % 2, bs=4) + b = _captured_step(ref, None, [1], slot=t % 2, bs=4) + _same(a, b, 1) + assert a.pages[1:] == [SINK_PAGE] * 3 + assert a.offsets[1:] == [0] * 3 + assert a.pool_slots[1:] == [SINK_SLOT] * 3 + assert a.positions[1:] == [0] * 3 + tables = live["kv"]._resident + assert int(tables.kv_len.view[TRASH_SLOT, 0]) == 0 + assert tables.block_table.view[TRASH_SLOT].eq(SINK_PAGE).all() + assert int(live["pool"]._resident_table.view[TRASH_SLOT, 0]) == SINK_SLOT + + +def test_a_mixed_step_leaves_the_next_decode_step_a_ready_row(): + live, ref, planner = _setup({1: 3}) + for t in range(3): + _captured_step(live, planner, [1], slot=t % 2, bs=1) + _captured_step(ref, None, [1], slot=t % 2, bs=1) + # a request's prompt beside the decode row, on the host plan; the decode row crosses here + for res in (live, ref): + _ingest(res, 2) + step, ctx = _host_step(res, [1, 2], [1, 7], walk="mixed") + if res is live: + planner.commit(step, ctx) + # and again, so the decode row is two columns behind + for _ in range(5): + for res in (live, ref): + step, ctx = _host_step(res, [1, 2], [1, 1], walk="mixed") + if res is live: + planner.commit(step, ctx) + a = _captured_step(live, planner, [1, 2], slot=0, bs=2) + b = _captured_step(ref, None, [1, 2], slot=0, bs=2) + _same(a, b, 2) + _assert_row(planner, live, 1) + _assert_row(planner, live, 2) + + +def test_an_abandoned_plan_that_rewrote_a_row_behind_by_columns_writes_it_again(): + live, ref, planner = _setup({1: 3}) + a = _captured_step(live, planner, [1], slot=0, bs=1) + b = _captured_step(ref, None, [1], slot=0, bs=1) + _same(a, b, 1) + # two page boundaries go by on the host + for res in (live, ref): + for _ in range(6): + step, ctx = _host_step(res, [1], [1], walk="mixed") + if res is live: + planner.commit(step, ctx) + assert 1 in planner._moved + staged = _captured_step(live, planner, [1], slot=1, bs=1, preplan=True, run=False) + assert 1 in staged.ctx.resident_plan.taken + for key in ("kv", "pool", "rope"): + live[key].clear_preplan() + c = _captured_step(live, planner, [1], slot=0, bs=1) + d = _captured_step(ref, None, [1], slot=0, bs=1) + _same(c, d, 1) + _assert_row(planner, live, 1) + + +def test_a_position_offset_carries_through_a_row(): + """A vision prompt advances positions by less than its tokens.""" + live, ref, planner = _setup({}) + for res in (live, ref): + _ingest(res, 1) + _host_step(res, [1], [10], walk="prefill_vision", advance=(6,)) + a = _captured_step(live, planner, [1], slot=0, bs=1) + b = _captured_step(ref, None, [1], slot=0, bs=1) + _same(a, b, 1) + assert a.positions[0] == 6 + + +def test_a_chunked_prefill_row_is_whole_at_its_first_decode_step(): + live, ref, planner = _setup({}) + for res in (live, ref): + _ingest(res, 1) + for chunk in (8, 8, 3): + step, ctx = _host_step(res, [1], [chunk], walk="prefill") + if res is live: + planner.commit(step, ctx) + a = _captured_step(live, planner, [1], slot=0, bs=1) + b = _captured_step(ref, None, [1], slot=0, bs=1) + _same(a, b, 1) + _assert_row(planner, live, 1) + + +def test_more_chunks_after_decode_rewrite_the_row(): + live, ref, planner = _setup({1: 2}) + _captured_step(live, planner, [1], slot=0, bs=1) + _captured_step(ref, None, [1], slot=0, bs=1) + for res in (live, ref): + step, ctx = _host_step(res, [1], [9], walk="prefill") + if res is live: + planner.commit(step, ctx) + assert 1 in planner._stale + a = _captured_step(live, planner, [1], slot=1, bs=1) + b = _captured_step(ref, None, [1], slot=1, bs=1) + _same(a, b, 1) + _assert_row(planner, live, 1) + assert not planner._stale + + +def test_a_fork_marks_the_row_stale(): + live, ref, planner = _setup({1: 9}) + del ref + _captured_step(live, planner, [1], slot=0, bs=1) + kv = live["kv"] + kv._touch_plan_cache(1) + assert 1 in planner._stale + + +def test_a_row_whose_pages_moved_is_rewritten(): + """What a reload leaves: the stream on other pages. Its hook is `_touch_plan_cache(rid)`; + `test_real_offload_and_reload_mark_the_row` runs the real one on a GPU.""" + live, ref, planner = _setup({1: 9}) + del ref + _captured_step(live, planner, [1], slot=0, bs=1) + kv = live["kv"] + stream = kv._streams[1]["main"] + moved = kv._arena.acquire(len(stream.page_indices)) + kv._arena.release(stream.page_indices) + stream.page_indices = moved + kv._touch_plan_cache(1) + assert 1 in planner._stale + s = _captured_step(live, planner, [1], slot=1, bs=1) + assert s.pages[0] == moved[-1] + _assert_row(planner, live, 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the host page pool is page-locked") +def test_real_offload_and_reload_mark_the_row(): + kv = KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=64, max_num_pages=16, page_size=PAGE, + cpu_offload_pages=16, + ), + name="kv", joint_comm_group=None, transfer_engine_info=None, + device=torch.device("cuda"), dtype=torch.float32, + ) + planner = ResidentPlanner("n", {WALK}, {"kv": kv}, torch.device("cuda"), max_slots=2) + kv.ingest_request(1, KVReqConfig()) + segments = (Segment(1, "main", 9),) + ctx = StepContext(request_ids=(1,), graph_walk="prefill", slot=0, capture=False) + step = KVStep(segments=segments) + assert kv.admit(step, ctx).ok + ctx.plan_results["kv"] = kv.plan(step, ctx) + kv.commit(step, ctx) + planner._slot_of[1] = 1 + assert kv.offload(1) > 0 + assert 1 in planner._stale + planner._stale.clear() + torch.cuda.synchronize() + assert kv.reload(1) + assert 1 in planner._stale + + +def test_an_abandoned_preplan_writes_nothing_and_gives_its_rows_back(): + live, ref, planner = _setup({1: 3}) + a = _captured_step(live, planner, [1], slot=0, bs=1) + b = _captured_step(ref, None, [1], slot=0, bs=1) + _same(a, b, 1) + # a host step dirties the row, a pre-plan takes it, and the step never runs + for res in (live, ref): + step, ctx = _host_step(res, [1], [6], walk="prefill") + if res is live: + planner.commit(step, ctx) + table_before = planner.store.clone() + staged = _captured_step(live, planner, [1], slot=1, bs=1, preplan=True, run=False) + assert staged.ctx.resident_plan.taken == {1} + assert torch.equal(planner.store, table_before) + live["kv"].clear_preplan() + live["pool"].clear_preplan() + live["rope"].clear_preplan() + # the next plan is another step: the row comes back, and is written whole then + c = _captured_step(live, planner, [1], slot=0, bs=1) + d = _captured_step(ref, None, [1], slot=0, bs=1) + _same(c, d, 1) + _assert_row(planner, live, 1) + + +def test_a_new_row_abandoned_before_its_first_step_is_still_written(): + live, ref, planner = _setup({1: 6}) + staged = _captured_step(live, planner, [1], slot=0, bs=1, preplan=True, run=False) + assert planner.slot_of(1) is not None + for key in ("kv", "pool", "rope"): + live[key].clear_preplan() + del staged + c = _captured_step(live, planner, [1], slot=1, bs=1) + d = _captured_step(ref, None, [1], slot=1, bs=1) + _same(c, d, 1) + _assert_row(planner, live, 1) + + +def test_a_promoted_preplan_reuses_its_stage(): + live, ref, planner = _setup({1: 3}) + del ref + _captured_step(live, planner, [1], slot=0, bs=1) + s = _captured_step(live, planner, [1], slot=1, bs=1, preplan=True, run=False) + plan = s.ctx.resident_plan + assert planner._staged is plan + _replay(live, planner, s) + assert planner._staged is None + assert s.ctx.resident_plan is plan + _assert_row(planner, live, 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the kernels run on CUDA") +def test_the_captured_kernels_match_the_torch_ops(): + from mstar.engine.resources.kv.resident import KVResidentTables + from mstar.engine.resources.resident import advance_column, gather_column, scatter_cells + + def run(device): + store = torch.arange(100, 100 + 7 * 20, device=device, dtype=torch.long) + slots = torch.tensor([3, 1, 0, 0], device=device) + cells = torch.tensor([[7, 0, 15, 0], [55, 0, 66, 0]], device=device) + scatter_cells(store, cells) + out = torch.zeros(4, dtype=torch.long, device=device) + small = torch.zeros(4, dtype=torch.int32, device=device) + gather_column(store, slots, 7, 1, out) + gather_column(store, slots, 7, 2, small, plus=out) + advance_column(store, slots, 7, 1) + tables = SimpleNamespace( + planner=SimpleNamespace(store=store), page_size=4, + block_table=SimpleNamespace(stride=7, col=3), kv_len=SimpleNamespace(stride=7, col=1), + ) + store[slots * 7 + 1] = torch.tensor([9, 5, 0, 0], device=device) + pages = torch.zeros(4, dtype=torch.long, device=device) + offsets = torch.zeros(4, dtype=torch.long, device=device) + step = SimpleNamespace( + slots=slots, bucket=SimpleNamespace(bs=4), buffers=SimpleNamespace( + cells=torch.zeros((2, 8), dtype=torch.long, device=device), kv_len=out, + ), + ) + KVResidentTables.gather(tables, step, pages, offsets) + return [t.tolist() for t in (store, out, small, pages, offsets)] + + assert run(torch.device("cuda")) == run(CPU) + + +# ---- slots ---- + +def test_a_freed_slot_waits_for_the_last_step_that_read_it(): + live, ref, planner = _setup({1: 3}) + del ref + # step N reads the slot and is still running: planned, not committed + s = _captured_step(live, planner, [1], slot=0, bs=1, preplan=True, run=False) + _promote(live, planner, s) + slot = s.ctx.resident_plan.slots[0] + planner.release(1) + planner._free.clear() + assert (int(slot), s.ctx.resident_plan.seq) in planner._cooling + with pytest.raises(RuntimeError, match="slots are held"): + with planner._lock: + planner._take_slot() + # commit lets the slot go; its late advance landed on the old owner's row + _replay(live, planner, s, promote=False) + with planner._lock: + assert planner._take_slot() == slot + + +def test_a_late_commit_cannot_touch_the_next_owner(): + live, ref, planner = _setup({1: 3, 2: 5}) + del ref + s = _captured_step(live, planner, [1], slot=0, bs=1, preplan=True, run=False) + old = int(s.ctx.resident_plan.slots[0]) + planner.release(1) + for key in ("kv", "pool", "rope"): + live[key].clear_preplan() + a = _captured_step(live, planner, [2], slot=1, bs=1, preplan=True, run=False) + assert int(a.ctx.resident_plan.slots[0]) != old + + +def test_release_purges_the_handle(): + live, ref, planner = _setup({1: 3}) + del ref + _captured_step(live, planner, [1], slot=0, bs=1) + planner.mark_stale(1) + planner.release(1) + assert planner.slot_of(1) is None + assert not planner._stale + assert planner._last_rids is None + + +def test_a_stage_over_a_released_row_does_not_promote(): + live, ref, planner = _setup({1: 3}) + del ref + s = _captured_step(live, planner, [1], slot=0, bs=1, preplan=True, run=False) + old = s.ctx.resident_plan + planner.release(1) + assert planner._staged is old + for key in ("kv", "pool", "rope"): + live[key].clear_preplan() + _replay(live, planner, s) + assert s.ctx.resident_plan is not old + assert planner.slot_of(1) is not None + + +def test_a_recycled_handle_gets_a_fresh_row(): + live, ref, planner = _setup({1: 9}) + del ref + _captured_step(live, planner, [1], slot=0, bs=1) + kv = live["kv"] + kv.remove_request(1) + live["pool"].remove_request(1) + live["rope"].remove_request(1) + planner.release(1) + _ingest(live, 1) + _host_step(live, [1], [2]) + _captured_step(live, planner, [1], slot=1, bs=1) + _assert_row(planner, live, 1) + + +def test_running_out_of_slots_says_how_to_fix_it(): + live, ref, planner = _setup({}) + del ref + planner._free.clear() + _ingest(live, 1) + _host_step(live, [1], [2]) + with pytest.raises(RuntimeError, match="resident_slots"): + _captured_step(live, planner, [1], slot=0, bs=1) + + +# ---- refusals ---- + +class _Spare(Resource): + @classmethod + def build(cls, spec, info): + return cls() + + +def _build(resources, nodes=None, walks=(WALK,)): + return build_planner("n", walks, resources, nodes or {}, CPU, 8) + + +def test_every_resident_resource_passes(): + res = _resources() + planner, why = _build({k: res[k] for k in ("kv", "pool", "rope")}) + assert planner is not None and why == "" + + +def test_no_walks_no_planner(): + res = _resources() + planner, why = _build({"kv": res["kv"]}, walks=()) + assert planner is None and "no resident walk" in why + + +def test_a_resource_without_a_table_refuses(): + planner, why = _build({"other": _Spare()}) + assert planner is None and "_Spare keeps no resident table" in why + + +def test_a_pool_shared_by_nodes_refuses(): + res = _resources() + planner, why = _build({"kv": res["kv"]}, nodes={"kv": ("a", "b")}) + assert planner is None and "shared by nodes" in why + + +def test_a_pool_read_across_nodes_refuses(): + res = _resources() + res["kv"]._needs_remote_transfer = True + planner, why = _build({"kv": res["kv"]}) + assert planner is None and "across nodes" in why + + +def test_the_mla_layout_refuses(): + """GLM-5.2's latent pages have no K/V axis to write through.""" + kv = KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=512, max_num_pages=64, page_size=PAGE, + layout=KVLayout.MLA, kv_lora_rank=4, qk_rope_head_dim=4, + ), + name="kv", joint_comm_group=None, transfer_engine_info=None, device=CPU, dtype=torch.float32, + ) + planner, why = _build({"kv": kv}) + assert planner is None and "MLA" in why + + +def test_a_sliding_window_pool_refuses(): + cfg = PagedKVConfig(num_layers=3, num_kv_heads=1, head_dim=8, max_seq_len=512, max_num_pages=64, page_size=8) + resolve_spec_dependencies([ + KVSpec(resource_key="kv", nodes={"n"}, config=cfg), + AttentionSpec(resource_key="local", nodes={"n"}, config=AttentionConfig("kv", sliding_window=6, layers=(1, 2))), + AttentionSpec(resource_key="full", nodes={"n"}, config=AttentionConfig("kv")), + ]) + kv = KVManager( + cfg=cfg, name="kv", joint_comm_group=None, transfer_engine_info=None, device=CPU, dtype=torch.float32, + ) + planner, why = _build({"kv": kv}) + assert planner is None and "sliding-window" in why + + +def test_block_positions_refuse(): + rope = RopeManager(PositionConfig(kv_cache="kv", scheme=PosScheme.BLOCK), CPU, max_seq_len=512, head_dim=8) + assert "doesn't advance" in rope.resident_refusal() + + +def test_a_pool_without_a_sink_refuses(): + pool = RecurrentStatePool(device=CPU, config=RecurrentStateConfig( + num_layers=1, max_slots=4, blocks={"state": RecurrentBlockConfig(shape=(2,), dtype=torch.float32)}, + disable_sink_slot=True, + )) + assert "sink" in pool.resident_refusal() + + +@pytest.mark.parametrize("module,name", [ + ("mstar.engine.resources.attn.cross", "CrossAttentionManager"), + ("mstar.engine.resources.kv.bounded.manager", "BoundedKVManager"), + ("mstar.engine.resources.attn.dense", "DenseAttentionManager"), + ("mstar.engine.resources.attn.flashinfer_mla", "FlashInferMLAManager"), + ("mstar.engine.resources.linear_attn.mamba2", "Mamba2Manager"), + ("mstar.engine.resources.kv.ring.manager", "RingKVManager"), +]) +def test_resources_with_no_table_keep_the_base_refusal(module, name): + """These stay on the host plan: none overrides the base's refusal.""" + import importlib + + cls = getattr(importlib.import_module(module), name) + assert cls.resident_refusal is Resource.resident_refusal + + +def test_flashinfer_refuses_a_sliding_window(): + from mstar.engine.resources.attn.flashinfer import FlashInferManager + + manager = FlashInferManager.__new__(FlashInferManager) + manager._window_left = 5 + assert "sliding window" in manager.resident_refusal() + manager._window_left = -1 + assert manager.resident_refusal() is None + + +@pytest.mark.parametrize("field,value", [ + ("pre_forks", (("a", "b"),)), ("post_forks", (("a", "b"),)), ("sink", frozenset({(1, "main")})), +]) +def test_a_resident_walk_can_not_declare_forks_or_a_sink(field, value): + _, _, planner = _setup({}) + segments = [Segment(1, "main", 1)] + kv = KVStep(segments=tuple(segments), **{field: value}) + with pytest.raises(ValueError, match=field): + planner.capture_step(BucketKey(WALK, 1, 1), 0, SubmoduleStep(steps={"kv": kv}, segments=segments)) + + +def test_a_resident_walk_is_one_token_a_row_in_one_stream(): + _, _, planner = _setup({}) + two = [Segment(1, "main", 2)] + with pytest.raises(ValueError, match="one token per row"): + planner.capture_step(BucketKey(WALK, 1, 2), 0, SubmoduleStep(steps={}, segments=two)) + labels = [Segment(1, "a", 1), Segment(1, "b", 1)] + with pytest.raises(ValueError, match="one token per row"): + planner.capture_step(BucketKey(WALK, 2, 2), 0, SubmoduleStep(steps={}, segments=labels)) + + +# ---- the engine ---- + +def _engine(walks, resources, device=CPU): + from mstar.engine.engine import Engine + from mstar.engine.resources.runner import StepRunner + + engine = Engine.__new__(Engine) + engine._device = device + engine._runner = StepRunner(resources) + engine._resource_nodes = {key: frozenset({"n"}) for key in resources} + submodule = SimpleNamespace(resident_walks=lambda: frozenset(walks)) + engine._submodules = {"n": SimpleNamespace(submodule=submodule, resources=resources, resident=None)} + return engine, engine._submodules["n"] + + +def test_the_engine_builds_a_planner_for_a_capable_node(monkeypatch): + monkeypatch.setattr(knobs, "resident_slots", lambda: 8) + res = _resources() + engine, mgmt = _engine({WALK}, {k: res[k] for k in ("kv", "pool", "rope")}) + engine._build_resident_planner("n") + assert isinstance(mgmt.resident, ResidentPlanner) + assert mgmt.resident.covers(BucketKey(WALK, 2, 2)) + assert not mgmt.resident.covers(BucketKey("prefill_text", 2, 64)) + + +def test_the_knob_turns_it_off(monkeypatch): + monkeypatch.setenv("MSTAR_RESIDENT_PLANNING", "off") + res = _resources() + engine, mgmt = _engine({WALK}, {"kv": res["kv"]}) + engine._build_resident_planner("n") + assert mgmt.resident is None + assert res["kv"]._resident is None + + +def test_a_refused_node_stays_on_the_host_and_says_why(monkeypatch, caplog): + monkeypatch.setattr(knobs, "resident_slots", lambda: 8) + res = _resources() + engine, mgmt = _engine({WALK}, {"kv": res["kv"], "other": _Spare()}) + with caplog.at_level("WARNING"): + engine._build_resident_planner("n") + assert mgmt.resident is None + assert "_Spare keeps no resident table" in caplog.text + assert res["kv"]._resident is None + + +def test_a_node_with_no_resident_walk_is_left_alone(caplog): + res = _resources() + engine, mgmt = _engine((), {"kv": res["kv"]}) + with caplog.at_level("WARNING"): + engine._build_resident_planner("n") + assert mgmt.resident is None + assert not caplog.text + + +def test_an_eager_step_marks_its_rows(): + live, ref, planner = _setup({1: 3}) + del ref + _captured_step(live, planner, [1], slot=0, bs=1) + step, ctx = _host_step(live, [1], [1], walk="decode") + planner.commit(step, ctx) + assert 1 in planner._moved + assert _captured_step(live, planner, [1], slot=1, bs=1).ctx.resident_plan.moved == {1} + _assert_row(planner, live, 1) + + +def test_qwen35_decode_needs_the_rope_inside_the_graph(monkeypatch): + from mstar.model.qwen3_5 import submodules + + monkeypatch.setattr(submodules, "ingraph_decode_rope", lambda: False) + llm = next( + cls for cls in vars(submodules).values() + if isinstance(cls, type) and cls.__module__ == submodules.__name__ and "resident_walks" in vars(cls) + ) + assert llm.resident_walks(None) == frozenset() + monkeypatch.setattr(submodules, "ingraph_decode_rope", lambda: True) + assert llm.resident_walks(None) == {"decode"} + + +# ---- the knob ---- + +def test_the_knob_takes_auto_and_off(): + assert knobs.check({"resident_planning": "auto"}) == {"resident_planning": "auto"} + assert knobs.check({"resident_planning": False}) == {"resident_planning": "off"} + with pytest.raises(ValueError, match="auto or off"): + knobs.check({"resident_planning": "maybe"}) + with pytest.raises(ValueError, match="auto or off"): + knobs.check({"resident_planning": True}) + + +def test_the_knob_defaults_to_auto(monkeypatch): + monkeypatch.delenv("MSTAR_RESIDENT_PLANNING", raising=False) + knobs.configure(None) + assert knobs.resident_planning() == "auto" + monkeypatch.setenv("MSTAR_RESIDENT_PLANNING", "off") + assert knobs.resident_planning() == "off" From dc2293637ba888fdbe80fc977ab3ef9850c76c0f Mon Sep 17 00:00:00 2001 From: merceod Date: Sun, 1 Nov 2026 19:55:22 +0000 Subject: [PATCH 02/15] resident: laps follow #852's plan phase, tests where the two meet The planner's stage runs before runner.plan sets plan_phase, so a fresh stage's laps landed under res.plan. A resident step doesn't read or write the recurrent addressing memo; the cached KV plan's np.insert still feeds fa2's attention plan on a resident step. --- mstar/engine/engine.py | 11 ++++-- mstar/engine/resources/recurrent/pool.py | 3 +- mstar/engine/resources/resident.py | 2 +- test/modular/test_resident_planning.py | 50 ++++++++++++++++++++++++ 4 files changed, 60 insertions(+), 6 deletions(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index fc619597b9..49212b49ae 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -1710,8 +1710,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: @@ -1722,10 +1727,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 @@ -2596,6 +2598,7 @@ def _pre_plan_for_batch( ctx = batch.step_context with stream if stream is not None else nullcontext(): 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) diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index cee3099658..cc011671d1 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -584,7 +584,8 @@ def plan(self, step: RecurrentStep, ctx: StepContext): "planning a different step ahead" ) if ctx.resident is not None: - # the replay's kernels read the buffers its graph fills from the table + # the replay's kernels read the buffers its graph fills from the + # table; the addressing memo is left to host-planned steps return self._current if self._preplanned: # Promotion: apply the fork copies staging left undone; see diff --git a/mstar/engine/resources/resident.py b/mstar/engine/resources/resident.py index 94ed1488fd..b155d0def6 100644 --- a/mstar/engine/resources/resident.py +++ b/mstar/engine/resources/resident.py @@ -354,7 +354,7 @@ def _restore(self, plan: ResidentPlan) -> None: def _build(self, rids: tuple[int, ...], key: tuple, label: str, ctx) -> ResidentPlan: lease = ctx.slot_lease - laps = PhaseLaps("res.pre_plan.resident." if ctx.is_preplan else "res.plan.resident.") + laps = PhaseLaps(f"{ctx.plan_phase}.resident.") with self._lock: if rids == self._last_rids: slots = self._last_slots diff --git a/test/modular/test_resident_planning.py b/test/modular/test_resident_planning.py index c98c73f4be..ab09200eb9 100644 --- a/test/modular/test_resident_planning.py +++ b/test/modular/test_resident_planning.py @@ -16,6 +16,7 @@ sys.path.insert(0, ".") +from mstar.engine.resources import resident as resident_mod # noqa: E402 from mstar.engine.resources.attn.config import AttentionConfig, AttentionSpec # noqa: E402 from mstar.engine.resources.base import CGSlotKey, Resource # noqa: E402 from mstar.engine.resources.kv import manager as manager_mod # noqa: E402 @@ -284,6 +285,55 @@ def test_a_mixed_step_leaves_the_next_decode_step_a_ready_row(): _assert_row(planner, live, 2) +def test_a_resident_step_leaves_the_addressing_memo_to_host_steps(): + live, ref, planner = _setup({1: 3, 2: 5}) + del ref + pool = live["pool"] + for slot in (0, 1): + _captured_step(live, planner, [1, 2], slot=slot, bs=2) + _host_step(live, [1, 2], [1, 1], walk="mixed") + memo, version = pool._addressing_memo["main"], pool._version + for t in range(6): + s = _captured_step(live, planner, [1, 2], slot=t % 2, bs=2) + assert s.pool_slots == [pool._slots[r]["main"].index for r in (1, 2)] + assert pool._addressing_memo["main"] is memo + assert pool._version == version + _host_step(live, [1, 2], [1, 1], walk="mixed") + addressing = pool._current["main"] + assert addressing.host_slot_indices is memo[1] + assert addressing.host_slot_indices == [pool._slots[r]["main"].index for r in (1, 2)] + + +def test_a_resident_step_keeps_the_cached_kv_plan_attention_reads(): + live, ref, planner = _setup({1: 3, 2: 5, 3: 8}) + for res in (live, ref): + res["kv"]._plan_cache_mode = 1 + rids = [1, 2, 3] + pages = set() + for t in range(40): + a = _captured_step(live, planner, rids, slot=t % 2, bs=4) + b = _captured_step(ref, None, rids, slot=t % 2, bs=4) + _same(a, b, 3) + x, y = (s.ctx.plan_results["kv"]["main"].cpu_indptrs for s in (a, b)) + for p, q in zip(x[:5], y[:5], strict=True): + assert torch.equal(torch.as_tensor(p), torch.as_tensor(q)) + pages.add(a.pages[0]) + assert len(pages) > 1 + assert live["kv"].plan_cache_stats["hit"] > 20 + + +def test_a_stage_is_timed_under_the_step_plan_phase(monkeypatch): + live, ref, planner = _setup({1: 3}) + del ref + seen = [] + laps = resident_mod.PhaseLaps + monkeypatch.setattr(resident_mod, "PhaseLaps", lambda prefix: seen.append(prefix) or laps(prefix)) + s = _captured_step(live, planner, [1], slot=0, bs=1, run=False) + s.ctx.plan_phase = "res.plan_fresh" + _replay(live, planner, s) + assert seen == ["res.plan_fresh.resident."] + + def test_an_abandoned_plan_that_rewrote_a_row_behind_by_columns_writes_it_again(): live, ref, planner = _setup({1: 3}) a = _captured_step(live, planner, [1], slot=0, bs=1) From 147c6f70b6a96bad5ee7ceedc760266a4680f11f Mon Sep 17 00:00:00 2001 From: merceod Date: Sun, 1 Nov 2026 22:03:27 +0000 Subject: [PATCH 03/15] tests: the chunked prefill engine stub knows about the planner --- test/modular/test_chunked_prefill.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/modular/test_chunked_prefill.py b/test/modular/test_chunked_prefill.py index 238d03bd24..df06999139 100644 --- a/test/modular/test_chunked_prefill.py +++ b/test/modular/test_chunked_prefill.py @@ -206,7 +206,7 @@ def _engine_with(submodule, policies=None): from mstar.engine.engine import Engine engine = Engine.__new__(Engine) - engine._submodules = {"LLM": SimpleNamespace(submodule=submodule, resources={})} + engine._submodules = {"LLM": SimpleNamespace(submodule=submodule, resources={}, resident=None)} engine._keyed_walks = {} engine._resources = {} engine._chunking = ChunkingManagement(engine._submodules) From fa20064631c3e09d3081ebbfa1fb0aa7df4b3782 Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 2 Nov 2026 01:22:02 +0000 Subject: [PATCH 04/15] resident: a refused admit drops the stage, as the runner's does #980 has the runner drop its stage when an admit is refused, since the unwind gave back the pages it planned into. The planner's stage names the same pages in its crossing cells, so with no commit in between a retry over the same rows promoted cells pointing at a page the row no longer held. The KV pool now drops it wherever it gives a stage's pages back. --- mstar/engine/resources/kv/manager.py | 8 ++++++++ mstar/engine/resources/resident.py | 8 ++++++++ test/modular/test_resident_planning.py | 23 +++++++++++++++++++++++ 3 files changed, 39 insertions(+) diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index 116f723991..fd09f9fd6b 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -2128,6 +2128,7 @@ def rollback_admit(self, step: KVStep, ctx: StepContext) -> None: since, the tail can no longer be safely removed. """ self._touch_plan_cache() + self._drop_resident_stage() with self._lock: # No owner check: every rollback finishes before `commit_done` is # released, and the only pre-admit (`Worker._preplan_spec`) waits on @@ -2496,6 +2497,12 @@ def _touch_plan_cache(self, rid: str | None = None, row: bool = True) -> None: if row and rid is not None and self._resident is not None: self._resident.planner.mark_stale(rid) + def _drop_resident_stage(self) -> None: + """Pages a staged step reserved go back: the resident stage that names + them goes too, as the runner's own does.""" + if self._resident is not None: + self._resident.planner.drop_stage() + def _split_padding(self, step: KVStep, ctx: StepContext): """``(n_real, (n_pad, span))`` of a step whose real segments come first and whose padding segments (a replay's dummy rows) all share @@ -2974,6 +2981,7 @@ def _on_sink(ctx: StepContext, rid) -> bool: def clear_preplan(self, promoted: bool = False): if not promoted: self._touch_plan_cache() + self._drop_resident_stage() # the staged step is not going to run, so undo what it did to live # state. the pre-forks need no rewinding — staging only read them — # but admit's reservations are real pages and real streams diff --git a/mstar/engine/resources/resident.py b/mstar/engine/resources/resident.py index b155d0def6..20da91f2a4 100644 --- a/mstar/engine/resources/resident.py +++ b/mstar/engine/resources/resident.py @@ -346,6 +346,14 @@ def plan(self, step: SubmoduleStep, ctx) -> None: if ctx.is_preplan: self._staged = plan + def drop_stage(self) -> None: + """The pages a pre-plan's stage names were given back (a refused admit, + an abandoned pre-plan): a retry over the same rows reserves its own, so + it plans afresh instead of promoting cells that name the old ones.""" + staged, self._staged = self._staged, None + if staged is not None: + self._restore(staged) + def _restore(self, plan: ResidentPlan) -> None: """Hand back the rows of a plan that won't run: it wrote nothing.""" with self._lock: diff --git a/test/modular/test_resident_planning.py b/test/modular/test_resident_planning.py index ab09200eb9..db9496cdc9 100644 --- a/test/modular/test_resident_planning.py +++ b/test/modular/test_resident_planning.py @@ -476,6 +476,29 @@ def test_an_abandoned_preplan_writes_nothing_and_gives_its_rows_back(): _assert_row(planner, live, 1) +def test_a_refused_admit_drops_the_stage_naming_its_pages(): + """The unwind gives back the page a pre-plan staged a crossing into. With no + commit in between, a retry over the same rows reserves another page, and it + must not promote the stage's cell naming the old one.""" + live, ref, planner = _setup({1: 3, 2: 4}) + for res, p in ((live, planner), (ref, None)): + _captured_step(res, p, [1], slot=0, bs=1) + # row 1 holds 4 tokens: the next one opens a page, which the pre-plan stages + s = _captured_step(res, p, [1], slot=1, bs=1, preplan=True, run=False) + # refused, as `StepRunner.admit` unwinds it + res["pool"].rollback_admit(s.pool, s.ctx) + res["kv"].rollback_admit(s.kv, s.ctx) + for key in ("kv", "pool", "rope"): + res[key].clear_preplan() + # a teardown frees pages before the retry, so its admit draws another + for key in ("kv", "pool", "rope"): + res[key].remove_request(2) + a = _captured_step(live, planner, [1], slot=1, bs=1) + b = _captured_step(ref, None, [1], slot=1, bs=1) + _same(a, b, 1) + _assert_row(planner, live, 1) + + def test_a_new_row_abandoned_before_its_first_step_is_still_written(): live, ref, planner = _setup({1: 6}) staged = _captured_step(live, planner, [1], slot=0, bs=1, preplan=True, run=False) From de9824c3228e13ebcaf4961a8877ea9c032df4d2 Mon Sep 17 00:00:00 2001 From: merceod Date: Sun, 1 Nov 2026 19:57:13 +0000 Subject: [PATCH 05/15] fmha reads a row's used columns off the resident table --- mstar/engine/resources/attn/flashinfer.py | 13 ++++++++++ mstar/engine/resources/attn/wrappers.py | 30 +++++++++++++++++++++++ mstar/engine/resources/kv/manager.py | 3 +++ mstar/engine/resources/kv/resident.py | 1 + mstar/engine/resources/resident.py | 9 +++++++ 5 files changed, 56 insertions(+) diff --git a/mstar/engine/resources/attn/flashinfer.py b/mstar/engine/resources/attn/flashinfer.py index 33c6b7b267..273b0e8550 100644 --- a/mstar/engine/resources/attn/flashinfer.py +++ b/mstar/engine/resources/attn/flashinfer.py @@ -1,6 +1,7 @@ """Paged attention through FlashInfer's prefill/decode wrappers.""" from collections.abc import Mapping +from typing import TYPE_CHECKING import torch @@ -27,6 +28,9 @@ from mstar.engine.resources.step import SlotLease, StepContext from mstar.utils.profiler import PhaseLaps +if TYPE_CHECKING: + from mstar.engine.resources.resident import CapturedResidentStep + class FlashInferManager(AttentionManager): def __init__( @@ -99,6 +103,13 @@ def resident_refusal(self) -> str | None: return "a sliding window reads pages a resident row no longer holds" return None + def resident_gather(self, step: "CapturedResidentStep") -> None: + wrapper = self._cg_plan_states.get(CGSlotKey(bucket=step.bucket, slot=step.slot, label=step.label)) + if isinstance(wrapper, FmhaDecodeWrapper): + table = step.block_table + wrapper.gather_resident(step.store, step.slots, step.buffers.kv_len, table.stride, table.col) + step.host_plan_free = True + def _check_layer(self, layer_idx: int) -> None: """Over a pool of more than one group, ``layer_idx`` must be a layer this attention runs, since its plans address only that group's pages. Checked in Python, so at eager and @@ -234,6 +245,8 @@ def force_double_buffer(self): return True def plan(self, step: AttentionStep, ctx: StepContext): + if ctx.resident is not None and ctx.slot_lease.bucket in ctx.resident.host_plan_free: + return refuse_context_only(step, type(self).__name__) if self._window_left >= 0 and not step.causal: raise ValueError("sliding_window requires causal attention") diff --git a/mstar/engine/resources/attn/wrappers.py b/mstar/engine/resources/attn/wrappers.py index e052d3b31f..fdf40bbcc1 100644 --- a/mstar/engine/resources/attn/wrappers.py +++ b/mstar/engine/resources/attn/wrappers.py @@ -19,6 +19,8 @@ import numpy as np import torch +import triton +import triton.language as tl from mstar.platforms import current_platform @@ -626,6 +628,23 @@ def run(self, q: torch.Tensor, kv_cache_layer: torch.Tensor) -> torch.Tensor: return self.attn_wrapper.run(q.to(self.dtype), kv_cache_layer) +@triton.jit +def _gather_block_rows_kernel( + slots_ptr, kv_len_ptr, store_ptr, table_ptr, lens_ptr, + STRIDE: tl.constexpr, BLOCK: tl.constexpr, PAGE: tl.constexpr, WIDTH: tl.constexpr, +): + i = tl.program_id(0) + n = tl.load(kv_len_ptr + i) + 1 + tl.store(lens_ptr + i, n.to(tl.int32)) + row = tl.load(slots_ptr + i) * STRIDE + BLOCK + pages = (n + PAGE - 1) // PAGE + for start in tl.range(0, pages, 32): + cols = start + tl.arange(0, 32) + keep = cols < pages + values = tl.load(store_ptr + row + cols, mask=keep, other=0) + tl.store(table_ptr + i * WIDTH + cols, values.to(tl.int32), mask=keep) + + class FmhaDecodeWrapper: """Captured decode attention on FlashInfer's trtllm-gen fmha kernel. @@ -712,6 +731,17 @@ def plan( self.dtype = dtype self._window_left = window_left + def gather_resident( + self, store: torch.Tensor, slots: torch.Tensor, kv_len: torch.Tensor, stride: int, block_col: int, + ) -> None: + """Inside a resident step's capture: this step's rows off the tables, + into the table and lengths the kernel holds the addresses of. Only the + columns a row reads; the rest of the table is never looked at.""" + _gather_block_rows_kernel[(self.batch_size,)]( + slots, kv_len, store, self._table, self._lens, + STRIDE=stride, BLOCK=block_col, PAGE=self._page_size, WIDTH=self._table.shape[1], num_warps=1, + ) + @torch.compiler.disable def run(self, q: torch.Tensor, kv_cache_layer: torch.Tensor) -> torch.Tensor: """``q`` is ``[batch, num_qo_heads, head_dim]``, ``kv_cache_layer`` diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index fd09f9fd6b..6115cef584 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -2830,6 +2830,9 @@ def plan(self, step: KVStep, ctx: StepContext) -> dict[str, KVPlanOutput]: ) laps = PhaseLaps(f"{ctx.plan_phase}.kv_cache.") self.reset_default_cursors() + if ctx.resident is not None and ctx.slot_lease.bucket in ctx.resident.host_plan_free: + # nothing reads a plan of this step, and its graph fills the addressing + return KVPlanOutputs({}) if self._preplanned: if self._preplan_key == self._plan_key(step, ctx): self._current_plan_states = self._preplan_states diff --git a/mstar/engine/resources/kv/resident.py b/mstar/engine/resources/kv/resident.py index 4533898f2c..00c2c90a2d 100644 --- a/mstar/engine/resources/kv/resident.py +++ b/mstar/engine/resources/kv/resident.py @@ -90,6 +90,7 @@ def advance(self, plan: ResidentPlan) -> None: def gather(self, step: CapturedResidentStep, token_to_page: torch.Tensor, token_to_cache: torch.Tensor) -> None: store, table = self.planner.store, self.block_table stride, page_size = table.stride, self.page_size + step.store, step.block_table = store, table # before the reads: a row that crossed writes into its new page scatter_cells(store, step.buffers.cells) kv_len = step.buffers.kv_len diff --git a/mstar/engine/resources/resident.py b/mstar/engine/resources/resident.py index 20da91f2a4..c2a15b2921 100644 --- a/mstar/engine/resources/resident.py +++ b/mstar/engine/resources/resident.py @@ -182,6 +182,11 @@ class CapturedResidentStep: label: str slots: torch.Tensor buffers: ResidentBuffers + # the store and the block table, for attention that reads a row's pages itself + store: torch.Tensor | None = None + block_table: ResidentTable | None = None + # set by an owner whose kernel runs off the tables alone: nothing plans this bucket + host_plan_free: bool = False class ResidentPlanner: @@ -203,6 +208,8 @@ def __init__( self.narrow_cols = 0 self._stager = PinnedStager(torch.long, numel=1 << 16) self._buffers: dict[tuple[BucketKey, int], ResidentBuffers] = {} + # buckets whose steps need no host plan at all + self.host_plan_free: set[BucketKey] = set() self._lock = threading.Lock() # rid -> slot. Handle-keyed: `release` purges it. @@ -495,6 +502,8 @@ def capture_step(self, bucket: BucketKey, slot: int, step: SubmoduleStep) -> Cap def gather(self, step: CapturedResidentStep) -> None: for resource in self._owners: resource.resident_gather(step) + if step.host_plan_free: + self.host_plan_free.add(step.bucket) def commit_in_graph(self, step: CapturedResidentStep) -> None: for resource in self._owners: From bfcf640805fdf5d7c609c4b0e47ffe01ced6f414 Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 2 Nov 2026 02:28:45 -0800 Subject: [PATCH 06/15] tests: the LoRA graph engine stub knows about the planner --- test/modular/test_lora_graphs.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/modular/test_lora_graphs.py b/test/modular/test_lora_graphs.py index d54b974403..3585586226 100644 --- a/test/modular/test_lora_graphs.py +++ b/test/modular/test_lora_graphs.py @@ -219,6 +219,7 @@ def lease_slot(graph_walk, bs, slot, num_tokens=None, cg_key_info=None): engine._submodules = {"LLM": SimpleNamespace( submodule=SimpleNamespace(cg_key_info=lambda walk, info, **kw: model_key), lora_pools=(pool,), accelerator_graph_runner=SimpleNamespace(lease_slot=lease_slot, is_packed=lambda lease: False), + resident=None, )} return engine, asked From e0cee953b8d1491f8dbe6fc5a884490962f2d5c8 Mon Sep 17 00:00:00 2001 From: merceod Date: Tue, 3 Nov 2026 01:15:11 +0000 Subject: [PATCH 07/15] Resident hooks go in a mixin Resource inherits --- mstar/engine/resources/base.py | 44 +-------------------- mstar/engine/resources/resident_hooks.py | 50 ++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 42 deletions(-) create mode 100644 mstar/engine/resources/resident_hooks.py diff --git a/mstar/engine/resources/base.py b/mstar/engine/resources/base.py index 2f78f17e78..a2571de0e0 100644 --- a/mstar/engine/resources/base.py +++ b/mstar/engine/resources/base.py @@ -8,6 +8,7 @@ import torch from mstar.distributed.communication import JointGroups +from mstar.engine.resources.resident_hooks import ResidentHooks from mstar.engine.resources.spec import NodeResourceSpec, ResourceReqConfig from mstar.engine.resources.step import ADMIT_OK, AdmitOutcome, BucketKey, ResourceStep, StepContext @@ -19,12 +20,6 @@ PiecewiseAcceleratorGraphConfig, ) from mstar.engine.resources.kv.transfer import TransferEngineInfo - from mstar.engine.resources.resident import ( - CapturedResidentStep, - ResidentPlan, - ResidentPlanner, - ResidentWrites, - ) @dataclass(frozen=True) @@ -122,7 +117,7 @@ class KeyedWalks(NamedTuple): growth: tuple[str, ...] = () -class Resource(ABC): +class Resource(ResidentHooks, ABC): @classmethod @abstractmethod def build( @@ -136,41 +131,6 @@ def depends_on(self) -> set[str]: def can_capture(self) -> bool: return True - # Resident planning (`resources/resident.py`). A captured decode step reads - # per-request state off tables on the device; each owner of that state - # keeps its own table. A resource overrides these only once it has one. - - def resident_refusal(self) -> str | None: - """Why a node holding this resource can't run on resident tables, or None if it can.""" - return f"{type(self).__name__} keeps no resident table" - - def bind_resident(self, planner: "ResidentPlanner") -> None: - """Declare this resource's tables on the node's planner, once, at load.""" - return - - def resident_row(self, rid: int, slot: int, row: list[int], writes: "ResidentWrites") -> None: - """Set this resource's cells of ``rid``'s row: it is new or stale. ``row`` - has the planner's narrow columns; the owner of the wide table appends.""" - return - - def resident_stage(self, plan: "ResidentPlan") -> None: - """Queue this step's per-row changes on ``plan.writes``. Host work only: - this can run beside a live graph.""" - return - - def resident_gather(self, step: "CapturedResidentStep") -> None: - """Record, inside the capture, the reads that fill what the step's forward - addresses. Runs at warmup and capture only; a replay repeats the kernels.""" - return - - def resident_commit(self, step: "CapturedResidentStep") -> None: - """Record, inside the capture and after the forward, the step's table updates.""" - return - - def resident_advance(self, plan: "ResidentPlan") -> None: - """Advance the host's copy of the table by the rule ``resident_commit`` recorded.""" - return - # stays correct when a hit skips the front of a request; required of every resource on a keyed node prefix_skip_safe: bool = False diff --git a/mstar/engine/resources/resident_hooks.py b/mstar/engine/resources/resident_hooks.py new file mode 100644 index 0000000000..b70d44a2e6 --- /dev/null +++ b/mstar/engine/resources/resident_hooks.py @@ -0,0 +1,50 @@ +"""What a resource overrides to take part in resident planning (`resident.py`).""" +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + # `resident.py` imports `Resource`, which inherits from here + from mstar.engine.resources.resident import ( + CapturedResidentStep, + ResidentPlan, + ResidentPlanner, + ResidentWrites, + ) + + +class ResidentHooks: + """A captured decode step reads per-request state off tables on the device; + each owner of that state keeps its own table. A resource overrides these + only once it has one.""" + + def resident_refusal(self) -> str | None: + """Why a node holding this resource can't run on resident tables, or None if it can.""" + return f"{type(self).__name__} keeps no resident table" + + def bind_resident(self, planner: ResidentPlanner) -> None: + """Declare this resource's tables on the node's planner, once, at load.""" + return + + def resident_row(self, rid: int, slot: int, row: list[int], writes: ResidentWrites) -> None: + """Set this resource's cells of ``rid``'s row: it is new or stale. ``row`` + has the planner's narrow columns; the owner of the wide table appends.""" + return + + def resident_stage(self, plan: ResidentPlan) -> None: + """Queue this step's per-row changes on ``plan.writes``. Host work only: + this can run beside a live graph.""" + return + + def resident_gather(self, step: CapturedResidentStep) -> None: + """Record, inside the capture, the reads that fill what the step's forward + addresses. Runs at warmup and capture only; a replay repeats the kernels.""" + return + + def resident_commit(self, step: CapturedResidentStep) -> None: + """Record, inside the capture and after the forward, the step's table updates.""" + return + + def resident_advance(self, plan: ResidentPlan) -> None: + """Advance the host's copy of the table by the rule ``resident_commit`` recorded.""" + return From 4d7148959967143c9f581146198d96a096d23fb5 Mon Sep 17 00:00:00 2001 From: merceod Date: Tue, 3 Nov 2026 01:15:11 +0000 Subject: [PATCH 08/15] Qwen3.5 decode is resident on every config #859 put decode's rope in the graph whatever the knob says, so the knob no longer guards anything. --- mstar/model/qwen3_5/submodules.py | 5 ++-- test/modular/test_resident_planning.py | 34 ++++++++++++++++++++++---- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index d05512f267..8b00516b6b 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -380,8 +380,9 @@ def prepare_inputs( return ARNodeInputs(input_seq_len=input_ids.shape[0], input_ids=input_ids) def resident_walks(self) -> frozenset[str]: - # with the tables staged outside the graph, `preprocess` reads positions - return frozenset({"decode"}) if ingraph_decode_rope() else frozenset() + # a captured decode step builds its rope inside the graph on every config + # (`_rope_in_graph`), so its `preprocess` reads nothing the tables fill + return frozenset({"decode"}) @staticmethod def _rope_in_graph(engine_inputs: ModelInputsFromEngine, graph_walk: str) -> bool: diff --git a/test/modular/test_resident_planning.py b/test/modular/test_resident_planning.py index db9496cdc9..6b754b619d 100644 --- a/test/modular/test_resident_planning.py +++ b/test/modular/test_resident_planning.py @@ -835,17 +835,41 @@ def test_an_eager_step_marks_its_rows(): _assert_row(planner, live, 1) -def test_qwen35_decode_needs_the_rope_inside_the_graph(monkeypatch): +def _qwen35_llm(): from mstar.model.qwen3_5 import submodules - monkeypatch.setattr(submodules, "ingraph_decode_rope", lambda: False) - llm = next( + return submodules, next( cls for cls in vars(submodules).values() if isinstance(cls, type) and cls.__module__ == submodules.__name__ and "resident_walks" in vars(cls) ) - assert llm.resident_walks(None) == frozenset() - monkeypatch.setattr(submodules, "ingraph_decode_rope", lambda: True) + + +@pytest.mark.parametrize("knob", [False, True]) +def test_qwen35_decode_builds_its_rope_in_the_graph_on_every_config(monkeypatch, knob): + submodules, llm = _qwen35_llm() + monkeypatch.setattr(submodules, "ingraph_decode_rope", lambda: knob) assert llm.resident_walks(None) == {"decode"} + assert llm._rope_in_graph(SimpleNamespace(captured=True), "decode") + assert not llm._rope_in_graph(SimpleNamespace(captured=False), "decode") + + +@pytest.mark.parametrize("rows", ["loopback", "own_ids"]) +def test_qwen35_captured_decode_preprocess_reads_nothing_the_tables_fill(rows): + """The graph fills positions and write locations, after `preprocess`.""" + submodules, llm = _qwen35_llm() + token = torch.tensor([5, 6]) + inputs = [ + SimpleNamespace(input_embeds=None, input_ids=None if rows == "loopback" else token[:1], custom_pos_ids=None, + input_seq_len=0 if rows == "loopback" else 1) + for _ in range(2) + ] + stub = SimpleNamespace( + _rope_in_graph=llm._rope_in_graph, _position_ids_3d=lambda inputs: pytest.fail("read positions"), + ) + sampler = SimpleNamespace(loopback_tokens=lambda rids: token) + engine_inputs = SimpleNamespace(captured=True, request_ids=[1, 2], resources={submodules.SAMPLER: sampler}) + out = llm.preprocess(stub, "decode", engine_inputs, inputs) + assert set(out) == {"input_ids"} # ---- the knob ---- From 00ab09ad7e747247c5b26c5e2b878dc46e532f51 Mon Sep 17 00:00:00 2001 From: merceod Date: Tue, 3 Nov 2026 01:15:11 +0000 Subject: [PATCH 09/15] Resident slots are a cache sized by the batch cap 2B+1 rows, B the widest capture of a resident walk. A step short of slots takes the least recently run ones, and an offload gives its slot up, so nothing here can refuse or raise, and the knob goes. Rows stop at the pool's pages, and the stager is sized at build so it never grows on the GPU thread. --- .claude/skills/engine-resources/SKILL.md | 2 +- docs/environment_variables.rst | 15 +- docs/serving.rst | 4 +- mstar/engine/accelerator_graph_runner.py | 14 +- mstar/engine/engine.py | 34 +- mstar/engine/resources/kv/manager.py | 11 +- mstar/engine/resources/kv/resident.py | 4 + mstar/engine/resources/resident.py | 132 +++-- mstar/utils/knobs.py | 8 - mstar/worker/rid_table.py | 3 +- test/modular/test_resident_planning.py | 703 +++++++++++++++++++++-- 11 files changed, 812 insertions(+), 118 deletions(-) diff --git a/.claude/skills/engine-resources/SKILL.md b/.claude/skills/engine-resources/SKILL.md index bdb2788a6b..46885d5029 100644 --- a/.claude/skills/engine-resources/SKILL.md +++ b/.claude/skills/engine-resources/SKILL.md @@ -34,7 +34,7 @@ Prototype a new kernel standalone against a dense reference before writing the r ## 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. 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 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. 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 diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index de3b028f9f..99a46b6a58 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -640,14 +640,13 @@ Worker scheduling 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. - Read by ``knobs.resident_planning``. - * - ``MSTAR_RESIDENT_SLOTS`` - - config's ``knobs: resident_slots``, else ``4096`` - - Request rows of those tables, allocated at startup because a captured - graph holds their addresses; a step that needs a row past the last - fails. Each row costs a cell per page of ``max_seq_len`` plus a few, 8 - bytes each. Read by ``knobs.resident_slots``. + 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: : 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 diff --git a/docs/serving.rst b/docs/serving.rst index 7cdd4646be..42cba6d172 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -388,8 +388,8 @@ 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``, ``compile_deterministic``, - ``resident_planning`` and ``resident_slots`` (``true``/``false`` or an integer; + ``gc_gen0_threshold``, ``pre_plan``, ``step_thread_core``, ``compile_deterministic`` + and ``resident_planning`` (``true``/``false`` or an integer; ``step_thread_core`` takes ``off``, the default, ``auto`` or a CPU number, ``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 diff --git a/mstar/engine/accelerator_graph_runner.py b/mstar/engine/accelerator_graph_runner.py index a8f33556d6..dc3c777bff 100644 --- a/mstar/engine/accelerator_graph_runner.py +++ b/mstar/engine/accelerator_graph_runner.py @@ -693,12 +693,11 @@ def __init__( num_slots: int, enable_nvtx: bool=False, device_capability: DeviceCapability | None = None, - resident: ResidentPlanner | None = None, ): self._submodule_name = submodule_name self._submodule = submodule self._comm_group = joint_comm_group - self._resident = resident + self._resident: ResidentPlanner | None = None self._capture_configs: list[AcceleratorGraphConfig] = ( submodule.get_accelerator_graph_configs( @@ -779,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]: diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index cdf18f0791..a3192cdb3b 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -695,12 +695,13 @@ def load_model( # the sampler keeps its slot masters only for a node that # reads them back on the device (see DeviceLoopback) resources[loopback.sampler].enable_device_loopback_reader() - self._build_resident_planner(node_name) - def _build_resident_planner(self, node_name: str) -> None: + 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.""" + 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] @@ -711,10 +712,11 @@ def _build_resident_planner(self, node_name: str) -> None: 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, knobs.resident_slots(), + 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.""" @@ -828,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, @@ -839,8 +841,8 @@ def warmup(self) -> None: num_slots=submodule_mgmt.num_slots, enable_nvtx=self._enable_nvtx, device_capability=self.device_capability, - resident=submodule_mgmt.resident, ) + self._build_resident_planner(node_name, runner) piecewise[node_name] = self._build_piecewise_runners( node_name, submodule_mgmt ) @@ -2960,6 +2962,9 @@ def _pre_plan_for_batch( 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() @@ -3049,12 +3054,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.""" @@ -3186,6 +3199,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, diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index 7990cf0fdf..767624398e 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -3460,7 +3460,8 @@ def _commit_offload( have a torn host copy, and that fork's source is one of these streams, so the whole request stays on device rather than half of it. """ - self._touch_plan_cache(rid) + # its resident slot goes with the offload (`Engine.offload_request`) + self._touch_plan_cache(rid, row=False) with self._lock: streams = self._streams.get(rid, {}) for claim in moved: @@ -3539,7 +3540,8 @@ def reload(self, rid: str) -> bool: from the host. False when the device can't fit them right now; nothing moves in that case, so the caller can evict further and try again. """ - self._touch_plan_cache(rid) + # an offload gave its resident slot up, so its next step writes the row whole + self._touch_plan_cache(rid, row=False) if self._cpu_pool is None or not self._cpu_pool.is_offloaded(rid): return False with self._lock: @@ -4514,7 +4516,10 @@ def bind_resident(self, planner: "ResidentPlanner") -> None: from mstar.engine.resources.kv.resident import KVResidentTables page_size = self._groups[0].page_size - self._resident = KVResidentTables(planner, page_size, -(-self.config.max_seq_len // page_size)) + # the pool's own context or override (not the serving config's top-level + # `max_seq_len`), and a stream's pages are distinct pages of the pool + width = min(-(-self.config.max_seq_len // page_size), self.config.request_pages) + self._resident = KVResidentTables(planner, page_size, width) def resident_row(self, rid: int, slot: int, row: list[int], writes: "ResidentWrites") -> None: stream = self._streams[rid][writes.label] diff --git a/mstar/engine/resources/kv/resident.py b/mstar/engine/resources/kv/resident.py index 00c2c90a2d..7589ce01c2 100644 --- a/mstar/engine/resources/kv/resident.py +++ b/mstar/engine/resources/kv/resident.py @@ -49,6 +49,8 @@ def __init__(self, planner: ResidentPlanner, page_size: int, max_pages: int): def row(self, stored_len: int, pages: list[int], slot: int, row: list[int]) -> None: """The cells of a stream about to append one token. Its host mirrors move here, not at the flush: a plan that never runs leaves the row stale.""" + if len(pages) > self.block_table.width: + raise ValueError(f"{len(pages)} pages of a stream don't fit a block table row of {self.block_table.width}") cols = stored_len // self.page_size + 1 row[self.kv_len.col] = stored_len row.extend(pages[:cols]) @@ -74,6 +76,8 @@ def stage_step(self, plan: ResidentPlan, stream_pages, rewrite) -> None: return for i in rows.tolist(): rid, slot, last = plan.rids[i], int(slots[i]), int(need[i]) - 1 + if last >= self.block_table.width: + raise ValueError(f"page {last} of a stream doesn't fit a block table row of {self.block_table.width}") if last > cols[i]: plan.taken.add(rid) rewrite(rid, slot, plan.writes) diff --git a/mstar/engine/resources/resident.py b/mstar/engine/resources/resident.py index c2a15b2921..ac4c80b888 100644 --- a/mstar/engine/resources/resident.py +++ b/mstar/engine/resources/resident.py @@ -18,7 +18,8 @@ The host stays the source of truth. A row is written whole the first time a step reads it and whenever something other than a resident commit changed it -(a host-planned step, an offload, a reload, a fork); see `mark_stale`. +(a host-planned step, a fork); see `mark_stale`. An offload gives the slot up, +so a reload starts over as a new request does. """ from __future__ import annotations @@ -190,23 +191,33 @@ class CapturedResidentStep: class ResidentPlanner: - """One node's request slots, tables and per-step rows.""" + """One node's request slots, tables and per-step rows. + + Slots are a cache, not something a request can run out of. Every cell of a + row mirrors host state, so a request can lose its slot whenever no step needs + it, and it takes another with a whole-row rewrite when it next runs, as a new + request does. ``2 * max_bs`` slots, ``max_bs`` being the widest capture of a + resident walk, always hold a step, on every rank and whatever the commit + timing: a slot is free, cooling or held, and a step can take any of them but + its own rows, at most ``max_bs``, and those of the one earlier step that may + still commit, at most ``max_bs``. Slot numbers are private to a rank's + planner, and no collective reads one. + """ def __init__( self, node: str, walks: Collection[str], resources: Mapping[str, Resource], - device: torch.device, max_slots: int, + device: torch.device, max_bs: int, ): self.node = node self.walks = frozenset(walks) self.device = device # request rows plus the trash row - self.max_slots = max_slots + 1 + self.max_slots = 2 * max_bs + 1 self._owners = tuple(resources.values()) self._tables: list[ResidentTable] = [] self._store: torch.Tensor | None = None # columns of the one-cell tables, which a rewritten row always covers self.narrow_cols = 0 - self._stager = PinnedStager(torch.long, numel=1 << 16) self._buffers: dict[tuple[BucketKey, int], ResidentBuffers] = {} # buckets whose steps need no host plan at all self.host_plan_free: set[BucketKey] = set() @@ -214,7 +225,7 @@ def __init__( self._lock = threading.Lock() # rid -> slot. Handle-keyed: `release` purges it. self._slot_of: dict[int, int] = {} - self._free: list[int] = list(range(max_slots, TRASH_SLOT, -1)) + self._free: list[int] = list(range(self.max_slots - 1, TRASH_SLOT, -1)) # (slot, seq of the last step that read it), freed under a step not yet committed self._cooling: list[tuple[int, int]] = [] self._last_read = np.zeros(self.max_slots, dtype=np.int64) @@ -233,6 +244,9 @@ def __init__( for resource in self._owners: resource.bind_resident(self) self._seal() + # the largest flush up front, plus a cells row: growing the stager waits + # out its copies, which on the GPU thread would be a sync in the step path + self._stager = PinnedStager(torch.long, numel=max_bs * self._tables[0].stride + 4 * max_bs) # Tables @@ -277,27 +291,54 @@ def covers(self, bucket: BucketKey) -> bool: def slot_of(self, rid: int) -> int | None: return self._slot_of.get(rid) - def _take_slot(self) -> int: - """Caller holds the lock. A freed slot comes back once the last step - that read it has committed: until then that step's commit still - advances the slot's host lengths.""" + def _take_slots(self, n: int, seq: int) -> list[int]: + """``n`` slots for the new rows of plan ``seq``, whose held rows carry + its stamp already. Caller holds the lock. A slot is taken only once the + last step that read it has committed, or can't: until then that step's + commit still advances the slot's host lengths. A step is planned behind + the commit of the one before it, so only the newest earlier plan can + still be waiting on one, however many before it never ran.""" + settled = max(self._committed_seq, seq - 2) if self._cooling: - done = self._committed_seq - ready = [slot for slot, seq in self._cooling if seq <= done] + ready = [slot for slot, read in self._cooling if read <= settled] if ready: - self._cooling = [(slot, seq) for slot, seq in self._cooling if seq > done] + self._cooling = [(slot, read) for slot, read in self._cooling if read > settled] self._free.extend(ready) - if not self._free: + free = self._free + taken = [free.pop() for _ in range(min(n, len(free)))] + if len(taken) < n: + taken += self._displace(n - len(taken), settled) + return taken + + def _displace(self, n: int, settled: int) -> list[int]: + """The slots of the ``n`` least recently run requests. Each loses its + slot and takes another when it next runs; the slot's host mirrors go to + the new tenant. Ordered by plan, not by clock or slot number, and a + step's rows in the order they took their slots, so a rank with the same + history picks the same requests. Nothing depends on that: a displaced + row is host state and one H2D, and no collective sees it.""" + held = self._slot_of + reads = self._last_read[np.fromiter(held.values(), dtype=np.int64, count=len(held))] + open_rows = np.flatnonzero(reads <= settled) + if open_rows.size < n: raise RuntimeError( - f"resident planning on {self.node}: all {self.max_slots - 1} slots are held; " - "raise `knobs: resident_slots`" + f"resident planning on {self.node}: {n} rows found no slot among {self.max_slots - 1}, " + f"{open_rows.size} of the {len(held)} held open. A step's rows and those of the one step " + "behind it can't hold them all, so the slot accounting is wrong" ) - return self._free.pop() + rids = list(held) + slots = [] + for i in open_rows[np.argsort(reads[open_rows], kind="stable")][:n].tolist(): + rid = rids[i] + slots.append(held.pop(rid)) + self._stale.discard(rid) + self._moved.discard(rid) + return slots def release(self, rid: int) -> None: - """The request is gone. Its slot cools until the last step that read it - has committed; on the device the row is only read again behind that - step, by stream order.""" + """The request is gone, or offloaded: a reload is written whole anyway. + Its slot cools until the last step that read it has committed; on the + device the row is only read again behind that step, by stream order.""" with self._lock: self._stale.discard(rid) self._moved.discard(rid) @@ -371,16 +412,20 @@ def _build(self, rids: tuple[int, ...], key: tuple, label: str, ctx) -> Resident lease = ctx.slot_lease laps = PhaseLaps(f"{ctx.plan_phase}.resident.") with self._lock: + self._plan_seq += 1 + seq = self._plan_seq if rids == self._last_rids: slots = self._last_slots else: get = self._slot_of.get row = [get(rid) for rid in rids] if None in row: - for i, slot in enumerate(row): - if slot is None: - row[i] = self._slot_of[rids[i]] = self._take_slot() - self._stale.add(rids[i]) + # stamped first: a row of this step is no one's victim + self._last_read[[slot for slot in row if slot is not None]] = seq + new = [i for i, slot in enumerate(row) if slot is None] + for i, slot in zip(new, self._take_slots(len(new), seq), strict=True): + row[i] = self._slot_of[rids[i]] = slot + self._stale.add(rids[i]) slots = np.array(row, dtype=np.int64) self._last_rids, self._last_slots = rids, slots stale = self._stale @@ -390,25 +435,28 @@ def _build(self, rids: tuple[int, ...], key: tuple, label: str, ctx) -> Resident refreshed = moved.intersection(rids) if moved else set() moved -= refreshed refreshed -= taken - self._plan_seq += 1 - seq = self._plan_seq self._last_read[slots] = seq rows = [(rid, self._slot_of[rid]) for rid in taken] plan = ResidentPlan(key, rids, slots, seq, self._commits, ResidentWrites(label), taken, refreshed) laps.lap("rows") - for rid, slot in rows: - self.rewrite(rid, slot, plan.writes) - laps.lap("rewrites") - buffers = self.buffers(lease.bucket, lease.slot) - if buffers.staged is not slots: - self._stager.copy_(buffers.slots, slots, pad_value=TRASH_SLOT) - buffers.staged = slots - laps.lap("slot_row") - for resource in self._owners: - resource.resident_stage(plan) - laps.lap("stage") - self._stage_cells(plan, buffers) - laps.lap("cells") + try: + for rid, slot in rows: + self.rewrite(rid, slot, plan.writes) + laps.lap("rewrites") + buffers = self.buffers(lease.bucket, lease.slot) + if buffers.staged is not slots: + self._stager.copy_(buffers.slots, slots, pad_value=TRASH_SLOT) + buffers.staged = slots + laps.lap("slot_row") + for resource in self._owners: + resource.resident_stage(plan) + laps.lap("stage") + self._stage_cells(plan, buffers) + laps.lap("cells") + except BaseException: + # nothing wrote these rows, so the next step that reads them must + self._restore(plan) + raise return plan def rewrite(self, rid: int, slot: int, writes: ResidentWrites) -> None: @@ -512,7 +560,7 @@ def commit_in_graph(self, step: CapturedResidentStep) -> None: def build_planner( node: str, walks: Iterable[str], resources: Mapping[str, Resource], - resource_nodes: Mapping[str, Collection[str]], device: torch.device, max_slots: int, + resource_nodes: Mapping[str, Collection[str]], device: torch.device, max_bs: int, ) -> tuple[ResidentPlanner | None, str]: """A planner for ``node``, or None and why not. Every resource of the node must be able to, since a capture is one graph: one host-planned resource @@ -520,10 +568,12 @@ def build_planner( walks = frozenset(walks) if not walks: return None, "its submodule names no resident walk" + if max_bs < 1: + return None, f"nothing captures its resident walks {sorted(walks)}" for key, resource in resources.items(): if len(resource_nodes.get(key, ())) > 1: return None, f"resource {key!r} is shared by nodes {sorted(resource_nodes[key])}" refusal = resource.resident_refusal() if refusal is not None: return None, f"resource {key!r}: {refusal}" - return ResidentPlanner(node, walks, resources, device, max_slots), "" + return ResidentPlanner(node, walks, resources, device, max_bs), "" diff --git a/mstar/utils/knobs.py b/mstar/utils/knobs.py index 2d2da1b980..48c10d294a 100644 --- a/mstar/utils/knobs.py +++ b/mstar/utils/knobs.py @@ -22,7 +22,6 @@ "step_thread_core": "MSTAR_STEP_THREAD_CORE", "compile_deterministic": "MSTAR_COMPILE_DETERMINISTIC", "resident_planning": "MSTAR_RESIDENT_PLANNING", - "resident_slots": "MSTAR_RESIDENT_SLOTS", } # this process's ``knobs:`` block, as strings in the env vars' format @@ -227,10 +226,3 @@ def resident_planning() -> str: resident-capable, and says why at startup when one is not; ``off`` keeps the host plan everywhere.""" return _resident_planning_value(_setting("resident_planning", "auto")) - - -def resident_slots() -> int: - """``resident_slots`` / ``MSTAR_RESIDENT_SLOTS`` (default 4096): request - rows of the resident tables, allocated once at startup since a captured - graph holds their addresses.""" - return int(_setting("resident_slots", "4096")) diff --git a/mstar/worker/rid_table.py b/mstar/worker/rid_table.py index 99afb088b0..ac11bae6f7 100644 --- a/mstar/worker/rid_table.py +++ b/mstar/worker/rid_table.py @@ -18,7 +18,8 @@ the tensor store's rid -> uuids ownership, held in the Rust bookkeeper under RustTensorStore; purged by force_cleanup_request), RecurrentStatePool.{_keyed, _leases, _restores, _snap_writes, _decode, _unkeyed, _admit_restores, -_admit_snaps} (purged by remove_request), and WorkerProfileInfo. +_admit_snaps} (purged by remove_request), ResidentPlanner.{_slot_of, _stale, +_moved} (purged by release), and WorkerProfileInfo. MicroScheduler.pending_tp_follow_count is the one deliberate exception: a ScheduleTPNode can arrive before the NEW_REQUEST that mints the handle, so it diff --git a/test/modular/test_resident_planning.py b/test/modular/test_resident_planning.py index 6b754b619d..ef1c384ce7 100644 --- a/test/modular/test_resident_planning.py +++ b/test/modular/test_resident_planning.py @@ -8,6 +8,7 @@ """ from __future__ import annotations +import random import sys from types import SimpleNamespace @@ -33,8 +34,17 @@ from mstar.engine.resources.recurrent.pool import SINK_SLOT, RecurrentStatePool # noqa: E402 from mstar.engine.resources.resident import TRASH_SLOT, ResidentPlanner, build_planner # noqa: E402 from mstar.engine.resources.spec import resolve_spec_dependencies # noqa: E402 -from mstar.engine.resources.step import BucketKey, Segment, SlotLease, StepContext, SubmoduleStep # noqa: E402 +from mstar.engine.resources.step import ( # noqa: E402 + AllocationFailed, + BucketKey, + Segment, + SlotLease, + StepContext, + SubmoduleStep, +) from mstar.utils import knobs # noqa: E402 +from mstar.utils.ipc_format import OffloadDelta # noqa: E402 +from mstar.worker.worker import Worker # noqa: E402 PAGE = 4 WALK = "decode" @@ -60,7 +70,7 @@ def _stub_transfer(monkeypatch): monkeypatch.setattr(manager_mod, "KVTransferManager", _StubTransfer) -def _resources(cpu_offload_pages: int = 0, max_num_pages: int = 64): +def _resources(cpu_offload_pages: int = 0, max_num_pages: int = 64, pool_slots: int = 9): kv = KVManager( cfg=PagedKVConfig( num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=512, @@ -70,7 +80,7 @@ def _resources(cpu_offload_pages: int = 0, max_num_pages: int = 64): ) kv._cg_max_seq_len = 64 pool = RecurrentStatePool(device=CPU, config=RecurrentStateConfig( - num_layers=1, max_slots=9, blocks={"state": RecurrentBlockConfig(shape=(2,), dtype=torch.float32)}, + num_layers=1, max_slots=pool_slots, blocks={"state": RecurrentBlockConfig(shape=(2,), dtype=torch.float32)}, )) rope = RopeManager(PositionConfig(kv_cache="kv"), CPU, max_seq_len=512, head_dim=8) # page-locked memory needs a GPU, and nothing here copies out of it asynchronously @@ -171,10 +181,11 @@ def _replay(res, planner, s, promote=True): return s -def _setup(prompts: dict[int, int], cpu_offload_pages: int = 0, max_num_pages: int = 64): +def _setup(prompts: dict[int, int], cpu_offload_pages: int = 0, max_num_pages: int = 64, pool_slots: int = 9): """The same requests prefilled on resident resources and on a host-planned reference.""" - live, ref = _resources(cpu_offload_pages, max_num_pages), _resources(cpu_offload_pages, max_num_pages) - planner = ResidentPlanner("n", {WALK}, {k: live[k] for k in ("kv", "pool", "rope")}, CPU, max_slots=8) + live = _resources(cpu_offload_pages, max_num_pages, pool_slots) + ref = _resources(cpu_offload_pages, max_num_pages, pool_slots) + planner = ResidentPlanner("n", {WALK}, {k: live[k] for k in ("kv", "pool", "rope")}, CPU, max_bs=4) for res in (live, ref): for rid, n in prompts.items(): _ingest(res, rid) @@ -408,8 +419,7 @@ def test_a_fork_marks_the_row_stale(): def test_a_row_whose_pages_moved_is_rewritten(): - """What a reload leaves: the stream on other pages. Its hook is `_touch_plan_cache(rid)`; - `test_real_offload_and_reload_mark_the_row` runs the real one on a GPU.""" + """A stream on other pages, as a fork or reset can leave it. Its hook is `_touch_plan_cache(rid)`.""" live, ref, planner = _setup({1: 9}) del ref _captured_step(live, planner, [1], slot=0, bs=1) @@ -426,16 +436,18 @@ def test_a_row_whose_pages_moved_is_rewritten(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="the host page pool is page-locked") -def test_real_offload_and_reload_mark_the_row(): +def test_a_real_offload_gives_the_slot_up_and_its_reload_is_written_whole(): + device = torch.device("cuda") kv = KVManager( cfg=PagedKVConfig( num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=64, max_num_pages=16, page_size=PAGE, cpu_offload_pages=16, ), - name="kv", joint_comm_group=None, transfer_engine_info=None, - device=torch.device("cuda"), dtype=torch.float32, + name="kv", joint_comm_group=None, transfer_engine_info=None, device=device, dtype=torch.float32, ) - planner = ResidentPlanner("n", {WALK}, {"kv": kv}, torch.device("cuda"), max_slots=2) + kv._cg_max_seq_len = 64 + planner = ResidentPlanner("n", {WALK}, {"kv": kv}, device, max_bs=2) + engine = _offload_engine({"kv": kv}, planner) kv.ingest_request(1, KVReqConfig()) segments = (Segment(1, "main", 9),) ctx = StepContext(request_ids=(1,), graph_walk="prefill", slot=0, capture=False) @@ -443,13 +455,15 @@ def test_real_offload_and_reload_mark_the_row(): assert kv.admit(step, ctx).ok ctx.plan_results["kv"] = kv.plan(step, ctx) kv.commit(step, ctx) - planner._slot_of[1] = 1 - assert kv.offload(1) > 0 - assert 1 in planner._stale - planner._stale.clear() + with planner._lock: + planner._slot_of[1] = planner._free.pop() + assert engine.offload_request("n", 1) > 0 + assert planner.slot_of(1) is None + assert not planner._stale torch.cuda.synchronize() - assert kv.reload(1) - assert 1 in planner._stale + assert engine.reload_request("n", 1) + assert not planner._stale + assert planner.slot_of(1) is None def test_an_abandoned_preplan_writes_nothing_and_gives_its_rows_back(): @@ -566,17 +580,16 @@ def test_a_freed_slot_waits_for_the_last_step_that_read_it(): # step N reads the slot and is still running: planned, not committed s = _captured_step(live, planner, [1], slot=0, bs=1, preplan=True, run=False) _promote(live, planner, s) - slot = s.ctx.resident_plan.slots[0] + slot = int(s.ctx.resident_plan.slots[0]) + seq = s.ctx.resident_plan.seq planner.release(1) - planner._free.clear() - assert (int(slot), s.ctx.resident_plan.seq) in planner._cooling - with pytest.raises(RuntimeError, match="slots are held"): - with planner._lock: - planner._take_slot() + assert (slot, seq) in planner._cooling + with planner._lock: + assert slot not in planner._take_slots(7, seq + 1) # commit lets the slot go; its late advance landed on the old owner's row _replay(live, planner, s, promote=False) with planner._lock: - assert planner._take_slot() == slot + assert planner._take_slots(1, seq + 2) == [slot] def test_a_late_commit_cannot_touch_the_next_owner(): @@ -631,14 +644,586 @@ def test_a_recycled_handle_gets_a_fresh_row(): _assert_row(planner, live, 1) -def test_running_out_of_slots_says_how_to_fix_it(): - live, ref, planner = _setup({}) +def _slot_ledger(planner: ResidentPlanner) -> None: + """Every request slot is exactly one of free, cooling or held.""" + held = list(planner._slot_of.values()) + cooling = [slot for slot, _ in planner._cooling] + assert sorted(held + cooling + planner._free) == list(range(1, planner.max_slots)) + + +def _assert_step_read_its_streams(live, s, rids) -> None: + """What the replay addressed against the host, which has committed the step since.""" + for i, rid in enumerate(rids): + stream = live["kv"]._streams[rid]["main"] + wrote = stream.stored_len - 1 + assert s.pages[i] == stream.page_indices[wrote // PAGE] + assert s.offsets[i] == wrote % PAGE + assert s.positions[i] == live["rope"].position(rid, "main") - 1 + assert s.pool_slots[i] == live["pool"]._slots[rid]["main"].index + + +def _assert_held_rows(planner, live) -> None: + """A row marked neither stale nor moved holds what its stream says.""" + for rid in planner._slot_of: + if rid not in planner._stale and rid not in planner._moved: + _assert_row(planner, live, rid) + + +def _recycle(live, planner, rid: int, tokens: int) -> None: + """The request ends, and its handle comes back as another's.""" + for key in ("kv", "pool", "rope"): + live[key].remove_request(rid) + planner.release(rid) + _ingest(live, rid) + step, ctx = _host_step(live, [rid], [tokens]) + planner.commit(step, ctx) + + +@pytest.mark.parametrize("seed", range(4)) +def test_more_live_requests_than_slots_run_off_the_cache(seed): + """B=4, so 8 slots and the trash row, under 20 live requests taking turns, with + pre-plans, abandoned ones, host steps and teardowns mixed in. Nothing raises, + a displaced row is written whole when it next runs, the victims are the least + recently run, and every row a step reads is its stream.""" + rng = random.Random(seed) + prompts = {rid: 1 + rid % 7 for rid in range(1, 21)} + live, ref, planner = _setup(prompts, max_num_pages=1024, pool_slots=32) del ref - planner._free.clear() - _ingest(live, 1) - _host_step(live, [1], [2]) - with pytest.raises(RuntimeError, match="resident_slots"): - _captured_step(live, planner, [1], slot=0, bs=1) + assert planner.max_slots == 9 + order, at, displaced, returned = list(prompts), 0, set(), 0 + for t in range(300): + rows = [order[(at + i) % len(order)] for i in range(rng.randint(1, 4))] + at += len(rows) + roll = rng.random() + if roll < 0.1: + recycled = rng.choice(order) + _recycle(live, planner, recycled, rng.randint(1, 9)) + displaced.discard(recycled) + elif roll < 0.2: + step, ctx = _host_step(live, rows, [1] * len(rows), walk="mixed") + planner.commit(step, ctx) + elif roll < 0.3: + _captured_step(live, planner, rows, slot=t % 2, bs=4, preplan=True, run=False) + for key in ("kv", "pool", "rope"): + live[key].clear_preplan() + before = {rid: int(planner._last_read[slot]) for rid, slot in planner._slot_of.items()} + s = _captured_step(live, planner, rows, slot=t % 2, bs=4, preplan=rng.random() < 0.5) + _assert_step_read_its_streams(live, s, rows) + for rid in rows: + _assert_row(planner, live, rid) + assert planner.slot_of(rid) != TRASH_SLOT + if rid in displaced: + assert rid in s.ctx.resident_plan.taken + returned += 1 + displaced.discard(rid) + gone = before.keys() - planner._slot_of.keys() + displaced |= gone + kept = before.keys() - gone - set(rows) + if gone and kept: + assert max(before[rid] for rid in gone) <= min(before[rid] for rid in kept) + _slot_ledger(planner) + _assert_held_rows(planner, live) + live["kv"].assert_pages_conserved() + assert returned > 100 + + +def test_a_displaced_row_is_written_whole_with_its_lengths(): + live, ref, planner = _setup({rid: 5 + rid for rid in range(1, 12)}, max_num_pages=256, pool_slots=16) + del ref + # eight rows fill the slots, the oldest step's first + _captured_step(live, planner, [1, 2, 3, 4], slot=0, bs=4) + _captured_step(live, planner, [5, 6, 7, 8], slot=1, bs=4) + assert not planner._free + victim_slot = planner.slot_of(1) + s = _captured_step(live, planner, [9], slot=0, bs=4) + assert planner.slot_of(1) is None + assert planner.slot_of(9) == victim_slot + assert s.ctx.resident_plan.taken == {9} + tables = live["kv"]._resident + assert tables.host_len[victim_slot] == live["kv"]._streams[9]["main"].stored_len + _assert_row(planner, live, 9) + # the next victims are the rest of the oldest step, in the order they took slots + _captured_step(live, planner, [10], slot=1, bs=4) + assert planner.slot_of(2) is None and planner.slot_of(3) is not None + # and the first comes back through the new-request path + s = _captured_step(live, planner, [1], slot=0, bs=4) + assert s.ctx.resident_plan.taken == {1} + _assert_step_read_its_streams(live, s, [1]) + _assert_row(planner, live, 1) + assert tables.host_len[planner.slot_of(1)] == live["kv"]._streams[1]["main"].stored_len + + +def test_a_step_does_not_displace_its_own_rows_or_the_step_behind_it(): + live, ref, planner = _setup({rid: 3 for rid in range(1, 16)}, max_num_pages=256, pool_slots=20) + del ref + _captured_step(live, planner, [1, 2, 3, 4], slot=0, bs=4) + _captured_step(live, planner, [5, 6, 7, 8], slot=1, bs=4) + # planned behind a step that hasn't committed: its rows are out of reach + behind = _captured_step(live, planner, [5, 6, 7, 8], slot=1, bs=4, preplan=True, run=False) + _promote(live, planner, behind) + s = _captured_step(live, planner, [9, 10, 11, 12], slot=0, bs=4, run=False) + planner.plan(s.step, s.ctx) + assert set(planner._slot_of) == {5, 6, 7, 8, 9, 10, 11, 12} + # the late commit advances the rows it read, none of them taken + _replay(live, planner, behind, promote=False) + _assert_row(planner, live, 5) + + +def test_a_step_keeps_its_own_oldest_row_and_a_victim_leaves_no_marks(): + live, ref, planner = _setup({rid: 3 for rid in range(1, 12)}, max_num_pages=256, pool_slots=16) + del ref + _captured_step(live, planner, [1, 2, 3, 4], slot=0, bs=4) + _captured_step(live, planner, [5, 6, 7, 8], slot=1, bs=4) + planner.mark_stale(2) + planner._moved.add(3) + # 1 is the oldest and in the step, so the victims come after it + s = _captured_step(live, planner, [1, 9, 10], slot=0, bs=4) + assert planner.slot_of(1) is not None + assert planner.slot_of(2) is None and planner.slot_of(3) is None and planner.slot_of(4) is not None + assert not planner._stale and not planner._moved + _slot_ledger(planner) + _assert_step_read_its_streams(live, s, [1, 9, 10]) + + +def test_plans_that_never_ran_hold_no_slots(): + """Refused retries and abandoned pre-plans, one after another with no commit + between: however many there are, only the newest can look uncommitted.""" + live, ref, planner = _setup({rid: 3 for rid in range(1, 26)}, max_num_pages=256, pool_slots=32) + del ref + rid = 1 + for t in range(6): + rows = list(range(rid, rid + 4)) + rid += 4 + s = _captured_step(live, planner, rows, slot=t % 2, bs=4, preplan=True, run=False) + assert set(s.ctx.resident_plan.rids) == set(rows) + for key in ("kv", "pool", "rope"): + live[key].clear_preplan() + _slot_ledger(planner) + + +def test_no_slot_is_an_invariant_violation_not_a_knob(): + live, ref, planner = _setup({1: 3}) + del ref + # every slot read by the step planned just now, which no step can do + with planner._lock: + planner._slot_of.update({100 + i: planner._free.pop() for i in range(8)}) + planner._last_read[1:] = 5 + planner._plan_seq = 5 + with pytest.raises(RuntimeError, match="accounting is wrong"): + planner._take_slots(1, 6) + + +def test_a_rank_whose_release_lands_on_the_other_side_of_a_commit_still_runs(): + """Ranks replay teardowns between the same steps, but a release can land + before or after a commit on the GPU thread, so the same slot is cooling on one + rank and free on the other. Neither may raise, and each rank's rows are its own.""" + ranks = [] + for release_before_commit in (False, True): + live, ref, planner = _setup({rid: 2 + rid for rid in range(1, 21)}, max_num_pages=1024, pool_slots=32) + del ref + ranks.append((live, planner, release_before_commit)) + next_rid = 21 + for t in range(60): + rows = [1 + (4 * t + i) % 20 for i in range(4)] + gone = rows[0] + for live, planner, early in ranks: + s = _captured_step(live, planner, rows, slot=t % 2, bs=4, preplan=True, run=False) + _promote(live, planner, s) + if early: + planner.release(gone) + _replay(live, planner, s, promote=False) + _assert_step_read_its_streams(live, s, rows) + if not early: + planner.release(gone) + for key in ("kv", "pool", "rope"): + live[key].remove_request(gone) + _ingest(live, gone) + step, ctx = _host_step(live, [gone], [next_rid % 5 + 1]) + planner.commit(step, ctx) + _slot_ledger(planner) + next_rid += 1 + + +def test_ranks_with_the_same_history_displace_the_same_requests(): + runs = [] + for shuffle in (0, 1): + live, ref, planner = _setup({rid: 3 for rid in range(1, 21)}, max_num_pages=1024, pool_slots=32) + del ref + # handles and slot numbers are per rank, so a tie must not depend on either + if shuffle: + planner._free.reverse() + seen = [] + for t in range(40): + rows = [1 + (3 * t + i) % 20 for i in range(4)] + _captured_step(live, planner, rows, slot=t % 2, bs=4) + seen.append(frozenset(planner._slot_of)) + runs.append(seen) + assert runs[0] == runs[1] + + +def test_the_store_is_sized_by_the_widest_capture(): + res = _resources() + planner = ResidentPlanner("n", {WALK}, {k: res[k] for k in ("kv", "pool", "rope")}, CPU, max_bs=8) + assert planner.max_slots == 17 + assert len(planner._free) == 16 + assert TRASH_SLOT not in planner._free + + +def test_a_flush_of_the_largest_step_stages_without_growing(monkeypatch): + res = _resources(max_num_pages=64) + planner = ResidentPlanner("n", {WALK}, {k: res[k] for k in ("kv", "pool", "rope")}, CPU, max_bs=4) + stride = planner._tables[0].stride + assert planner._stager._numel >= 4 * stride + 4 * 4 + monkeypatch.setattr(planner._stager, "_grow", lambda numel: pytest.fail(f"grew to {numel}")) + rows = [(slot, list(range(stride))) for slot in range(1, 5)] + plan = SimpleNamespace(writes=SimpleNamespace(rows=rows)) + planner.flush(plan) + assert planner.store.view(9, stride)[1:5].tolist() == [list(range(stride))] * 4 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the stager pins and copies on a GPU") +def test_a_flush_of_the_largest_step_stages_on_the_device_without_growing(monkeypatch): + device = torch.device("cuda") + kv = KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=512, max_num_pages=64, page_size=PAGE, + ), + name="kv", joint_comm_group=None, transfer_engine_info=None, device=device, dtype=torch.float32, + ) + planner = ResidentPlanner("n", {WALK}, {"kv": kv}, device, max_bs=4) + stride = planner._tables[0].stride + monkeypatch.setattr(planner._stager, "_grow", lambda numel: pytest.fail(f"grew to {numel}")) + rows = [(slot, list(range(100, 100 + stride))) for slot in range(1, 5)] + planner.flush(SimpleNamespace(writes=SimpleNamespace(rows=rows))) + planner._stager.copy_(planner.buffers(BucketKey(WALK, 4, 4), 0).cells, [0] * 16) + torch.cuda.synchronize() + assert planner.store.view(9, stride)[1:5].tolist() == [list(range(100, 100 + stride))] * 4 + + +# ---- the block table's width ---- + +def _kv(max_seq_len: int, max_num_pages: int) -> KVManager: + return KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=max_seq_len, + max_num_pages=max_num_pages, page_size=PAGE, + ), + name="kv", joint_comm_group=None, transfer_engine_info=None, device=CPU, dtype=torch.float32, + ) + + +def test_a_row_is_as_wide_as_the_pool_when_the_context_is_wider(): + kv = _kv(max_seq_len=1 << 20, max_num_pages=64) + ResidentPlanner("n", {WALK}, {"kv": kv}, CPU, max_bs=1) + assert kv._resident.block_table.width == kv.config.request_pages == 63 + + +def test_a_row_is_as_wide_as_the_context_when_the_pool_is_wider(): + kv = _kv(max_seq_len=4 * 10 + 1, max_num_pages=64) + ResidentPlanner("n", {WALK}, {"kv": kv}, CPU, max_bs=1) + assert kv._resident.block_table.width == 11 + + +def test_a_stream_wider_than_a_row_raises(): + kv = _kv(max_seq_len=8, max_num_pages=64) + planner = ResidentPlanner("n", {WALK}, {"kv": kv}, CPU, max_bs=1) + tables = kv._resident + assert tables.block_table.width == 2 + with pytest.raises(ValueError, match="fit a block table row"): + tables.row(9, [1, 2, 3], 1, [0] * planner.narrow_cols) + plan = SimpleNamespace( + slots=torch.tensor([1]).numpy(), rids=(7,), taken=set(), writes=resident_mod.ResidentWrites("main"), + crossed=None, + ) + tables.host_len[1], tables.host_cols[1] = 8, 2 + with pytest.raises(ValueError, match="fit a block table row"): + tables.stage_step(plan, lambda rid: [1, 2, 3], planner.rewrite) + assert not plan.writes.cells + + +# ---- what a raise in the plan leaves ---- + +def _preplan_worker(engine): + worker = Worker.__new__(Worker) + worker.worker_id = "w" + worker.enable_nvtx = False + worker._phase_period = 0 + worker.engine_manager = SimpleNamespace(get_engine=lambda node: engine) + return worker + + +def _preplan_engine(res, planner): + from mstar.engine.engine import Engine + from mstar.engine.resources.runner import StepRunner + + resources = {k: res[k] for k in ("kv", "pool", "rope")} + engine = Engine.__new__(Engine) + engine._enable_nvtx = False + engine._device = CPU + engine._runner = StepRunner(resources) + runner = SimpleNamespace( + plan_stream=lambda: None, pad_inputs=lambda lease, inputs: inputs, + step_ids=lambda lease, rids: list(rids), release=lambda lease, rows: None, + ) + + def declare(graph_walk, request_ids, inputs, **kwargs): + del graph_walk, inputs, kwargs + segments = [Segment(rid, "main", 1) for rid in request_ids] + keys = {"kv": KVStep(), "pool": RecurrentStep(), "rope": PositionStep()} + return SubmoduleStep(segments=segments, steps=keys) + + engine._submodules = {"n": SimpleNamespace( + submodule=SimpleNamespace(declare_step=declare, step_is_input_free=lambda walk: False), + accelerator_graph_runner=runner, piecewise_runners={}, resident=planner, + )} + engine.can_pre_plan = lambda node: True + engine.reserve_replay_slot = lambda batch: batch.step_context.slot_lease + return engine + + +def _preplan_batch(planner, rids, bs): + bucket = BucketKey(graph_walk=WALK, bs=bs, num_tokens=bs) + lease = SlotLease(slot=0, bucket=bucket) + ctx = StepContext(request_ids=tuple(rids), graph_walk=WALK, slot=0, capture=False, slot_lease=lease) + ctx.resident = planner + return SimpleNamespace( + node_name="n", inputs=[SimpleNamespace() for _ in rids], request_ids=list(rids), per_request_info={}, + step_context=ctx, preplan_event=None, preplanned_rids=None, step=None, + ) + + +@pytest.mark.parametrize("raiser", ["resident", "runner"]) +def test_a_raise_in_the_preplan_leaves_the_reservation_for_exec(monkeypatch, raiser): + """`pre_admit` reserved a crossing page and something raises before the stage + is done. The plan thread's wrapper resets the pre-plan: the page stays with the + stream, the in-flight mark goes, and exec's admit finds the page and takes no + other. The same for a raise from the planner and from the runner's own pre-plan.""" + res = _resources() + planner = ResidentPlanner("n", {WALK}, {k: res[k] for k in ("kv", "pool", "rope")}, CPU, max_bs=4) + _ingest(res, 1) + step, ctx = _host_step(res, [1], [3]) + planner.commit(step, ctx) + s = _captured_step(res, planner, [1], slot=0, bs=1) + _assert_step_read_its_streams(res, s, [1]) + kv, stream = res["kv"], res["kv"]._streams[1]["main"] + assert stream.stored_len == 4 and len(stream.page_indices) == 1 + + engine = _preplan_engine(res, planner) + batch = _preplan_batch(planner, [1], 1) + target = planner if raiser == "resident" else engine._runner + name = "plan" if raiser == "resident" else "pre_plan" + real = getattr(target, name) + + def boom(*args, **kwargs): + raise RuntimeError("raised mid pre-plan") + + monkeypatch.setattr(target, name, boom) + assert _preplan_worker(engine)._preplan_spec(None, SimpleNamespace(node_batch=batch)) is False + monkeypatch.setattr(target, name, real) + + assert batch.preplan_event is None and batch.preplanned_rids is None and batch.step is None + assert engine._runner._pre_admitted is None and not engine._runner.staged + assert not stream.step_in_flight + assert len(stream.page_indices) == 2 + pages, free = list(stream.page_indices), kv._arena.num_free + kv.assert_pages_conserved() + + # exec re-drives it: the page it finds is the one pre_admit took + s = _captured_step(res, planner, [1], slot=1, bs=1) + assert stream.page_indices == pages and kv._arena.num_free == free + assert s.pages[0] == pages[1] + _assert_step_read_its_streams(res, s, [1]) + _assert_row(planner, res, 1) + kv.assert_pages_conserved() + + +def test_a_plan_that_raises_leaves_its_rows_to_the_next_step(monkeypatch): + live, ref, planner = _setup({1: 3, 2: 5}) + del ref + kv = live["kv"] + real = kv.resident_row + + def boom(rid, slot, row, writes): + if rid == 2: + raise ValueError("a row that can't be written") + real(rid, slot, row, writes) + + monkeypatch.setattr(kv, "resident_row", boom) + s = _captured_step(live, planner, [1, 2], slot=0, bs=2, run=False) + with pytest.raises(ValueError, match="can't be written"): + planner.plan(s.step, s.ctx) + assert planner._stale == {1, 2} + monkeypatch.setattr(kv, "resident_row", real) + _replay(live, planner, s) + _assert_row(planner, live, 1) + _assert_row(planner, live, 2) + + +# ---- what an offload does to a slot ---- + +class _HostPagesKV: + """The offload hooks of a KV manager over a real one's streams and arena, for a + host with no page-locked memory: a stream's pages go back to the pool and the + reload takes others, which is all a row can see of the real move.""" + + def __init__(self, kv: KVManager): + self.kv = kv + self.away: dict[int, int] = {} + kv.__class__ = type("_Offloading", (KVManager,), { + "offload": lambda _, rid: self.offload(rid), + "reload": lambda _, rid: self.reload(rid), + "is_offloaded": lambda _, rid: rid in self.away, + "can_reload": lambda _, rid: True, + "supports_eviction": property(lambda _: True), + "reclaimable": lambda _, rid: len(kv._streams[rid]["main"].page_indices) if rid not in self.away else 0, + }) + + def offload(self, rid: int) -> int: + stream = self.kv._streams[rid]["main"] + if stream.step_in_flight: + return 0 + pages = stream.page_indices + self.kv._touch_plan_cache(rid, row=False) + stream.page_indices = [] + self.kv._arena.release(pages) + self.away[rid] = len(pages) + return len(pages) + + def reload(self, rid: int) -> bool: + self.kv._touch_plan_cache(rid, row=False) + stream = self.kv._streams[rid]["main"] + stream.page_indices = self.kv._arena.acquire(self.away.pop(rid)) + return True + + +def _offload_engine(resources, planner): + from mstar.engine.engine import Engine + + engine = Engine.__new__(Engine) + engine._tp_follower_nodes = set() + engine._resident_delta = {} + engine._resources = resources + engine._submodules = {"n": SimpleNamespace(resources=resources, resident=planner)} + engine._graph_runtime = SimpleNamespace(get_rid_string=str, get_rid_handle=int) + return engine + + +def _refused_worker(engine, line): + """The real handler and the real victim search over a real engine.""" + worker = Worker.__new__(Worker) + worker.parallel_nodes = set() + worker.parallel_leader_nodes = set() + worker._hold_logged = {} + worker.pushed = [] + worker.held = [] + worker._apply_removes_whose_step_landed = lambda: None + worker._apply_pending_removes_safe_to_drop = lambda rids: None + worker._graph_runtime = SimpleNamespace(push_back_node=lambda node, rids, graphs: worker.pushed.append(rids)) + worker.engine_manager = SimpleNamespace(get_engine=lambda node: engine) + worker.eviction_policy = SimpleNamespace(value="last_in_line") + worker.scheduler = SimpleNamespace( + arrival=list(line), hold_requests=worker.held.extend, could_refill=lambda *a: False, + last_in_line=lambda node, walk, rids: max(rids, key=line.index), + ) + return worker + + +def test_a_refused_admit_offloads_a_victim_and_its_reload_is_written_whole(): + """Through the worker's refusal path: the KV pool refuses, `_handle_allocation_failure` + offloads the row furthest from a step, its slot goes with it, the retry runs, and the + victim's reload is written whole at its own lengths.""" + from mstar.engine.resources.runner import StepRunner + + # seven pages to hand out, all held: each of 1 and 2 has room for one more token + live = _resources(max_num_pages=8) + owned = {k: live[k] for k in ("kv", "pool", "rope")} + planner = ResidentPlanner("n", {WALK}, owned, CPU, max_bs=4) + for rid, tokens in ((1, 7), (2, 7), (3, 11)): + _ingest(live, rid) + step, ctx = _host_step(live, [rid], [tokens]) + planner.commit(step, ctx) + kv = live["kv"] + _HostPagesKV(kv) + _captured_step(live, planner, [1, 2, 3], slot=0, bs=3) + assert planner.slot_of(3) is not None and kv._arena.num_free == 0 + engine = _offload_engine(owned, planner) + + # both cross a page now + segments = tuple(Segment(rid, "main", 1) for rid in (1, 2)) + step = SubmoduleStep( + segments=list(segments), steps={"kv": KVStep(), "pool": RecurrentStep(), "rope": PositionStep()}, + ) + step.set_ctx(StepContext(request_ids=(1, 2), graph_walk=WALK, slot=0, capture=False)) + outcome = StepRunner(owned).admit(step) + assert not outcome.ok and isinstance(outcome.reason, AllocationFailed) and outcome.failed_resource == "kv" + + worker = _refused_worker(engine, line=[1, 2, 3]) + batch = SimpleNamespace(node_name="n", graph_walk=WALK, request_to_worker_graph={1: "g", 2: "g"}) + worker._handle_allocation_failure(batch, SimpleNamespace( + node_name="n", admit_error=outcome.reason, failed_resource=outcome.failed_resource, + )) + assert worker.pushed == [[1, 2]] and worker.held == [3] + assert engine.is_offloaded("n", 3) + assert planner.slot_of(3) is None and 3 not in planner._stale + + # the retry runs on the victim's pages + s = _captured_step(live, planner, [1, 2], slot=0, bs=2) + _assert_step_read_its_streams(live, s, [1, 2]) + _assert_row(planner, live, 1) + + for key in ("kv", "pool", "rope"): + live[key].remove_request(2) + planner.release(2) + assert engine.reload_request("n", 3) + assert planner.slot_of(3) is None + s = _captured_step(live, planner, [3], slot=0, bs=1) + assert s.ctx.resident_plan.taken == {3} + stream = kv._streams[3]["main"] + assert stream.stored_len == 13 + assert kv._resident.host_len[planner.slot_of(3)] == stream.stored_len + _assert_step_read_its_streams(live, s, [3]) + _assert_row(planner, live, 3) + + +def test_a_recurrent_pool_refusal_holds_the_request_since_the_pool_cannot_evict(): + """The pool refuses with `AllocationFailed`, which is backpressure: it has no + eviction, so nothing is offloaded and the row it named is held for a smaller retry.""" + live = _resources(pool_slots=2) + planner = ResidentPlanner("n", {WALK}, {k: live[k] for k in ("kv", "pool", "rope")}, CPU, max_bs=4) + for rid in (1, 2): + _ingest(live, rid) + step, ctx = _host_step(live, [1], [3]) + planner.commit(step, ctx) + _HostPagesKV(live["kv"]) + pool_step = RecurrentStep(segments=(Segment(2, "main", 1),)) + pool_ctx = StepContext(request_ids=(2,), graph_walk="prefill", slot=0, capture=False) + live["kv"].admit(KVStep(segments=(Segment(2, "main", 1),)), pool_ctx) + outcome = live["pool"].admit(pool_step, pool_ctx) + assert not outcome.ok and isinstance(outcome.reason, AllocationFailed) + assert outcome.reason.request_id == 2 + + engine = _offload_engine({k: live[k] for k in ("kv", "pool", "rope")}, planner) + worker = _refused_worker(engine, line=[1, 2]) + batch = SimpleNamespace(node_name="n", graph_walk="prefill", request_to_worker_graph={2: "g"}) + worker._handle_allocation_failure(batch, SimpleNamespace( + node_name="n", admit_error=outcome.reason, failed_resource="pool", + )) + assert worker.pushed == [[2]] + assert worker.held == [2] + assert not engine.is_offloaded("n", 1) + + +def test_a_follower_replaying_an_offload_gives_the_slot_up_too(): + live, ref, planner = _setup({1: 9}) + del ref + _captured_step(live, planner, [1], slot=0, bs=1) + _HostPagesKV(live["kv"]) + engine = _offload_engine({k: live[k] for k in ("kv", "pool", "rope")}, planner) + engine._tp_follower_nodes = {"n"} + delta = OffloadDelta.new() + delta.add_offloaded("1") + assert engine.apply_resident_delta("n", delta) + assert planner.slot_of(1) is None # ---- refusals ---- @@ -649,8 +1234,8 @@ def build(cls, spec, info): return cls() -def _build(resources, nodes=None, walks=(WALK,)): - return build_planner("n", walks, resources, nodes or {}, CPU, 8) +def _build(resources, nodes=None, walks=(WALK,), max_bs=4): + return build_planner("n", walks, resources, nodes or {}, CPU, max_bs) def test_every_resident_resource_passes(): @@ -665,6 +1250,13 @@ def test_no_walks_no_planner(): assert planner is None and "no resident walk" in why +def test_no_capture_of_the_walks_no_planner(): + res = _resources() + planner, why = _build({"kv": res["kv"]}, max_bs=0) + assert planner is None and "nothing captures" in why + assert res["kv"]._resident is None + + def test_a_resource_without_a_table_refuses(): planner, why = _build({"other": _Spare()}) assert planner is None and "_Spare keeps no resident table" in why @@ -772,6 +1364,13 @@ def test_a_resident_walk_is_one_token_a_row_in_one_stream(): # ---- the engine ---- +def _capture_runner(max_bs=4): + bound = {} + return SimpleNamespace( + max_capture_bs=lambda walks: max_bs, bind_resident=lambda planner: bound.update(p=planner), bound=bound, + ) + + def _engine(walks, resources, device=CPU): from mstar.engine.engine import Engine from mstar.engine.resources.runner import StepRunner @@ -785,12 +1384,14 @@ def _engine(walks, resources, device=CPU): return engine, engine._submodules["n"] -def test_the_engine_builds_a_planner_for_a_capable_node(monkeypatch): - monkeypatch.setattr(knobs, "resident_slots", lambda: 8) +def test_the_engine_builds_a_planner_for_a_capable_node(): res = _resources() engine, mgmt = _engine({WALK}, {k: res[k] for k in ("kv", "pool", "rope")}) - engine._build_resident_planner("n") + runner = _capture_runner() + engine._build_resident_planner("n", runner) assert isinstance(mgmt.resident, ResidentPlanner) + assert runner.bound["p"] is mgmt.resident + assert mgmt.resident.max_slots == 9 assert mgmt.resident.covers(BucketKey(WALK, 2, 2)) assert not mgmt.resident.covers(BucketKey("prefill_text", 2, 64)) @@ -799,17 +1400,16 @@ def test_the_knob_turns_it_off(monkeypatch): monkeypatch.setenv("MSTAR_RESIDENT_PLANNING", "off") res = _resources() engine, mgmt = _engine({WALK}, {"kv": res["kv"]}) - engine._build_resident_planner("n") + engine._build_resident_planner("n", _capture_runner()) assert mgmt.resident is None assert res["kv"]._resident is None -def test_a_refused_node_stays_on_the_host_and_says_why(monkeypatch, caplog): - monkeypatch.setattr(knobs, "resident_slots", lambda: 8) +def test_a_refused_node_stays_on_the_host_and_says_why(caplog): res = _resources() engine, mgmt = _engine({WALK}, {"kv": res["kv"], "other": _Spare()}) with caplog.at_level("WARNING"): - engine._build_resident_planner("n") + engine._build_resident_planner("n", _capture_runner()) assert mgmt.resident is None assert "_Spare keeps no resident table" in caplog.text assert res["kv"]._resident is None @@ -819,7 +1419,7 @@ def test_a_node_with_no_resident_walk_is_left_alone(caplog): res = _resources() engine, mgmt = _engine((), {"kv": res["kv"]}) with caplog.at_level("WARNING"): - engine._build_resident_planner("n") + engine._build_resident_planner("n", _capture_runner()) assert mgmt.resident is None assert not caplog.text @@ -872,6 +1472,20 @@ def test_qwen35_captured_decode_preprocess_reads_nothing_the_tables_fill(rows): assert set(out) == {"input_ids"} +def test_a_capture_runner_answers_for_the_walks_it_captures(): + from mstar.engine.accelerator_graph_runner import AcceleratorGraphRunner as Runner + + runner = Runner.__new__(Runner) + runner._capture_configs = [ + SimpleNamespace(capture_graph_walk="decode", capture_batch_sizes=[1, 8, 128]), + SimpleNamespace(capture_graph_walk="prefill", capture_batch_sizes=[1, 512]), + SimpleNamespace(capture_graph_walk="other", capture_batch_sizes=None), + ] + assert runner.max_capture_bs({"decode"}) == 128 + assert runner.max_capture_bs({"other"}) == max(Runner.CAPTURE_BATCH_SIZES) + assert runner.max_capture_bs({"absent"}) == 0 + + # ---- the knob ---- def test_the_knob_takes_auto_and_off(): @@ -889,3 +1503,8 @@ def test_the_knob_defaults_to_auto(monkeypatch): assert knobs.resident_planning() == "auto" monkeypatch.setenv("MSTAR_RESIDENT_PLANNING", "off") assert knobs.resident_planning() == "off" + + +def test_the_slot_count_is_no_knob(): + with pytest.raises(ValueError): + knobs.check({"resident_slots": 8}) From d279b7a8d4622df57275ff75fc1ab0d0a1ff807c Mon Sep 17 00:00:00 2001 From: merceod Date: Tue, 3 Nov 2026 01:36:34 +0000 Subject: [PATCH 10/15] test_resident_delta's engine stub binds the slot release --- test/modular/test_resident_delta.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/test/modular/test_resident_delta.py b/test/modular/test_resident_delta.py index 284a4cb93a..162711c174 100644 --- a/test/modular/test_resident_delta.py +++ b/test/modular/test_resident_delta.py @@ -191,6 +191,7 @@ class _Engine: check_ready = Engine.check_ready is_offloaded = Engine.is_offloaded offload_request = Engine.offload_request + _release_resident_slot = Engine._release_resident_slot reclaimable = Engine.reclaimable reload_request = Engine.reload_request take_resident_delta = Engine.take_resident_delta @@ -209,7 +210,7 @@ def __init__(self, role: str = "leader"): self.resource = _Resource() self._resources = {"kv": self.resource} self._submodules = { - "node": SimpleNamespace(resources={"kv": self.resource}) + "node": SimpleNamespace(resources={"kv": self.resource}, resident=None) } self._runner = SimpleNamespace(admit_retrieve=lambda **kw: FULL_ADMIT_OK) self._resident_delta = ( @@ -317,7 +318,7 @@ def test_a_partial_reload_gives_back_what_it_reloaded(): first.offloaded = {"a"} second.offloaded = {"a"} engine._submodules = { - "node": SimpleNamespace(resources={"first": first, "second": second}) + "node": SimpleNamespace(resources={"first": first, "second": second}, resident=None) } assert engine.reload_request("node", "a") is False @@ -346,7 +347,7 @@ def test_a_resource_that_never_offloaded_the_request_does_not_block_its_reload() holder, bystander = _Unsure(), _Unsure() holder.offloaded = {"a"} engine._submodules = { - "node": SimpleNamespace(resources={"holder": holder, "bystander": bystander}) + "node": SimpleNamespace(resources={"holder": holder, "bystander": bystander}, resident=None) } assert engine.reload_request("node", "a") is True @@ -360,7 +361,7 @@ def test_a_resource_that_cannot_fit_stops_the_reload_before_anything_moves(): second.offloaded = {"a"} first.reload = second.reload = lambda rid: pytest.fail("reloaded past a refusal") engine._submodules = { - "node": SimpleNamespace(resources={"first": first, "second": second}) + "node": SimpleNamespace(resources={"first": first, "second": second}, resident=None) } assert engine.reload_request("node", "a") is False From 254e0fd7cbc117ea2ec971a672c77e33df93331f Mon Sep 17 00:00:00 2001 From: merceod Date: Tue, 3 Nov 2026 02:23:09 +0000 Subject: [PATCH 11/15] An abandoned stage is staged again, and capture refuses commit=False Its slot row went out on the plan stream with nothing ordering the next step behind it. Neither changes a steady step. --- mstar/engine/resources/resident.py | 17 ++++++++++----- test/modular/test_resident_planning.py | 29 +++++++++++++++++++++++++- 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/mstar/engine/resources/resident.py b/mstar/engine/resources/resident.py index ac4c80b888..a49b327229 100644 --- a/mstar/engine/resources/resident.py +++ b/mstar/engine/resources/resident.py @@ -246,7 +246,7 @@ def __init__( self._seal() # the largest flush up front, plus a cells row: growing the stager waits # out its copies, which on the GPU thread would be a sync in the step path - self._stager = PinnedStager(torch.long, numel=max_bs * self._tables[0].stride + 4 * max_bs) + self._stager = PinnedStager(torch.long, numel=max_bs * self._stride + 4 * max_bs) # Tables @@ -265,7 +265,7 @@ def _seal(self) -> None: raise ValueError("a resident row holds one wide table") narrow = [t for t in self._tables if t.width == 1] self.narrow_cols = len(narrow) - stride = sum(t.width for t in self._tables) + stride = self._stride = sum(t.width for t in self._tables) # int64: what the KV write addressing and the positions are, and index ops take self._store = torch.zeros(self.max_slots * stride, dtype=torch.long, device=self.device) rows = self._store.view(self.max_slots, stride) @@ -403,10 +403,16 @@ def drop_stage(self) -> None: self._restore(staged) def _restore(self, plan: ResidentPlan) -> None: - """Hand back the rows of a plan that won't run: it wrote nothing.""" + """Hand back the rows of a plan that won't run: it wrote nothing but its + staging. The next plan stages its row again, even an equal one: this + one's copy is on the plan stream, and the replacing step waits on no + event for it.""" with self._lock: self._stale |= {rid for rid in plan.taken if rid in self._slot_of} self._moved |= {rid for rid in plan.moved if rid in self._slot_of} + buffers = self._buffers.get(plan.key[1:]) + if buffers is not None: + buffers.staged = None def _build(self, rids: tuple[int, ...], key: tuple, label: str, ctx) -> ResidentPlan: lease = ctx.slot_lease @@ -485,7 +491,7 @@ def flush(self, plan: ResidentPlan) -> None: rows = plan.writes.rows if not rows: return - stride = self._tables[0].stride + stride = self._stride dsts = [self._store[slot * stride: slot * stride + len(row)] for slot, row in rows] self._stager.copy_runs_(dsts, np.fromiter(chain.from_iterable(r for _, r in rows), dtype=np.int64)) @@ -537,7 +543,8 @@ def capture_step(self, bucket: BucketKey, slot: int, step: SubmoduleStep) -> Cap name for name in ("pre_forks", "post_forks", "retention", "sink", "advance", "pos_ids", "fork_rids") if getattr(resource_step, name, None) ), - None, + # a step that keeps no tokens: the table's commit would advance past its stream + "commit=False" if getattr(resource_step, "commit", True) is False else None, ) if refusal is not None: raise ValueError( diff --git a/test/modular/test_resident_planning.py b/test/modular/test_resident_planning.py index ef1c384ce7..f15d045231 100644 --- a/test/modular/test_resident_planning.py +++ b/test/modular/test_resident_planning.py @@ -526,6 +526,28 @@ def test_a_new_row_abandoned_before_its_first_step_is_still_written(): _assert_row(planner, live, 1) +def test_an_abandoned_stage_of_a_steady_batch_is_staged_again(monkeypatch): + """Its slot row went out on the plan stream, and the step replacing it waits on no + event for that, so an equal row is no reason to skip the copy on its own stream.""" + live, ref, planner = _setup({1: 3, 2: 5}) + del ref + _captured_step(live, planner, [1, 2], slot=0, bs=2) + _captured_step(live, planner, [1, 2], slot=0, bs=2, preplan=True, run=False) + for key in ("kv", "pool", "rope"): + live[key].clear_preplan() + buffers = planner.buffers(BucketKey(WALK, 2, 2), 0) + copies = [] + real = planner._stager.copy_ + monkeypatch.setattr(planner._stager, "copy_", lambda dst, *a, **k: (copies.append(dst), real(dst, *a, **k))[1]) + s = _captured_step(live, planner, [1, 2], slot=0, bs=2) + assert any(dst is buffers.slots for dst in copies) + _assert_step_read_its_streams(live, s, [1, 2]) + # and a steady step after it still copies nothing + copies.clear() + _captured_step(live, planner, [1, 2], slot=0, bs=2) + assert not any(dst is buffers.slots for dst in copies) + + def test_a_promoted_preplan_reuses_its_stage(): live, ref, planner = _setup({1: 3}) del ref @@ -1238,6 +1260,11 @@ def _build(resources, nodes=None, walks=(WALK,), max_bs=4): return build_planner("n", walks, resources, nodes or {}, CPU, max_bs) +def test_a_node_with_no_table_owner_still_builds(): + planner = ResidentPlanner("n", {WALK}, {}, CPU, max_bs=2) + assert planner.max_slots == 5 + + def test_every_resident_resource_passes(): res = _resources() planner, why = _build({k: res[k] for k in ("kv", "pool", "rope")}) @@ -1342,7 +1369,7 @@ def test_flashinfer_refuses_a_sliding_window(): @pytest.mark.parametrize("field,value", [ - ("pre_forks", (("a", "b"),)), ("post_forks", (("a", "b"),)), ("sink", frozenset({(1, "main")})), + ("pre_forks", (("a", "b"),)), ("post_forks", (("a", "b"),)), ("sink", frozenset({(1, "main")})), ("commit", False), ]) def test_a_resident_walk_can_not_declare_forks_or_a_sink(field, value): _, _, planner = _setup({}) From 0af5053dd08e25a64f4d76a5707ca3e3054dbff1 Mon Sep 17 00:00:00 2001 From: merceod Date: Tue, 3 Nov 2026 04:39:57 +0000 Subject: [PATCH 12/15] Resident state per resource key, and pools a step only reads The planner held one stream per row: one wide table, one label, one stored length and one block table per capture, so a second KV pool on a node didn't work. Each of those is keyed by resource key now: - A row is the narrow columns, then a span per block table. A rewrite queues a run per span, and a run that starts where the last one ended joins it, so a node with one pool still writes a row in one copy. - A capture records each resource's stream (label, and whether it appends a token a row or only reads) once per bucket; every owner reads its own. Positions take their offset against, and gather with, the length of the pool their PositionConfig.kv_cache names. - Each KV pool gathers its lengths into a buffer of its own, and an fmha decode wrapper reads its own pool's table and lengths. A bucket skips the host plan only if every attention read its pages off the tables. - A pool declared at span 0 (a cross-attention context) is read only: the row holds its pages and length, with no write addressing, no advance and no crossings, and a row a host step moved is rewritten only if that stream changed. FlashInfer cross attention over a pool of its own runs on the planner, host-planned as before. --- .claude/skills/engine-resources/SKILL.md | 2 +- docs/adding_models.rst | 8 +- mstar/engine/resources/attn/cross.py | 14 +- mstar/engine/resources/attn/flashinfer.py | 14 +- mstar/engine/resources/kv/manager.py | 40 +- mstar/engine/resources/kv/resident.py | 65 +++- mstar/engine/resources/position/manager.py | 18 +- mstar/engine/resources/recurrent/pool.py | 10 +- mstar/engine/resources/resident.py | 170 ++++++--- mstar/engine/resources/resident_hooks.py | 13 +- test/modular/test_resident_planning.py | 418 ++++++++++++++++++++- 11 files changed, 657 insertions(+), 115 deletions(-) diff --git a/.claude/skills/engine-resources/SKILL.md b/.claude/skills/engine-resources/SKILL.md index 46885d5029..124a8201d9 100644 --- a/.claude/skills/engine-resources/SKILL.md +++ b/.claude/skills/engine-resources/SKILL.md @@ -34,7 +34,7 @@ Prototype a new kernel standalone against a dense reference before writing the r ## 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. 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. +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 diff --git a/docs/adding_models.rst b/docs/adding_models.rst index becf88534b..c9a42d7827 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -1638,9 +1638,11 @@ 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 one token per row in one stream, 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. Capture raises on a declaration that breaks this. The engine +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 diff --git a/mstar/engine/resources/attn/cross.py b/mstar/engine/resources/attn/cross.py index ae08485751..157649bf52 100644 --- a/mstar/engine/resources/attn/cross.py +++ b/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) + 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 diff --git a/mstar/engine/resources/attn/flashinfer.py b/mstar/engine/resources/attn/flashinfer.py index 273b0e8550..96d77bb1f8 100644 --- a/mstar/engine/resources/attn/flashinfer.py +++ b/mstar/engine/resources/attn/flashinfer.py @@ -104,11 +104,15 @@ def resident_refusal(self) -> str | None: return None def resident_gather(self, step: "CapturedResidentStep") -> None: - wrapper = self._cg_plan_states.get(CGSlotKey(bucket=step.bucket, slot=step.slot, label=step.label)) - if isinstance(wrapper, FmhaDecodeWrapper): - table = step.block_table - wrapper.gather_resident(step.store, step.slots, step.buffers.kv_len, table.stride, table.col) - step.host_plan_free = True + # planned per label of its KV cache's plan, so its wrapper is under that stream's label + pool, stream = self._kv_cache_name, step.streams[self._kv_cache_name] + wrapper = self._cg_plan_states.get(CGSlotKey(bucket=step.bucket, slot=step.slot, label=stream.label)) + if isinstance(wrapper, FmhaDecodeWrapper) and not stream.reads: + table = step.block_tables[pool] + wrapper.gather_resident(step.store, step.slots, step.buffers.kv_len[pool], table.stride, table.col) + step.reads_tables = True + else: + step.host_planned = True def _check_layer(self, layer_idx: int) -> None: """Over a pool of more than one group, ``layer_idx`` must be a layer this attention runs, diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index 767624398e..1e08784db0 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -2947,7 +2947,8 @@ def plan(self, step: KVStep, ctx: StepContext) -> dict[str, KVPlanOutput]: for g, outputs in enumerate(res.groups): self._setup_plan_states(outputs, ctx, ctx.slot_lease, g, eager_write=step.eager_write) else: - # the replay fills the slot's addressing, so no reuse may skip re-staging it + # the replay fills the slot's addressing (a stream the step only + # reads has none), so no reuse may skip re-staging it for label in res: self._staged_outputs.pop((ctx.slot_lease.slot, (label, 0)), None) laps.lap("setup_states") @@ -4512,39 +4513,50 @@ def resident_refusal(self) -> str | None: return "a pool read across nodes takes pages a resident row doesn't track" return None - def bind_resident(self, planner: "ResidentPlanner") -> None: + def bind_resident(self, planner: "ResidentPlanner", key: str) -> None: from mstar.engine.resources.kv.resident import KVResidentTables page_size = self._groups[0].page_size # the pool's own context or override (not the serving config's top-level # `max_seq_len`), and a stream's pages are distinct pages of the pool width = min(-(-self.config.max_seq_len // page_size), self.config.request_pages) - self._resident = KVResidentTables(planner, page_size, width) + self._resident = KVResidentTables(planner, key, page_size, width) def resident_row(self, rid: int, slot: int, row: list[int], writes: "ResidentWrites") -> None: - stream = self._streams[rid][writes.label] + tables = self._resident + label = writes.streams[tables.key].label + stream = self._streams[rid][label] if stream.released: raise ValueError( - f"KV {self.name}: {rid}/{writes.label} released {stream.released} tokens, " + f"KV {self.name}: {rid}/{label} released {stream.released} tokens, " "which a block table row can't address" ) - writes.kv_len = stream.stored_len - self._resident.row(stream.stored_len, stream.page_indices, slot, row) + tables.row(stream.stored_len, stream.page_indices, slot, row, writes) def resident_stage(self, plan: "ResidentPlan") -> None: - tables, streams, label = self._resident, self._streams, plan.writes.label - planner = tables.planner - for rid in plan.moved: - tables.refresh_len(streams[rid][label].stored_len, planner.slot_of(rid), plan.writes) - tables.stage_step(plan, lambda rid: streams[rid][label].page_indices, planner.rewrite) + tables, streams = self._resident, self._streams + stream = plan.writes.streams[tables.key] + label, planner = stream.label, tables.planner + if not stream.reads: + for rid in plan.moved: + tables.refresh_len(streams[rid][label].stored_len, planner.slot_of(rid), plan.writes) + tables.stage_step( + plan, lambda rid: streams[rid][label].stored_len, lambda rid: streams[rid][label].page_indices, + planner.rewrite, + ) if _DEBUG_ASSERTS: for rid, slot in zip(plan.rids, plan.slots.tolist(), strict=True): held, stored = int(tables.host_len[slot]), streams[rid][label].stored_len assert held == stored, f"KV {self.name}: resident row of {rid} at {held}, its stream at {stored}" def resident_gather(self, step: "CapturedResidentStep") -> None: - state = self._static_plan_state(step.slot, (step.label, 0)) - self._resident.gather(step, state.token_to_page, state.token_to_cache) + tables = self._resident + stream = step.streams[tables.key] + if stream.reads: + tables.gather(step, None, None) + return + state = self._static_plan_state(step.slot, (stream.label, 0)) + tables.gather(step, state.token_to_page, state.token_to_cache) def resident_commit(self, step: "CapturedResidentStep") -> None: self._resident.commit(step) diff --git a/mstar/engine/resources/kv/resident.py b/mstar/engine/resources/kv/resident.py index 7589ce01c2..3aad74d9e7 100644 --- a/mstar/engine/resources/kv/resident.py +++ b/mstar/engine/resources/kv/resident.py @@ -13,7 +13,6 @@ ResidentWrites, advance_column, gather_column, - scatter_cells, ) @@ -33,27 +32,37 @@ class KVResidentTables: stored length: what a step appending one token per row reads to find where it writes and what it attends to. + A step that only reads the stream (declared at span 0, as a context is) + leaves it where it is: the row holds the stream's pages and length, and + has no write addressing, no advance and no crossings. + The stream stays the truth. ``host_len`` advances by the rule the graph does, so nothing comes back from the device, and ``host_cols`` counts the columns a row holds, which is how the host sees a crossing. """ - def __init__(self, planner: ResidentPlanner, page_size: int, max_pages: int): + def __init__(self, planner: ResidentPlanner, key: str, page_size: int, max_pages: int): self.planner = planner + self.key = key self.page_size = page_size self.block_table = planner.table(max_pages) - self.kv_len = planner.table() + self.kv_len = planner.length_table(key) self.host_len = np.zeros(planner.max_slots, dtype=np.int64) self.host_cols = np.zeros(planner.max_slots, dtype=np.int64) - def row(self, stored_len: int, pages: list[int], slot: int, row: list[int]) -> None: - """The cells of a stream about to append one token. Its host mirrors - move here, not at the flush: a plan that never runs leaves the row stale.""" + def row(self, stored_len: int, pages: list[int], slot: int, row: list[int], writes: ResidentWrites) -> None: + """The cells of a stream about to append one token, or to be read. Its + host mirrors move here, not at the flush: a plan that never runs leaves + the row stale.""" if len(pages) > self.block_table.width: raise ValueError(f"{len(pages)} pages of a stream don't fit a block table row of {self.block_table.width}") - cols = stored_len // self.page_size + 1 + page = self.page_size + # a stream that appends holds the page its next token goes to, too + cols = -(-stored_len // page) if writes.streams[self.key].reads else stored_len // page + 1 row[self.kv_len.col] = stored_len - row.extend(pages[:cols]) + writes.kv_len[self.key] = stored_len + if cols: + writes.run(slot, self.block_table.col, pages[:cols]) self.host_len[slot] = stored_len self.host_cols[slot] = cols @@ -62,12 +71,24 @@ def refresh_len(self, stored_len: int, slot: int, writes: ResidentWrites) -> Non writes.put(self.kv_len.cell(slot), stored_len) self.host_len[slot] = stored_len - def stage_step(self, plan: ResidentPlan, stream_pages, rewrite) -> None: + def stage_step(self, plan: ResidentPlan, stream_len, stream_pages, rewrite) -> None: """Bring each row's columns up to this step. The usual case, one new column, rides the graph's scatter, ahead of its gather: only columns past the row's length change, so a step still replaying reads none of them. A row further behind (a host-planned step took its pages) is + rewritten whole. + + A stream the step only reads moves only under a host-planned step, so + just the rows one moved are compared, and one that changed is rewritten whole.""" + if plan.writes.streams[self.key].reads: + slot_of = self.planner.slot_of + for rid in plan.moved: + slot = slot_of(rid) + if self.host_len[slot] != stream_len(rid): + plan.taken.add(rid) + rewrite(rid, slot, plan.writes) + return slots = plan.slots cols = self.host_cols[slots] need = self.host_len[slots] // self.page_size + 1 @@ -83,22 +104,29 @@ def stage_step(self, plan: ResidentPlan, stream_pages, rewrite) -> None: rewrite(rid, slot, plan.writes) else: plan.writes.put(self.block_table.cell(slot, last), stream_pages(rid)[last]) - plan.crossed = (rows, need[rows]) + plan.crossed[self.key] = (rows, need[rows]) def advance(self, plan: ResidentPlan) -> None: + if plan.writes.streams[self.key].reads: + return self.host_len[plan.slots] += 1 - if plan.crossed is not None: - rows, need = plan.crossed + crossed = plan.crossed.get(self.key) + if crossed is not None: + rows, need = crossed self.host_cols[plan.slots[rows]] = need - def gather(self, step: CapturedResidentStep, token_to_page: torch.Tensor, token_to_cache: torch.Tensor) -> None: + def gather( + self, step: CapturedResidentStep, token_to_page: torch.Tensor | None, token_to_cache: torch.Tensor | None, + ) -> None: + """The rows' stored lengths, and where each writes unless the step only + reads. The planner has scattered the step's single cells already.""" store, table = self.planner.store, self.block_table stride, page_size = table.stride, self.page_size - step.store, step.block_table = store, table - # before the reads: a row that crossed writes into its new page - scatter_cells(store, step.buffers.cells) - kv_len = step.buffers.kv_len + step.block_tables[self.key] = table + kv_len = step.buffers.kv_len[self.key] gather_column(store, step.slots, stride, self.kv_len.col, kv_len) + if step.streams[self.key].reads: + return if store.is_cuda: _gather_write_slots_kernel[(step.bucket.bs,)]( step.slots, kv_len, store, token_to_page, token_to_cache, @@ -111,4 +139,5 @@ def gather(self, step: CapturedResidentStep, token_to_page: torch.Tensor, token_ token_to_cache[:bs].copy_(kv_len - col * page_size) def commit(self, step: CapturedResidentStep) -> None: - advance_column(self.planner.store, step.slots, self.kv_len.stride, self.kv_len.col) + if not step.streams[self.key].reads: + advance_column(self.planner.store, step.slots, self.kv_len.stride, self.kv_len.col) diff --git a/mstar/engine/resources/position/manager.py b/mstar/engine/resources/position/manager.py index 595802d80e..5524e1e880 100644 --- a/mstar/engine/resources/position/manager.py +++ b/mstar/engine/resources/position/manager.py @@ -367,19 +367,29 @@ def resident_refusal(self) -> str | None: return f"position scheme {self._config.scheme} doesn't advance with the stream" return None - def bind_resident(self, planner: "ResidentPlanner") -> None: + def bind_resident(self, planner: "ResidentPlanner", key: str) -> None: self._resident_store = planner + self._resident_key = key self._resident_delta = planner.table() def resident_row(self, rid: int, slot: int, row: list[int], writes: "ResidentWrites") -> None: - row[self._resident_delta.col] = self.position(rid, writes.label) - writes.kv_len + # the KV pool binds ahead of this, its dependent, and sets its length first + label = writes.streams[self._resident_key].label + row[self._resident_delta.col] = self.position(rid, label) - writes.kv_len[self._kv_cache_name] def resident_gather(self, step: "CapturedResidentStep") -> None: - key = CGSlotKey(bucket=step.bucket, slot=step.slot, label=step.label) + own, pool = step.streams[self._resident_key], step.streams[self._kv_cache_name] + if own.label != pool.label or own.reads or pool.reads: + # the offset holds only while both advance a token a step, as one stream + raise ValueError( + f"positions on stream {own.label!r} over {self._kv_cache_name!r}'s {pool.label!r}: " + "a resident step's positions advance with the stream their KV cache appends to" + ) + key = CGSlotKey(bucket=step.bucket, slot=step.slot, label=own.label) table = self._resident_delta gather_column( self._resident_store.store, step.slots, table.stride, table.col, - self._static_buffer(key, step.bucket.num_tokens), plus=step.buffers.kv_len, + self._static_buffer(key, step.bucket.num_tokens), plus=step.buffers.kv_len[self._kv_cache_name], ) def commit(self, step: PositionStep, ctx: StepContext): diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index 3ca3c3e6ba..a406e5e4de 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -1125,15 +1125,19 @@ def resident_refusal(self) -> str | None: return "a pool with no sink slot has nowhere for a padding row" return None - def bind_resident(self, planner: "ResidentPlanner") -> None: + def bind_resident(self, planner: "ResidentPlanner", key: str) -> None: self._resident = planner + self._resident_key = key self._resident_table = planner.table() def resident_row(self, rid: int, slot: int, row: list[int], writes: "ResidentWrites") -> None: - row[self._resident_table.col] = self._slots[rid][writes.label].index + row[self._resident_table.col] = self._slots[rid][writes.streams[self._resident_key].label].index def resident_gather(self, step: "CapturedResidentStep") -> None: - key = CGSlotKey(bucket=step.bucket, slot=step.slot, label=step.label) + stream = step.streams[self._resident_key] + if stream.reads: + raise ValueError(f"recurrent pool {self._resident_key!r}: a resident step advances its state, not reads it") + key = CGSlotKey(bucket=step.bucket, slot=step.slot, label=stream.label) table = self._resident_table # a padding row reads the trash row's zero, which is SINK_SLOT gather_column(self._resident.store, step.slots, table.stride, table.col, self._cg_addressing[key].slot_indices) diff --git a/mstar/engine/resources/resident.py b/mstar/engine/resources/resident.py index a49b327229..8cfd483d9e 100644 --- a/mstar/engine/resources/resident.py +++ b/mstar/engine/resources/resident.py @@ -26,7 +26,7 @@ import logging import threading from collections.abc import Collection, Iterable, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, field from itertools import chain import numpy as np @@ -124,36 +124,61 @@ def cell(self, slot: int, col: int = 0) -> int: return slot * self.stride + self.col + col -class ResidentWrites: - """What one step changes. ``rows`` are rewritten whole by `flush`, from - their first column; ``cells`` and ``values`` are single cells, which ride - the graph's own scatter.""" +class ResidentStream: + """One resource's stream on a resident bucket's steps, fixed by its capture.""" - __slots__ = ("label", "rows", "cells", "values", "kv_len") + __slots__ = ("label", "reads") - def __init__(self, label: str): + def __init__(self, label: str, reads: bool): self.label = label - self.rows: list[tuple[int, list[int]]] = [] + # declared at span 0: the step reads the stream and appends nothing + self.reads = reads + + +class ResidentWrites: + """What one step changes. ``rows`` are runs (slot, first column, cells) + that `flush` writes whole; ``cells`` and ``values`` are single cells, which + ride the graph's own scatter. ``streams`` is each resource's stream, by + resource key.""" + + __slots__ = ("streams", "rows", "cells", "values", "kv_len") + + def __init__(self, streams: Mapping[str, ResidentStream]): + self.streams = streams + self.rows: list[tuple[int, int, list[int]]] = [] self.cells: list[int] = [] self.values: list[int] = [] - # the stored length of the row being rewritten, for the owners after the KV cache - self.kv_len = 0 + # each KV pool's stored length of the row being rewritten, for the owners after it + self.kv_len: dict[str, int] = {} def put(self, cell: int, value: int) -> None: self.cells.append(cell) self.values.append(value) + def run(self, slot: int, col: int, cells: list[int]) -> None: + """``slot``'s row from column ``col``. A run that starts where the one + before it ends joins it, so a row of narrow cells and one block table + is still one copy. Takes ``cells``, which a later run may extend.""" + rows = self.rows + if rows: + last_slot, last_col, last = rows[-1] + if last_slot == slot and last_col + len(last) == col: + last.extend(cells) + return + rows.append((slot, col, cells)) + @dataclass(slots=True) class ResidentBuffers: - """What one (bucket, slot) capture reads per step. ``cells`` is `[2, 2 * bs]`, - flat cells then values, so a step's single-cell writes are one copy: a row - has at most a length and a new column to write. Unused entries write zero + """What one (bucket, slot) capture reads per step. ``cells`` is flat cells + then values, so a step's single-cell writes are one copy: a row has at most + a length and a new column to write per KV pool. Unused entries write zero to cell zero, a trash cell, and ``dirty`` says whether any entry is in use.""" slots: torch.Tensor cells: torch.Tensor - # the rows' stored lengths this step: the graph writes them, for what reads them after - kv_len: torch.Tensor + # each KV pool's stored lengths of the rows this step, by resource key: the + # graph writes them, for what reads them after + kv_len: dict[str, torch.Tensor] # the host row ``slots`` holds: a steady batch keeps one array, so no H2D and no compare staged: np.ndarray | None = None dirty: bool = False @@ -171,8 +196,8 @@ class ResidentPlan: # rows written whole by this plan, and rows whose length alone it refreshes taken: set[int] moved: set[int] - # (rows, pages after the step) of the step's crossings, for `resident_advance` - crossed: tuple[np.ndarray, np.ndarray] | None = None + # each KV pool's (rows, pages after the step) of its crossings, for `resident_advance` + crossed: dict[str, tuple[np.ndarray, np.ndarray]] = field(default_factory=dict) @dataclass(slots=True) @@ -180,14 +205,17 @@ class CapturedResidentStep: """What one capture's kernels share, at warmup and capture only.""" bucket: BucketKey slot: int - label: str slots: torch.Tensor buffers: ResidentBuffers - # the store and the block table, for attention that reads a row's pages itself - store: torch.Tensor | None = None - block_table: ResidentTable | None = None - # set by an owner whose kernel runs off the tables alone: nothing plans this bucket - host_plan_free: bool = False + store: torch.Tensor + # each resource's stream, by resource key + streams: Mapping[str, ResidentStream] + # each KV pool's block table, by resource key, for attention that reads a row's pages itself + block_tables: dict[str, ResidentTable] = field(default_factory=dict) + # How each attention read its pages: off the tables alone, or through a + # host plan of the step. Nothing plans the bucket only if none needed one. + reads_tables: bool = False + host_planned: bool = False class ResidentPlanner: @@ -213,12 +241,17 @@ def __init__( self.device = device # request rows plus the trash row self.max_slots = 2 * max_bs + 1 + self._keys = tuple(resources) self._owners = tuple(resources.values()) self._tables: list[ResidentTable] = [] + # the KV pools whose stored lengths a capture gathers, as they bind + self._length_keys: list[str] = [] self._store: torch.Tensor | None = None # columns of the one-cell tables, which a rewritten row always covers self.narrow_cols = 0 self._buffers: dict[tuple[BucketKey, int], ResidentBuffers] = {} + # each resource's stream on a bucket's steps, by bucket then resource key, from its capture + self._streams: dict[BucketKey, dict[str, ResidentStream]] = {} # buckets whose steps need no host plan at all self.host_plan_free: set[BucketKey] = set() @@ -241,12 +274,12 @@ def __init__( self._last_rids: tuple[int, ...] | None = None self._last_slots: np.ndarray | None = None - for resource in self._owners: - resource.bind_resident(self) + for key, resource in resources.items(): + resource.bind_resident(self, key) self._seal() # the largest flush up front, plus a cells row: growing the stager waits # out its copies, which on the GPU thread would be a sync in the step path - self._stager = PinnedStager(torch.long, numel=max_bs * self._stride + 4 * max_bs) + self._stager = PinnedStager(torch.long, numel=max_bs * self._stride + 2 * self._cells_width(max_bs)) # Tables @@ -258,11 +291,16 @@ def table(self, width: int = 1) -> ResidentTable: self._tables.append(table) return table + def length_table(self, key: str) -> ResidentTable: + """KV pool ``key``'s table of stored lengths, which each capture + gathers into a buffer of its own (`ResidentBuffers.kv_len`).""" + self._length_keys.append(key) + return self.table() + def _seal(self) -> None: - # the one wide table last, so a row's used cells are contiguous from column 0 + # The narrow tables first, then a span per wide table, so a row's used + # cells are a run per span and the narrow cells join the first one. wide = [t for t in self._tables if t.width > 1] - if len(wide) > 1: - raise ValueError("a resident row holds one wide table") narrow = [t for t in self._tables if t.width == 1] self.narrow_cols = len(narrow) stride = self._stride = sum(t.width for t in self._tables) @@ -363,15 +401,21 @@ def mark_stale(self, rid: int) -> None: # The step + def _cells_width(self, bs: int) -> int: + """Single-cell writes a step of ``bs`` rows can queue: a length and a + new column a row per KV pool.""" + return 2 * bs * max(1, len(self._length_keys)) + def buffers(self, bucket: BucketKey, slot: int) -> ResidentBuffers: """The per-step rows of one capture, every slot at the trash row until staged.""" key = (bucket, slot) found = self._buffers.get(key) if found is None: + bs, device = bucket.bs, self.device found = self._buffers[key] = ResidentBuffers( - slots=torch.full((bucket.bs,), TRASH_SLOT, dtype=torch.long, device=self.device), - cells=torch.zeros((2, 2 * bucket.bs), dtype=torch.long, device=self.device), - kv_len=torch.zeros(bucket.bs, dtype=torch.long, device=self.device), + slots=torch.full((bs,), TRASH_SLOT, dtype=torch.long, device=device), + cells=torch.zeros((2, self._cells_width(bs)), dtype=torch.long, device=device), + kv_len={pool: torch.zeros(bs, dtype=torch.long, device=device) for pool in self._length_keys}, ) return found @@ -389,7 +433,7 @@ def plan(self, step: SubmoduleStep, ctx) -> None: return # the step it was for never ran, or something ran first self._restore(staged) - plan = self._build(rids, key, step.segments[0].label, ctx) + plan = self._build(rids, key, self._streams[lease.bucket], ctx) ctx.resident_plan = plan if ctx.is_preplan: self._staged = plan @@ -414,7 +458,7 @@ def _restore(self, plan: ResidentPlan) -> None: if buffers is not None: buffers.staged = None - def _build(self, rids: tuple[int, ...], key: tuple, label: str, ctx) -> ResidentPlan: + def _build(self, rids: tuple[int, ...], key: tuple, streams: Mapping[str, ResidentStream], ctx) -> ResidentPlan: lease = ctx.slot_lease laps = PhaseLaps(f"{ctx.plan_phase}.resident.") with self._lock: @@ -443,7 +487,7 @@ def _build(self, rids: tuple[int, ...], key: tuple, label: str, ctx) -> Resident refreshed -= taken self._last_read[slots] = seq rows = [(rid, self._slot_of[rid]) for rid in taken] - plan = ResidentPlan(key, rids, slots, seq, self._commits, ResidentWrites(label), taken, refreshed) + plan = ResidentPlan(key, rids, slots, seq, self._commits, ResidentWrites(streams), taken, refreshed) laps.lap("rows") try: for rid, slot in rows: @@ -466,12 +510,13 @@ def _build(self, rids: tuple[int, ...], key: tuple, label: str, ctx) -> Resident return plan def rewrite(self, rid: int, slot: int, writes: ResidentWrites) -> None: - """Queue ``rid``'s row whole. The narrow columns come first, then the - wide table's used cells, which its owner appends.""" + """Queue ``rid``'s row whole: the narrow columns as one run, which the + owners fill, and each wide table's used cells as a run its owner adds.""" row = [0] * self.narrow_cols + if row: + writes.run(slot, 0, row) for resource in self._owners: resource.resident_row(rid, slot, row, writes) - writes.rows.append((slot, row)) def _stage_cells(self, plan: ResidentPlan, buffers: ResidentBuffers) -> None: """One copy, on the plan stream: the graph scatters these first thing.""" @@ -491,9 +536,9 @@ def flush(self, plan: ResidentPlan) -> None: rows = plan.writes.rows if not rows: return - stride = self._stride - dsts = [self._store[slot * stride: slot * stride + len(row)] for slot, row in rows] - self._stager.copy_runs_(dsts, np.fromiter(chain.from_iterable(r for _, r in rows), dtype=np.int64)) + stride, store = self._stride, self._store + dsts = [store[slot * stride + col: slot * stride + col + len(cells)] for slot, col, cells in rows] + self._stager.copy_runs_(dsts, np.fromiter(chain.from_iterable(c for _, _, c in rows), dtype=np.int64)) def commit(self, step: SubmoduleStep, ctx) -> None: """After the resources' commit, before the step's commit event: the @@ -525,18 +570,13 @@ def _note_host_step(self, step: SubmoduleStep, ctx) -> None: def capture_step(self, bucket: BucketKey, slot: int, step: SubmoduleStep) -> CapturedResidentStep: """What a capture of ``bucket`` records around its forward. Refuses a - walk whose declared step isn't one token per row in one stream, each - row advancing alone: a captured step can't hold anything else.""" + walk whose declared step isn't, for each resource, one stream that it + appends one token a row to or only reads, each row advancing alone: a + captured step can't hold anything else.""" if step is None: raise ValueError(f"resident walk {bucket.graph_walk!r} on {self.node} declares no step") - segments = step.segments or next(iter(step.steps.values())).segments - labels = {seg.label for seg in segments} - if len(labels) != 1 or any(seg.span != 1 for seg in segments): - raise ValueError( - f"resident walk {bucket.graph_walk!r} on {self.node} declares " - f"{len(segments)} segments over labels {sorted(labels)}; a resident step is " - "one token per row in one stream" - ) + # once per bucket: each of its captures, and every step it replays, reads these + streams = self._streams.get(bucket) or self._declared_streams(bucket, step) for key, resource_step in step.steps.items(): refusal = next( ( @@ -551,13 +591,37 @@ def capture_step(self, bucket: BucketKey, slot: int, step: SubmoduleStep) -> Cap f"resident walk {bucket.graph_walk!r} on {self.node}: resource {key!r} declares " f"{refusal}, which a resident step can't carry" ) + self._streams[bucket] = streams buffers = self.buffers(bucket, slot) - return CapturedResidentStep(bucket, slot, labels.pop(), buffers.slots, buffers) + return CapturedResidentStep(bucket, slot, buffers.slots, buffers, self._store, streams) + + def _declared_streams(self, bucket: BucketKey, step: SubmoduleStep) -> dict[str, ResidentStream]: + """Each resource's stream on ``bucket``'s steps: the label of its own + segments, or of the step's when it names none, and whether it appends + one token a row (span 1) or only reads (span 0).""" + streams = {} + for key in self._keys: + resource_step = step.steps.get(key) + segments = (resource_step.segments if resource_step is not None else None) or step.segments or () + labels = {seg.label for seg in segments} + spans = {seg.span for seg in segments} + if len(labels) != 1 or spans not in ({1}, {0}): + raise ValueError( + f"resident walk {bucket.graph_walk!r} on {self.node}: resource {key!r} declares " + f"{len(segments)} segments over labels {sorted(labels)} and spans {sorted(spans)}; " + "a resident step is one token per row in one stream a resource, or a stream it only reads" + ) + streams[key] = ResidentStream(labels.pop(), reads=spans == {0}) + return streams def gather(self, step: CapturedResidentStep) -> None: + if self._length_keys: + # before the owners' reads: a row that crossed writes into its new + # page. Only a KV pool queues single cells. + scatter_cells(self._store, step.buffers.cells) for resource in self._owners: resource.resident_gather(step) - if step.host_plan_free: + if step.reads_tables and not step.host_planned: self.host_plan_free.add(step.bucket) def commit_in_graph(self, step: CapturedResidentStep) -> None: diff --git a/mstar/engine/resources/resident_hooks.py b/mstar/engine/resources/resident_hooks.py index b70d44a2e6..fc404f642e 100644 --- a/mstar/engine/resources/resident_hooks.py +++ b/mstar/engine/resources/resident_hooks.py @@ -22,13 +22,16 @@ def resident_refusal(self) -> str | None: """Why a node holding this resource can't run on resident tables, or None if it can.""" return f"{type(self).__name__} keeps no resident table" - def bind_resident(self, planner: ResidentPlanner) -> None: - """Declare this resource's tables on the node's planner, once, at load.""" + def bind_resident(self, planner: ResidentPlanner, key: str) -> None: + """Declare this resource's tables on the node's planner, once, at load. + ``key`` is its resource key, under which each step and capture name its + stream (`ResidentWrites.streams`, `CapturedResidentStep.streams`).""" return def resident_row(self, rid: int, slot: int, row: list[int], writes: ResidentWrites) -> None: """Set this resource's cells of ``rid``'s row: it is new or stale. ``row`` - has the planner's narrow columns; the owner of the wide table appends.""" + has the planner's narrow columns; the owner of a wide table queues its + used cells as a run of their own (`ResidentWrites.run`).""" return def resident_stage(self, plan: ResidentPlan) -> None: @@ -38,7 +41,9 @@ def resident_stage(self, plan: ResidentPlan) -> None: def resident_gather(self, step: CapturedResidentStep) -> None: """Record, inside the capture, the reads that fill what the step's forward - addresses. Runs at warmup and capture only; a replay repeats the kernels.""" + addresses. Runs at warmup and capture only; a replay repeats the kernels. + An attention sets `CapturedResidentStep.reads_tables` or ``host_planned``, + for whether its kernel reads its pages off the tables or a host plan.""" return def resident_commit(self, step: CapturedResidentStep) -> None: diff --git a/test/modular/test_resident_planning.py b/test/modular/test_resident_planning.py index f15d045231..902420e45c 100644 --- a/test/modular/test_resident_planning.py +++ b/test/modular/test_resident_planning.py @@ -577,7 +577,7 @@ def run(device): gather_column(store, slots, 7, 2, small, plus=out) advance_column(store, slots, 7, 1) tables = SimpleNamespace( - planner=SimpleNamespace(store=store), page_size=4, + planner=SimpleNamespace(store=store), page_size=4, key="kv", block_table=SimpleNamespace(stride=7, col=3), kv_len=SimpleNamespace(stride=7, col=1), ) store[slots * 7 + 1] = torch.tensor([9, 5, 0, 0], device=device) @@ -585,8 +585,8 @@ def run(device): offsets = torch.zeros(4, dtype=torch.long, device=device) step = SimpleNamespace( slots=slots, bucket=SimpleNamespace(bs=4), buffers=SimpleNamespace( - cells=torch.zeros((2, 8), dtype=torch.long, device=device), kv_len=out, - ), + cells=torch.zeros((2, 8), dtype=torch.long, device=device), kv_len={"kv": out}, + ), streams={"kv": resident_mod.ResidentStream("main", reads=False)}, block_tables={}, ) KVResidentTables.gather(tables, step, pages, offsets) return [t.tolist() for t in (store, out, small, pages, offsets)] @@ -897,7 +897,7 @@ def test_a_flush_of_the_largest_step_stages_without_growing(monkeypatch): stride = planner._tables[0].stride assert planner._stager._numel >= 4 * stride + 4 * 4 monkeypatch.setattr(planner._stager, "_grow", lambda numel: pytest.fail(f"grew to {numel}")) - rows = [(slot, list(range(stride))) for slot in range(1, 5)] + rows = [(slot, 0, list(range(stride))) for slot in range(1, 5)] plan = SimpleNamespace(writes=SimpleNamespace(rows=rows)) planner.flush(plan) assert planner.store.view(9, stride)[1:5].tolist() == [list(range(stride))] * 4 @@ -915,7 +915,7 @@ def test_a_flush_of_the_largest_step_stages_on_the_device_without_growing(monkey planner = ResidentPlanner("n", {WALK}, {"kv": kv}, device, max_bs=4) stride = planner._tables[0].stride monkeypatch.setattr(planner._stager, "_grow", lambda numel: pytest.fail(f"grew to {numel}")) - rows = [(slot, list(range(100, 100 + stride))) for slot in range(1, 5)] + rows = [(slot, 0, list(range(100, 100 + stride))) for slot in range(1, 5)] planner.flush(SimpleNamespace(writes=SimpleNamespace(rows=rows))) planner._stager.copy_(planner.buffers(BucketKey(WALK, 4, 4), 0).cells, [0] * 16) torch.cuda.synchronize() @@ -951,15 +951,15 @@ def test_a_stream_wider_than_a_row_raises(): planner = ResidentPlanner("n", {WALK}, {"kv": kv}, CPU, max_bs=1) tables = kv._resident assert tables.block_table.width == 2 + writes = resident_mod.ResidentWrites({"kv": resident_mod.ResidentStream("main", reads=False)}) with pytest.raises(ValueError, match="fit a block table row"): - tables.row(9, [1, 2, 3], 1, [0] * planner.narrow_cols) + tables.row(9, [1, 2, 3], 1, [0] * planner.narrow_cols, writes) plan = SimpleNamespace( - slots=torch.tensor([1]).numpy(), rids=(7,), taken=set(), writes=resident_mod.ResidentWrites("main"), - crossed=None, + slots=torch.tensor([1]).numpy(), rids=(7,), taken=set(), writes=writes, crossed={}, moved=set(), ) tables.host_len[1], tables.host_cols[1] = 8, 2 with pytest.raises(ValueError, match="fit a block table row"): - tables.stage_step(plan, lambda rid: [1, 2, 3], planner.rewrite) + tables.stage_step(plan, lambda rid: 8, lambda rid: [1, 2, 3], planner.rewrite) assert not plan.writes.cells @@ -1248,6 +1248,406 @@ def test_a_follower_replaying_an_offload_gives_the_slot_up_too(): assert planner.slot_of(1) is None +# ---- two KV pools on a node ---- + +def _pool(name: str, max_seq_len: int = 512, max_num_pages: int = 64) -> KVManager: + kv = KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=max_seq_len, + max_num_pages=max_num_pages, page_size=PAGE, + ), + name=name, joint_comm_group=None, transfer_engine_info=None, device=CPU, dtype=torch.float32, + ) + kv._cg_max_seq_len = 64 + return kv + + +def _decoder(positions: str = "kv"): + """A decoder's resources: the self-attention pool ``kv``, a second pool + ``ctx`` holding each request's context, and positions over ``positions``.""" + rope = RopeManager(PositionConfig(kv_cache=positions), CPU, max_seq_len=512, head_dim=8) + rope._pinned_pos_ids_buffer = lambda key, n: torch.empty(n, dtype=torch.long) + return {"kv": _pool("kv"), "ctx": _pool("ctx", max_seq_len=64), "rope": rope} + + +def _decoder_setup(prompts: dict[int, tuple[int, int]], positions: str = "kv"): + """Each request's prompt and context prefilled on resident resources and a + host-planned reference. ``prompts`` is rid -> (prompt tokens, context tokens).""" + live, ref = _decoder(positions), _decoder(positions) + planner = ResidentPlanner("n", {WALK}, {k: live[k] for k in ("kv", "ctx", "rope")}, CPU, max_bs=4) + for res in (live, ref): + for rid, (prompt, context) in prompts.items(): + res["kv"].ingest_request(rid, KVReqConfig()) + res["ctx"].ingest_request(rid, KVReqConfig()) + res["rope"].ingest_request(rid, None) + _decoder_host_step(res, [rid], [prompt], [context]) + return live, ref, planner + + +def _decoder_host_step(res, rids, spans, context_spans, walk="prefill"): + ctx = StepContext(request_ids=tuple(rids), graph_walk=walk, slot=0, capture=False) + segments = tuple(Segment(r, "main", n) for r, n in zip(rids, spans, strict=True)) + context = tuple(Segment(r, "main", n) for r, n in zip(rids, context_spans, strict=True)) + kv, cx, rope = KVStep(segments=segments), KVStep(segments=context), PositionStep(segments=segments) + assert res["kv"].admit(kv, ctx).ok + assert res["ctx"].admit(cx, ctx).ok + ctx.plan_results["kv"] = res["kv"].plan(kv, ctx) + ctx.plan_results["ctx"] = res["ctx"].plan(cx, ctx) + res["rope"].plan(rope, ctx) + for key, resource_step in (("kv", kv), ("ctx", cx), ("rope", rope)): + res[key].commit(resource_step, ctx) + return SubmoduleStep(steps={"kv": kv, "ctx": cx, "rope": rope}, segments=list(segments)), ctx + + +def _decoder_step(res, planner, rids, slot, bs, context_span=0): + """One captured decode step: ``kv`` appends a token a row, and ``ctx`` is read + at span 0, or appended to as well at span 1. With a planner, the replay of a + resident capture; without one, the host plan of the same step.""" + padded = tuple(rids) + tuple(-(i + 1) for i in range(bs - len(rids))) + bucket = BucketKey(graph_walk=WALK, bs=bs, num_tokens=bs) + lease = SlotLease(slot=slot, bucket=bucket) + ctx = StepContext(request_ids=tuple(rids), graph_walk=WALK, slot=slot, capture=False, slot_lease=lease) + ctx.set_padded_rids(padded) + segments = tuple(Segment(r, "main", 1) for r in padded) + context = tuple(Segment(r, "main", context_span) for r in padded) + kv, cx, rope = KVStep(segments=segments), KVStep(segments=context), PositionStep(segments=segments) + step = SubmoduleStep(steps={"kv": kv, "ctx": cx, "rope": rope}, segments=list(segments)) + assert res["kv"].admit(kv, ctx).ok + assert res["ctx"].admit(cx, ctx).ok + if (bucket, slot) not in res.setdefault("captured", set()): + res["captured"].add((bucket, slot)) + first = StepContext(request_ids=padded, graph_walk=WALK, slot=slot, capture=True, slot_lease=lease) + first.plan_results["kv"] = res["kv"].plan(kv, first) + first.plan_results["ctx"] = res["ctx"].plan(cx, first) + res["rope"].plan(rope, first) + captured = None + if planner is not None: + captured = planner.capture_step(bucket, slot, step) + ctx.resident = planner + planner.plan(step, ctx) + ctx.plan_results["kv"] = res["kv"].plan(kv, ctx) + ctx.plan_results["ctx"] = res["ctx"].plan(cx, ctx) + res["rope"].plan(rope, ctx) + if planner is not None: + planner.flush(ctx.resident_plan) + planner.gather(captured) + out = SimpleNamespace(plan=ctx.resident_plan, captured=captured, lens={}) + ck = CGSlotKey(bucket=bucket, slot=slot, label="main") + out.positions = res["rope"]._static_pos_ids[ck][:bs].tolist() + for key in ("kv", "ctx") if context_span else ("kv",): + state = res[key]._static_plan_state(slot, ("main", 0)) + out.lens[key] = None + setattr(out, f"{key}_pages", state.token_to_page[:bs].tolist()) + setattr(out, f"{key}_offsets", state.token_to_cache[:bs].tolist()) + if planner is not None: + out.lens = {key: lens[:bs].tolist() for key, lens in captured.buffers.kv_len.items()} + planner.commit_in_graph(captured) + for key, resource_step in (("kv", kv), ("ctx", cx), ("rope", rope)): + res[key].commit(resource_step, ctx) + if planner is not None: + planner.commit(step, ctx) + return out + + +def _assert_pool_row(planner, live, key: str, rid: int) -> None: + """``key``'s cells of ``rid``'s row hold its stream: the stored length, and + every page holding a token.""" + tables, stream = live[key]._resident, live[key]._streams[rid]["main"] + slot = planner.slot_of(rid) + n = int(tables.kv_len.view[slot, 0]) + assert n == stream.stored_len + assert tables.block_table.view[slot, : -(-n // PAGE)].tolist() == stream.page_indices[: -(-n // PAGE)] + + +def _assert_trash_row(live) -> None: + for key in ("kv", "ctx"): + tables = live[key]._resident + assert int(tables.kv_len.view[TRASH_SLOT, 0]) == 0 + assert tables.block_table.view[TRASH_SLOT].eq(SINK_PAGE).all() + + +def test_two_pools_lay_out_a_span_each_after_the_narrow_columns(): + live, _, planner = _decoder_setup({}) + kv, ctx = live["kv"]._resident, live["ctx"]._resident + assert planner.narrow_cols == 3 + assert kv.block_table.col == 3 and kv.block_table.width == 63 + assert ctx.block_table.col == 3 + 63 and ctx.block_table.width == 16 + assert planner._stride == 3 + 63 + 16 + assert {kv.kv_len.col, ctx.kv_len.col, live["rope"]._resident_delta.col} == {0, 1, 2} + for table in (kv.block_table, ctx.block_table, kv.kv_len, ctx.kv_len): + assert table.view.data_ptr() == planner.store.data_ptr() + 8 * table.col + + +def test_a_self_attention_pool_and_a_context_pool_run_through_captured_steps(): + """``kv`` appends a token a step, ``ctx`` is read: every step addresses what + the host plan does and gathers both pools' lengths, a padding row reads the + trash row, and after every step both rows hold their streams.""" + prompts = {1: (3, 9), 2: (6, 13), 3: (8, 1)} + live, ref, planner = _decoder_setup(prompts) + rids, contexts = [1, 2, 3], [9, 13, 1] + crossed = 0 + for t in range(14): + before = [live["kv"]._streams[rid]["main"].stored_len for rid in rids] + a = _decoder_step(live, planner, rids, slot=t % 2, bs=4) + b = _decoder_step(ref, None, rids, slot=t % 2, bs=4) + assert a.kv_pages[:3] == b.kv_pages[:3] and a.kv_offsets[:3] == b.kv_offsets[:3] + assert a.positions[:3] == b.positions[:3] + assert a.lens == {"kv": before + [0], "ctx": contexts + [0]} + # the padding row + assert (a.kv_pages[3], a.kv_offsets[3], a.positions[3]) == (SINK_PAGE, 0, 0) + assert set(a.plan.crossed) <= {"kv"} + crossed += "kv" in a.plan.crossed + for rid in rids: + _assert_pool_row(planner, live, "kv", rid) + _assert_pool_row(planner, live, "ctx", rid) + _assert_trash_row(live) + ctx = live["ctx"]._resident + assert [int(ctx.host_len[planner.slot_of(rid)]) for rid in rids] == contexts + assert crossed >= 3 + + +def test_a_crossing_on_the_appending_pool_writes_its_table_alone(): + live, ref, planner = _decoder_setup({1: (2, 9)}) + del ref + for t in range(2): + _decoder_step(live, planner, [1], slot=t % 2, bs=1) + kv, ctx = live["kv"]._resident, live["ctx"]._resident + slot = planner.slot_of(1) + held = planner.store.view(planner.max_slots, planner._stride)[slot].clone() + # stored 4 of a page of 4: this step's token goes to the row's second page + s = _decoder_step(live, planner, [1], slot=0, bs=1) + assert list(s.plan.crossed) == ["kv"] + assert s.plan.writes.cells == [kv.block_table.cell(slot, 1)] + row = planner.store.view(planner.max_slots, planner._stride)[slot] + changed = (row != held).nonzero().flatten().tolist() + assert changed == [kv.kv_len.col, kv.block_table.col + 1] + assert row[ctx.block_table.col: ctx.block_table.col + 3].tolist() == live["ctx"]._streams[1]["main"].page_indices + _assert_pool_row(planner, live, "kv", 1) + _assert_pool_row(planner, live, "ctx", 1) + + +def test_a_stale_row_is_rewritten_whole_with_both_spans(): + live, ref, planner = _decoder_setup({1: (3, 9), 2: (6, 13)}) + del ref + for t in range(3): + _decoder_step(live, planner, [1, 2], slot=t % 2, bs=2) + slot = planner.slot_of(2) + kv, ctx = live["kv"]._resident, live["ctx"]._resident + # whatever the device row held, a stale row comes back from the streams + planner.store.view(planner.max_slots, planner._stride)[slot].fill_(77) + planner.mark_stale(2) + s = _decoder_step(live, planner, [1, 2], slot=1, bs=2) + runs = [(col, len(cells)) for row_slot, col, cells in s.plan.writes.rows if row_slot == slot] + assert runs == [(0, 3 + 9 // PAGE + 1), (ctx.block_table.col, 4)] + assert kv.block_table.col == 3 + _assert_pool_row(planner, live, "kv", 2) + _assert_pool_row(planner, live, "ctx", 2) + assert int(live["rope"]._resident_delta.view[slot, 0]) == 0 + _assert_pool_row(planner, live, "kv", 1) + + +def test_a_context_moved_by_a_host_step_is_rewritten_and_an_unmoved_one_is_not(): + """A host step that ran a row one token leaves it to the next resident step + to refresh; the context pool compares its own stream, which nothing but a + host step moves, and rewrites the row only if it changed.""" + live, ref, planner = _decoder_setup({1: (3, 9), 2: (6, 13)}) + del ref + for t in range(2): + _decoder_step(live, planner, [1, 2], slot=t % 2, bs=2) + # one token on each self stream; request 2's context grows by a page's worth, 1's doesn't + step, ctx = _decoder_host_step(live, [1, 2], [1, 1], [0, 5], walk="mixed") + planner.commit(step, ctx) + assert planner._moved == {1, 2} + s = _decoder_step(live, planner, [1, 2], slot=0, bs=2) + assert s.plan.taken == {2} + assert s.lens["ctx"] == [9, 18] + for rid in (1, 2): + _assert_pool_row(planner, live, "kv", rid) + _assert_pool_row(planner, live, "ctx", rid) + + +def test_a_context_the_step_only_reads_has_no_advance_and_no_write_addressing(): + live, ref, planner = _decoder_setup({1: (3, 7)}) + del ref + ctx = live["ctx"]._resident + for t in range(6): + s = _decoder_step(live, planner, [1], slot=t % 2, bs=2) + slot = planner.slot_of(1) + assert int(ctx.kv_len.view[slot, 0]) == int(ctx.host_len[slot]) == 7 + assert int(ctx.host_cols[slot]) == 2 + assert s.captured.streams["ctx"].reads and not s.captured.streams["kv"].reads + assert "ctx" in s.captured.block_tables and "kv" in s.captured.block_tables + + +@pytest.mark.parametrize("positions", ["kv", "ctx"]) +def test_positions_follow_the_pool_their_config_names(positions): + """Two pools that both append, at different lengths: the position offset is + taken against, and gathered with, the length of the pool positions name.""" + live, ref, planner = _decoder_setup({1: (3, 11), 2: (9, 2)}, positions=positions) + for t in range(8): + a = _decoder_step(live, planner, [1, 2], slot=t % 2, bs=4, context_span=1) + b = _decoder_step(ref, None, [1, 2], slot=t % 2, bs=4, context_span=1) + assert a.positions[:2] == b.positions[:2] + for key in ("kv", "ctx"): + assert getattr(a, f"{key}_pages")[:2] == getattr(b, f"{key}_pages")[:2] + assert getattr(a, f"{key}_offsets")[:2] == getattr(b, f"{key}_offsets")[:2] + assert planner.buffers(BucketKey(WALK, 4, 4), 0).cells.shape == (2, 2 * 4 * 2) + + +def test_positions_over_a_pool_the_step_only_reads_refuse_at_capture(): + live, ref, planner = _decoder_setup({1: (3, 7)}, positions="ctx") + del ref + with pytest.raises(ValueError, match="advance with the stream"): + _decoder_step(live, planner, [1], slot=0, bs=1) + + +def test_capture_takes_a_label_for_each_resource(): + live, _, planner = _decoder_setup({}) + del live + rows = [Segment(1, "self", 1)] + step = SubmoduleStep( + steps={"kv": KVStep(), "ctx": KVStep(segments=(Segment(1, "context", 0),)), "rope": PositionStep()}, + segments=rows, + ) + captured = planner.capture_step(BucketKey(WALK, 1, 1), 0, step) + streams = {key: (s.label, s.reads) for key, s in captured.streams.items()} + assert streams == {"kv": ("self", False), "ctx": ("context", True), "rope": ("self", False)} + assert planner._streams[BucketKey(WALK, 1, 1)] is captured.streams + + +def test_capture_refuses_two_labels_on_one_resource(): + live, _, planner = _decoder_setup({}) + del live + rows = [Segment(1, "main", 1)] + two = (Segment(1, "main", 1), Segment(1, "aux", 1)) + step = SubmoduleStep(steps={"kv": KVStep(segments=two), "ctx": KVStep(), "rope": PositionStep()}, segments=rows) + with pytest.raises(ValueError, match="resource 'kv' declares 2 segments over labels"): + planner.capture_step(BucketKey(WALK, 1, 1), 0, step) + mixed = (Segment(1, "main", 1), Segment(2, "main", 0)) + step = SubmoduleStep(steps={"kv": KVStep(), "ctx": KVStep(segments=mixed), "rope": PositionStep()}, segments=rows) + with pytest.raises(ValueError, match="resource 'ctx'"): + planner.capture_step(BucketKey(WALK, 2, 2), 0, step) + + +class _Fmha: + """A captured fmha decode wrapper's gather, recorded.""" + + def __init__(self): + self.calls = [] + + def gather_resident(self, store, slots, kv_len, stride, block_col): + self.calls.append((kv_len, block_col)) + + +def _fmha_manager(pool: str, wrapper): + from mstar.engine.resources.attn.flashinfer import FlashInferManager + + manager = FlashInferManager.__new__(FlashInferManager) + manager._kv_cache_name = pool + manager._cg_plan_states = {CGSlotKey(bucket=BucketKey(WALK, 2, 2), slot=0, label="main"): wrapper} + return manager + + +@pytest.mark.parametrize("pool", ["kv", "ctx"]) +def test_the_fmha_gather_reads_its_own_pools_table_and_lengths(monkeypatch, pool): + from mstar.engine.resources.attn import flashinfer as flashinfer_mod + + monkeypatch.setattr(flashinfer_mod, "FmhaDecodeWrapper", _Fmha) + live, _, planner = _decoder_setup({}) + wrapper = _Fmha() + manager = _fmha_manager(pool, wrapper) + segments = [Segment(1, "main", 1), Segment(2, "main", 1)] + # both pools append here, so either can be read off the tables + step = SubmoduleStep(steps={"kv": KVStep(), "ctx": KVStep(), "rope": PositionStep()}, segments=segments) + captured = planner.capture_step(BucketKey(WALK, 2, 2), 0, step) + for key in ("kv", "ctx"): + live[key].resident_gather(captured) + manager.resident_gather(captured) + assert wrapper.calls == [(captured.buffers.kv_len[pool], live[pool]._resident.block_table.col)] + assert captured.reads_tables and not captured.host_planned + + +def test_an_fmha_wrapper_over_a_pool_the_step_only_reads_keeps_its_host_plan(monkeypatch): + from mstar.engine.resources.attn import flashinfer as flashinfer_mod + + monkeypatch.setattr(flashinfer_mod, "FmhaDecodeWrapper", _Fmha) + _, _, planner = _decoder_setup({}) + wrapper = _Fmha() + manager = _fmha_manager("ctx", wrapper) + segments = [Segment(1, "main", 1), Segment(2, "main", 1)] + context = (Segment(1, "main", 0), Segment(2, "main", 0)) + steps = {"kv": KVStep(), "ctx": KVStep(segments=context), "rope": PositionStep()} + captured = planner.capture_step(BucketKey(WALK, 2, 2), 0, SubmoduleStep(steps=steps, segments=segments)) + manager.resident_gather(captured) + assert not wrapper.calls and captured.host_planned + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the kernel runs on CUDA") +def test_the_fmha_gather_kernel_reads_the_second_pools_table(): + from mstar.engine.resources.attn.wrappers import FmhaDecodeWrapper + + device = torch.device("cuda") + store = torch.zeros(5 * 20, dtype=torch.long, device=device) + rows = store.view(5, 20) + # the first pool's table at columns 4-9 and the second's at 10-19, one row of each per slot + rows[:, 4:10] = torch.arange(5 * 6, device=device).view(5, 6) + 1000 + rows[:, 10:20] = torch.arange(5 * 10, device=device).view(5, 10) + 2000 + wrapper = FmhaDecodeWrapper.__new__(FmhaDecodeWrapper) + wrapper.batch_size, wrapper._page_size = 3, 4 + wrapper._table = torch.zeros((3, 10), dtype=torch.int32, device=device) + wrapper._lens = torch.zeros(3, dtype=torch.int32, device=device) + slots = torch.tensor([2, 4, 0], device=device) + kv_len = torch.tensor([5, 12, 0], device=device) + wrapper.gather_resident(store, slots, kv_len, 20, 10) + torch.cuda.synchronize() + assert wrapper._lens.tolist() == [6, 13, 1] + assert wrapper._table[0, :2].tolist() == rows[2, 10:12].tolist() + assert wrapper._table[1, :4].tolist() == rows[4, 10:14].tolist() + assert wrapper._table[2, :1].tolist() == rows[0, 10:11].tolist() + + +class _Attention(Resource): + @classmethod + def build(cls, spec, info): + return cls() + + def __init__(self, reads_tables: bool): + self._reads = reads_tables + + def resident_refusal(self): + return None + + def resident_gather(self, step): + if self._reads: + step.reads_tables = True + else: + step.host_planned = True + + +@pytest.mark.parametrize("attentions,free", [ + ((True,), True), ((True, True), True), ((True, False), False), ((False, True), False), ((), False), +]) +def test_a_bucket_is_host_plan_free_only_when_every_attention_reads_the_tables(attentions, free): + resources = {f"attn{i}": _Attention(reads) for i, reads in enumerate(attentions)} + planner = ResidentPlanner("n", {WALK}, resources, CPU, max_bs=2) + step = SubmoduleStep(steps={}, segments=[Segment(1, "main", 1)]) + planner.gather(planner.capture_step(BucketKey(WALK, 1, 1), 0, step)) + assert (BucketKey(WALK, 1, 1) in planner.host_plan_free) is free + + +def test_cross_attention_over_its_own_context_pool_runs_on_the_planner(): + from mstar.engine.resources.attn.cross import FlashInferCrossManager + + manager = FlashInferCrossManager.__new__(FlashInferCrossManager) + manager._kv_cache_name, manager._query_kv_cache_name = "ctx", "kv" + assert manager.resident_refusal() is None + manager._query_kv_cache_name = "ctx" + assert "second stream" in manager.resident_refusal() + step = SimpleNamespace(reads_tables=False, host_planned=False) + FlashInferCrossManager.resident_gather(manager, step) + assert step.host_planned and not step.reads_tables + + # ---- refusals ---- class _Spare(Resource): From 8c3af2a17d74355e0bc0c059c28ca80b54f58e6b Mon Sep 17 00:00:00 2001 From: merceod Date: Tue, 3 Nov 2026 04:45:56 +0000 Subject: [PATCH 13/15] tests: name the read-only context test for what it checks --- test/modular/test_resident_planning.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/modular/test_resident_planning.py b/test/modular/test_resident_planning.py index 902420e45c..26603c7b34 100644 --- a/test/modular/test_resident_planning.py +++ b/test/modular/test_resident_planning.py @@ -1336,7 +1336,6 @@ def _decoder_step(res, planner, rids, slot, bs, context_span=0): out.positions = res["rope"]._static_pos_ids[ck][:bs].tolist() for key in ("kv", "ctx") if context_span else ("kv",): state = res[key]._static_plan_state(slot, ("main", 0)) - out.lens[key] = None setattr(out, f"{key}_pages", state.token_to_page[:bs].tolist()) setattr(out, f"{key}_offsets", state.token_to_cache[:bs].tolist()) if planner is not None: @@ -1466,7 +1465,7 @@ def test_a_context_moved_by_a_host_step_is_rewritten_and_an_unmoved_one_is_not() _assert_pool_row(planner, live, "ctx", rid) -def test_a_context_the_step_only_reads_has_no_advance_and_no_write_addressing(): +def test_a_context_the_step_only_reads_never_advances(): live, ref, planner = _decoder_setup({1: (3, 7)}) del ref ctx = live["ctx"]._resident From f4b14f21d94c9e16649b5a91b47b03c8d9b444cb Mon Sep 17 00:00:00 2001 From: merceod Date: Tue, 3 Nov 2026 06:07:53 +0000 Subject: [PATCH 14/15] resident_walks: a stream a resource, or one it only reads --- mstar/model/submodule_base.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 9a938a60be..914482edb4 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -1103,10 +1103,11 @@ def uniform_row_inputs( def resident_walks(self) -> frozenset[str]: """Walks whose captured steps may read their rows off the resident - tables (`mstar.engine.resources.resident`). A walk here runs one token - per row in one stream, with no forks or retention, and its `preprocess` - reads no planned position or KV location: the forward's graph fills - those in. The engine still refuses if a resource can't.""" + tables (`mstar.engine.resources.resident`). A walk here runs, for each + resource, one stream that takes one token a row or one it only reads + (span 0), with no forks or retention, and its `preprocess` reads no + planned position or KV location: the forward's graph fills those in. + The engine still refuses if a resource can't.""" return frozenset() def step_is_input_free(self, graph_walk: str) -> bool: From cf7704024fbf0922a63049ba6f3ea2d6f8ca4cb8 Mon Sep 17 00:00:00 2001 From: merceod Date: Tue, 3 Nov 2026 10:19:37 +0000 Subject: [PATCH 15/15] Stub resident_walks in the workspace seal test The engine asks each submodule for its resident walks at warmup, before sizing; #1142's test builds its submodule as a SimpleNamespace. --- test/modular/test_attn_workspace.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/modular/test_attn_workspace.py b/test/modular/test_attn_workspace.py index 1117215833..545da5daa3 100644 --- a/test/modular/test_attn_workspace.py +++ b/test/modular/test_attn_workspace.py @@ -449,7 +449,8 @@ def post_warmup_validate(self): engine._resources = {"attn": _Resource()} engine._runner = None engine._submodules = {"dit": SimpleNamespace( - submodule=SimpleNamespace(warmup_compiles=lambda: None), resources={}, joint_comm_group=None, num_slots=2, + submodule=SimpleNamespace(warmup_compiles=lambda: None, resident_walks=frozenset), resources={}, + joint_comm_group=None, num_slots=2, fp32_precision=None, node_dtype=None, lora_pools=(), )} monkeypatch.setattr(engine_mod, "AcceleratorGraphRunner", _Runner)