from __future__ import annotations
from abc import ABC, abstractmethod
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, NamedTuple
import torch
from mstar.distributed.communication import JointGroups
from mstar.engine.resources.spec import NodeResourceSpec, ResourceReqConfig
from mstar.engine.resources.step import ADMIT_OK, AdmitOutcome, BucketKey, ResourceStep, StepContext
if TYPE_CHECKING:
# the config reaches back here through the submodule base, so keep the
# import out of module exec
from mstar.engine.accelerator_graph_config import (
AcceleratorGraphConfig,
PiecewiseAcceleratorGraphConfig,
)
from mstar.engine.resources.kv.transfer import TransferEngineInfo
[docs]
@dataclass(frozen=True)
class CGSlotSpec:
bucket: BucketKey
slot: int
# a whole forward's capture, or one piecewise region's
config: AcceleratorGraphConfig | PiecewiseAcceleratorGraphConfig
config_idx: int | None = None
@property
def bs(self):
return self.bucket.bs
@property
def num_tokens(self):
return self.bucket.num_tokens
def __str__(self) -> str:
return f"{self.bucket} slot={self.slot}"
[docs]
@dataclass(frozen=True)
class CGSlotKey:
bucket: BucketKey
slot: int
label: str
[docs]
@dataclass(frozen=True)
class EngineResourceInfo:
"""What the engine has to offer a resource at build time.
One struct rather than per-kind keyword arguments: a resource takes what it
needs and ignores the rest, and a name that does not exist here is a
TypeError rather than something silently swallowed by a ``**kwargs``.
"""
device: torch.device
joint_comm_group: JointGroups | None = None
transfer_engine_info: "TransferEngineInfo | None" = None
kv_dtype: torch.dtype = torch.bfloat16
# Whether this logical resource has consumers in another worker instance.
needs_remote_transfer: bool = True
# the specs this one named in `depends_on`, by resource key
dependencies: "Mapping[str, NodeResourceSpec]" = field(
default_factory=dict
)
# walk -> rows of it a step of its combined walk may hold (None uncapped),
# for every walk of this resource's local nodes
max_rows_in_combined: "Mapping[str, int | None]" = field(
default_factory=dict
)
[docs]
def dependency(self, key: str) -> NodeResourceSpec:
spec = self.dependencies.get(key)
if spec is None:
raise KeyError(
f"resource {key!r} was not resolved; declare it in the "
"spec's `depends_on`"
)
return spec
[docs]
class CachedPrefix(NamedTuple):
"""How much of a walk's input a resource already holds for ``label``.
``tokens`` counts the walk's leading inputs held, and ``position`` is the
stream's position counter once they are skipped.
"""
label: str
tokens: int
position: int
[docs]
class KeyedWalks(NamedTuple):
"""A keyed label's walks, as a resource's prefix cache takes them: the walks whose
writes are keyed, and the growth walks whose sampled ids extend the keys."""
writes: tuple[str, ...]
growth: tuple[str, ...] = ()
[docs]
class Resource(ABC):
[docs]
@classmethod
@abstractmethod
def build(
cls, spec: NodeResourceSpec, info: EngineResourceInfo,
) -> "Resource":
...
[docs]
def depends_on(self) -> set[str]:
return set()
[docs]
def can_capture(self) -> bool:
return True
# stays correct when a hit skips the front of a request; required of every resource on a keyed node
prefix_skip_safe: bool = False
[docs]
def enable_prefix_cache(self, root: bytes, walks: dict[str, KeyedWalks] | None = None) -> bool:
"""Open whatever this resource keeps across requests, under ``root``;
True if anything opened."""
return False
[docs]
def resolve_cached_prefix(
self, rid: str, node_name: str, graph_walk: str, inputs: Any,
) -> CachedPrefix | None:
"""How much of this walk's leading inputs this resource already holds.
None leaves the answer to other resources; 0 tokens means nothing can be skipped.
"""
return None
[docs]
def apply_cached_prefix(
self, rid: str, node_name: str, graph_walk: str,
inputs: Any, prefix: CachedPrefix | None,
) -> None:
"""Take on the prefix every resource agreed to skip, or None if there is none.
The inputs are untrimmed, so a resource that needs the tokens being
skipped can still read them.
"""
return
[docs]
def extend_prefix_chain(
self, rid: str, node_name: str, graph_walk: str, outputs: Any,
) -> None:
"""Take this step's sampled tokens, so what was generated can be keyed."""
return
@property
def keys_prefix_chains(self) -> bool:
"""Whether ``extend_prefix_chain`` can do anything right now (a prefix
cache is open). The runner skips the per-request sweep otherwise. A
resource that overrides ``extend_prefix_chain`` but not this is swept
every step."""
return False
[docs]
def complete_cached_walk(self, rid: str, node_name: str, graph_walk: str) -> None:
"""Called for a walk the cache served whole, since that walk never admits or commits."""
return
[docs]
def lease_cached_prefix(self, rid: str, node_name: str, graph_walk: str) -> list[int] | None:
"""On rank 0 of a TP node, lease the request's cached prefix as its step is sent and
return the pages, or None."""
return None
[docs]
def adopt_cached_prefix(
self, rid: str, node_name: str, graph_walk: str, pages: list[int], caught_up: bool,
) -> bool:
"""On a TP follower, lease ``pages`` as rank 0 did. False until this rank holds them too, and
raises if it still lacks them once ``caught_up``."""
return True
[docs]
def release_cached_prefix(self, rid: str, node_name: str, graph_walk: str) -> None:
"""Give back what `lease_cached_prefix` or `adopt_cached_prefix` took, for a step that will not run."""
return
[docs]
def page_trail(self) -> tuple[int, int] | None:
"""How many page allocations this resource has made and a hash of their pages in
order, for a TP follower to check against rank 0's. None where it keeps no trail."""
return None
[docs]
def fingerprint(self) -> bytes | None:
"""What this resource contributes to the prefix cache's root.
Everything a stored page depends on beyond its own tokens. None leaves
the resource out of the root, which is right only where it cannot
change what a page holds.
"""
return None
# Request lifecycle
[docs]
def ingest_request(self, rid: str, overrides: ResourceReqConfig | None):
"""Must be idempotent: the conductor sends one NewRequest per partition,
so a worker serving several partitions of a request ingests it once
for each, possibly after a node has already run a step for it."""
return
[docs]
def ingest_capture_row(self, rid: str):
"""Open a CUDA graph capture padding row; by default, as a request."""
self.ingest_request(rid, None)
[docs]
def remove_request(self, rid: str):
return
[docs]
def admit_retrieve(
self, rid: str,
node_name: str,
graph_walk: str,
published: "PublishedInfo | None"
) -> AdmitOutcome:
"""
Takes the output of publish, possibly from another device, and kicks
of a retrieval if needed (e.g., PD disaggregation KV transfer).
Returns whether the retrieve has completed.
"""
return ADMIT_OK
# Step lifecycle
[docs]
def admit(self, step: ResourceStep, ctx: StepContext) -> AdmitOutcome:
"""
Reserve space for the given step. In the case where requests in a batch must
be executed sequentially, this may be called for all requests in a loop
before the per-request plan -> forward -> commit cycle.
"""
return ADMIT_OK
[docs]
def rollback_admit(self, step: ResourceStep, ctx: StepContext) -> None:
"""Give back whatever the live reservation for this step took.
Called on a refused admit, for every resource the step reached including
the one that refused.
Only what was freshly taken. A page retained from a shared prefix, or a
lease converted onto a stream, belongs to whoever else holds it.
"""
return
[docs]
def plan(self, step: ResourceStep, ctx: StepContext) -> Any:
"""ret is immutable and opaque to runner; only gives to `ctx.plan_results`"""
return None
[docs]
def commit(self, step: ResourceStep, ctx: StepContext) -> None:
"""record step consumption"""
return
[docs]
def publish(self, request_id: str) -> "PublishedInfo | None":
return None
[docs]
def publish_for_step(
self,
request_id: str,
node_name: str | None,
graph_walk: str | None,
) -> "PublishedInfo | None":
"""Publish after one node step, with context for selective exporters.
Existing resources remain compatible through ``publish``; resources
that need node/walk context can override this method.
"""
del node_name, graph_walk
return self.publish(request_id)
[docs]
def publish_after_stop(
self,
request_id: str,
node_name: str | None,
graph_walk: str | None,
) -> "PublishedInfo | None":
"""Publish state that is useful only when a dynamic loop stops."""
del request_id, node_name, graph_walk
return None
[docs]
def publish_snapshot_batch(
self,
request_ids: list[str],
node_name: str | None,
graph_walk: str | None,
) -> list:
"""``publish_snapshot_for_step`` for every request of a step, one
entry per request in order (None for nothing to publish). A resource
with per-request locking or lookups overrides this to do them once."""
return [
self.publish_snapshot_for_step(
rid, node_name=node_name, graph_walk=graph_walk,
)
for rid in request_ids
]
[docs]
def publish_snapshot_for_step(
self,
request_id: str,
node_name: str | None,
graph_walk: str | None,
) -> Any:
"""A cheap record of what ``publish_for_step`` would export right now.
The worker finishes it with ``publish_from_snapshot`` only for the
requests whose frame or completion carries the publication, which in a
steady decode loop is none of them. The default takes the eager
publication itself, so a resource is correct without overriding; one
whose publish builds something per request (a page list, a transfer
descriptor) overrides both to record a few scalars here instead.
"""
return self.publish_for_step(
request_id, node_name=node_name, graph_walk=graph_walk,
)
[docs]
def publish_from_snapshot(
self,
request_id: str,
snapshot: Any,
node_name: str | None,
graph_walk: str | None,
) -> "PublishedInfo | None":
"""Finish a ``publish_snapshot_for_step`` record into the published info."""
del request_id, node_name, graph_walk
return snapshot
[docs]
def reset_request(self, rid: str, free: bool=False):
"""For clearing dummy RIDs during cuda graph capture"""
return
# Pre-planning
@property
def supports_preplan(self):
return False
@property
def force_double_buffer(self):
"""Whether this resource needs double-buffering even off the pre-plan path.
A resource that stages a step's layout into a reused host buffer and
issues a non-blocking H2D into a graph-read device buffer has a race the
moment the CPU runs ahead of the GPU: plan(N+1) can overwrite the buffer
before step N's DMA has retired, and the replay attends with N+1's data.
The buffer need not be ours — FlashInfer's ``plan`` holds one per
wrapper, which is why the attention managers key theirs per slot.
The main runner already double-buffers whenever a resource
``supports_preplan``; this flag extends that to a runner that has no
pre-plan path but still replays such a resource — notably the piecewise
runner — and makes ``_exec_per_request`` fence before it reuses a slot,
since its sub-steps all run inside one step. Off by default; a resource
opts in only if it has this hazard.
"""
return False
[docs]
def clear_preplan(self):
return
# Eviction
#
# A resource that holds enough per-request state to be worth reclaiming
# (the KV cache, today) opts in here; the worker picks victims and drives
# the move. Which requests exist and how recently they ran is the
# scheduler's knowledge, so the resource only answers for its own state.
@property
def supports_eviction(self):
return False
[docs]
def is_offloaded(self, rid: str) -> bool:
return False
[docs]
def offload(self, rid: str) -> int:
"""Move the request's state off-device. Returns what was reclaimed."""
return 0
[docs]
def reload(self, rid: str) -> bool:
"""Bring it back. False when it doesn't fit on device yet."""
return True
[docs]
def can_reload(self, rid: str) -> bool:
"""Whether the request can be reloaded. If this is False, the Engine
will not attempt to reload. If True, it will attempt to reload but still
check that the reload was successful."""
return True
[docs]
def reclaimable(self, rid: str) -> int:
"""What `offload` would free, in whatever this resource counts.
0 means the request holds nothing worth taking, so it is not an
eviction candidate however cold it is. Distinct from
`get_offload_priority`, which orders candidates rather than sizing them.
"""
return 0
[docs]
def get_offload_priority(self, rid: str) -> float:
"""How much this resource wants ``rid`` gone, higher being more.
Only consulted under a PRIORITY eviction policy, which names the
resource to ask; LRU never calls it.
"""
return 0.0
# Engine lifecycle
[docs]
def build_accelerator_graph_buffers(
self, slots: list[CGSlotSpec],
max_bs: int, max_seq_len: int
) -> None:
"""Size whatever the captured replays will read.
Called once per runner that captures against this node — the whole
forward's, and one per piecewise region — so it must tolerate repeated
calls: grow to the largest shape asked for, never clobber what an
earlier call already sized.
"""
# NOTE @nsagan: this should probably be refined; it was just the first
# thing that came to mind
return
[docs]
def post_warmup_validate(self):
"""
For, e.g., the KV cache to check that num_free_pages is identical
across TP ranks after cuda graph capture.
Raises an error (fails loudly) if invalid.
"""
return
[docs]
def cleanup(self):
return
[docs]
class AttentionResource(Resource):
"""A resource a layer stack calls per layer, under one plan label.
Adds the label / layer-index cursors: a caller running the whole stack sets
them once instead of threading them through every call, and an explicit
argument still supersedes. Class-level defaults so a subclass picks them up
without touching its __init__; subclasses that read them clear them in
`plan`, so a step that never binds cannot inherit the previous step's.
The KV, attention and cross-attention resources; not the sampler or the
position resource, which are called once per step rather than per layer.
"""
_default_label: str = "main"
_default_layer_idx: int | None = None
@property
def default_label(self) -> str:
return self._default_label
[docs]
@torch.compiler.disable
def set_default_label(self, label: str) -> None:
self._default_label = label
[docs]
@torch.compiler.disable
def set_default_layer_idx(self, layer_idx: int) -> None:
self._default_layer_idx = layer_idx
[docs]
def reset_default_cursors(self) -> None:
self._default_label = "main"
self._default_layer_idx = None
[docs]
class PublishedInfo(ABC):
[docs]
@abstractmethod
def update(self, other: "PublishedInfo") -> None:
...
[docs]
@abstractmethod
def clone(self) -> "PublishedInfo":
"""Return independent publication metadata for deferred sending."""
...
[docs]
def build_resource(spec: NodeResourceSpec, info: EngineResourceInfo) -> Resource:
return spec.resource_class.build(spec, info)