[Preemption 1/7] Slim Admission Control (A1)

docs/adding_models.rst

@@ -334,7 +334,10 @@ A layout has five fields:
 The footprint adds the spans, the growth at each loop's iterations and the scratch, and
 rounds up to whole pages per label. A span's own scratch counts only with the spans before
 it, so a window's denoise steps and its commit count once. Under a ``RetentionPolicy`` it
-takes the window's peak instead.
+takes the window's peak instead. ``KVReqConfig.validate`` refuses a request whose footprint
+is bigger than its pool with a 400 at intake, and ``KVReqConfig.loop_room`` cuts the default
+length of a request that asks for none until it fits. With every loop at 0 iterations, the
+same sum is the prompt a request reserves when its pool admits it.
 
 ``prompt_spans(plan, texts, text_walk, item_spans)`` builds the common case: a prompt's
 text and attachments in ``prefill_plan`` order. Qwen3.5 lays out its one label with it:
@@ -524,6 +527,9 @@ The spec types are:
        ``kv_lora_rank`` and ``qk_rope_head_dim``, and sets ``num_kv_heads`` to 1 and
        ``head_dim`` to their sum. Every rank holds the whole latent; only
        ``num_qo_heads`` shards, and MLA attention needs it set to the model's query heads.
+       A spec whose nodes can run on several workers, as Bagel's LLM nodes do under CFG
+       parallel, sets ``leader`` to the node whose pool admits each request. Left unset, each
+       worker's pool admits requests on its own.
    * - ``KVSpec(config=RingKVConfig(...))``
      - A fixed-capacity frame ring. It adds ``tokens_per_frame``, one
        ``RingKVLayerConfig`` per layer, and the deployment-tunable ``num_sessions``.
@@ -632,8 +638,10 @@ The spec types are:
    * - ``RecurrentStateSpec(config=RecurrentStateConfig(...))``
      - Fixed-size per-request state that carries between steps (an LSTM's ``h`` and
        ``c``, a cursor, a conv tail, a Mamba or gated-DeltaNet layer's conv window and SSM
-       state, a streaming encoder's caches): one slot per request and label, ``blocks``
-       naming a shape and dtype for each piece (dtype None: the engine's compute dtype,
+       state, a streaming encoder's caches): one slot per request and label, ``labels``
+       naming every label a request's steps address or fork into (``("main",)`` by default)
+       so that admission reserves a slot for each, ``blocks`` naming a shape and dtype for
+       each piece (dtype None: the engine's compute dtype,
        what its KV caches store), ``num_layers`` copies of each, ``max_slots`` slots
        (deployment-tunable) with slot 0 the sink that padding rows address. A recurrent
        kernel reads the step's ``addressing``. ``declare_step`` gives each row a
@@ -2025,7 +2033,7 @@ that its block loop can be CUDA-graph captured:
        return [
            RaggedAttentionSpec(resource_key=VIT_ATTN, nodes={"vit_encoder"},
                                config=RaggedAttentionConfig(...)),
-           KVSpec(resource_key="kv", nodes=nodes, config=self._kv_config()),
+           KVSpec(resource_key="kv", nodes=nodes, config=self._kv_config(), leader="LLM"),
            AttentionSpec(resource_key="attn", nodes=nodes,
                          config=AttentionConfig(kv_cache="kv")),
            PositionSpec(resource_key="rope", nodes=nodes,

docs/environment_variables.rst

@@ -616,11 +616,12 @@ Worker scheduling
        worker main loop (speculate, await_gpu, submit_spec, ...).
    * - ``MSTAR_KV_DEBUG_ASSERTS``
      - ``0``
-     - ``1``: after every ``admit``, ``commit``, ``reset_request`` and
-       ``remove_request``, check the KV page bookkeeping against the streams
-       holding the pages, and each commit against its request's layout. A recurrent
-       snapshot's key is checked against the KV chain's key at its boundary. Logs a
-       label left under its footprint at removal. Tests and debugging only.
+     - ``1``: after every ``admit``, ``commit``, ``reset_request``, ``remove_request``
+       and change to an admitted request's pages, check the KV page bookkeeping against
+       the streams holding the pages, the reservations against the free and evictable
+       pages, and each commit against its request's layout. A recurrent snapshot's key is
+       checked against the KV chain's key at its boundary. Logs a label left under its
+       footprint at removal. Tests and debugging only.
 
 Compilation
 -----------

docs/serving.rst

@@ -380,9 +380,14 @@ ranks adopt it. ``prefix_cache_tp`` can't be combined with ``cpu_offload_pages``
 cache's nodes must run on one TP group, with no replicas and no walks placed on other
 ranks. Loading fails otherwise.
 
-A cached prompt no longer reserves its pages, so the pool no longer limits how many
-requests are admitted. It has to hold the decode of every request allowed to run at once,
-and a decode step that finds no free page holds its requests until they time out.
+On a single rank, a request waits at the head of its pool's queue until its prompt fits
+beside the prompts already admitted, and no request behind it starts first. A cache hit on
+a page only the prefix index holds counts as a page the request takes. A request that
+can't fit in its pool at its full output length gets a 400 if it asked for that length,
+and otherwise has the model's default length cut until it fits. The pool
+reserves only prompts, so it still has to hold the decode of every request allowed to run
+at once, and a decode step that finds no page outside the admitted prompts holds its
+requests until they time out.
 ``max_concurrent_requests`` prevents that. It caps how many requests run at once, and the
 pool's pages divided by the pages one request needs, prompt and output together, is a
 safe value. ``cpu_offload_pages`` lets a full pool move a request to the host. An offload

mstar/conductor/conductor.py

@@ -9,7 +9,7 @@
 import time
 from collections import defaultdict, deque
 from collections.abc import Collection
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from typing import NamedTuple
 
 import numpy as np
@@ -652,32 +652,42 @@ def _loop_caps(self, model_kwargs: dict) -> dict[str, int]:
                 )
         return {name: caps.get(name, n) for name, n in self._loop_max_iters.items()}
 
-    def _resource_specs(self) -> dict:
-        """The model's resource specs by key, after this deployment's YAML, for admission-time checks."""
+    def _resource_specs(self) -> list[dict]:
+        """Each worker's resource specs by key, after this deployment's YAML, for admission-time checks.
+        A spec's ``nodes`` is cut to the worker's own, as the worker's engine does when it builds the resource."""
         specs = getattr(self, "_resource_specs_cache", None)
         if specs is None:
-            built = self.model.get_node_resources() or []
-            apply_yaml_overrides(built, self.model_config)
-            specs = {s.resource_key: s for s in built}
-            self._resource_specs_cache = specs
+            built = self._node_resources()
+            specs = self._resource_specs_cache = [
+                {s.resource_key: replace(s, nodes=s.nodes & nodes) for s in built if s.nodes & nodes}
+                for nodes in {frozenset(self._worker_nodes(worker_id)) for worker_id in self._per_worker_graphs}
+            ]
         return specs
 
+    def _node_resources(self) -> list:
+        """The model's resource specs after this deployment's YAML, each with all of its nodes."""
+        built = getattr(self, "_node_resources_cache", None)
+        if built is None:
+            built = self._node_resources_cache = self.model.get_node_resources() or []
+            apply_yaml_overrides(built, self.model_config)
+        return built
+
     def _streamed_partitions(self) -> frozenset[str]:
         streamed = getattr(self, "_streamed_partitions_cache", None)
         if streamed is None:
-            streamed = streamed_partitions(list(self._resource_specs().values()), self.model)
+            streamed = streamed_partitions(self._node_resources(), self.model)
             self._streamed_partitions_cache = streamed
         return streamed
 
     def _fit_loop_caps(self, model_kwargs: dict, configs: dict, cap: int) -> tuple[int, dict[str, int]]:
         """The largest ``max_output_tokens`` up to ``cap``, and its loop caps, that keeps every loop
         within what the request's resources hold. A loop that ``max_output_tokens`` doesn't drive is
         left as it is, and ``validate`` then refuses the request if it doesn't fit."""
-        specs = self._resource_specs()
         room: dict[str, int] = {}
-        for key, cfg in configs.items():
-            for name, n in cfg.loop_room(specs.get(key)).items():
-                room[name] = min(n, room.get(name, n))
+        for specs in self._resource_specs():
+            for key, cfg in configs.items():
+                for name, n in cfg.loop_room(specs.get(key)).items():
+                    room[name] = min(n, room.get(name, n))
 
         def overrun(caps: dict[str, int]) -> int:
             return max((caps[name] - n for name, n in room.items() if name in caps), default=0)
@@ -1379,9 +1389,9 @@ def _do_ingest_request(
             )
             for key, cfg in request_data.resource_configs.items():
                 cfg.apply_conductor_config(seed=seed, resource_key=key, loop_caps=loop_caps)
-        specs = self._resource_specs()
-        for key, cfg in request_data.resource_configs.items():
-            cfg.validate(specs.get(key))
+        for specs in self._resource_specs():
+            for key, cfg in request_data.resource_configs.items():
+                cfg.validate(specs.get(key))
 
         # Send NewRequest to each worker with the appropriate partition's inputs
         request_data.dispatched = True

mstar/engine/engine.py

@@ -548,6 +548,7 @@ def load_model(
                             spec.nodes, relevant_nodes
                         )
                     ),
+                    nodes=frozenset(relevant_nodes),
                     dependencies={
                         key: specs_by_key[key] for key in spec.depends_on()
                     },
@@ -2341,6 +2342,13 @@ def check_ready(
             published=request_info.resource_publish_info,
         )
 
+    def admit_waiting(
+        self, published_of: Callable[[str, str], Mapping[str, PublishedInfo] | None],
+    ) -> list[tuple[str, str]]:
+        """The (node, rid) pairs admitted since `check_ready` answered them ``queued``; see
+        `Admission.admit_waiting`."""
+        return self._runner.admit_waiting(published_of)
+
     def reserve_replay_slot(self, batch: ExecutingBatch) -> SlotLease | None:
         """Lease the slot this batch will replay on, before it is dispatched.
 

mstar/engine/resources/admission.py

@@ -0,0 +1,187 @@
+"""Admission across the resources of one runner, see `Resource.ask`."""
+
+from __future__ import annotations
+
+import bisect
+import itertools
+import logging
+from collections.abc import Callable, Collection, Mapping
+from dataclasses import dataclass, field
+
+from mstar.engine.resources.base import PublishedInfo, Resource
+from mstar.engine.resources.step import FULL_ADMIT_QUEUED, FullAdmitOutcome
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class _Waiting:
+    """A waiting request: the node and walk that asked each resource it waits for, and every node
+    whose readiness check waits on it, with what the request had published for that node then."""
+    asked: dict[str, tuple[str, str]] = field(default_factory=dict)
+    nodes: set[str] = field(default_factory=set)
+    published: dict[str, Mapping[str, PublishedInfo] | None] = field(default_factory=dict)
+
+
+class Admission:
+    """The queues of the resources that ask, and the nodes each request was admitted on.
+
+    A waiting request is admitted once it is first in every queue it is in and each resource it
+    asked reserves for it, all or nothing.
+    """
+
+    def __init__(self, resources: Mapping[str, Resource], ask_order: Collection[str]) -> None:
+        self._resources = resources
+        # `_arrival` numbers each request in the order it first asked.
+        # `_waiting` holds what each waiting request asked.
+        # `_queues` holds, per resource, the requests waiting for it in arrival order.
+        # `_short` maps each resource whose first waiting request didn't fit to the mark it didn't fit
+        # against, until that request leaves the front of the queue.
+        # `_untried` is set when a queue has a new first request that `admit_waiting` hasn't tried.
+        # `_admitted` holds the nodes each request was admitted on.
+        # `_started` holds, for a request admitted during one node's check, the other nodes waiting on
+        # it, which `admit_waiting` reports so their parked checks resume.
+        self._arrivals = itertools.count()
+        self._arrival: dict[str, int] = {}
+        self._waiting: dict[str, _Waiting] = {}
+        self._queues: dict[str, list[str]] = {key: [] for key in ask_order}
+        self._short: dict[str, object] = {}
+        self._untried = False
+        self._admitted: dict[str, set[str]] = {}
+        self._started: dict[str, set[str]] = {}
+
+    def admitted(self, rid: str, node_name: str) -> bool:
+        return node_name in self._admitted.get(rid, ())
+
+    def admit(
+        self, rid: str, node_name: str, graph_walk: str, published: Mapping[str, PublishedInfo] | None,
+        keys: Collection[str],
+    ) -> FullAdmitOutcome | None:
+        """None once ``rid`` is admitted on ``node_name``, else why it waits or fails. ``keys`` are the
+        resources of ``node_name`` that ask."""
+        entry = self._waiting.get(rid)
+        waits = False
+        for key in keys:
+            if entry is not None and key in entry.asked:
+                waits = True
+                continue
+            asked = self._resources[key].ask(rid, node_name, graph_walk)
+            if asked is None:
+                continue
+            if not asked.ok:
+                logger.warning("Admission for resource %s failed with error: %s", key, asked.reason.message)
+                if entry is not None:
+                    # a request a resource can't serve waits in no queue
+                    del self._waiting[rid]
+                    self._leave(rid, entry.asked)
+                return FullAdmitOutcome(asked, key)
+            entry = self._park_request(rid, key, node_name, graph_walk)
+            waits = True
+        if not waits:
+            self._admitted.setdefault(rid, set()).add(node_name)
+            return None
+        entry.nodes.add(node_name)
+        entry.published[node_name] = published
+        tried = all(self._queues[key][0] == rid and key not in self._short for key in entry.asked)
+        if tried and (nodes := self._reserve(rid, entry.published.get)):
+            if nodes - {node_name}:
+                self._started.setdefault(rid, set()).update(nodes - {node_name})
+            return None
+        return FULL_ADMIT_QUEUED
+
+    def admit_waiting(
+        self, published_of: Callable[[str, str], Mapping[str, PublishedInfo] | None],
+    ) -> list[tuple[str, str]]:
+        """Admit each waiting request that is now first in every queue it is in and fits, and return
+        the (node, rid) pairs that may start. ``published_of(rid, node)`` gives what the request
+        published for that node's resources.
+
+        Nothing is tried until a resource whose first waiting request didn't fit has grown, or a
+        queue has a new first request, so a pass that changed neither costs one
+        `Resource.grew_since` for each such resource.
+        """
+        started = []
+        if self._started:
+            started = [(node, rid) for rid, nodes in self._started.items() for node in nodes]
+            self._started.clear()
+        if not self._waiting:
+            return started
+        grown = [key for key, mark in self._short.items() if self._resources[key].grew_since(mark)]
+        if not grown and not self._untried:
+            return started
+        for key in grown:
+            del self._short[key]
+        admitted = True
+        while admitted:
+            admitted = False
+            for queue in self._queues.values():
+                if not queue:
+                    continue
+                rid = queue[0]
+                entry = self._waiting[rid]
+                if any(self._queues[key][0] != rid or key in self._short for key in entry.asked):
+                    continue
+                nodes = self._reserve(rid, lambda node, rid=rid: published_of(rid, node))
+                if nodes:
+                    started.extend((node, rid) for node in nodes)
+                    admitted = True
+        # the last sweep admitted nothing, so it tried every queue's first request
+        self._untried = False
+        return started
+
+    def remove_request(self, rid: str) -> None:
+        entry = self._waiting.pop(rid, None)
+        if entry is not None:
+            self._leave(rid, entry.asked)
+        self._admitted.pop(rid, None)
+        self._started.pop(rid, None)
+        self._arrival.pop(rid, None)
+
+    def _park_request(self, rid: str, key: str, node_name: str, graph_walk: str) -> _Waiting:
+        """Put ``rid`` in ``key``'s queue at its arrival, and return what it waits for."""
+        entry = self._waiting.get(rid)
+        if entry is None:
+            entry = self._waiting[rid] = _Waiting()
+            self._arrival.setdefault(rid, next(self._arrivals))
+        entry.asked[key] = (node_name, graph_walk)
+        queue = self._queues[key]
+        bisect.insort(queue, rid, key=self._arrival.__getitem__)
+        if queue[0] == rid:
+            # ahead of the request the mark measured, so `admit` tries it unless another queue holds it back
+            self._short.pop(key, None)
+        return entry
+
+    def _reserve(
+        self, rid: str, published_of: Callable[[str], Mapping[str, PublishedInfo] | None],
+    ) -> set[str] | None:
+        """Take everything waiting ``rid`` asked for and return the nodes waiting on it, or take
+        nothing and return None. ``published_of`` gives what ``rid`` published for a node's
+        resources."""
+        entry = self._waiting[rid]
+        taken: list[str] = []
+        for key, (node_name, graph_walk) in entry.asked.items():
+            published = published_of(node_name)
+            mark = self._resources[key].reserve(
+                rid, node_name, graph_walk, None if published is None else published.get(key),
+            )
+            if mark is not None:
+                for done in taken:
+                    self._resources[done].unreserve(rid)
+                self._short[key] = mark
+                return None
+            taken.append(key)
+        del self._waiting[rid]
+        self._leave(rid, entry.asked)
+        self._admitted.setdefault(rid, set()).update(entry.nodes)
+        return entry.nodes
+
+    def _leave(self, rid: str, keys: Collection[str]) -> None:
+        """Take ``rid`` out of the queues of ``keys``. A queue it was first in drops its mark,
+        which measured ``rid``, and its next request is left for `admit_waiting` to try."""
+        for key in keys:
+            queue = self._queues[key]
+            if queue[0] == rid:
+                self._short.pop(key, None)
+                if len(queue) > 1:
+                    self._untried = True
+            queue.remove(rid)